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
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347/* qxmx_fa.cpp: M1.7b.7 causal prefill attention, 3-kernel split per K-tile. * See qxmx_fa.h for the design. * * Per K-tile (TILE_T keys): * 1. dequant K-tile -> Kt [NQ][HD][TILE_T] (transposed for Q@K^T; each * q-head's slot holds its kv-head's K-tile, replicated across GQA=6) * 2. QK GEMM (gemm_tf32_batched, DPAS): S[hh][n_tok][TILE_T] * = Qsep[hh][n_tok][HD] @ Kt[hh][HD][TILE_T] * 3. online-softmax merge (plain SYCL): causal-mask S, exp, merge into * P + (m, l) running state * 4. dequant V-tile -> Vt [NQ][TILE_T][HD] (TurboQuant rotated domain, T13) * 5. PV GEMM (gemm_tf32_batched): O[hh][n_tok][HD] += P[hh][n_tok][TILE_T] * @ Vt[hh][TILE_T][HD] (accumulates across tiles via Oacc) * 6. final: butterfly each 32-elem block of O (T13 inverse WHT) x NRM -> out * * Why the split (vs M1.7b.6 fused scalar): the scalar HD=256 dot per * (token,head,key) capped prefill at 230 tok/s. The QK and PV GEMMs use * DPAS at high occupancy (gemm_tf32_batched, M1.7a's proven 4.2 TFLOPs * primitive). Materializes S/P per tile (~3 MB at CHUNK=256, TILE_T=128), * bounded regardless of context length. The K-tile loop walks * 0..pos_start+n_tok-1 in TILE_T-aligned blocks; the final partial tile and * the per-token causal cap are masked in the softmax-merge step. * * M1.7b.1-6 history: M1.7b.4 ESIMD (176 tok/s) and M1.7b.6 plain-SYCL scalar * (230 tok/s) both did the QK dot as a scalar loop over HD. The inline-asm * DPAS path (M1.7b.5) was blocked by a vISA JIT asm-dialect wall -- see * qxmx_m1_7b_5_dpas_attempt.md. This split reuses the WORKING ESIMD * gemm_tf32_batched (which the M1.7b.5 attempt confirmed REQUIRES * [[intel::sycl_explicit_simd]] -- the GEMM kernels have it, the softmax * merge kernel does not need it). */#include "qxmx_fa.h"
#include <cmath>#include <cstdlib>
#include <sycl/sycl.hpp>
#include "qxmx_fp8.h"#include "qxmx_gemm_tf32.h"
namespace qx {
static constexpr int HD = QX_HEAD_DIM; /* 256 */static constexpr int NQ = QX_N_HEAD; /* 24 */static constexpr int NKV = QX_N_HEAD_KV; /* 4 */static constexpr int GQA = NQ / NKV; /* 6 */static constexpr int KV_DIM = QX_KV_DIM; /* 1024 */static constexpr int VBLK_L = VBLK; /* 32 */static constexpr int BLK_PER_HEAD = HD / VBLK_L; /* 8 */static constexpr int TILE_T = 128; /* K-tokens per block */static constexpr float NRM = 0.17677669529663687f; /* 1/sqrt(32) */
/* ---- Step 1: dequant Q [n_tok][NQ*HD] -> Qsep [NQ][n_tok][HD] ------------- * One WG per (q-head, query-token), LWS=HD lanes cooperatively copy + scale. * Also initializes m_state=-inf, l_state=0, Oacc=0 for the chunk. */static void fa_prep_q(sycl::queue& q, const float* Q, int n_tok, fa_split_workspace* ws) { /* Copy struct members to locals -- the device kernel can't dereference * the host struct pointer. */ float* Qsep = ws->Qsep; float* Oacc = ws->Oacc; float* m_state = ws->m_state; float* l_state = ws->l_state; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)NQ * n_tok * HD), sycl::range<1>(HD)), [=](sycl::nd_item<1> it) { int gid = it.get_group(0); int hh = gid / n_tok; int qi = gid % n_tok; int d = it.get_local_id(0); float v = Q[(size_t)qi * (NQ * HD) + (size_t)hh * HD + d]; Qsep[(size_t)hh * n_tok * HD + (size_t)qi * HD + d] = v; Oacc[(size_t)hh * n_tok * HD + (size_t)qi * HD + d] = 0.0f; if (d == 0) { m_state[(size_t)hh * n_tok + qi] = -1e30f; l_state[(size_t)hh * n_tok + qi] = 0.0f; } }); });}
/* ---- Step 2: dequant K-tile [t0..t0+TILE_T-1] -> Kt [NQ][HD][TILE_T] ------ * Transposed for Q@K^T: Kt[hh][d][t] = K[hk][t][d]. Each q-head's slot holds * its kv-head's K-tile (replicated across GQA). One WG per (q-head, key-tile * element), LWS=HD lanes. */template <typename Kct>static void fa_dequant_kt(sycl::queue& q, const uint8_t* kv_k, int64_t t0, int tile_len, fa_split_workspace* ws) { int n_wgs = NQ * TILE_T; float* Kt = ws->Kt; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD), sycl::range<1>(HD)), [=](sycl::nd_item<1> it) { int gid = it.get_group(0); int hh = gid / TILE_T; int tloc = gid % TILE_T; int hk = hh / GQA; int d = it.get_local_id(0); int64_t kglob = t0 + tloc; float kval = 0.0f; if (tloc < tile_len) { if constexpr (std::is_same_v<Kct, KFp16>) { const sycl::half* kp = reinterpret_cast<const sycl::half*>(kv_k); kval = float(kp[(size_t)kglob * KV_DIM + (size_t)hk * HD + d]); } else { const uint8_t* kb = kv_k + (size_t)kglob * Kct::row_bytes + (size_t)(hk * BLK_PER_HEAD + (d >> 5)) * KBLK_BYTES; uint16_t sbits = (uint16_t)kb[0] | ((uint16_t)kb[1] << 8); float dd = float(sycl::bit_cast<sycl::half>(sbits)); uint8_t code = kb[2 + (d & 31)]; if constexpr (std::is_same_v<Kct, KQ8_0>) kval = (float)(int8_t)code * dd; else kval = e4m3_to_fp32(code) * dd; } } /* Kt[hh][d][tloc] = Kt[hh * HD * TILE_T + d * TILE_T + tloc] */ Kt[(size_t)hh * HD * TILE_T + (size_t)d * TILE_T + tloc] = kval; }); });}
/* ---- Step 3: QK GEMM via gemm_tf32_batched ------------------------------- * S[hh][n_tok][TILE_T] = Qsep[hh][n_tok][HD] @ Kt[hh][HD][TILE_T] * batch=NQ, M=n_tok, K=HD, N=TILE_T. Requires n_tok % 8, HD % 8, TILE_T % 16. */static void fa_qk_gemm(sycl::queue& q, int n_tok, fa_split_workspace* ws) { gemm_tf32_batched(q, ws->Qsep, ws->Kt, ws->S, NQ, n_tok, HD, TILE_T);}
/* ---- Step 4: online-softmax merge (causal-masked) ----------------------- * One WG per (q-head, query-token), LWS=TILE_T. For each query token qi, the * valid keys in this tile are tloc in [0..tile_len-1] with t0+tloc <= pos_start+qi. * S is scaled (1/sqrt(HD)); apply causal mask, merge exp(S - m_new) into P, * update m_state/l_state. */static void fa_softmax_merge(sycl::queue& q, int64_t pos_start, int64_t t0, int tile_len, int n_tok, float scale, fa_split_workspace* ws) { int n_wgs = NQ * n_tok; const float* S = ws->S; float* P = ws->P; float* m_state = ws->m_state; float* l_state = ws->l_state; float* Oalpha = ws->O; /* temp alpha store */ q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * TILE_T), sycl::range<1>(TILE_T)), [=](sycl::nd_item<1> it) { int gid = it.get_group(0); int hh = gid / n_tok; int qi = gid % n_tok; int tloc = it.get_local_id(0); auto grp = it.get_group();
int64_t qpos = pos_start + qi; int64_t kglob = t0 + tloc; bool valid = (tloc < tile_len) && (kglob <= qpos); float s = valid ? S[(size_t)hh * n_tok * TILE_T + (size_t)qi * TILE_T + tloc] * scale : -1e30f; float tile_max = sycl::reduce_over_group(grp, s, sycl::maximum<float>()); float prev_m = m_state[(size_t)hh * n_tok + qi]; float new_m = sycl::fmax(prev_m, tile_max); float alpha = sycl::native::exp(prev_m - new_m); float e = (s > -1e29f) ? sycl::native::exp(s - new_m) : 0.0f; float tile_sum = sycl::reduce_over_group(grp, e, sycl::plus<float>()); float prev_l = l_state[(size_t)hh * n_tok + qi]; float new_l = alpha * prev_l + tile_sum; if (tloc == 0) { m_state[(size_t)hh * n_tok + qi] = new_m; l_state[(size_t)hh * n_tok + qi] = new_l; Oalpha[(size_t)hh * n_tok + qi] = alpha; } P[(size_t)hh * n_tok * TILE_T + (size_t)qi * TILE_T + tloc] = e; }); });}
/* ---- Step 5: dequant V-tile -> Vt [NQ][TILE_T][HD] (rotated domain, T13) - * V_LM_LUT[code] * vscale * vnorm_corr (NO per-token butterfly -- T13 defers * the inverse WHT to the final O). One WG per (q-head, key-tile element), * LWS=HD lanes. */static void fa_dequant_vt(sycl::queue& q, const uint32_t* kv_v, int64_t t0, int tile_len, fa_split_workspace* ws) { int n_wgs = NQ * TILE_T; float* Vt = ws->Vt; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD), sycl::range<1>(HD)), [=](sycl::nd_item<1> it) { int gid = it.get_group(0); int hh = gid / TILE_T; int tloc = gid % TILE_T; int hk = hh / GQA; int d = it.get_local_id(0); int64_t kglob = t0 + tloc; float vval = 0.0f; if (tloc < tile_len) { int blk = hk * BLK_PER_HEAD + (d >> 5); int vlane = d & 31; const uint32_t* db = kv_v + (size_t)kglob * VROW_U32 + blk * 5; uint16_t sc_u = (uint16_t)(db[0] & 0xFFFFu); uint16_t nc_u = (uint16_t)((db[0] >> 16) & 0xFFFFu); float vscale = float(sycl::bit_cast<sycl::half>(sc_u)); float vnorm_corr = float(sycl::bit_cast<sycl::half>(nc_u)); uint32_t c0 = db[1], c1 = db[2], c2 = db[3], c3 = db[4]; int w = vlane >> 3; int off = (vlane & 7) * 4; uint32_t word = (w == 0) ? c0 : (w == 1) ? c1 : (w == 2) ? c2 : c3; uint32_t code = (word >> off) & 0xFu; vval = V_LM_LUT[code] * vscale * vnorm_corr; } Vt[(size_t)hh * TILE_T * HD + (size_t)tloc * HD + d] = vval; }); });}
/* ---- Step 6: PV GEMM + accumulate --------------------------------------- * Oacc[hh][n_tok][HD] = alpha * Oacc + P[hh][n_tok][TILE_T] @ Vt[hh][TILE_T][HD] * We fold the alpha rescale by pre-multiplying P by alpha (per query token), * then call gemm_tf32_batched with accum=Oacc. batch=NQ, M=n_tok, K=TILE_T, * N=HD. Requires TILE_T % 8, HD % 16. */static void fa_pv_gemm(sycl::queue& q, int n_tok, fa_split_workspace* ws) { /* Online-softmax rescale: Oacc = alpha * Oacc, then Oacc += P @ Vt. * alpha is stashed in ws->O[hh*n_tok+qi] by fa_softmax_merge (one value per * (q-head, query-token)). For the first tile alpha=0 (prev_m=-inf), which * correctly zeroes Oacc (already 0 from prep) and the tile's P@Vt sets the * initial accumulator. Pre-scale Oacc (not P -- alpha rescales the OLD * state, the new P@Vt is already at the new_m scale). */ int n_wgs = NQ * n_tok; const float* Oalpha = ws->O; float* Oacc = ws->Oacc; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD), sycl::range<1>(HD)), [=](sycl::nd_item<1> it) { int gid = it.get_group(0); int hh = gid / n_tok; int qi = gid % n_tok; int d = it.get_local_id(0); float alpha = Oalpha[(size_t)hh * n_tok + qi]; Oacc[(size_t)hh * n_tok * HD + (size_t)qi * HD + d] *= alpha; }); }); /* Oacc += P @ Vt (accum=Oacc, in-place). M=n_tok, K=TILE_T, N=HD. */ gemm_tf32_batched(q, ws->P, ws->Vt, ws->Oacc, NQ, n_tok, TILE_T, HD, ws->Oacc);}
/* ---- Step 7: final -- Oacc / l, T13 inverse WHT, store ----------------- * One WG per (q-head, query-token), LWS=HD. Normalize by l_state, butterfly * each 32-elem block (inverse WHT) x NRM -> value domain, write to out. */static void fa_finalize(sycl::queue& q, float* out, int n_tok, fa_split_workspace* ws) { int n_wgs = NQ * n_tok; const float* Oacc = ws->Oacc; const float* l_state = ws->l_state; q.submit([&](sycl::handler& h) { sycl::local_accessor<float, 1> bfly0(sycl::range<1>(HD), h); sycl::local_accessor<float, 1> bfly1(sycl::range<1>(HD), h); h.parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t)n_wgs * HD), sycl::range<1>(HD)), [=](sycl::nd_item<1> it) { int gid = it.get_group(0); int hh = gid / n_tok; int qi = gid % n_tok; int lid = it.get_local_id(0); auto grp = it.get_group(); float l = l_state[(size_t)hh * n_tok + qi]; float inv_l = 1.0f / l; bfly0[lid] = Oacc[(size_t)hh * n_tok * HD + (size_t)qi * HD + lid] * inv_l; sycl::group_barrier(grp); int blk = lid >> 5; int lane = lid & 31; int base = blk * VBLK_L; auto* cur = &bfly0; auto* nxt = &bfly1; for (int bit = 1; bit < VBLK_L; bit <<= 1) { float a = (*cur)[base + lane]; float b = (*cur)[base + (lane ^ bit)]; sycl::group_barrier(grp); (*nxt)[base + lane] = (lane & bit) ? (b - a) : (a + b); sycl::group_barrier(grp); std::swap(cur, nxt); } float oval = (*cur)[base + lane] * NRM; out[(size_t)qi * (NQ * HD) + (size_t)hh * HD + lid] = oval; }); });}
/* ---- Entry: fused kernel by default; QXMX_FA_SPLIT=1 keeps the M1.7b.7 * 3-kernel split as an A/B oracle. The fused path fills Oacc + l_state and * reuses this file's fa_finalize for /l + T13 inverse WHT. */template <typename Kct>void gpu_flash_attn(sycl::queue& q, const float* Q, const uint8_t* kv_k, const uint32_t* kv_v, float* out, int64_t pos_start, int n_tok, fa_split_workspace* ws) { static const bool use_split = [] { const char* e = std::getenv("QXMX_FA_SPLIT"); return e && std::atoi(e) != 0; }(); if (!use_split) { gpu_flash_attn_fused<Kct>(q, Q, kv_k, kv_v, pos_start, n_tok, ws); fa_finalize(q, out, n_tok, ws); return; } const float scale = 1.0f / sycl::sqrt((float)HD); /* Prep Q + init state. */ fa_prep_q(q, Q, n_tok, ws); /* Walk K in TILE_T-aligned blocks from 0 to pos_start+n_tok-1. */ int64_t total = pos_start + n_tok; int n_tiles = (int)((total + TILE_T - 1) / TILE_T); for (int ti = 0; ti < n_tiles; ti++) { int64_t t0 = (int64_t)ti * TILE_T; int tile_len = (int)(total - t0); if (tile_len > TILE_T) tile_len = TILE_T; fa_dequant_kt<Kct>(q, kv_k, t0, tile_len, ws); fa_qk_gemm(q, n_tok, ws); fa_softmax_merge(q, pos_start, t0, tile_len, n_tok, scale, ws); fa_dequant_vt(q, kv_v, t0, tile_len, ws); fa_pv_gemm(q, n_tok, ws); } fa_finalize(q, out, n_tok, ws);}
template void gpu_flash_attn<KFp16>(sycl::queue&, const float*, const uint8_t*, const uint32_t*, float*, int64_t, int, fa_split_workspace*);template void gpu_flash_attn<KQ8_0>(sycl::queue&, const float*, const uint8_t*, const uint32_t*, float*, int64_t, int, fa_split_workspace*);template void gpu_flash_attn<KFp8 >(sycl::queue&, const float*, const uint8_t*, const uint32_t*, float*, int64_t, int, fa_split_workspace*);
} // namespace qx