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
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226/* gemm_engine_test: M1.6 go/no-go — port gemm_gpu_v4 into the engine and * validate (a) correctness vs a CPU fp32 reference and (b) in-engine sustained * TFLOPs with a real activation producer (gpu_rmsnorm_quant_b) feeding B. * * Shapes from the real model (FFN gate/up/down, attn qkv): * ffn_gate/up: M=17408, K=5120 attn_q: M=12288, K=5120 * ffn_down: M=5120, K=17408 * * We can't load the real weights here, so we synthesize ternary Q2 weights * (codes {-1,0,1}, random f16 scale per 64) in the exact GGUF layout, run * rmsnorm_quant_b to produce the B operand, run the GEMM, and compare. * * Build: meson compile -C build gemm_engine_test && ./build/gemm_engine_test */#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(0x6E6E6EULL);
/* f32 -> f16 bits (software; matches the GGUF/ggml_half layout). */static uint16_t f32_to_f16_bits(float f) { sycl::half h = sycl::half(f); return sycl::bit_cast<uint16_t>(h);}
/* f16 bits -> f32 (software; matches qxmx_gpu.cpp::f16_to_f32). */static float f16_to_f32_local(uint16_t bits) { return float(sycl::bit_cast<sycl::half>(bits));}
/* Synthesize a Q2_0_g64 weight matrix in GGUF layout (qx_block_q2 array): * M rows, K cols (K multiple of 64). codes in {-1,0,1} -> q in {0,1,2}. */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 Q2 weight -> engine A_codes layout [M/8][ngrp][2][64B] + A_sc_t [ngrp][M]. * Identical to qxmx_gpu.cpp::repack_q2 (kept local so the test is standalone). */static void repack_q2_local(const qx_block_q2* src, int M, int K, uint8_t* codes, float* scales_t) { int ngrp = K / QX_QK2; 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 + (((size_t)(m8 * ngrp + gg) * 2 + half) * 64 + r * 8), src[(m8 * 8 + r) * ngrp + gg].qs + half * 8, 8); for (int m = 0; m < M; m++) for (int gg = 0; gg < ngrp; gg++) scales_t[gg * M + m] = f16_to_f32_local(src[m * ngrp + gg].d);}
/* CPU fp32 reference: C = W @ rmsnorm(x_in), matching the GPU path which * feeds rmsnorm_quant_b(x_in) into the GEMM. rmsnorm is fp32 here; the GPU * adds s8 quantization noise on top (~1e-3 mean|d| expected). */static void rmsnorm_row(const float* x, const float* w, float* out, int K) { double ss = 0; for (int k = 0; k < K; k++) ss += (double)x[k] * x[k]; float inv = 1.0f / sqrtf((float)(ss / K) + QX_RMS_EPS); for (int k = 0; k < K; k++) out[k] = x[k] * inv * w[k];}
static void cpu_ref_row(const qx_block_q2* W, const float* b, float* acc, int M, int K) { /* For one output row m of W (M weights), accumulate into acc[m] += W[m]·b. */ int ngrp = K / QX_QK2; for (int m = 0; m < M; m++) { const qx_block_q2* row = W + (size_t)m * ngrp; float p = 0; for (int gg = 0; gg < ngrp; gg++) { float d = f16_to_f32_local(row[gg].d); float s = 0; for (int k = 0; k < QX_QK2; k++) { int q = (row[gg].qs[k / 4] >> ((k % 4) * 2)) & 3; s += (float)(q - 1) * b[gg * QX_QK2 + k]; } p += d * s; } acc[m] += p; }}
struct Shape { int M, K; const char* name; };
int main() { sycl::queue q(sycl::gpu_selector_v, sycl::property::queue::in_order{}); printf("gemm_engine_test (M1.6 go/no-go): gemm_gpu_v4 in-engine, real producer feeding B\n");
Shape shapes[] = { {17408, 5120, "ffn_gate/up"}, {5120, 17408, "ffn_down"}, {12288, 5120, "attn_q"}, }; int N = 256; /* n_tok = chunk size */
int bad_total = 0; for (auto sh : shapes) { int M = sh.M, K = sh.K; int ngrp = K / QX_QK2; printf("\n=== %s M=%d K=%d N=%d ===\n", sh.name, M, K, N);
/* ---- synthesize weight + activation ---- */ auto W = synth_q2_weight(M, K); std::vector<float> x_in((size_t)N * K); std::normal_distribution<float> g(0.0f, 1.0f); for (int t = 0; t < N; t++) { float sc = powf(10.0f, -2.0f + 3.0f * (t % 7) / 6.0f); for (int k = 0; k < K; k++) x_in[(size_t)t * K + k] = g(rng) * sc; } /* rmsnorm weight (K entries) = 1+noise so x*inv*w is non-trivial */ std::vector<float> rms_w(K); for (int k = 0; k < K; k++) rms_w[k] = 1.0f + 0.1f * g(rng);
/* ---- upload + repack weight to A layout ---- */ size_t codes_sz = (size_t)(M / 8) * ngrp * 128; size_t sc_sz = (size_t)ngrp * M; uint8_t* A_codes = (uint8_t*)aligned_alloc(32, codes_sz); float* A_sc = (float*)aligned_alloc(32, sc_sz * sizeof(float)); repack_q2_local(W.data(), M, K, A_codes, A_sc); uint8_t* d_A_codes = sycl::malloc_shared<uint8_t>(codes_sz, q); float* d_A_sc = sycl::malloc_shared<float>(sc_sz, q); q.memcpy(d_A_codes, A_codes, codes_sz).wait(); q.memcpy(d_A_sc, A_sc, sc_sz * sizeof(float)).wait(); free(A_codes); free(A_sc);
/* ---- B operand: produced by gpu_rmsnorm_quant_b (the real producer) ---- */ float* d_x = sycl::malloc_shared<float>((size_t)N * K, q); float* d_rmsw = sycl::malloc_shared<float>(K, q); memcpy(d_x, x_in.data(), (size_t)N * K * sizeof(float)); memcpy(d_rmsw, rms_w.data(), K * sizeof(float)); /* B_vnni: [K/4][N] int32 = (K/4 * N * 4) bytes. sum_a: [ngrp][N]. */ int8_t* d_B = sycl::malloc_shared<int8_t>((size_t)(K / 4) * N * 4, q); int32_t* d_sa = sycl::malloc_shared<int32_t>((size_t)ngrp * N, q); float* d_asc = sycl::malloc_shared<float>(N, q); gpu_rmsnorm_quant(q, d_x, d_rmsw, K, d_B, d_sa, d_asc, N); q.wait();
/* ---- run GEMM ---- */ float* d_C = sycl::malloc_shared<float>((size_t)M * N, q); gemm_gpu_v4(q, d_A_codes, d_A_sc, d_B, d_sa, d_asc, d_C, M, K, N, nullptr, true);
/* ---- correctness: CPU ref applies rmsnorm (fp32) then W@b per sampled row. * GPU path = rmsnorm_quant_b (fp32 rmsnorm + s8 quant) + GEMM. Expected * mean|d| ~ 1e-3 from the s8 quantization noise. ---- */ int M_SAMPLE = std::min(M, 64); std::vector<int> sample_rows(M_SAMPLE); for (int i = 0; i < M_SAMPLE; i++) sample_rows[i] = (int)((uint64_t)rng() % M); std::vector<float> normed(K); double md = 0; float wd = 0; int ncmp = 0; for (int n = 0; n < N; n++) { rmsnorm_row(x_in.data() + (size_t)n * K, rms_w.data(), normed.data(), K); for (int i = 0; i < M_SAMPLE; i++) { int m = sample_rows[i]; const qx_block_q2* row = W.data() + (size_t)m * ngrp; float acc = 0; for (int gg = 0; gg < ngrp; gg++) { float d = f16_to_f32_local(row[gg].d); float p = 0; for (int k = 0; k < QX_QK2; k++) { int qq = (row[gg].qs[k / 4] >> ((k % 4) * 2)) & 3; p += (float)(qq - 1) * normed[gg * QX_QK2 + k]; } acc += d * p; } double dd = fabs(acc - d_C[(size_t)m * N + n]); md += dd; ncmp++; if (dd > wd) wd = (float)dd; } } printf(" correctness vs CPU fp32+rmsnorm (sampled %d rows x %d toks): mean|d|=%.3e worst=%.3e\n", M_SAMPLE, N, md / ncmp, wd);
/* ---- perf: sustained TFLOPs over 50 back-to-back calls ---- */ int REPS = 50; q.wait(); auto t0 = std::chrono::steady_clock::now(); for (int r = 0; r < REPS; r++) gemm_gpu_v4(q, d_A_codes, d_A_sc, d_B, d_sa, d_asc, d_C, M, K, N, nullptr, false); q.wait(); auto t1 = std::chrono::steady_clock::now(); double ms = std::chrono::duration<double, std::milli>(t1 - t0).count() / REPS; double flops = 2.0 * (double)M * K * N; double tflops = flops / (ms / 1e3) / 1e12; /* tok/s-eq: per-token GEMM cost across all 64 layers is ~25.7B MACs * (M1.2 convention). tok/s = TFLOPs * 1e12 / 2 / 25.7e9. */ double tokseq = tflops * 1e12 / 2.0 / 25.7e9; printf(" perf: %.2f ms/call, %.1f TFLOPs (%.0f tok/s-eq)\n", ms, tflops, tokseq);
sycl::free(d_A_codes, q); sycl::free(d_A_sc, q); sycl::free(d_x, q); sycl::free(d_rmsw, q); sycl::free(d_B, q); sycl::free(d_sa, q); sycl::free(d_asc, q); sycl::free(d_C, q); } return bad_total ? 1 : 0;}