Something went wrong. Try again.
A single-model (Bonsai-27B) inference engine for Intel Arc GPUs
Something went wrong. Try again.
C++
at extra-debug
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232/* ffn_batch_test: M1.6 S1 — ffn_forward_b (batched FFN over M tokens via * gemm_gpu_v4) vs the decode per-token oracle (ffn_forward called M times). * * Synthesizes a dev_ffn_blk with random ternary Q2 weights in the exact GGUF * layout (same pattern as gemm_engine_test), runs both paths on the same * random x[M][D_MODEL], and compares the residual-stream output. * * Gate (per plan.md S1): mean|d| of FFN output <= ~0.05 vs per-token decode. * NOT bit-exact — the batched GEMM reassociates fp32 differently than the * per-token gemv reduction (ground rule 2: greedy-token gate, drift OK). * * Build: meson compile -C build ffn_batch_test && ./build/ffn_batch_test */#include <algorithm>#include <chrono>#include <cmath>#include <cstdint>#include <cstdio>#include <cstdlib>#include <cstring>#include <random>#include <vector>
#include <sycl/sycl.hpp>
#include "qxmx.h"#include "qxmx_gpu.h"#include "qxmx_gemm.h"
using namespace qx;
static std::mt19937_64 rng(0xF1FFULL);
/* ---- Q2 weight synthesis + repack (matches gemm_engine_test) ------------ */static uint16_t f32_to_f16_bits(float f) { return sycl::bit_cast<uint16_t>(sycl::half(f));}static float f16_to_f32_local(uint16_t bits) { return float(sycl::bit_cast<sycl::half>(bits));}
static std::vector<qx_block_q2> synth_q2_weight(int M, int K) { int ngrp = K / QX_QK2; std::vector<qx_block_q2> W((size_t)M * ngrp); std::normal_distribution<float> g(0.0f, 1.0f); for (int m = 0; m < M; m++) for (int gg = 0; gg < ngrp; gg++) { float sc = 0.05f * std::fabs(g(rng)) + 0.001f; W[(size_t)m * ngrp + gg].d = f32_to_f16_bits(sc); for (int b = 0; b < QX_QK2 / 4; b++) { uint8_t pack = 0; for (int k = 0; k < 4; k++) { int q = rng() % 3; /* {0,1,2} -> value {-1,0,+1} */ pack |= (uint8_t)(q << (k * 2)); } W[(size_t)m * ngrp + gg].qs[b] = pack; } } return W;}
/* Repack to A_codes [M/8][ngrp][2][64B] + A_sc_t [ngrp][M] (decode upload_q2 * layout, shared with gemm_gpu_v4 per M1.6 layout discovery). */static dev_q2 upload_synth(sycl::queue& q, const std::vector<qx_block_q2>& W, int M, int K) { int ngrp = K / QX_QK2; size_t codes_sz = (size_t)(M / 8) * ngrp * 128; size_t sc_sz = (size_t)ngrp * M; uint8_t* codes_h = (uint8_t*)aligned_alloc(32, codes_sz); float* sc_h = (float*)aligned_alloc(32, sc_sz * sizeof(float)); for (int m8 = 0; m8 < M / 8; m8++) for (int gg = 0; gg < ngrp; gg++) for (int half = 0; half < 2; half++) for (int r = 0; r < 8; r++) memcpy(codes_h + (((size_t)(m8 * ngrp + gg) * 2 + half) * 64 + r * 8), W[(m8 * 8 + r) * ngrp + gg].qs + half * 8, 8); for (int m = 0; m < M; m++) for (int gg = 0; gg < ngrp; gg++) sc_h[gg * M + m] = f16_to_f32_local(W[m * ngrp + gg].d); dev_q2 d; d.codes = sycl::malloc_shared<uint8_t>(codes_sz, q); d.scales = sycl::malloc_shared<float>(sc_sz, q); q.memcpy(d.codes, codes_h, codes_sz).wait(); q.memcpy(d.scales, sc_h, sc_sz * sizeof(float)).wait(); free(codes_h); free(sc_h); d.M = M; d.K = K; d.src = nullptr; return d;}
/* Fill an [n_tok][dim] tile with varied row scales (matches the * gemm_engine_test distribution: rows pick scales across decades, with a * zero-row and an outlier to exercise amax edge cases). */static void fill_tile(float* p, int n_tok, int dim, int seed_off) { std::normal_distribution<float> g(0.0f, 1.0f); for (int t = 0; t < n_tok; t++) { float sc = (((t + seed_off) % 17) == 3) ? 0.0f : powf(10.0f, -2.0f + 3.0f * (((t + seed_off) % 13) / 12.0f)); for (int i = 0; i < dim; i++) p[(size_t)t * dim + i] = g(rng) * sc; if (((t + seed_off) % 31) == 5) p[(size_t)t * dim + 7] = 123.0f; }}
struct Cfg { int M; const char* note; };
int main() { sycl::queue q(sycl::gpu_selector_v, sycl::property::queue::in_order{}); printf("ffn_batch_test (M1.6 S1): ffn_forward_b vs decode ffn_forward oracle\n");
/* ---- synthesize FFN block 0 shapes ---- */ const int Md = QX_D_MODEL; /* 5120 */ const int Mf = QX_FF_INTER; /* 17408 */ auto W_gate = synth_q2_weight(Mf, Md); auto W_up = synth_q2_weight(Mf, Md); auto W_down = synth_q2_weight(Md, Mf);
dev_ffn_blk f; f.gate = upload_synth(q, W_gate, Mf, Md); f.up = upload_synth(q, W_up, Mf, Md); f.down = upload_synth(q, W_down, Md, Mf); std::vector<float> post_norm(Md); { std::normal_distribution<float> g(0.0f, 1.0f); for (int i = 0; i < Md; i++) post_norm[i] = 1.0f + 0.1f * g(rng); } f.post_norm = sycl::malloc_shared<float>(Md, q); q.memcpy(f.post_norm, post_norm.data(), Md * sizeof(float)).wait();
/* ---- batched workspace (sized for the largest M we test) ---- */ const int Mmax = 256; /* bvnni_d layout = [K/4][M] int32 = K*M bytes; max K = FF_INTER = 17408. */ int8_t* bvn_d = sycl::malloc_shared<int8_t>((size_t)Mf * Mmax, q); int32_t* sum_d = sycl::malloc_shared<int32_t>((size_t)(Mf / QX_QK2) * Mmax, q); float* asc_d = sycl::malloc_shared<float>(Mmax, q); float* g_b = sycl::malloc_shared<float>((size_t)Mmax * Mf, q); float* u_b = sycl::malloc_shared<float>((size_t)Mmax * Mf, q); float* t_b = sycl::malloc_shared<float>((size_t)Mmax * Mf, q); float* x_b = sycl::malloc_shared<float>((size_t)Mmax * Md, q);
/* ---- decode oracle workspace (single-token) ---- */ /* The decode path uses globals via engine_set_quant_stage. Device-sized * (malloc_device would suffice; shared is fine and simpler to inspect). */ int8_t* dec_bvn = sycl::malloc_shared<int8_t>(Mf, q); int32_t* dec_sum = sycl::malloc_shared<int32_t>(Mf / QX_QK2, q); float* dec_asc = sycl::malloc_shared<float>(1, q); engine_set_quant_stage(dec_bvn, dec_sum, dec_asc); float* dec_part = sycl::malloc_shared<float>((size_t)8 * Mf, q); float* x_dec = sycl::malloc_shared<float>(Md, q); float* dec_g = sycl::malloc_shared<float>(Mf, q); float* dec_u = sycl::malloc_shared<float>(Mf, q); float* dec_t = sycl::malloc_shared<float>(Mf, q);
Cfg cfgs[] = { {16, "smallest valid GEMM block"}, {64, "mid chunk"}, {256, "full chunk"}, };
int fail = 0; for (auto c : cfgs) { int M = c.M; printf("\n=== M=%d (%s) ===\n", M, c.note);
/* seed x identically for both paths. */ fill_tile(x_b, M, Md, 1); std::vector<float> x_save((size_t)M * Md); memcpy(x_save.data(), x_b, (size_t)M * Md * sizeof(float));
/* ---- batched: ffn_forward_b on the M tokens ---- */ ffn_forward_b(q, &f, x_b, M, bvn_d, sum_d, asc_d, g_b, u_b, t_b); q.wait();
/* ---- oracle: decode ffn_forward M times, row by row ---- */ std::vector<float> x_oracle((size_t)M * Md); for (int tok = 0; tok < M; tok++) { memcpy(x_dec, x_save.data() + (size_t)tok * Md, Md * sizeof(float)); ffn_forward(q, &f, x_dec, dec_part, dec_g, dec_u, dec_t); q.wait(); memcpy(x_oracle.data() + (size_t)tok * Md, x_dec, Md * sizeof(float)); }
/* ---- compare ---- */ double sum_abs = 0; float worst = 0; size_t n = (size_t)M * Md; for (size_t i = 0; i < n; i++) { double d = fabs(x_b[i] - x_oracle[i]); sum_abs += d; if (d > worst) worst = (float)d; } double mean_abs = sum_abs / n; /* reference magnitude for context */ double ref_mag = 0; for (size_t i = 0; i < n; i++) ref_mag += fabs(x_oracle[i]); ref_mag /= n; bool pass = mean_abs <= 0.05; printf(" mean|d|=%.3e worst=%.3e (mean|x|=%.3e) %s\n", mean_abs, worst, ref_mag, pass ? "PASSED" : "FAILED (gate 0.05)"); if (!pass) fail++; }
/* ---- perf: M=256 back-to-back to amortize clock ramp ---- */ { int M = 256, REPS = 30; q.wait(); auto t0 = std::chrono::steady_clock::now(); for (int r = 0; r < REPS; r++) ffn_forward_b(q, &f, x_b, M, bvn_d, sum_d, asc_d, g_b, u_b, t_b); q.wait(); auto t1 = std::chrono::steady_clock::now(); double ms = std::chrono::duration<double, std::milli>(t1 - t0).count() / REPS; /* 3 GEMMs: 2x (Mf*Md) + 1x (Md*Mf), per token. */ double macs = (double)(2 * Mf * Md + Md * Mf); double tflops = macs * M / (ms / 1e3) / 1e12; double tokseq = M / (ms / 1e3); printf("\nperf M=256: %.2f ms/FFN-block, %.1f TFLOPs, %.0f tok/s (FFN only)\n", ms, tflops, tokseq); }
/* cleanup */ auto free_devq2 = [&](dev_q2& d){ sycl::free(d.codes, q); sycl::free(d.scales, q); }; free_devq2(f.gate); free_devq2(f.up); free_devq2(f.down); sycl::free(f.post_norm, q); sycl::free(bvn_d, q); sycl::free(sum_d, q); sycl::free(asc_d, q); sycl::free(g_b, q); sycl::free(u_b, q); sycl::free(t_b, q); sycl::free(x_b, q); sycl::free(dec_bvn, q); sycl::free(dec_sum, q); sycl::free(dec_asc, q); sycl::free(dec_part, q); sycl::free(x_dec, q); sycl::free(dec_g, q); sycl::free(dec_u, q); sycl::free(dec_t, q);
printf(fail ? "\nFAILED: %d config(s)\n" : "\nALL PASSED — mean|d| within gate vs decode oracle\n", fail); return fail ? 1 : 0;}