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
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248/* attn_forward_batch_test: M1.6 S3 — attn_forward_b (batched attention block * via gemm_gpu_v4 + S2 elementwise _b + M1.4 attn_prefill) vs the decode * per-token oracle (attn_forward called M times, one per position). * * Synthesizes a dev_attn_blk (random ternary Q2 weights in GGUF layout) and * runs both paths on the same x[M][D_MODEL], comparing the residual-stream * output. Exercises the full attn pipeline: rmsnorm_quant -> q/k/v GEMMs -> * split_qgate -> group_rmsnorm -> RoPE -> KV-cache quant writes -> * attn_prefill -> sigmoid_mul -> o-proj residual fold. * * Gate (per plan.md S3): mean|d| of attn output <= ~0.1 (GEMM reassociation * compounds through 4 GEMMs + attention); greedy-token match. * * Build: meson compile -C build attn_forward_batch_test && ./build/attn_forward_batch_test */#include <algorithm>#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(0xA77A04ULL);
/* ---- Q2 synthesis + repack (shared with ffn_batch_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++) pack |= (uint8_t)((rng() % 3) << (k * 2)); W[(size_t)m * ngrp + gg].qs[b] = pack; } } return W;}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;}
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; }}
/* Run one (M, k_ctk) config; return 0 on pass, 1 on fail. */static int run_cfg(sycl::queue& q, const dev_attn_blk* a, int M, engine::KCT k_ctk, const char* kname, uint8_t* kvkA, uint8_t* kvkB, uint32_t* kvvA, uint32_t* kvvB, int max_ctx, int Mmax) { /* zero caches so writes are byte-clean */ size_t krow = (k_ctk == engine::KCT::fp16) ? KROW_BYTES_F16 : KROW_BYTES_Q8; memset(kvkA, 0, (size_t)max_ctx * krow); memset(kvkB, 0, (size_t)max_ctx * krow); memset(kvvA, 0, (size_t)max_ctx * VROW_U32 * 4); memset(kvvB, 0, (size_t)max_ctx * VROW_U32 * 4);
/* ---- batched workspace ---- * bvnni_d / suma_d must cover the largest projection K used here. For * attn that's max(D_MODEL=5120 for q/k/v, Q_DIM=6144 for o-proj) = 6144; * size for FF_INTER=17408 to match the engine's uniform sizing (S5). */ int8_t* bvn_d = sycl::malloc_shared<int8_t>((size_t)QX_FF_INTER * Mmax, q); int32_t* sum_d = sycl::malloc_shared<int32_t>((size_t)(QX_FF_INTER / QX_QK2) * Mmax, q); float* asc_d = sycl::malloc_shared<float>(Mmax, q); float* proj_b = sycl::malloc_shared<float>((size_t)Mmax * 2 * QX_Q_DIM, q); float* qv_b = sycl::malloc_shared<float>((size_t)Mmax * QX_Q_DIM, q); float* gv_b = sycl::malloc_shared<float>((size_t)Mmax * QX_Q_DIM, q); float* kv_b = sycl::malloc_shared<float>((size_t)Mmax * QX_KV_DIM, q); float* vv_b = sycl::malloc_shared<float>((size_t)Mmax * QX_KV_DIM, q); float* out_b = sycl::malloc_shared<float>((size_t)Mmax * QX_Q_DIM, q); float* rc = sycl::malloc_shared<float>((size_t)Mmax * (QX_ROT_DIM / 2), q); float* rs = sycl::malloc_shared<float>((size_t)Mmax * (QX_ROT_DIM / 2), q); float* x_b = sycl::malloc_shared<float>((size_t)Mmax * QX_D_MODEL, q); fill_tile(x_b, M, QX_D_MODEL, 1); engine_compute_rope_table(0, M, rc, rs);
/* ---- decode oracle workspace ---- */ int8_t* dec_bvn = sycl::malloc_shared<int8_t>(QX_D_MODEL, q); int32_t* dec_sum = sycl::malloc_shared<int32_t>(QX_D_MODEL / 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 * QX_Q_DIM, q); float* dec_proj = sycl::malloc_shared<float>(2 * QX_Q_DIM, q); float* dec_q = sycl::malloc_shared<float>(QX_Q_DIM, q); float* dec_g = sycl::malloc_shared<float>(QX_Q_DIM, q); float* dec_k = sycl::malloc_shared<float>(QX_KV_DIM, q); float* dec_v = sycl::malloc_shared<float>(QX_KV_DIM, q); float* dec_out = sycl::malloc_shared<float>(QX_Q_DIM, q); float* dec_x = sycl::malloc_shared<float>(QX_D_MODEL, q); float* dec_rc = sycl::malloc_shared<float>(QX_ROT_DIM / 2, q); float* dec_rs = sycl::malloc_shared<float>(QX_ROT_DIM / 2, q);
/* save x for both paths */ std::vector<float> x_save((size_t)M * QX_D_MODEL); memcpy(x_save.data(), x_b, (size_t)M * QX_D_MODEL * sizeof(float));
/* ---- batched ---- */ fa_split_workspace fa_ws = fa_split_alloc(q, Mmax); attn_forward_b(q, a, x_b, M, kvkB, kvvB, /*pos_start=*/0, k_ctk, bvn_d, sum_d, asc_d, proj_b, qv_b, gv_b, kv_b, vv_b, out_b, rc, rs, &fa_ws); q.wait(); fa_split_free(q, fa_ws);
/* ---- oracle: per-token decode at pos=0..M-1 ---- */ std::vector<float> x_oracle((size_t)M * QX_D_MODEL); for (int tok = 0; tok < M; tok++) { memcpy(dec_x, x_save.data() + (size_t)tok * QX_D_MODEL, QX_D_MODEL * sizeof(float)); engine_compute_rope_table(tok, 1, dec_rc, dec_rs); attn_forward(q, a, dec_x, kvkA, kvvA, /*pos=*/tok, k_ctk, dec_part, dec_proj, dec_q, dec_g, dec_k, dec_v, dec_out, dec_rc, dec_rs); q.wait(); memcpy(x_oracle.data() + (size_t)tok * QX_D_MODEL, dec_x, QX_D_MODEL * sizeof(float)); }
/* ---- compare residual stream ---- */ double sum_abs = 0; float worst = 0; size_t n = (size_t)M * QX_D_MODEL; 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; 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.1; printf(" K=%-5s M=%-4d : mean|d|=%.3e worst=%.3e (mean|x|=%.3e) %s\n", kname, M, mean_abs, worst, ref_mag, pass ? "PASSED" : "FAILED (gate 0.1)");
sycl::free(bvn_d, q); sycl::free(sum_d, q); sycl::free(asc_d, q); sycl::free(proj_b, q); sycl::free(qv_b, q); sycl::free(gv_b, q); sycl::free(kv_b, q); sycl::free(vv_b, q); sycl::free(out_b, q); sycl::free(rc, q); sycl::free(rs, 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(dec_proj, q); sycl::free(dec_q, q); sycl::free(dec_g, q); sycl::free(dec_k, q); sycl::free(dec_v, q); sycl::free(dec_out, q); sycl::free(dec_x, q); sycl::free(dec_rc, q); sycl::free(dec_rs, q); return pass ? 0 : 1;}
int main() { sycl::queue q(sycl::gpu_selector_v, sycl::property::queue::in_order{}); printf("attn_forward_batch_test (M1.6 S3): attn_forward_b vs decode attn_forward oracle\n");
/* ---- synthesize attn block shapes ---- */ dev_attn_blk a; a.attn_norm = sycl::malloc_shared<float>(QX_D_MODEL, q); { std::normal_distribution<float> g(0.0f, 1.0f); std::vector<float> w(QX_D_MODEL); for (int i = 0; i < QX_D_MODEL; i++) w[i] = 1.0f + 0.1f * g(rng); q.memcpy(a.attn_norm, w.data(), QX_D_MODEL * sizeof(float)).wait(); } a.q = upload_synth(q, synth_q2_weight(2 * QX_Q_DIM, QX_D_MODEL), 2 * QX_Q_DIM, QX_D_MODEL); a.k = upload_synth(q, synth_q2_weight(QX_KV_DIM, QX_D_MODEL), QX_KV_DIM, QX_D_MODEL); a.v = upload_synth(q, synth_q2_weight(QX_KV_DIM, QX_D_MODEL), QX_KV_DIM, QX_D_MODEL); a.o = upload_synth(q, synth_q2_weight(QX_D_MODEL, QX_Q_DIM), QX_D_MODEL, QX_Q_DIM); a.q_norm = sycl::malloc_shared<float>(QX_HEAD_DIM, q); a.k_norm = sycl::malloc_shared<float>(QX_HEAD_DIM, q); for (int t = 0; t < 2; t++) { std::vector<float> w(QX_HEAD_DIM); std::normal_distribution<float> g(0.0f, 1.0f); for (int i = 0; i < QX_HEAD_DIM; i++) w[i] = 1.0f + 0.1f * g(rng); float* dst = t ? a.k_norm : a.q_norm; q.memcpy(dst, w.data(), QX_HEAD_DIM * sizeof(float)).wait(); }
/* KV caches: one per K format (different row strides), each per-path * (decode A vs batched B). Use a small max_ctx since we test pos_start=0. */ int Mmax = 256, max_ctx = Mmax; size_t krow_max = KROW_BYTES_F16; /* fp16 is the larger */ uint8_t* kvkA = sycl::malloc_shared<uint8_t>((size_t)max_ctx * krow_max, q); uint8_t* kvkB = sycl::malloc_shared<uint8_t>((size_t)max_ctx * krow_max, q); uint32_t* kvvA = sycl::malloc_shared<uint32_t>((size_t)max_ctx * VROW_U32, q); uint32_t* kvvB = sycl::malloc_shared<uint32_t>((size_t)max_ctx * VROW_U32, q);
int fail = 0; struct Cfg { int M; engine::KCT k; const char* name; }; for (auto c : std::vector<Cfg>{ {16, engine::KCT::q8_0, "q8_0"}, {64, engine::KCT::q8_0, "q8_0"}, {256, engine::KCT::q8_0, "q8_0"}, {256, engine::KCT::fp16, "fp16"}, {256, engine::KCT::fp8, "fp8"}, }) { fail += run_cfg(q, &a, c.M, c.k, c.name, kvkA, kvkB, kvvA, kvvB, max_ctx, Mmax); }
auto free_devq2 = [&](dev_q2& d){ sycl::free(d.codes, q); sycl::free(d.scales, q); }; free_devq2(a.q); free_devq2(a.k); free_devq2(a.v); free_devq2(a.o); sycl::free(a.attn_norm, q); sycl::free(a.q_norm, q); sycl::free(a.k_norm, q); sycl::free(kvkA, q); sycl::free(kvkB, q); sycl::free(kvvA, q); sycl::free(kvvB, q);
printf(fail ? "\nFAILED: %d config(s)\n" : "\nALL PASSED — mean|d| within gate vs decode oracle\n", fail); return fail ? 1 : 0;}