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
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247/* ssm_batch_test: M1.6 S4 — ssm_forward_b (batched SSM block over M tokens) * vs the decode per-token oracle (ssm_forward called M times, one per token). * * Synthesizes a dev_ssm_blk (random ternary Q2 weights + conv1d/a_log/dt_bias) * and runs both paths on the same x[M][D_MODEL], comparing the residual-stream * output AND the final ssm_state/ssm_conv (the recurrent state must also match * since gpu_deltanet_rmsnorm_b and gpu_conv1d_step_b are bit-exact per row). * * Gate (per plan.md S4): mean|d| of ssm output <= ~0.1; conv1d bit-exact; * greedy-token match. * * Build: meson compile -C build ssm_batch_test && ./build/ssm_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(0x554DULL);
/* ---- Q2 synthesis + repack ---- */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; /* Pad M to 256 (gemm_gpu_v4 WGM8=32 * 8 = 256-row A-load tile; see * upload_q2 in qxmx_gpu.cpp). Without this, alpha/beta (M=48) cause * OOB A_codes/A_sc_t reads -> DEVICE_LOST. */ int M_pad = (M + 255) & ~255; size_t codes_sz = (size_t)M_pad * ngrp * 16; size_t sc_sz = ngrp * M + 256; uint8_t* codes_h = (uint8_t*)calloc(1, codes_sz); float* sc_h = (float*)calloc(1, 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 float* upload_f32(sycl::queue& q, const std::vector<float>& v) { float* d = sycl::malloc_shared<float>(v.size(), q); q.memcpy(d, v.data(), v.size() * sizeof(float)).wait(); 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; }}
int main() { sycl::queue q(sycl::gpu_selector_v, sycl::property::queue::in_order{}); printf("ssm_batch_test (M1.6 S4): ssm_forward_b vs decode ssm_forward oracle\n");
/* ---- synthesize ssm block ---- */ dev_ssm_blk s; { 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); s.attn_norm = upload_f32(q, w); } s.qkv = upload_synth(q, synth_q2_weight(QX_SSM_CONVD, QX_D_MODEL), QX_SSM_CONVD, QX_D_MODEL); s.alpha = upload_synth(q, synth_q2_weight(QX_SSM_DT_RANK, QX_D_MODEL), QX_SSM_DT_RANK, QX_D_MODEL); s.beta = upload_synth(q, synth_q2_weight(QX_SSM_DT_RANK, QX_D_MODEL), QX_SSM_DT_RANK, QX_D_MODEL); s.gate = upload_synth(q, synth_q2_weight(QX_SSM_VALD, QX_D_MODEL), QX_SSM_VALD, QX_D_MODEL); s.out = upload_synth(q, synth_q2_weight(QX_D_MODEL, QX_SSM_VALD), QX_D_MODEL, QX_SSM_VALD); { std::normal_distribution<float> g(0.0f, 1.0f); std::vector<float> conv(QX_SSM_CONV_K * QX_SSM_CONVD); for (auto& v : conv) v = 0.1f * g(rng); s.conv1d = upload_f32(q, conv); std::vector<float> al(QX_SSM_DT_RANK); for (int h = 0; h < QX_SSM_DT_RANK; h++) al[h] = -0.5f - 0.1f * (float)(h % 5); s.a_log = upload_f32(q, al); std::vector<float> dtb(QX_SSM_DT_RANK); for (auto& v : dtb) v = 0.5f * g(rng); s.dt_bias = upload_f32(q, dtb); std::vector<float> nw(QX_SSM_STATE); for (int i = 0; i < QX_SSM_STATE; i++) nw[i] = 1.0f + 0.1f * g(rng); s.ssm_norm = upload_f32(q, nw); }
const int Mmax = 256; const size_t state_n = (size_t)QX_SSM_DT_RANK * QX_SSM_STATE * QX_SSM_STATE; const size_t conv_n = (size_t)(QX_SSM_CONV_K - 1) * QX_SSM_CONVD;
/* initial recurrent state (same for both paths) */ std::vector<float> s0_init(state_n), c0_init(conv_n); { std::normal_distribution<float> g(0.0f, 1.0f); for (auto& v : s0_init) v = 0.05f * g(rng); for (auto& v : c0_init) v = 0.1f * g(rng); }
/* batched + decode workspaces */ 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 * QX_SSM_CONVD, q); float* alpha_b = sycl::malloc_shared<float>((size_t)Mmax * QX_SSM_DT_RANK, q); float* beta_b = sycl::malloc_shared<float>((size_t)Mmax * QX_SSM_DT_RANK, q); float* ssm_out_b = sycl::malloc_shared<float>((size_t)Mmax * QX_SSM_VALD, q); float* zg_b = sycl::malloc_shared<float>((size_t)Mmax * QX_SSM_VALD, q); float* x_b = sycl::malloc_shared<float>((size_t)Mmax * QX_D_MODEL, q);
int8_t* dec_bvn = sycl::malloc_shared<int8_t>(QX_FF_INTER, q); int32_t* dec_sum = sycl::malloc_shared<int32_t>(QX_FF_INTER / 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_SSM_CONVD, q); float* dec_proj = sycl::malloc_shared<float>(QX_SSM_CONVD, q); float* dec_alpha = sycl::malloc_shared<float>(QX_SSM_DT_RANK, q); float* dec_beta = sycl::malloc_shared<float>(QX_SSM_DT_RANK, q); float* dec_ssm_out = sycl::malloc_shared<float>(QX_SSM_VALD, q); float* dec_zg = sycl::malloc_shared<float>(QX_SSM_VALD, q); float* dec_x = sycl::malloc_shared<float>(QX_D_MODEL, q);
float* state_b = sycl::malloc_shared<float>(state_n, q); float* state_d = sycl::malloc_shared<float>(state_n, q); float* conv_b = sycl::malloc_shared<float>(conv_n, q); float* conv_d = sycl::malloc_shared<float>(conv_n, q);
int fail = 0; for (int M : {16, 64, 256}) { printf("\n=== M=%d ===\n", M); fill_tile(x_b, M, QX_D_MODEL, 1); 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)); /* reset recurrent state */ memcpy(state_b, s0_init.data(), state_n * sizeof(float)); memcpy(state_d, s0_init.data(), state_n * sizeof(float)); memcpy(conv_b, c0_init.data(), conv_n * sizeof(float)); memcpy(conv_d, c0_init.data(), conv_n * sizeof(float));
/* ---- batched ---- */ ssm_forward_b(q, &s, state_b, conv_b, x_b, M, bvn_d, sum_d, asc_d, proj_b, alpha_b, beta_b, ssm_out_b, zg_b); q.wait();
/* ---- oracle: per-token decode ---- */ 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)); ssm_forward(q, &s, state_d, conv_d, dec_x, dec_part, dec_proj, dec_alpha, dec_beta, dec_ssm_out, dec_zg); q.wait(); memcpy(x_oracle.data() + (size_t)tok * QX_D_MODEL, dec_x, QX_D_MODEL * sizeof(float)); }
/* ---- compare x, state, conv ---- */ size_t n = (size_t)M * QX_D_MODEL; double sx = 0; float wx = 0; for (size_t i = 0; i < n; i++) { double d = fabs(x_b[i] - x_oracle[i]); sx += d; if (d > wx) wx = (float)d; } double mx = sx / n; double ss = 0; float ws = 0; for (size_t i = 0; i < state_n; i++) { double d = fabs(state_b[i] - state_d[i]); ss += d; if (d > ws) ws = (float)d; } double ms = ss / state_n; double sc = 0; float wc = 0; for (size_t i = 0; i < conv_n; i++) { double d = fabs(conv_b[i] - conv_d[i]); sc += d; if (d > wc) wc = (float)d; } double mc = sc / conv_n; bool pass = (mx <= 0.1) && (ms <= 0.1) && (mc <= 1e-3); printf(" x: mean|d|=%.3e worst=%.3e\n", mx, wx); printf(" state: mean|d|=%.3e worst=%.3e\n", ms, ws); printf(" conv: mean|d|=%.3e worst=%.3e (gate: <=1e-3; bit-exact in isolation)\n", mc, wc); printf(" %s\n", pass ? "PASSED" : "FAILED"); if (!pass) fail++; }
auto free_devq2 = [&](dev_q2& d){ sycl::free(d.codes, q); sycl::free(d.scales, q); }; free_devq2(s.qkv); free_devq2(s.alpha); free_devq2(s.beta); free_devq2(s.gate); free_devq2(s.out); sycl::free(s.attn_norm, q); sycl::free(s.conv1d, q); sycl::free(s.a_log, q); sycl::free(s.dt_bias, q); sycl::free(s.ssm_norm, q); sycl::free(bvn_d, q); sycl::free(sum_d, q); sycl::free(asc_d, q); sycl::free(proj_b, q); sycl::free(alpha_b, q); sycl::free(beta_b, q); sycl::free(ssm_out_b, q); sycl::free(zg_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(dec_proj, q); sycl::free(dec_alpha, q); sycl::free(dec_beta, q); sycl::free(dec_ssm_out, q); sycl::free(dec_zg, q); sycl::free(dec_x, q); sycl::free(state_b, q); sycl::free(state_d, q); sycl::free(conv_b, q); sycl::free(conv_d, q);
printf(fail ? "\nFAILED: %d config(s)\n" : "\nALL PASSED — mean|d| within gate vs decode oracle\n", fail); return fail ? 1 : 0;}