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
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788#include "qxmx_gpu.h"#include "qxmx_fp8.h"#include "qxmx_gemm.h"#include "qxmx_fa.h"#include "qxmx_prefix_cache.h"#include "qxmx_device.h" /* qx::device_name + chk_shared/chk_device */
#include <algorithm>#include <cmath>#include <cstdio>#include <cstdlib>#include <cstring>#include <vector>#include <type_traits>
#include <sycl/ext/intel/esimd.hpp>
#include <chrono>
namespace esimd = sycl::ext::intel::esimd;namespace xmx = sycl::ext::intel::esimd::xmx;using xmx::dpas_argument_type;
namespace qx {
/* File-static model pointer, set in engine::init. Forward-declared here so * forward()/run_chunk_b() (defined below) can reach it; the actual * dequant_embd helper is defined after f16_to_f32 (which it calls). The * CLIs (qxmx_run, qxmx_diff) reach this via their own `extern` decl to set * it before init — keep it non-static. */qx_model* g_model = nullptr;static void dequant_embd(const qx_model* m, int token_id, float* out);
static float f16_to_f32(uint16_t bits) { uint32_t sign = (uint32_t)(bits & 0x8000) << 16; uint32_t exp = (bits >> 10) & 0x1f, mant = bits & 0x3ff; union { uint32_t u; float f; } cv; if (exp == 0) { cv.f = (float)mant * 0.00006103515625f; cv.u |= sign; return cv.f; } if (exp == 31) { cv.u = sign | 0x7f800000u | (mant ? 0x00400000u : 0); return cv.f; } cv.u = sign | ((exp + 112) << 23) | (mant << 13); return cv.f;}
/* Dequantize one embedding row (token_id) into out[D_MODEL]. Host-side; * used by forward() (single-token) — run_chunk_b() open-codes the same loop * over M rows. Handles both g64 and g128 Q2 formats. */static void dequant_embd(const qx_model* m, int token_id, float* out) { int K = m->token_embd.dims[0]; int qk2 = (int)m->qk2; int ngrp = K / qk2; if (qk2 == 64) { const qx_block_q2* row = (const qx_block_q2*)m->token_embd.data + (size_t)token_id * ngrp; for (int g = 0; g < ngrp; g++) { float d = f16_to_f32(row[g].d); for (int k = 0; k < 64; k++) { int q = (row[g].qs[k/4] >> ((k%4)*2)) & 3; out[g*64+k] = d*(float)(q-1); } } } else { const qx_block_q2_g128* row = (const qx_block_q2_g128*)m->token_embd.data + (size_t)token_id * ngrp; for (int g = 0; g < ngrp; g++) { float d = f16_to_f32(row[g].d); for (int k = 0; k < 128; k++) { int q = (row[g].qs[k/4] >> ((k%4)*2)) & 3; out[g*128+k] = d*(float)(q-1); } } }}
/* Repack g64 Q2 weights: GGUF layout -> DPAS A_codes + transposed scales. */static void repack_q2_g64(const qx_block_q2* src, int M, int K, uint8_t* codes, float* scales_t) { int ngrp = K / 64; for (int m8 = 0; m8 < M / 8; m8++) for (int g = 0; g < ngrp; g++) for (int half = 0; half < 2; half++) for (int r = 0; r < 8; r++) memcpy(codes + (((size_t)(m8 * ngrp + g) * 2 + half) * 64 + r * 8), src[(m8 * 8 + r) * ngrp + g].qs + half * 8, 8); for (int m = 0; m < M; m++) for (int g = 0; g < ngrp; g++) scales_t[g * M + m] = f16_to_f32(src[m * ngrp + g].d);}
/* Repack g128 Q2 weights: split each g128 block into two g64 blocks sharing * the same scale. Output is identical to g64 repack (ngrp*2 groups). */static void repack_q2_g128(const qx_block_q2_g128* src, int M, int K, uint8_t* codes, float* scales_t) { int ngrp128 = K / 128; int ngrp64 = ngrp128 * 2; for (int m8 = 0; m8 < M / 8; m8++) { for (int g = 0; g < ngrp128; g++) { for (int half = 0; half < 2; half++) { for (int r = 0; r < 8; r++) { const qx_block_q2_g128* blk = &src[(m8 * 8 + r) * ngrp128 + g]; /* First g64 half: bytes 0-7 of the 32-byte qs */ memcpy(codes + (((size_t)(m8 * ngrp64 + g*2) * 2 + half) * 64 + r * 8), blk->qs + half * 8, 8); /* Second g64 half: bytes 16-23 of the 32-byte qs */ memcpy(codes + (((size_t)(m8 * ngrp64 + g*2 + 1) * 2 + half) * 64 + r * 8), blk->qs + 16 + half * 8, 8); } } } } for (int m = 0; m < M; m++) { for (int g = 0; g < ngrp128; g++) { float d = f16_to_f32(src[m * ngrp128 + g].d); scales_t[(g*2) * M + m] = d; scales_t[(g*2 + 1) * M + m] = d; } }}
static constexpr int GRP = QX_QK2;
/* GPU gemv_dpas8 (ESIMD DPAS) — proven 240-468 GB/s in benchmark. */static constexpr int RC8 = 8, SD = 8;
static void gemv_reduce(sycl::queue& q, const float* part, float* y, const float* act_scale_d, int M, int S, const float* accum = nullptr) { q.parallel_for(sycl::range<1>(M), [=](sycl::id<1> i) { float a = 0; for (int s = 0; s < S; s++) a += part[(size_t)s * M + i]; y[i] = a * (*act_scale_d) + (accum ? accum[i] : 0.f); });}
/* Split-K tuning knobs for the decode GEMVs (read once per process). * QXMX_GEMV_TGTWGS: target WGs before the split-K heuristic stops doubling * S (default 512 = 2 WGs/EU on B70). Raise to dispatch more WGs/EU. * QXMX_GEMV_SMAX: hard cap on S (default 16). Bounded by part_buf * (2*SMAX*FF_INTER floats <= 2*16*17408 = 557K << 8*VOCAB = 1.98M). * See latency-hunt-plan.md Subtask 2. */static int gemv_smax() { static int v = []{ int x = 16; if (const char* e = std::getenv("QXMX_GEMV_SMAX")) x = std::atoi(e); return x; }(); return v;}static int gemv_tgtwgs() { static int v = []{ int x = 512; if (const char* e = std::getenv("QXMX_GEMV_TGTWGS")) x = std::atoi(e); return x; }(); return v;}
static void gemv_dpas8(sycl::queue& q, const uint8_t* A_dpas, const float* A_scale_t, const int8_t* B_vnni, const int32_t* sum_a, const float* act_scale_d, float* y, float* part, int M, int K, const float* accum = nullptr) { int ngrp = K / GRP, M8 = M / 8; constexpr int WGM8 = 32, GPP = 4; constexpr int SLM_A_SZ = WGM8 * 128 * GPP; constexpr int SLM_B_SZ = 1024 * GPP; constexpr int LWS = WGM8; int NGM = (M8 + WGM8 - 1) / WGM8; int S = 1; /* Raise split-K to hide latency: the decode GEMVs are 0.27-0.31 WGs/EU at * the old `NGM*S<64` cap, starving the EUs (50% idle @ 100% occupancy). * Target ~1 WG/EU (256). Only double while ngrp stays divisible by 2*S, * which guarantees gps=ngrp/S is integral and [0,ngrp) is fully covered * (no dropped groups -> qxmx_diff stable). Partial last GPP-tile is * handled by the gmax guard below. See latency-hunt-plan.md. */ const int GV_SMAX = gemv_smax(); const int GV_TGTWGS = gemv_tgtwgs(); while (S < GV_SMAX && ngrp % (S * 2) == 0 && NGM * S * 2 <= GV_TGTWGS) S *= 2; int G = NGM * S * LWS; int gps = ngrp / S; auto props = esimd::properties{esimd::cache_hint_L1<esimd::cache_hint::uncached>, esimd::cache_hint_L2<esimd::cache_hint::cached>}; q.parallel_for(sycl::nd_range<1>(G, LWS), [=](sycl::nd_item<1> it) [[intel::sycl_explicit_simd]] { int group = it.get_group(0); int gm = group % NGM, ks = group / NGM; int lid = it.get_local_id(0); int m8 = gm * WGM8 + lid, m0 = m8 * 8; bool valid = m8 < M8; float asc0 = *act_scale_d; /* read scale from device USM (broadcast) */ esimd::slm_init<SLM_A_SZ + SLM_B_SZ>(); esimd::simd<float, 128> acc(0.f); int gend = (ks + 1) * gps; for (int gb = ks * gps; gb < gend; gb += GPP) { /* Cooperative A-load: chunk = pass*LWS + lid, so am8 = lid * (pass-invariant), ag = pass. Hoist the HBM base * (gm8*ngrp + gb)*128 and the SLM base lid*128 out of the * pass loop; step +pass*128 (HBM) and +WGM8*128 (SLM). */ { const int gm8 = gm * WGM8 + lid; const bool gm8_ok = gm8 < M8; const uint32_t* ap = reinterpret_cast<const uint32_t*>( A_dpas + (size_t)(gm8 * ngrp + gb) * 128); int slm_a = lid * 128; for (int pass = 0; pass < (WGM8 * GPP) / LWS; pass++) { int g = gb + pass; if (gm8_ok && g < gend) { esimd::simd<uint32_t, 32> av; av.copy_from(ap + (size_t)pass * 32, props); esimd::slm_block_store(slm_a, av.template bit_cast_view<uint8_t>()); } slm_a += WGM8 * 128; } } for (int pass = 0; pass < (GPP * 16) / LWS; pass++) { int chunk = pass * LWS + lid; int bgg = chunk / 16, bkk = chunk % 16, g = gb + bgg; if (g < gend) { int k4 = g * 16 + bkk; /* T8: load 1 int32 (pack), splat to 16 SLM slots (was 16 replicated int32s in global). */ esimd::simd<int32_t, 1> bv1; bv1.copy_from(reinterpret_cast<const int32_t*>(B_vnni + (size_t)k4 * 4), props); esimd::simd<int32_t, 16> bv(bv1[0]); esimd::slm_block_store(SLM_A_SZ + bgg * 1024 + bkk * 64, bv.template bit_cast_view<int8_t>()); } } if (gb + GPP < gend) { /* Prefetch next A tile: same hoist as A-load but g = gb+GPP+pass. */ const int gm8 = gm * WGM8 + lid; const bool gm8_ok = gm8 < M8; const uint32_t* ap = reinterpret_cast<const uint32_t*>( A_dpas + (size_t)(gm8 * ngrp + gb + GPP) * 128); for (int pass = 0; pass < (WGM8 * GPP) / LWS; pass++) { int g = gb + GPP + pass; if (gm8_ok && g < gend) esimd::prefetch<uint32_t, 32>(ap + (size_t)pass * 32, props); } } esimd::barrier(); if (valid) { int gmax = (gb + GPP < gend) ? GPP : (gend - gb); /* Hoist loop-invariant address bases out of the ag loop: * IGC recomputes ag*WGM8*128, ag*M, ag*1024 from scratch * each iteration (shl/or) instead of strength-reducing * to a running offset. Compute the bases once and step * them per ag so IGC keeps a single add per iter. */ const float* dfp = A_scale_t + (size_t)gb * M + m0; const int32_t* sap = sum_a + gb; int slm_a_off = lid * 128; /* +ag*WGM8*128 each iter */ int slm_b_off = SLM_A_SZ; /* +ag*1024 each iter */ for (int ag = 0; ag < gmax; ag++) { esimd::simd<uint8_t, 128> av = esimd::slm_block_load<uint8_t, 128>(slm_a_off); esimd::simd<int, 128> p(0); for (int half = 0; half < 2; half++) { esimd::simd<int8_t, 512> vb = esimd::slm_block_load<int8_t, 512>(slm_b_off + half * 512); esimd::simd<uint8_t, 64> va_half = av.select<64, 1>(half * 64); p = xmx::dpas<SD, RC8, int, int, int8_t, uint8_t, dpas_argument_type::s8, dpas_argument_type::u2>(p, vb, va_half); } esimd::simd<float, 8> df; df.copy_from(dfp); int32_t sa = *sap; esimd::simd<int, 128> corr = p - sa; esimd::simd<float, 128> corrf = corr, df128; for (int r = 0; r < 8; r++) df128.select<16, 1>(r * 16) = esimd::simd<float, 16>(df[r]); acc = acc + df128 * corrf; dfp += M; ++sap; slm_a_off += WGM8 * 128; slm_b_off += 1024; } } esimd::barrier(); } if (valid) { if (S == 1) for (int r = 0; r < 8; r++) y[m0 + r] = acc[r * 16] * asc0 + (accum ? accum[m0 + r] : 0.f); else for (int r = 0; r < 8; r++) part[(size_t)ks * M + m0 + r] = acc[r * 16]; } }); if (S > 1) gemv_reduce(q, part, y, act_scale_d, M, S, accum);}
/* FFN gate+up fusion: two A matrices (gate, up) share one quantized B. * Stage two A tiles + one shared B per group-block; compute runs two DPAS * accumulators against the shared B (loaded once from SLM). SLM_A holds * A0||A1 (2x). Bit-exact vs two separate gemv_dpas8 calls (validated in * gemv.cpp dual_test). Reads act_scale from device USM (same convention). */static void gemv_dpas8_dual(sycl::queue& q, const uint8_t* A0_dpas, const float* A0_sc, const uint8_t* A1_dpas, const float* A1_sc, const int8_t* B_vnni, const int32_t* sum_a, const float* act_scale_d, float* y0, float* y1, float* part, int M, int K) { int ngrp = K / GRP, M8 = M / 8; constexpr int WGM8 = 32, GPP = 4; constexpr int SLM_A0 = WGM8 * 128 * GPP; constexpr int SLM_A_SZ = 2 * SLM_A0; constexpr int SLM_B_SZ = 1024 * GPP; constexpr int LWS = WGM8; int NGM = (M8 + WGM8 - 1) / WGM8; int S = 1; /* Same raised split-K heuristic as gemv_dpas8 (see comment there). */ const int GV_SMAX = gemv_smax(); const int GV_TGTWGS = gemv_tgtwgs(); while (S < GV_SMAX && ngrp % (S * 2) == 0 && NGM * S * 2 <= GV_TGTWGS) S *= 2; int G = NGM * S * LWS; int gps = ngrp / S; auto props = esimd::properties{esimd::cache_hint_L1<esimd::cache_hint::uncached>, esimd::cache_hint_L2<esimd::cache_hint::cached>}; q.parallel_for(sycl::nd_range<1>(G, LWS), [=](sycl::nd_item<1> it) [[intel::sycl_explicit_simd]] { int group = it.get_group(0); int gm = group % NGM, ks = group / NGM; int lid = it.get_local_id(0); int m8 = gm * WGM8 + lid, m0 = m8 * 8; bool valid = m8 < M8; float asc0 = *act_scale_d; esimd::slm_init<SLM_A_SZ + SLM_B_SZ>(); esimd::simd<float, 128> acc0(0.f), acc1(0.f); int gend = (ks + 1) * gps; for (int gb = ks * gps; gb < gend; gb += GPP) { /* Cooperative A0-load: am8 = lid (pass-invariant), ag = pass. * Hoist (gm8*ngrp + gb)*128 (HBM base) and lid*128 (SLM base). */ { const int gm8 = gm * WGM8 + lid; const bool gm8_ok = gm8 < M8; const uint32_t* ap = reinterpret_cast<const uint32_t*>( A0_dpas + (size_t)(gm8 * ngrp + gb) * 128); int slm_a = lid * 128; for (int pass = 0; pass < (WGM8 * GPP) / LWS; pass++) { int g = gb + pass; if (gm8_ok && g < gend) { esimd::simd<uint32_t, 32> av; av.copy_from(ap + (size_t)pass * 32, props); esimd::slm_block_store(slm_a, av.template bit_cast_view<uint8_t>()); } slm_a += WGM8 * 128; } } /* Cooperative A1-load (same hoist, SLM_A0 offset). */ { const int gm8 = gm * WGM8 + lid; const bool gm8_ok = gm8 < M8; const uint32_t* ap = reinterpret_cast<const uint32_t*>( A1_dpas + (size_t)(gm8 * ngrp + gb) * 128); int slm_a = SLM_A0 + lid * 128; for (int pass = 0; pass < (WGM8 * GPP) / LWS; pass++) { int g = gb + pass; if (gm8_ok && g < gend) { esimd::simd<uint32_t, 32> av; av.copy_from(ap + (size_t)pass * 32, props); esimd::slm_block_store(slm_a, av.template bit_cast_view<uint8_t>()); } slm_a += WGM8 * 128; } } for (int pass = 0; pass < (GPP * 16) / LWS; pass++) { int chunk = pass * LWS + lid; int bgg = chunk / 16, bkk = chunk % 16, g = gb + bgg; if (g < gend) { int k4 = g * 16 + bkk; /* T8: load 1 int32 (pack), splat to 16 SLM slots. */ esimd::simd<int32_t, 1> bv1; bv1.copy_from(reinterpret_cast<const int32_t*>(B_vnni + (size_t)k4 * 4), props); esimd::simd<int32_t, 16> bv(bv1[0]); esimd::slm_block_store(SLM_A_SZ + bgg * 1024 + bkk * 64, bv.template bit_cast_view<int8_t>()); } } if (gb + GPP < gend) { /* Prefetch next A0/A1 tile (same hoist as the loads). */ const int gm8 = gm * WGM8 + lid; const bool gm8_ok = gm8 < M8; const uint32_t* ap0 = reinterpret_cast<const uint32_t*>( A0_dpas + (size_t)(gm8 * ngrp + gb + GPP) * 128); const uint32_t* ap1 = reinterpret_cast<const uint32_t*>( A1_dpas + (size_t)(gm8 * ngrp + gb + GPP) * 128); for (int pass = 0; pass < (WGM8 * GPP) / LWS; pass++) { int g = gb + GPP + pass; if (gm8_ok && g < gend) { esimd::prefetch<uint32_t, 32>(ap0 + (size_t)pass * 32, props); esimd::prefetch<uint32_t, 32>(ap1 + (size_t)pass * 32, props); } } } esimd::barrier(); if (valid) { int gmax = (gb + GPP < gend) ? GPP : (gend - gb); /* Hoist loop-invariant address bases out of the ag loop * (same as gemv_dpas8 -- running pointers so IGC's LICM * keeps a single add per iter instead of recomputing * ag*WGM8*128, ag*M, ag*1024 from scratch). */ const float* df0p = A0_sc + (size_t)gb * M + m0; const float* df1p = A1_sc + (size_t)gb * M + m0; const int32_t* sap = sum_a + gb; int slm_a0_off = lid * 128; /* +ag*WGM8*128 each iter */ int slm_a1_off = SLM_A0 + lid * 128; int slm_b_off = SLM_A_SZ; /* +ag*1024 each iter */ for (int ag = 0; ag < gmax; ag++) { esimd::simd<uint8_t, 128> av0 = esimd::slm_block_load<uint8_t, 128>(slm_a0_off); esimd::simd<uint8_t, 128> av1 = esimd::slm_block_load<uint8_t, 128>(slm_a1_off); int32_t sa = *sap; esimd::simd<int, 128> p0(0), p1(0); for (int half = 0; half < 2; half++) { esimd::simd<int8_t, 512> vbh = esimd::slm_block_load<int8_t, 512>(slm_b_off + half * 512); esimd::simd<uint8_t, 64> a0h = av0.select<64, 1>(half * 64); esimd::simd<uint8_t, 64> a1h = av1.select<64, 1>(half * 64); p0 = xmx::dpas<SD, RC8, int, int, int8_t, uint8_t, dpas_argument_type::s8, dpas_argument_type::u2>(p0, vbh, a0h); p1 = xmx::dpas<SD, RC8, int, int, int8_t, uint8_t, dpas_argument_type::s8, dpas_argument_type::u2>(p1, vbh, a1h); } esimd::simd<float, 8> df0, df1; df0.copy_from(df0p); df1.copy_from(df1p); esimd::simd<int, 128> c0 = p0 - sa, c1 = p1 - sa; esimd::simd<float, 128> cf0 = c0, cf1 = c1, d0, d1; for (int r = 0; r < 8; r++) { d0.select<16, 1>(r * 16) = esimd::simd<float, 16>(df0[r]); d1.select<16, 1>(r * 16) = esimd::simd<float, 16>(df1[r]); } acc0 = acc0 + d0 * cf0; acc1 = acc1 + d1 * cf1; df0p += M; df1p += M; ++sap; slm_a0_off += WGM8 * 128; slm_a1_off += WGM8 * 128; slm_b_off += 1024; } } esimd::barrier(); } if (valid) { if (S == 1) { for (int r = 0; r < 8; r++) { y0[m0 + r] = acc0[r * 16] * asc0; y1[m0 + r] = acc1[r * 16] * asc0; } } else { for (int r = 0; r < 8; r++) { part[(size_t)ks * M + m0 + r] = acc0[r * 16]; part[(size_t)(S + ks) * M + m0 + r] = acc1[r * 16]; } } } }); if (S > 1) { gemv_reduce(q, part, y0, act_scale_d, M, S); gemv_reduce(q, part + (size_t)S * M, y1, act_scale_d, M, S); }}
static int8_t* g_bvnni_d = nullptr;static int32_t* g_suma_d = nullptr;static float* g_act_scale_d = nullptr;
/* Per-head RMSNorm: same head_dim weights applied to each of n_heads groups. * Exposed (drop-static) so the M1.6 S2 batched variant can be tested against * it as the per-token oracle. */void gpu_group_rmsnorm(sycl::queue& q, const float* x, const float* w, float* out, int n_heads, int head_dim) { int LWS = head_dim; q.submit([&](sycl::handler& h) { sycl::local_accessor<double,1> lss(sycl::range<1>(LWS), h); h.parallel_for(sycl::nd_range<1>(n_heads * LWS, LWS), [=](sycl::nd_item<1> it) { int hd = it.get_group(0); int lid = it.get_local_id(0); const float* xh = x + (size_t)hd * head_dim; lss[lid] = (double)xh[lid] * xh[lid]; sycl::group_barrier(it.get_group()); for (int s = LWS / 2; s > 0; s >>= 1) { if (lid < s) lss[lid] += lss[lid + s]; sycl::group_barrier(it.get_group()); } float inv = 1.0f / sycl::sqrt((float)(lss[0] / head_dim) + QX_RMS_EPS); out[(size_t)hd * head_dim + lid] = xh[lid] * inv * w[lid]; }); });}
/* M1.6 S2: batched per-(token, head) RMSNorm. One WG per (tok, head); same * per-row arithmetic as the decode oracle (validated bit-exact per row). */void gpu_group_rmsnorm_b(sycl::queue& q, const float* x, const float* w, float* out, int n_heads, int head_dim, int n_tok) { int LWS = head_dim; q.submit([&](sycl::handler& h) { sycl::local_accessor<double,1> lss(sycl::range<1>(LWS), h); h.parallel_for(sycl::nd_range<1>((size_t)n_tok * n_heads * LWS, LWS), [=](sycl::nd_item<1> it) { int tok = it.get_group(0) / n_heads; int hd = it.get_group(0) % n_heads; int lid = it.get_local_id(0); const float* xh = x + (size_t)tok * n_heads * head_dim + (size_t)hd * head_dim; lss[lid] = (double)xh[lid] * xh[lid]; sycl::group_barrier(it.get_group()); for (int s = LWS / 2; s > 0; s >>= 1) { if (lid < s) lss[lid] += lss[lid + s]; sycl::group_barrier(it.get_group()); } float inv = 1.0f / sycl::sqrt((float)(lss[0] / head_dim) + QX_RMS_EPS); out[(size_t)tok * n_heads * head_dim + (size_t)hd * head_dim + lid] = xh[lid] * inv * w[lid]; }); });}
void gpu_silu_inplace(sycl::queue& q, float* x, int n) { q.parallel_for(sycl::range<1>(n), [=](sycl::id<1> i) { x[i] = x[i] / (1.0f + sycl::exp(-x[i])); });}
/* M1.6 S2: batched silu_inplace over n_tok rows. */void gpu_silu_inplace_b(sycl::queue& q, float* x, int n, int n_tok) { q.parallel_for(sycl::range<2>(n_tok, n), [=](sycl::id<2> idx) { size_t off = (size_t)idx[0] * n + idx[1]; x[off] = x[off] / (1.0f + sycl::exp(-x[off])); });}
/* Apply rotary embedding using a precomputed cos/sin table (one table per * decode step, computed on host with exact double/float math to avoid sycl::pow * precision drift). One thread per (half-)frequency, iterates over all heads. * In-place on q and k. Exposed (drop-static) so the M1.6 S2 batched variant * can be tested against it as the per-token oracle. */void gpu_apply_rope(sycl::queue& q, float* q_vec, float* k_vec, int n_q, int n_kv, int head_dim, int rot_dim, const float* rope_c, const float* rope_s) { int half = rot_dim / 2; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<1>(half), [=](sycl::id<1> idx) { int i = idx; float c = rope_c[i], s = rope_s[i]; for (int hh = 0; hh < n_q; hh++) { float* qp = q_vec + (size_t)hh * head_dim; float a = qp[i], b = qp[i + half]; qp[i] = a * c - b * s; qp[i + half] = a * s + b * c; } for (int hh = 0; hh < n_kv; hh++) { float* kp = k_vec + (size_t)hh * head_dim; float a = kp[i], b = kp[i + half]; kp[i] = a * c - b * s; kp[i + half] = a * s + b * c; } }); });}
/* M1.6 S2: batched RoPE. One thread per (tok, half-frequency); rope_c/s are * [n_tok][rot_dim/2] tables the host precomputes (one row per position * pos_start..pos_start+n_tok-1, same double/float math as decode so per-position * output is bit-exact vs decode). q_vec: [n_tok][n_q*head_dim], k_vec similar. */void gpu_apply_rope_b(sycl::queue& q, float* q_vec, float* k_vec, int n_q, int n_kv, int head_dim, int rot_dim, const float* rope_c, const float* rope_s, int n_tok) { int half = rot_dim / 2; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<2>(n_tok, half), [=](sycl::id<2> idx) { int tok = idx[0], i = idx[1]; float c = rope_c[(size_t)tok * half + i]; float s = rope_s[(size_t)tok * half + i]; float* qb = q_vec + (size_t)tok * n_q * head_dim; float* kb = k_vec + (size_t)tok * n_kv * head_dim; for (int hh = 0; hh < n_q; hh++) { float* qp = qb + (size_t)hh * head_dim; float a = qp[i], b = qp[i + half]; qp[i] = a * c - b * s; qp[i + half] = a * s + b * c; } for (int hh = 0; hh < n_kv; hh++) { float* kp = kb + (size_t)hh * head_dim; float a = kp[i], b = kp[i + half]; kp[i] = a * c - b * s; kp[i + half] = a * s + b * c; } }); });}
/* Split interleaved [q|gate] projection into separate q and gate buffers. * One thread per (head, dim). Exposed for the S2 batched oracle. */void gpu_attn_split_qgate(sycl::queue& q, const float* proj, float* q_vec, float* gate, int n_heads, int head_dim) { int n = n_heads * head_dim; q.parallel_for(sycl::range<1>(n), [=](sycl::id<1> idx) { int hh = idx / head_dim, d = idx % head_dim; q_vec[(size_t)hh * head_dim + d] = proj[(size_t)hh * 2 * head_dim + d]; gate[(size_t)hh * head_dim + d] = proj[(size_t)hh * 2 * head_dim + head_dim + d]; });}
/* M1.6 S2: batched q/gate split over n_tok rows. proj: [n_tok][2*n_heads*head_dim]. */void gpu_attn_split_qgate_b(sycl::queue& q, const float* proj, float* q_vec, float* gate, int n_heads, int head_dim, int n_tok) { int n = n_heads * head_dim; q.parallel_for(sycl::range<2>(n_tok, n), [=](sycl::id<2> idx) { int tok = idx[0]; int hh = idx[1] / head_dim, d = idx[1] % head_dim; size_t qo = (size_t)tok * n + (size_t)hh * head_dim + d; size_t po = (size_t)tok * 2 * n + (size_t)hh * 2 * head_dim + d; q_vec[qo] = proj[po]; gate[qo] = proj[po + head_dim]; });}
/* M4 FlashDecoding: max KV splits for split-KV decode attention. Partial * acc [S][NQ][HD] + (m,l) [S][NQ][2] live in engine-hoisted USM. */constexpr int FD_MAX_SPLIT = 64;
/* Cross-split merge: per (head, dim), * out = sum_s exp(m_s-M)*acc_s / sum_s exp(m_s-M)*l_s, M = max_s m_s. */static void attn_fd_merge(sycl::queue& q, const float* fd_acc, const float* fd_ml, float* out, int S) { q.submit([&](sycl::handler& h) { h.parallel_for( sycl::nd_range<1>((size_t)QX_N_HEAD * QX_HEAD_DIM, QX_HEAD_DIM), [=](sycl::nd_item<1> it) { int hh = it.get_group(0); int lid = it.get_local_id(0); float M = -1e30f; for (int s = 0; s < S; s++) M = sycl::fmax(M, fd_ml[((size_t)s * QX_N_HEAD + hh) * 2]); float num = 0.0f, den = 0.0f; for (int s = 0; s < S; s++) { float w = sycl::exp(fd_ml[((size_t)s * QX_N_HEAD + hh) * 2] - M); num += w * fd_acc[((size_t)s * QX_N_HEAD + hh) * QX_HEAD_DIM + lid]; den += w * fd_ml[((size_t)s * QX_N_HEAD + hh) * 2 + 1]; } out[(size_t)hh * QX_HEAD_DIM + lid] = num / den; }); });}
/* Online-softmax attention (FlashAttention-style). One work-group per * (query head, KV split): HD=256 threads = head_dim. Each step does a * cooperative HD-dim dot product (group reduction), then updates the running * max/sum/accumulator. No scores buffer needed — everything in registers/ * SLM. Writes attn_out [NQ*HD] (S==1) or partials for attn_fd_merge (S>1). * Reads q [NQ*HD], kv_k [per K-cache format], kv_v [Phase 1 V]. * GQA: hk=hh/GQA. Templated on K cache format (fp16 | q8_0 | fp8); the K * dequant is selected at compile time via if constexpr — no runtime branch * in the hot per-token loop. V dequant (TurboQuant WHT+4-bit) is identical * across K types. */template <typename Kct>void gpu_attn_softmax(sycl::queue& q, const float* q_vec, const uint8_t* kv_k, const uint32_t* kv_v, float* out, int64_t pos, float* fd_acc, float* fd_ml, int fd_max_split) { constexpr int HD = QX_HEAD_DIM; constexpr int NQ = QX_N_HEAD; constexpr int GQA = NQ / QX_N_HEAD_KV; const float scale = 1.0f / sqrtf((float)HD); const int kv_dim = QX_KV_DIM; /* Per-head block coverage: HD lanes → HD/VBLK = 8 blocks/head. K and V * share the same 32-elem block structure, so blk_row/lane mapping is * identical for both. */ constexpr int BLK_PER_HEAD = HD / VBLK; constexpr float NRM = 0.17677669529663687f; /* 1/sqrt(32) for inverse WHT */ /* FlashDecoding split count: one WG per (head, split) of ~QXMX_FD_CHUNK * tokens turns 24 serial WG scans into NQ*S parallel ones. S==1 keeps * the exact single-pass oracle path (shallow ctx, no scratch, QXMX_FD=0). */ static const int fd_chunk = [] { const char* e = std::getenv("QXMX_FD_CHUNK"); return e ? std::atoi(e) : 256; }(); static const bool fd_disabled = [] { const char* e = std::getenv("QXMX_FD"); return e && e[0] == '0'; }(); const int64_t n_kv = pos + 1; if (fd_max_split < 1) fd_max_split = 1; /* GQA-folded scan (default): one WG per (kv-head, split) covers all GQA * sibling query heads in a single pass — K dequant and V dequant are done * once per token instead of once per sibling WG (the llama.cpp fattn-vec * ncols=gqa_ratio pattern). QXMX_FOLD=0 reverts to one WG per * (query-head, split). Same fd_acc/fd_ml layout + merge kernel; per-head * math is unchanged apart from the dot reduction tree (SG butterfly + one * SLM combine vs reduce_over_group) — reassociation-only drift. */ static const bool fold_disabled = [] { const char* e = std::getenv("QXMX_FOLD"); return e && e[0] == '0'; }(); static const bool vec_disabled = [] { const char* e = std::getenv("QXMX_VEC"); return e && e[0] == '0'; }(); int S = 1; if (fd_acc && fd_ml && !fd_disabled && fd_chunk > 0) { /* Folded/vec WGs do GQA heads' work per pass — scale the split target * by GQA to keep the device saturated (NKV*S folded ~= NQ*S unfolded). */ int64_t s64 = (fold_disabled && vec_disabled) ? (n_kv + fd_chunk - 1) / fd_chunk : (n_kv * GQA + fd_chunk - 1) / fd_chunk; S = (int)(s64 < fd_max_split ? s64 : (int64_t)fd_max_split); } const int64_t fd_span = (n_kv + S - 1) / S; /* v3 llama-vec scan (default): FD-in-WG. One WG per (kv-head, head-group * of JG=3, split); each of the WG's 8 SGs scans its own strided token * stream with SG-LOCAL online-softmax state — zero barriers in the scan * (fold kernel: 2/token) and V dequant runs once per token (v1 vec: 8x * redundant). Lane owns 8 consecutive elems (4 lanes per 32-block): K = * byte loads (uint32 miscompiles at 2-mod-4 alignment on icpx 2026.1), * V = 1 uint32. acc[3][8] = 24 regs — sized to avoid the v1 spill * (VTune: 13KB). Scan ends with an 8-way SG merge through SLM (same * math as attn_fd_merge, once per WG), then the deferred inverse WHT. * QXMX_VEC=0 falls back to the GQA-fold kernel. */ if (!vec_disabled) { constexpr int NKV = QX_N_HEAD_KV; constexpr int NSG = HD / VBLK; /* 8 SGs/WG = token streams */ constexpr int JG = GQA / 2; /* heads per WG */ static_assert(GQA == 2 * JG, "v3 needs GQA = 2*JG"); q.submit([&](sycl::handler& h) { /* sl: ml[NSG*JG][2] | a_s[NSG*JG][HD] | q_s[JG][HD] */ sycl::local_accessor<float,1> sl(sycl::range<1>( (size_t)NSG*JG*2 + (size_t)NSG*JG*HD + (size_t)JG*HD), h); h.parallel_for(sycl::nd_range<1>((size_t)NKV * 2 * S * HD, HD), [=](sycl::nd_item<1> it) [[sycl::reqd_sub_group_size(VBLK)]] { float* ml = &sl[0]; float* a_s = &sl[NSG * JG * 2]; float* q_s = &sl[NSG * JG * 2 + NSG * JG * HD]; int g = it.get_group(0); int hk = g % NKV; int hg = (g / NKV) & 1; int sp = g / (NKV * 2); int lid = it.get_local_id(0); int lane = lid & 31; int sg_id = lid >> 5; auto sg = it.get_sub_group(); /* lane -> 8 elems: block lane>>2, offset (lane&3)*8 */ const int blk_row = hk * BLK_PER_HEAD + (lane >> 2); const int dim0 = (lane >> 2) * VBLK + (lane & 3) * 8; const int o_l = (lane & 3) * 8;#pragma unroll for (int j = 0; j < JG; j++) q_s[j * HD + lid] = q_vec[(size_t)(hk * GQA + hg * JG + j) * HD + lid]; sycl::group_barrier(it.get_group()); float acc[JG][8], max_s[JG], sum_s[JG];#pragma unroll for (int j = 0; j < JG; j++) { max_s[j] = -1e30f; sum_s[j] = 0.0f;#pragma unroll for (int i = 0; i < 8; i++) acc[j][i] = 0.0f; } const int64_t t0 = (int64_t)sp * fd_span; const int64_t t1 = (t0 + fd_span < n_kv) ? t0 + fd_span : n_kv; for (int64_t t = t0 + sg_id; t < t1; t += NSG) { /* K dequant: 8 elems per lane (t always valid). */ float k8[8]; if constexpr (std::is_same_v<Kct, KFp16>) { const sycl::half* kp = reinterpret_cast<const sycl::half*>(kv_k) + (size_t)t * kv_dim + hk * HD + dim0; const uint32_t* k32 = reinterpret_cast<const uint32_t*>(kp);#pragma unroll for (int w = 0; w < 4; w++) { uint32_t u = k32[w]; k8[2*w] = float(sycl::bit_cast<sycl::half>((uint16_t)(u & 0xFFFFu))); k8[2*w + 1] = float(sycl::bit_cast<sycl::half>((uint16_t)(u >> 16))); } } else { const uint8_t* kb = kv_k + (size_t)t * Kct::row_bytes + blk_row * KBLK_BYTES; uint16_t sbits = (uint16_t)kb[0] | ((uint16_t)kb[1] << 8); float d = float(sycl::bit_cast<sycl::half>(sbits));#pragma unroll for (int i = 0; i < 8; i++) { uint8_t code = kb[2 + o_l + i]; if constexpr (std::is_same_v<Kct, KQ8_0>) { k8[i] = (float)(int8_t)code * d; } else { /* KFp8 */ k8[i] = e4m3_to_fp32(code) * d; } } } /* V dequant once for this SG's own token. */ const uint32_t* db = kv_v + (size_t)t * VROW_U32 + blk_row * 5; uint16_t sc_u = (uint16_t)(db[0] & 0xFFFFu); uint16_t nc_u = (uint16_t)((db[0] >> 16) & 0xFFFFu); float vsn = float(sycl::bit_cast<sycl::half>(sc_u)) * float(sycl::bit_cast<sycl::half>(nc_u)); /* 32 codes x 4bit = one uint32 per (block, lane). */ uint32_t w0 = db[1 + (lane & 3)]; float v8[8];#pragma unroll for (int i = 0; i < 8; i++) v8[i] = V_LM_LUT[(w0 >> (i * 4)) & 0xFu] * vsn; /* Dots + SG-local online softmax per head. */#pragma unroll for (int j = 0; j < JG; j++) { float v = 0.0f;#pragma unroll for (int i = 0; i < 8; i++) v += q_s[j * HD + dim0 + i] * k8[i];#pragma unroll for (int off = VBLK >> 1; off > 0; off >>= 1) v += sycl::permute_group_by_xor(sg, v, off); float score = v * scale; float m_new = sycl::fmax(max_s[j], score); float f = sycl::native::exp(max_s[j] - m_new); float e = sycl::native::exp(score - m_new); sum_s[j] = f * sum_s[j] + e;#pragma unroll for (int i = 0; i < 8; i++) acc[j][i] = f * acc[j][i] + e * v8[i]; max_s[j] = m_new; } } /* 8-way SG partial merge through SLM (once per WG). Empty * SGs carry (m=-1e30, l=0, acc=0) → zero weight. */ if (lane == 0)#pragma unroll for (int j = 0; j < JG; j++) { ml[(sg_id * JG + j) * 2] = max_s[j]; ml[(sg_id * JG + j) * 2 + 1] = sum_s[j]; } sycl::group_barrier(it.get_group()); float mmax[JG], lsum[JG];#pragma unroll for (int j = 0; j < JG; j++) { float m = -1e30f;#pragma unroll for (int s = 0; s < NSG; s++) m = sycl::fmax(m, ml[(s * JG + j) * 2]); mmax[j] = m; float ls = 0.0f;#pragma unroll for (int s = 0; s < NSG; s++) ls += sycl::native::exp(ml[(s * JG + j) * 2] - m) * ml[(s * JG + j) * 2 + 1]; lsum[j] = ls; }#pragma unroll for (int j = 0; j < JG; j++) { float wj = sycl::native::exp(max_s[j] - mmax[j]);#pragma unroll for (int i = 0; i < 8; i++) a_s[(sg_id * JG + j) * HD + dim0 + i] = acc[j][i] * wj; } sycl::group_barrier(it.get_group()); /* Epilogue per head: cross-SG sum, deferred inverse WHT — * in-lane bits {1,2,4}, cross-lane bits {8,16} — then NRM. */#pragma unroll for (int j = 0; j < JG; j++) { float a[8];#pragma unroll for (int i = 0; i < 8; i++) a[i] = 0.0f;#pragma unroll for (int s = 0; s < NSG; s++)#pragma unroll for (int i = 0; i < 8; i++) a[i] += a_s[(s * JG + j) * HD + dim0 + i];#pragma unroll for (int bit = 1; bit < 8; bit <<= 1) { float na[8];#pragma unroll for (int i = 0; i < 8; i++) na[i] = (i & bit) ? a[i ^ bit] - a[i] : a[i] + a[i ^ bit];#pragma unroll for (int i = 0; i < 8; i++) a[i] = na[i]; }#pragma unroll for (int bit = 1; bit <= 2; bit <<= 1) { bool hi = (lane & bit) != 0;#pragma unroll for (int i = 0; i < 8; i++) { float p = sycl::permute_group_by_xor(sg, a[i], bit); a[i] = hi ? p - a[i] : a[i] + p; } } int hh = hk * GQA + hg * JG + j; if (S == 1) { float inv = lsum[j] > 0.0f ? NRM / lsum[j] : 0.0f;#pragma unroll for (int i = 0; i < 8; i++) out[(size_t)hh * HD + dim0 + i] = a[i] * inv; } else {#pragma unroll for (int i = 0; i < 8; i++) fd_acc[((size_t)sp * NQ + hh) * HD + dim0 + i] = a[i] * NRM; } } if (S > 1 && lid == 0) {#pragma unroll for (int j = 0; j < JG; j++) { int hh = hk * GQA + hg * JG + j; fd_ml[((size_t)sp * NQ + hh) * 2] = mmax[j]; fd_ml[((size_t)sp * NQ + hh) * 2 + 1] = lsum[j]; } } }); }); if (S > 1) attn_fd_merge(q, fd_acc, fd_ml, out, S); return; } if (!fold_disabled) { constexpr int NKV = QX_N_HEAD_KV; constexpr int NSG = HD / VBLK; /* SGs per WG */ static_assert(NQ == NKV * GQA, "GQA fold requires exact grouping"); q.submit([&](sycl::handler& h) { /* red: per-SG dot partials [NSG][GQA] + combined scores [GQA]. */ sycl::local_accessor<float,1> red(sycl::range<1>(NSG * GQA + GQA), h); h.parallel_for(sycl::nd_range<1>((size_t)NKV * S * HD, HD), [=](sycl::nd_item<1> it) [[sycl::reqd_sub_group_size(VBLK)]] { int g = it.get_group(0); int hk = g % NKV; int sp = g / NKV; int lid = it.get_local_id(0); int lane = lid & 31; int blk_row = hk * BLK_PER_HEAD + (lid >> 5); auto sg = it.get_sub_group(); const int sg_id = lid >> 5; float q_r[GQA], acc[GQA], max_s[GQA], sum_s[GQA];#pragma unroll for (int j = 0; j < GQA; j++) { q_r[j] = q_vec[(size_t)(hk * GQA + j) * HD + lid]; acc[j] = 0.0f; max_s[j] = -1e30f; sum_s[j] = 0.0f; } const int64_t t0 = (int64_t)sp * fd_span; const int64_t t1 = (t0 + fd_span < n_kv) ? t0 + fd_span : n_kv; for (int64_t t = t0; t < t1; t++) { /* K dequant once for the whole GQA group (per-Kct layout). */ float k_val; if constexpr (std::is_same_v<Kct, KFp16>) { const sycl::half* kp = reinterpret_cast<const sycl::half*>(kv_k); k_val = float(kp[(size_t)t * kv_dim + hk * HD + lid]); } else { const uint8_t* kb = kv_k + (size_t)t * Kct::row_bytes + blk_row * KBLK_BYTES; uint16_t sbits = (uint16_t)kb[0] | ((uint16_t)kb[1] << 8); float d = float(sycl::bit_cast<sycl::half>(sbits)); uint8_t code = kb[2 + lane]; if constexpr (std::is_same_v<Kct, KQ8_0>) { k_val = (float)(int8_t)code * d; } else { /* KFp8 */ k_val = e4m3_to_fp32(code) * d; } } /* GQA dots: SG xor-butterfly per head (all lanes end with * the SG sum, no barrier), then ONE SLM exchange combines * all GQA partials at once — 2 barriers/token total, vs 2 * per head per WG across 6 sibling WGs unfolded. */#pragma unroll for (int j = 0; j < GQA; j++) { float v = q_r[j] * k_val;#pragma unroll for (int off = VBLK >> 1; off > 0; off >>= 1) v += sycl::permute_group_by_xor(sg, v, off); if (lane == 0) red[sg_id * GQA + j] = v; } sycl::group_barrier(it.get_group()); if (lid < GQA) { float s6 = 0.0f;#pragma unroll for (int b = 0; b < NSG; b++) s6 += red[b * GQA + lid]; red[NSG * GQA + lid] = s6; } sycl::group_barrier(it.get_group()); /* V dequant once; feeds all GQA sibling accumulators. */ const uint32_t* db = kv_v + (size_t)t * VROW_U32 + blk_row * 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 word = db[1 + (lane >> 3)]; uint32_t code = (word >> ((lane & 7) * 4)) & 0xFu; float v_pre = V_LM_LUT[code] * (vscale * vnorm_corr);#pragma unroll for (int j = 0; j < GQA; j++) { float score = red[NSG * GQA + j] * scale; float m_new = sycl::fmax(max_s[j], score); float f = sycl::native::exp(max_s[j] - m_new); float e = sycl::native::exp(score - m_new); sum_s[j] = f * sum_s[j] + e; acc[j] = f * acc[j] + e * v_pre; max_s[j] = m_new; } } /* Per-head deferred inverse WHT + write (same as unfolded). */#pragma unroll for (int j = 0; j < GQA; j++) { float a = acc[j]; for (int bit = 1; bit < VBLK; bit <<= 1) { float partner = sycl::permute_group_by_xor(sg, a, bit); a = (lane & bit) ? (partner - a) : (a + partner); } a *= NRM; int hh = hk * GQA + j; if (S == 1) { out[(size_t)hh * HD + lid] = a / sum_s[j]; } else { fd_acc[((size_t)sp * NQ + hh) * HD + lid] = a; } } if (S > 1 && lid == 0) {#pragma unroll for (int j = 0; j < GQA; j++) { int hh = hk * GQA + j; fd_ml[((size_t)sp * NQ + hh) * 2] = max_s[j]; fd_ml[((size_t)sp * NQ + hh) * 2 + 1] = sum_s[j]; } } }); }); if (S > 1) attn_fd_merge(q, fd_acc, fd_ml, out, S); return; } q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>((size_t)NQ * S * HD, HD), [=](sycl::nd_item<1> it) [[sycl::reqd_sub_group_size(VBLK)]] { int g = it.get_group(0); int hh = g % NQ; int sp = g / NQ; int hk = hh / GQA; int lid = it.get_local_id(0); int blk_local = lid >> 5; /* 0..7 within the head */ int lane = lid & 31; /* 0..31 within the block */ int blk_row = hk * BLK_PER_HEAD + blk_local; auto sg = it.get_sub_group(); /* Lane owns exactly one q element for the whole scan -- register, * not SLM (was a per-token SLM read after a staging barrier). */ const float q_l = q_vec[(size_t)hh * HD + lid]; float acc_pre = 0.0f, max_s = -1e30f, sum_s = 0.0f; const int64_t t0 = (int64_t)sp * fd_span; const int64_t t1 = (t0 + fd_span < n_kv) ? t0 + fd_span : n_kv; for (int64_t t = t0; t < t1; t++) { /* K dequant — compile-time selected per Kct. Thread lid reads * its own K element (hk*HD + lid) from the per-format layout. */ float k_val; if constexpr (std::is_same_v<Kct, KFp16>) { const sycl::half* kp = reinterpret_cast<const sycl::half*>(kv_k); k_val = float(kp[(size_t)t * kv_dim + hk * HD + lid]); } else { /* q8_0 / fp8 share the 34-B block layout (fp16 scale + 32 * codes). blk_row maps to the block; lane to the code. */ const uint8_t* kb = kv_k + (size_t)t * Kct::row_bytes + blk_row * KBLK_BYTES; uint16_t sbits = (uint16_t)kb[0] | ((uint16_t)kb[1] << 8); float d = float(sycl::bit_cast<sycl::half>(sbits)); uint8_t code = kb[2 + lane]; if constexpr (std::is_same_v<Kct, KQ8_0>) { k_val = (float)(int8_t)code * d; } else { /* KFp8 */ k_val = e4m3_to_fp32(code) * d; } } float dot = q_l * k_val; float score = sycl::reduce_over_group(it.get_group(), dot, sycl::plus<float>()) * scale; float m_new = sycl::fmax(max_s, score); /* native::exp: the M1.7c fused FA kernel (qxmx_fa_fused.cpp:284) * uses the fast hardware exp for the softmax here; this decode * path is the same recurrence and the inputs (score - m_new, * max_s - m_new) are bounded softmax diffs, so native is safe. * Was sycl::exp -- 2 of these per token, ~12.7M calls/step at * 33k ctx, so the cheaper intrinsic matters on the scan's * critical path. */ float f = sycl::native::exp(max_s - m_new); float e = sycl::native::exp(score - m_new); sum_s = f * sum_s + e; /* Phase 1 V dequant — T13 deferred WHT: fold vscale and * vnorm_corr into the pre-WHT dequant; skip the butterfly * here (it would run 5 permute_xor per (token, block)). * Apply both WHT and NRM once per output row after the scan. * Valid because WHT is linear and vscale/vnorm_corr are * per-block scalars that commute through it (see * qxmx_fa.cpp:185 for the prefill T13 reference). */ const uint32_t* db = kv_v + (size_t)t * VROW_U32 + blk_row * 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)); int w = lane >> 3; /* 0..3 */ uint32_t word = db[1 + w]; /* one load, not four + select */ uint32_t code = (word >> ((lane & 7) * 4)) & 0xFu; float v_pre = V_LM_LUT[code] * (vscale * vnorm_corr); acc_pre = f * acc_pre + e * v_pre; max_s = m_new; } /* Inverse WHT once per (block, lane) at scan end — T13 deferral. * Same self-inverse permute_xor butterfly as before, but now it * runs ~n_kv times less often. NRM (1/sqrt(32)) applied after. */ for (int bit = 1; bit < VBLK; bit <<= 1) { float partner = sycl::permute_group_by_xor(sg, acc_pre, bit); acc_pre = (lane & bit) ? (partner - acc_pre) : (acc_pre + partner); } acc_pre = acc_pre * NRM; if (S == 1) { out[(size_t)hh * HD + lid] = acc_pre / sum_s; return; } fd_acc[((size_t)sp * NQ + hh) * HD + lid] = acc_pre; if (lid == 0) { fd_ml[((size_t)sp * NQ + hh) * 2] = max_s; fd_ml[((size_t)sp * NQ + hh) * 2 + 1] = sum_s; } }); }); if (S > 1) attn_fd_merge(q, fd_acc, fd_ml, out, S);}
/* Causal conv1d with in-place history update. Safe for x_new == out: each * thread reads x_new[c] into a register before writing out[c]. hist is * updated in-place: each thread owns column c and reads its k-1 history * values into registers before writing any. */static void gpu_conv1d_step(sycl::queue& q, const float* wc, float* hist, const float* x_new, float* out, int k, int channels) { q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<1>(channels), [=](sycl::id<1> idx) { int c = idx; const float* wcc = wc + c * k; float xn = x_new[c]; float hvals[QX_SSM_CONV_K - 1]; float acc = wcc[k - 1] * xn; for (int j = 0; j < k - 1; j++) { hvals[j] = hist[(size_t)j * channels + c]; acc += wcc[j] * hvals[j]; } out[c] = acc; for (int j = 0; j < k - 2; j++) hist[(size_t)j * channels + c] = hvals[j + 1]; hist[(size_t)(k - 2) * channels + c] = xn; }); });}
/* M1.6 S4: batched conv1d over n_tok tokens. One thread per channel, loops * t=0..n_tok-1 sequentially (token t's history depends on token t-1's input). * History is carried in registers across the chunk, written back once. Same * fp ops per (channel, token) as the decode oracle -> bit-exact. hist is the * pre-chunk state on input, post-chunk state on output (in-place RMW; must be * malloc_shared per the SYCL gotchas doc). x_new/out are [n_tok][channels]. */void gpu_conv1d_step_b(sycl::queue& q, const float* wc, float* hist, const float* x_new, float* out, int k, int channels, int n_tok, int valid_count /*=-1*/) { if (valid_count < 0) valid_count = n_tok; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<1>(channels), [=](sycl::id<1> idx) { int c = idx; const float* wcc = wc + c * k; float hvals[QX_SSM_CONV_K - 1]; for (int j = 0; j < k - 1; j++) hvals[j] = hist[(size_t)j * channels + c]; /* Only the first valid_count tokens advance the conv state. * The output for tokens >= valid_count is left as whatever was * in `out` before the call (caller discards those rows). */ for (int t = 0; t < valid_count; t++) { float xn = x_new[(size_t)t * channels + c]; float acc = wcc[k - 1] * xn; for (int j = 0; j < k - 1; j++) acc += wcc[j] * hvals[j]; out[(size_t)t * channels + c] = acc; for (int j = 0; j < k - 2; j++) hvals[j] = hvals[j + 1]; hvals[k - 2] = xn; } for (int j = 0; j < k - 1; j++) hist[(size_t)j * channels + c] = hvals[j]; }); });}
/* Like tmatvec but skips quant_pack — producer already filled g_bvnni_d / * g_suma_d / g_act_scale_d (e.g. gpu_rmsnorm_quant). */static void tmatvec_q(sycl::queue& q, const dev_q2* w, float* y, float* part) { gemv_dpas8(q, w->codes, w->scales, g_bvnni_d, g_suma_d, g_act_scale_d, y, part, w->M, w->K);}
/* Like tmatvec_q but accumulates into y (residual add folded into gemv * epilogue): y[i] = proj_out[i] + accum[i]. Bit-exact vs tmatvec_q followed * by gpu_residual_add (same two fp roundings, same order). */static void tmatvec_accum(sycl::queue& q, const dev_q2* w, float* y, const float* accum, float* part) { gemv_dpas8(q, w->codes, w->scales, g_bvnni_d, g_suma_d, g_act_scale_d, y, part, w->M, w->K, accum);}
/* Like tmatvec_gate_up but skips quant_pack — producer already quantized. */static void tmatvec_gate_up_q(sycl::queue& q, const dev_q2* w0, const dev_q2* w1, float* y0, float* y1, float* part) { gemv_dpas8_dual(q, w0->codes, w0->scales, w1->codes, w1->scales, g_bvnni_d, g_suma_d, g_act_scale_d, y0, y1, part, w0->M, w0->K);}
void engine_set_quant_stage(int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d) { g_bvnni_d = bvnni_d; g_suma_d = suma_d; g_act_scale_d = act_scale_d;}
/* Host-side RoPE table precompute for prefill (see qxmx_gpu.h). Matches * decode's exact double-then-float math so per-position output is bit-exact * vs decode at the same position. */void engine_compute_rope_table(int64_t pos_start, int n_tok, float* rope_c, float* rope_s) { int half = QX_ROT_DIM / 2; for (int tok = 0; tok < n_tok; tok++) { int64_t pos = pos_start + tok; float* cr = rope_c + (size_t)tok * half; float* sr = rope_s + (size_t)tok * half; for (int i = 0; i < half; i++) { double freq = 1.0 / pow((double)QX_ROPE_FREQ_BASE, (double)(2 * i) / QX_ROT_DIM); float a = (float)(pos * freq); cr[i] = cosf(a); sr[i] = sinf(a); } }}
void ffn_forward(sycl::queue& q, const dev_ffn_blk* f, float* x, float* part, float* g, float* u, float* t) { /* rmsnorm producer packs once; gate+up dual gemv reuses that pack. */ gpu_rmsnorm_quant(q, x, f->post_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_gate_up_q(q, &f->gate, &f->up, g, u, part); /* swiglu producer packs; down matvec skips quant_pack. */ gpu_swiglu_quant(q, g, u, t, QX_FF_INTER, g_bvnni_d, g_suma_d, g_act_scale_d); /* T6: down-proj writes into x with accum=x (residual add folded into epilogue). */ tmatvec_accum(q, &f->down, x, x, part);}
/* M1.6 S1: batched FFN over M tokens via gemm_gpu_v4 (see qxmx_gpu.h). */void ffn_forward_b(sycl::queue& q, const dev_ffn_blk* f, float* x, int M, int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d, float* g, float* u, float* t) { static bool prof = []{ return std::getenv("QXMX_PROFILE") != nullptr; }(); auto phase = [&](const char* tag) { if (!prof) return; q.wait(); static auto last = std::chrono::steady_clock::now(); auto t = std::chrono::steady_clock::now(); if (std::getenv("QXMX_FFN_DEBUG")) fprintf(stderr, " [ffn %s] %.2f ms\n", tag, std::chrono::duration<double, std::milli>(t - last).count()); last = t; }; /* transpose_out=true: gemm_gpu_v4 writes [N_tok][M_dim] (token-major), so * the engine's existing token-major producers and residual fold consume * the GEMM output directly — no layout adapter needed. One rmsnorm_quant * pack serves both gate and up GEMMs (decode fuses them via gemv_dpas8_dual; * the batched GEMM has no dual variant, so two launches at M=256). */ phase("start"); gpu_rmsnorm_quant(q, x, f->post_norm, QX_D_MODEL, bvnni_d, suma_d, act_scale_d, M); phase("rmsnorm"); gemm_gpu_v4_dual(q, f->gate.codes, f->gate.scales, f->up.codes, f->up.scales, bvnni_d, suma_d, act_scale_d, g, u, f->gate.M, f->gate.K, M, /*accum0=*/nullptr, /*accum1=*/nullptr, /*sync=*/false, /*transpose_out=*/true); phase("gate+up"); gpu_swiglu_quant(q, g, u, t, QX_FF_INTER, bvnni_d, suma_d, act_scale_d, M); phase("swiglu"); /* T6: down-proj folds residual via gemm_gpu_v4 accum=x (token-major both). */ gemm_gpu_v4(q, f->down.codes, f->down.scales, bvnni_d, suma_d, act_scale_d, x, f->down.M, f->down.K, M, /*accum=*/x, false, /*transpose_out=*/true); phase("down");}
void attn_forward(sycl::queue& q, const dev_attn_blk* a, float* x, uint8_t* kv_k, uint32_t* kv_v, int64_t pos, engine::KCT k_ctk, float* part, float* proj_buf, float* qvec, float* gate, float* kvec, float* vvec, float* attn_out_buf, const float* rope_c, const float* rope_s, float* fd_acc, float* fd_ml, int fd_max_split) { /* One rmsnorm_quant; q/k/v all share the packed activation. */ gpu_rmsnorm_quant(q, x, a->attn_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &a->q, proj_buf, part); gpu_attn_split_qgate(q, proj_buf, qvec, gate, QX_N_HEAD, QX_HEAD_DIM); tmatvec_q(q, &a->k, kvec, part); tmatvec_q(q, &a->v, vvec, part); gpu_group_rmsnorm(q, qvec, a->q_norm, qvec, QX_N_HEAD, QX_HEAD_DIM); gpu_group_rmsnorm(q, kvec, a->k_norm, kvec, QX_N_HEAD_KV, QX_HEAD_DIM); gpu_apply_rope(q, qvec, kvec, QX_N_HEAD, QX_N_HEAD_KV, QX_HEAD_DIM, QX_ROT_DIM, rope_c, rope_s); /* K cache write — per-format quant. Row stride differs per format. */ size_t k_row_off; if (k_ctk == engine::KCT::fp16) { k_row_off = (size_t)pos * KROW_BYTES_F16; gpu_fp32_to_fp16(q, kvec, reinterpret_cast<sycl::half*>(kv_k + k_row_off), QX_KV_DIM); } else if (k_ctk == engine::KCT::q8_0) { k_row_off = (size_t)pos * KROW_BYTES_Q8; gpu_quant_k_row_q8(q, kvec, kv_k + k_row_off); } else { k_row_off = (size_t)pos * KROW_BYTES_Q8; gpu_quant_k_row_fp8(q, kvec, kv_k + k_row_off); } /* Phase 1: V cache stored as packed 4-bit TurboQuant blocks. */ gpu_quant_v_row(q, vvec, kv_v + (size_t)pos*VROW_U32); /* Attn dequant — compile-time-selected template per K format. */ if (k_ctk == engine::KCT::fp16) gpu_attn_softmax<KFp16>(q, qvec, kv_k, kv_v, attn_out_buf, pos, fd_acc, fd_ml, fd_max_split); else if (k_ctk == engine::KCT::q8_0) gpu_attn_softmax<KQ8_0>(q, qvec, kv_k, kv_v, attn_out_buf, pos, fd_acc, fd_ml, fd_max_split); else gpu_attn_softmax<KFp8>(q, qvec, kv_k, kv_v, attn_out_buf, pos, fd_acc, fd_ml, fd_max_split); /* sigmoid_mul producer packs; o-proj writes into x with accum=x (T6). */ gpu_sigmoid_mul_quant(q, attn_out_buf, gate, QX_Q_DIM, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_accum(q, &a->o, x, x, part);}
/* M1.6 S3: batched attention block over M tokens. Mirrors decode attn_forward * but M-token-shaped: q/k/v via gemm_gpu_v4 (transpose_out=true, token-major), * S2 elementwise _b kernels, the validated M1.4 KV-quant + gpu_attn_prefill, * and the o-proj residual fold via gemm_gpu_v4 accum=x. KV cache rows * [pos_start..pos_start+M-1] are written; gpu_attn_prefill reads the whole * cache (history + chunk) with per-token causal cap. Greedy-token gate: * GEMM epilogue reassociation drift is OK (per ground rule 2). */void attn_forward_b(sycl::queue& q, const dev_attn_blk* a, float* x, int M, uint8_t* kv_k, uint32_t* kv_v, int64_t pos_start, engine::KCT k_ctk, int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d, float* proj_buf, float* qvec, float* gate, float* kvec, float* vvec, float* attn_out_buf, const float* rope_c, const float* rope_s, fa_split_workspace* fa_ws) { /* One rmsnorm_quant pack serves q/k/v GEMMs (all K=D_MODEL=5120). */ gpu_rmsnorm_quant(q, x, a->attn_norm, QX_D_MODEL, bvnni_d, suma_d, act_scale_d, M); /* q projection: [M][Q_DIM*2=12288] (q+gate interleaved), then split. */ gemm_gpu_v4(q, a->q.codes, a->q.scales, bvnni_d, suma_d, act_scale_d, proj_buf, a->q.M, a->q.K, M, nullptr, false, /*transpose_out=*/true); gpu_attn_split_qgate_b(q, proj_buf, qvec, gate, QX_N_HEAD, QX_HEAD_DIM, M); gemm_gpu_v4(q, a->k.codes, a->k.scales, bvnni_d, suma_d, act_scale_d, kvec, a->k.M, a->k.K, M, nullptr, false, /*transpose_out=*/true); gemm_gpu_v4(q, a->v.codes, a->v.scales, bvnni_d, suma_d, act_scale_d, vvec, a->v.M, a->v.K, M, nullptr, false, /*transpose_out=*/true); gpu_group_rmsnorm_b(q, qvec, a->q_norm, qvec, QX_N_HEAD, QX_HEAD_DIM, M); gpu_group_rmsnorm_b(q, kvec, a->k_norm, kvec, QX_N_HEAD_KV, QX_HEAD_DIM, M); gpu_apply_rope_b(q, qvec, kvec, QX_N_HEAD, QX_N_HEAD_KV, QX_HEAD_DIM, QX_ROT_DIM, rope_c, rope_s, M); /* KV cache writes — batched quant at pos_start..pos_start+M-1. * Row stride differs per K format; V is always TurboQuant 4-bit. */ if (k_ctk == engine::KCT::fp16) { gpu_fp32_to_fp16(q, kvec, reinterpret_cast<sycl::half*>(kv_k + (size_t)pos_start * KROW_BYTES_F16), QX_KV_DIM, M); } else if (k_ctk == engine::KCT::q8_0) { gpu_quant_k_row_q8(q, kvec, kv_k + (size_t)pos_start * KROW_BYTES_Q8, M); } else { gpu_quant_k_row_fp8(q, kvec, kv_k + (size_t)pos_start * KROW_BYTES_Q8, M); } gpu_quant_v_row(q, vvec, kv_v + (size_t)pos_start * VROW_U32, M); /* M1.7b causal attention (history + chunk, per-token causal cap). * FlashAttention-2 fused kernel (qxmx_fa): one WG per (q-head, query * token), so any M works (no BLK_Q divisibility requirement). This * replaces the M1.4 naive O(M*P) gpu_attn_prefill entirely for the * chunk path; the <16 tail still uses decode (see engine::prefill). */ if (k_ctk == engine::KCT::fp16) gpu_flash_attn<KFp16>(q, qvec, kv_k, kv_v, attn_out_buf, pos_start, M, fa_ws); else if (k_ctk == engine::KCT::q8_0) gpu_flash_attn<KQ8_0>(q, qvec, kv_k, kv_v, attn_out_buf, pos_start, M, fa_ws); else gpu_flash_attn<KFp8 >(q, qvec, kv_k, kv_v, attn_out_buf, pos_start, M, fa_ws); /* sigmoid_mul producer packs; o-proj residual fold via accum=x. */ gpu_sigmoid_mul_quant(q, attn_out_buf, gate, QX_Q_DIM, bvnni_d, suma_d, act_scale_d, M); gemm_gpu_v4(q, a->o.codes, a->o.scales, bvnni_d, suma_d, act_scale_d, x, a->o.M, a->o.K, M, /*accum=*/x, false, /*transpose_out=*/true);}
void ssm_forward(sycl::queue& q, const dev_ssm_blk* s, float* ssm_state, float* ssm_conv, float* x, float* part, float* proj_buf, float* alpha_buf, float* beta_buf, float* ssm_out_buf, float* zg_buf) { /* One rmsnorm_quant; qkv/alpha/beta/gate all share the packed activation. */ gpu_rmsnorm_quant(q, x, s->attn_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &s->qkv, proj_buf, part); /* GPU conv1d: x_new == out == proj_buf is safe (kernel reads xn into a * register before writing out[c]). ssm_conv must be malloc_shared for * correct in-place history RMW (malloc_host gives stale reads on Xe2). */ gpu_conv1d_step(q, s->conv1d, ssm_conv, proj_buf, proj_buf, QX_SSM_CONV_K, QX_SSM_CONVD); gpu_silu_inplace(q, proj_buf, QX_SSM_CONVD); /* qk = [q || k] = proj_buf[0 .. 2*KEYD); v = proj_buf + 2*KEYD. * Pass proj_buf directly (no separate qk_comb copy). */ tmatvec_q(q, &s->alpha, alpha_buf, part); tmatvec_q(q, &s->beta, beta_buf, part); /* T11: fused DeltaNet (single-sweep S update in SLM) + per-V-head * group_rmsnorm, one launch. Replaces gpu_deltanet + gpu_group_rmsnorm. */ gpu_deltanet_rmsnorm(q, proj_buf, proj_buf + (size_t)2*QX_SSM_KEYD, alpha_buf, beta_buf, s->a_log, s->dt_bias, s->ssm_norm, ssm_state, ssm_out_buf); tmatvec_q(q, &s->gate, zg_buf, part); /* silu_mul producer packs; out-proj writes into x with accum=x (T6). */ gpu_silu_mul_quant(q, ssm_out_buf, zg_buf, QX_SSM_VALD, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_accum(q, &s->out, x, x, part);}
/* M1.6 S4: batched SSM block over M tokens. Mirrors decode ssm_forward but * M-token-shaped: qkv GEMM (transpose_out), S4 conv1d_step_b (sequential * in-chunk, history in registers), S2 silu_inplace_b, alpha/beta GEMMs, * validated M1.5 gpu_deltanet_rmsnorm_b (sequential SLM scan; qk/v read from * the fused proj_buf via stride params), gate GEMM, silu_mul_quant, and the * out-proj residual fold via gemm_gpu_v4 accum=x. * x [M][D_MODEL] residual stream (in/out) * proj_buf [M][CONVD=10240] qkv fused (qk=[0..2*KEYD), v=[2*KEYD..CONVD)) * alpha/beta [M][NVH=48]; zg_buf/ssm_out_buf [M][VALD=6144] * ssm_state [NVH*HD*HD] (in/out, sequential RMW across the chunk) * ssm_conv [(K-1)*CONVD] (in/out, sequential RMW; malloc_shared) * bvnni_d/suma_d: max K = max(D_MODEL, VALD, CONVD) = FF_INTER works * Requires M % 16 == 0. */void ssm_forward_b(sycl::queue& q, const dev_ssm_blk* s, float* ssm_state, float* ssm_conv, float* x, int M, int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d, float* proj_buf, float* alpha_buf, float* beta_buf, float* ssm_out_buf, float* zg_buf, wyut_workspace* wyut_ws, int valid_count /*=-1*/) { if (valid_count < 0) valid_count = M; static bool prof = []{ return std::getenv("QXMX_PROFILE") != nullptr; }(); double t_norm = 0, t_qkv = 0, t_conv = 0, t_ab = 0, t_wyut = 0, t_gate_out = 0; auto phase = [&](double& acc) { if (!prof) return; q.wait(); static auto last = std::chrono::steady_clock::now(); auto t = std::chrono::steady_clock::now(); acc += std::chrono::duration<double, std::milli>(t - last).count(); last = t; }; if (prof) { auto t = std::chrono::steady_clock::now(); (void)t; } /* mark */
gpu_rmsnorm_quant(q, x, s->attn_norm, QX_D_MODEL, bvnni_d, suma_d, act_scale_d, M); phase(t_norm); /* rmsnorm_quant alone */ gemm_gpu_v4(q, s->qkv.codes, s->qkv.scales, bvnni_d, suma_d, act_scale_d, proj_buf, s->qkv.M, s->qkv.K, M, nullptr, false, /*transpose_out=*/true); phase(t_qkv); /* qkv GEMM alone */ /* conv1d + silu in-place on proj_buf ([M][CONVD]). valid_count limits the * conv state advance to the first valid_count tokens (tail-batched prefill). */ gpu_conv1d_step_b(q, s->conv1d, ssm_conv, proj_buf, proj_buf, QX_SSM_CONV_K, QX_SSM_CONVD, M, valid_count); gpu_silu_inplace_b(q, proj_buf, QX_SSM_CONVD, M); phase(t_conv); /* alpha/beta GEMMs (M=NVH=48 each, tiny). */ gemm_gpu_v4(q, s->alpha.codes, s->alpha.scales, bvnni_d, suma_d, act_scale_d, alpha_buf, s->alpha.M, s->alpha.K, M, nullptr, false, /*transpose_out=*/true); gemm_gpu_v4(q, s->beta.codes, s->beta.scales, bvnni_d, suma_d, act_scale_d, beta_buf, s->beta.M, s->beta.K, M, nullptr, false, /*transpose_out=*/true); phase(t_ab); /* DeltaNet scan: read qk/v from the fused proj_buf (stride CONVD), write * ssm_out at stride VALD. Dispatched via gpu_deltanet_batched (meson- * selected: wyut engine source set -> WY/UT M1.7a, naive -> M1.7b.12). * Greedy-token gate accepts the WY/UT + tf32 reassociation drift (~0.01 * abs). ws is the wyut_workspace (engine member; ignored by naive). * valid_count: naive honors it (per-token loop stops early); wyut ignores * it (the WY/UT state is a T-wide batched matmul, can't be sliced). The * tail-batched prefill path is only used by the main engine (naive build). */ gpu_deltanet_batched(q, proj_buf, proj_buf + (size_t)2*QX_SSM_KEYD, alpha_buf, beta_buf, s->a_log, s->dt_bias, s->ssm_norm, ssm_state, ssm_out_buf, M, /*qk_stride=*/QX_SSM_CONVD, /*v_stride=*/QX_SSM_CONVD, /*out_stride=*/QX_SSM_VALD, /*ws=*/wyut_ws, /*valid_count=*/valid_count); phase(t_wyut); gemm_gpu_v4(q, s->gate.codes, s->gate.scales, bvnni_d, suma_d, act_scale_d, zg_buf, s->gate.M, s->gate.K, M, nullptr, false, /*transpose_out=*/true); gpu_silu_mul_quant(q, ssm_out_buf, zg_buf, QX_SSM_VALD, bvnni_d, suma_d, act_scale_d, M); gemm_gpu_v4(q, s->out.codes, s->out.scales, bvnni_d, suma_d, act_scale_d, x, s->out.M, s->out.K, M, /*accum=*/x, false, /*transpose_out=*/true); phase(t_gate_out);
if (prof) { static int call = 0; call++; if (call <= 2 || call % 48 == 0) fprintf(stderr, " [ssm L%d] norm=%.1f qkv=%.1f conv=%.1f ab=%.1f wyut=%.1f gate+out=%.1f\n", 0, t_norm, t_qkv, t_conv, t_ab, t_wyut, t_gate_out); }}
static dev_q2 upload_q2(sycl::queue& q, const qx_weight* w) { int K = w->dims[0], M = w->dims[1]; /* Detect group size from on-disk footprint: g64 = 18 bytes per 64 elems, * g128 = 34 bytes per 128 elems = 17 bytes per 64 elems. */ int qk2 = (w->nbytes * 64 / w->n_elem == 18) ? 64 : 128; int ngrp = K / 64; /* always g64 groups after repack */ /* Pad M to a multiple of 256 (gemm_gpu_v4's WGM8=32 * 8 = 256-row A-load * tile). The batched GEMM unconditionally loads A_codes/A_sc_t for a full * WGM8*8-row tile even when M isn't a multiple of it; the writes are * guarded (m8t >= M8 -> skip) but the loads are not. For the model's * alpha/beta projections (M=DT_RANK=48) the unpadded loads read unmapped * USM -> GPU page fault -> UR_RESULT_ERROR_DEVICE_LOST. Padding the * allocation (zeroed) makes the OOB reads safe; their outputs are * discarded by the write guard. * A_codes layout [M/8][ngrp][2][64B] has M/8 outermost, so end-padding * suffices. A_sc_t [ngrp][M] uses stride M: non-last-group OOB reads * land in the next group (in-bounds), only the last group needs a tail * pad of WGM8*8 = 256 floats. */ int M_pad = (M + 255) & ~255; size_t codes_sz = (size_t)M_pad * ngrp * 16; size_t sc_sz = (size_t)ngrp * M + 256; uint8_t* codes_h = (uint8_t*)calloc(1, codes_sz); float* sc_h = (float*)calloc(1, sc_sz * sizeof(float)); if (qk2 == 64) repack_q2_g64((const qx_block_q2*)w->data, M, K, codes_h, sc_h); else repack_q2_g128((const qx_block_q2_g128*)w->data, M, K, codes_h, sc_h); dev_q2 d; d.codes = chk_shared<uint8_t>(q, codes_sz, "q2.codes"); d.scales = chk_shared<float>(q, sc_sz, "q2.scales"); d.M = M; d.K = K; d.src = w; 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); return d;}
static float* upload_f32(sycl::queue& q, const qx_weight* w) { float* d = chk_shared<float>(q, w->n_elem, "upload_f32"); q.memcpy(d, w->data, w->nbytes).wait(); return d;}
/* USM allocation failure -> nullptr. A subsequent q.fill/q.memset on that * null pointer throws the cryptic "NULL pointer argument in memory fill * operation" from the oneAPI runtime, hiding the real cause (OOM, or shared * USM unsupported under a given backend — e.g. WSL Level-Zero). The checked * allocators (qx::chk_shared/chk_device, header-only in qxmx_device.h) abort * with the buffer name, size, and device instead. */static float* dev_shared_alloc(sycl::queue& q, size_t n, const char* what) { return qx::chk_shared<float>(q, n, what);}
bool engine::init(qx_model* m, int mc) { max_ctx = mc; { /* Enable profiling unconditionally (cheap): QXMX_PROFILE=2 reads * event timestamps; QXMX_PROFILE=1/0 never query them. */ using namespace sycl::property::queue; q = sycl::queue(sycl::gpu_selector_v, {in_order{}, enable_profiling{}}); } output = upload_q2(q, &m->output); output_norm = upload_f32(q, &m->output_norm); int ai = 0, si = 0; for (int L = 0; L < QX_N_LAYER; L++) { qx_block* b = &m->blocks[L]; if (b->is_attention) { dev_attn_blk* a = &attns[ai]; a->attn_norm = upload_f32(q, &b->attn.attn_norm); a->q = upload_q2(q, &b->attn.q); a->k = upload_q2(q, &b->attn.k); a->v = upload_q2(q, &b->attn.v); a->q_norm = upload_f32(q, &b->attn.q_norm); a->k_norm = upload_f32(q, &b->attn.k_norm); a->o = upload_q2(q, &b->attn.o); ai++; } else { dev_ssm_blk* s = &ssms[si]; s->attn_norm = upload_f32(q, &b->ssm.attn_norm); s->qkv = upload_q2(q, &b->ssm.qkv); s->conv1d = upload_f32(q, &b->ssm.conv1d); s->a_log = upload_f32(q, &b->ssm.a); s->dt_bias = upload_f32(q, &b->ssm.dt_bias); s->alpha = upload_q2(q, &b->ssm.alpha); s->beta = upload_q2(q, &b->ssm.beta); s->ssm_norm = upload_f32(q, &b->ssm.ssm_norm); s->gate = upload_q2(q, &b->ssm.gate); s->out = upload_q2(q, &b->ssm.out); si++; } dev_ffn_blk* f = &ffns[L]; f->post_norm = upload_f32(q, &b->ffn.post_norm); f->gate = upload_q2(q, &b->ffn.gate); f->up = upload_q2(q, &b->ffn.up); f->down = upload_q2(q, &b->ffn.down); } n_attn = ai; n_ssm = si; token_embd_raw.M = m->token_embd.dims[1]; token_embd_raw.K = m->token_embd.dims[0]; /* Per-slot KV/SSM state is allocated in open_slot (M2.1), not here. */ q.wait(); /* Hot activation buffers: sized for `chunk` tokens (the runtime chunk size, * env QXMX_CHUNK, default 256). Host touches only x (host writes * embeddings), logits (host reads argmax), and rope_c/s (host writes the * table) — those stay malloc_shared (coherent). Total ~120 MB at chunk=256, * trivial vs the 7.8 GB weights. Changing QXMX_CHUNK requires re-init. */ /* Runtime chunk size (env override; default 256, the gemm_gpu_v4 sweet * spot). Must be a multiple of 16 (gemm_gpu_v4 N-tile constraint). */ if (const char* cv = std::getenv("QXMX_CHUNK")) { int v = std::atoi(cv); if (v >= 16 && (v % 16) == 0) chunk = v; else std::fprintf(stderr, "QXMX_CHUNK=%s invalid (need >=16, %%16); using %d\n", cv, chunk); } int C = chunk; x = dev_shared_alloc(q, (size_t)C * QX_D_MODEL, "x"); proj_buf = dev_shared_alloc(q, (size_t)C * 2 * QX_Q_DIM, "proj_buf"); /* > VOCAB at C=256 */ part_buf = dev_shared_alloc(q, (size_t)8 * QX_N_VOCAB, "part_buf"); logits = dev_shared_alloc(q, QX_N_VOCAB, "logits"); alpha_buf = dev_shared_alloc(q, (size_t)C * QX_SSM_DT_RANK, "alpha_buf"); beta_buf = dev_shared_alloc(q, (size_t)C * QX_SSM_DT_RANK, "beta_buf"); ssm_out_buf = dev_shared_alloc(q, (size_t)C * QX_SSM_VALD, "ssm_out_buf"); zg_buf = dev_shared_alloc(q, (size_t)C * QX_SSM_VALD, "zg_buf"); ffn_g = dev_shared_alloc(q, (size_t)C * QX_FF_INTER, "ffn_g"); ffn_u = dev_shared_alloc(q, (size_t)C * QX_FF_INTER, "ffn_u"); ffn_t = dev_shared_alloc(q, (size_t)C * QX_FF_INTER, "ffn_t"); attn_q = dev_shared_alloc(q, (size_t)C * QX_Q_DIM, "attn_q"); attn_gate = dev_shared_alloc(q, (size_t)C * QX_Q_DIM, "attn_gate"); attn_k = dev_shared_alloc(q, (size_t)C * QX_KV_DIM, "attn_k"); attn_v = dev_shared_alloc(q, (size_t)C * QX_KV_DIM, "attn_v"); attn_out = dev_shared_alloc(q, (size_t)C * QX_Q_DIM, "attn_out"); fd_acc = chk_device<float>(q, (size_t)FD_MAX_SPLIT * QX_N_HEAD * QX_HEAD_DIM, "fd_acc"); fd_ml = chk_device<float>(q, (size_t)FD_MAX_SPLIT * QX_N_HEAD * 2, "fd_ml"); logits_batch_buf = dev_shared_alloc(q, (size_t)16 * QX_N_VOCAB, "logits_batch_buf"); bvnni_d = chk_device<int8_t>(q, (size_t)QX_FF_INTER * C, "bvnni_d"); suma_d = chk_device<int32_t>(q, (size_t)(QX_FF_INTER / QX_QK2) * C, "suma_d"); act_scale_d = chk_shared<float>(q, C, "act_scale_d"); rope_c = dev_shared_alloc(q, (size_t)C * (QX_ROT_DIM / 2), "rope_c"); rope_s = dev_shared_alloc(q, (size_t)C * (QX_ROT_DIM / 2), "rope_s"); /* M1.7a WY/UT workspace (sized for C=CHUNK=256, NVH=48, HD=128). */ { const int T = chunk, Hd = QX_SSM_STATE, Nvh = QX_SSM_DT_RANK; wy_ws.Kb = chk_shared<float>(q, (size_t)Nvh*T*Hd, "wy_ws.Kb"); wy_ws.Qb = chk_shared<float>(q, (size_t)Nvh*T*Hd, "wy_ws.Qb"); wy_ws.KbT = chk_shared<float>(q, (size_t)Nvh*Hd*T, "wy_ws.KbT"); wy_ws.QbT = chk_shared<float>(q, (size_t)Nvh*Hd*T, "wy_ws.QbT"); wy_ws.KK = chk_shared<float>(q, (size_t)Nvh*T*T, "wy_ws.KK"); wy_ws.KQ = chk_shared<float>(q, (size_t)Nvh*T*T, "wy_ws.KQ"); wy_ws.p = chk_shared<float>(q, (size_t)Nvh*T*Hd, "wy_ws.p"); wy_ws.u = chk_shared<float>(q, (size_t)Nvh*T*Hd, "wy_ws.u"); wy_ws.S0Q = chk_shared<float>(q, (size_t)Nvh*T*Hd, "wy_ws.S0Q"); wy_ws.gl = chk_shared<float>(q, (size_t)Nvh*T, "wy_ws.gl"); wy_ws.Ga = chk_shared<float>(q, (size_t)Nvh*T, "wy_ws.Ga"); wy_ws.bs = chk_shared<float>(q, (size_t)Nvh*T, "wy_ws.bs"); wy_ws.g_one = chk_shared<float>(q, (size_t)Nvh*T, "wy_ws.g_one"); wy_ws.out_acc = chk_shared<float>(q, (size_t)Nvh*T*Hd, "wy_ws.out_acc"); wy_ws.GKQ = chk_shared<float>(q, (size_t)Nvh*T*T, "wy_ws.GKQ"); wy_ws.S0_scaled = chk_shared<float>(q, (size_t)Nvh*Hd*Hd, "wy_ws.S0_scaled"); wy_ws.u_scaled = chk_shared<float>(q, (size_t)Nvh*T*Hd, "wy_ws.u_scaled"); } /* M1.7b.7 3-kernel FA-split workspace. Sized for `chunk` tokens (the * runtime chunk), TILE_T=128. Kt/Vt sized [NQ] (replicated across GQA -- * see fa_dequant_kt/vt). */ { const int T = chunk, Tt = 128; fa_ws.Kt = chk_shared<float>(q, (size_t)QX_N_HEAD * QX_HEAD_DIM * Tt, "fa_ws.Kt"); fa_ws.Vt = chk_shared<float>(q, (size_t)QX_N_HEAD * Tt * QX_HEAD_DIM, "fa_ws.Vt"); fa_ws.S = chk_shared<float>(q, (size_t)QX_N_HEAD * T * Tt, "fa_ws.S"); fa_ws.P = chk_shared<float>(q, (size_t)QX_N_HEAD * T * Tt, "fa_ws.P"); fa_ws.O = chk_shared<float>(q, (size_t)QX_N_HEAD * T * QX_HEAD_DIM, "fa_ws.O"); fa_ws.Oacc = chk_shared<float>(q, (size_t)QX_N_HEAD * T * QX_HEAD_DIM, "fa_ws.Oacc"); fa_ws.m_state = chk_shared<float>(q, (size_t)QX_N_HEAD * T, "fa_ws.m_state"); fa_ws.l_state = chk_shared<float>(q, (size_t)QX_N_HEAD * T, "fa_ws.l_state"); fa_ws.Qsep = chk_shared<float>(q, (size_t)QX_N_HEAD * T * QX_HEAD_DIM, "fa_ws.Qsep"); /* M1.7c fused-FA fp16 K/V shadow (grown on demand by * fa_shadow_ensure). 8192 tokens covers the bench prompt. */ fa_ws.shadow_tok = 8192; fa_ws.Kh = chk_shared<sycl::half>(q, (size_t)fa_ws.shadow_tok * QX_KV_DIM, "fa_ws.Kh"); fa_ws.Vh = chk_shared<sycl::half>(q, (size_t)fa_ws.shadow_tok * QX_KV_DIM, "fa_ws.Vh"); } g_bvnni_d = bvnni_d; g_suma_d = suma_d; g_act_scale_d = act_scale_d; return true;}
void engine::free() {}
/* ---------- slot lifecycle (M2.1) ---------- * open_slot/close_slot are concrete in engine_i; they manage the slot id and * the per-slot RNG vector, then call on_open_slot/on_close_slot here to * allocate/free the per-slot KV+SSM+conv state. */
void engine::on_open_slot(int slot) { /* slot is the just-reserved index; grow slots[] to match. The base * already advanced the id space, so slots.size() must reach slot+1. */ while ((int)slots.size() <= slot) slots.emplace_back(); slot_state& s = slots[slot]; size_t krow_bytes = (k_ctk == KCT::fp16) ? (size_t)KROW_BYTES_F16 : (size_t)KROW_BYTES_Q8; for (int i = 0; i < n_ssm; i++) { s.ssm_state[i] = chk_shared<float>(q, QX_SSM_DT_RANK * QX_SSM_STATE * QX_SSM_STATE, "slot.ssm_state"); s.ssm_conv[i] = chk_shared<float>(q, (QX_SSM_CONV_K-1) * QX_SSM_CONVD, "slot.ssm_conv"); q.memset(s.ssm_state[i], 0, QX_SSM_DT_RANK*QX_SSM_STATE*QX_SSM_STATE*sizeof(float)); q.memset(s.ssm_conv[i], 0, (QX_SSM_CONV_K-1)*QX_SSM_CONVD*sizeof(float)); } for (int i = 0; i < n_attn; i++) { s.kv_k[i] = chk_shared<uint8_t>(q, (size_t)max_ctx * krow_bytes, "slot.kv_k"); s.kv_v[i] = chk_shared<uint32_t>(q, (size_t)max_ctx * VROW_U32, "slot.kv_v"); q.memset(s.kv_k[i], 0, (size_t)max_ctx * krow_bytes); q.memset(s.kv_v[i], 0, (size_t)max_ctx*VROW_U32*sizeof(uint32_t)); } s.n_tokens = 0; q.wait();}
void engine::on_close_slot(int slot) { if (slot < 0 || slot >= (int)slots.size()) return; slot_state& s = slots[slot]; for (int i = 0; i < n_ssm; i++) { if (s.ssm_state[i]) sycl::free(s.ssm_state[i], q); if (s.ssm_conv[i]) sycl::free(s.ssm_conv[i], q); } for (int i = 0; i < n_attn; i++) { if (s.kv_k[i]) sycl::free(s.kv_k[i], q); if (s.kv_v[i]) sycl::free(s.kv_v[i], q); } /* The base close_slot erases the slot id; shrink slots[] to match. Slots * are usually opened in order and closed in reverse, so pop_back is the * common path. */ if (slot == (int)slots.size() - 1) slots.pop_back(); else slots.erase(slots.begin() + slot);}
void engine::reset(int slot) { slot_state& s = slots.at(slot); size_t krow_bytes = (k_ctk == KCT::fp16) ? (size_t)KROW_BYTES_F16 : (size_t)KROW_BYTES_Q8; for (int i = 0; i < n_ssm; i++) { q.memset(s.ssm_state[i], 0, QX_SSM_DT_RANK*QX_SSM_STATE*QX_SSM_STATE*sizeof(float)); q.memset(s.ssm_conv[i], 0, (QX_SSM_CONV_K-1)*QX_SSM_CONVD*sizeof(float)); } for (int i = 0; i < n_attn; i++) { q.memset(s.kv_k[i], 0, (size_t)max_ctx * krow_bytes); q.memset(s.kv_v[i], 0, (size_t)max_ctx*VROW_U32*sizeof(uint32_t)); } s.n_tokens = 0;}
void engine::forward(int slot, int token_id) { slot_state& st = slots.at(slot); /* Drain any pending GPU writes to the engine-wide `x` before the host * writes the new token's embedding. USM shared host writes are NOT * ordered w.r.t. in-flight GPU kernels on the in-order queue -- a prior * forward/prefill's tmatvec_accum (which writes x with accum=x) could * still be running and race the host dequant_embd write. The queue's * in-order property only orders GPU kernels against each other, not * against host memory writes. (M2.5 isolation fix.) */ q.wait(); dequant_embd(g_model, token_id, x); int64_t pos = st.n_tokens; int ai = 0, si = 0;
static int prof_lvl = []{ const char* e = std::getenv("QXMX_PROFILE"); return e ? std::atoi(e) : 0; }(); bool prof = prof_lvl >= 1; int half = QX_ROT_DIM / 2; for (int i = 0; i < half; i++) { double freq = 1.0 / pow((double)QX_ROPE_FREQ_BASE, (double)(2 * i) / QX_ROT_DIM); rope_c[i] = cosf((float)(pos * freq)); rope_s[i] = sinf((float)(pos * freq)); }
static int t_count = 0; auto t0 = std::chrono::steady_clock::now(); double t_ssm = 0, t_attn = 0, t_ffn = 0, t_res = 0, t_head = 0; /* QXMX_PROFILE=2: bracket the layer loop with 1-thread profiling markers * to measure total device-side busy time for the 64-layer stack. Delta * vs wall-clock = dispatch/idle overhead. The markers add ~28us each * (small vs the ~44ms layer loop). */ sycl::event dev_start, dev_end; if (prof_lvl >= 2) dev_start = q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<1>(1), [=](sycl::id<1>){}); }); auto phase = [&](double& acc) { if (!prof) return; q.wait(); auto t = std::chrono::steady_clock::now(); static auto last = t0; acc += std::chrono::duration<double, std::milli>(t - last).count(); last = t; };
for (int L = 0; L < QX_N_LAYER; L++) { bool is_attn = qx_layer_is_attention(L); if (is_attn) { attn_forward(q, &attns[ai], x, st.kv_k[ai], st.kv_v[ai], pos, k_ctk, part_buf, proj_buf, attn_q, attn_gate, attn_k, attn_v, attn_out, rope_c, rope_s, fd_acc, fd_ml, FD_MAX_SPLIT); ai++; phase(t_attn); } else { ssm_forward(q, &ssms[si], st.ssm_state[si], st.ssm_conv[si], x, part_buf, proj_buf, alpha_buf, beta_buf, ssm_out_buf, zg_buf); si++; phase(t_ssm); } ffn_forward(q, &ffns[L], x, part_buf, ffn_g, ffn_u, ffn_t); phase(t_ffn); } gpu_rmsnorm_quant(q, x, output_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &output, logits, part_buf); phase(t_head); if (prof_lvl >= 2) dev_end = q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<1>(1), [=](sycl::id<1>){}); }); /* THE ONLY SYNC (non-profiling): wait for the whole async pipeline before * the host reads logits. In profiling mode the last phase() already * waited. */ if (!prof) q.wait(); st.n_tokens++;
auto t1 = std::chrono::steady_clock::now(); t_count++; double total = std::chrono::duration<double,std::milli>(t1-t0).count(); if (t_count > 3 && t_count % 20 != 0) { return; }
if (!prof) { fprintf(stderr, " [decode %d: %.1f ms total]\r", t_count, total); return; }
fprintf(stderr, " [decode %d: %.1f ms] ssm=%.1f attn=%.1f ffn=%.1f res=%.1f head=%.1f", t_count, total, t_ssm, t_attn, t_ffn, t_res, t_head); if (prof_lvl >= 2) { q.wait(); auto s = dev_start.get_profiling_info<sycl::info::event_profiling::command_start>(); auto e = dev_end.get_profiling_info<sycl::info::event_profiling::command_end>(); double dev_ms = (e - s) / 1e6; fprintf(stderr, " | device=%.1f ms dispatch=%.1f ms (%.0f%%)", dev_ms, total - dev_ms, (total - dev_ms) * 100.0 / total); } fprintf(stderr, "\n");}
/* M4.1 multi-slot batched decode spike. Drives N slots one step each, but * runs the FFN block as a single M=M_pad gemm_gpu_v4 call (M_pad = N rounded * up to a multiple of 16 for transpose_out) instead of N separate gemv calls. * attn/ssm/head stay serialized per-slot (the next lever). Per-slot logits * are written to logits_batch[i]; the caller samples each. * * Correctness vs N separate forward() calls: the FFN path switches from * gemv_dpas8 (M=1) to gemm_gpu_v4 (M=N) -- same reassociation drift class as * prefill (ground rule 2). Greedy-token gate is the qxmx_diff / fa_kernel * reassociation gate, not bit-exact. */void engine::forward_batch(const int* slot_ids, const int* token_ids, int N) { /* Drain pending GPU writes to x before the host writes N embeddings. */ q.wait(); int M_pad = (N + 15) & ~15; /* gemm_gpu_v4 transpose_out wants M%16 */ for (int i = 0; i < N; i++) { slot_state& st = slots[slot_ids[i]]; dequant_embd(g_model, token_ids[i], x + (size_t)i * QX_D_MODEL); /* Per-slot rope tables at row i of the shared rope_c/s workspace. */ int64_t pos = st.n_tokens; int half = QX_ROT_DIM / 2; float* rc = rope_c + (size_t)i * half; float* rs = rope_s + (size_t)i * half; for (int j = 0; j < half; j++) { double freq = 1.0 / pow((double)QX_ROPE_FREQ_BASE, (double)(2 * j) / QX_ROT_DIM); rc[j] = cosf((float)(pos * freq)); rs[j] = sinf((float)(pos * freq)); } } /* Padded donor rows (M_pad > N): copy slot 0's x so the GEMM reads valid * memory; the padded FFN output is discarded. */ for (int i = N; i < M_pad; i++) q.memcpy(x + (size_t)i * QX_D_MODEL, x, (size_t)QX_D_MODEL * sizeof(float));
static int prof_lvl = []{ const char* e = std::getenv("QXMX_PROFILE"); return e ? std::atoi(e) : 0; }(); bool prof = prof_lvl >= 1;
static int t_count = 0; auto t0 = std::chrono::steady_clock::now(); double t_attnssm = 0, t_ffn = 0, t_head = 0; sycl::event dev_start, dev_end; if (prof_lvl >= 2) dev_start = q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<1>(1), [=](sycl::id<1>){}); }); auto phase = [&](double& acc) { if (!prof) return; q.wait(); auto t = std::chrono::steady_clock::now(); static auto last = t0; acc += std::chrono::duration<double, std::milli>(t - last).count(); last = t; };
/* Per-layer: run attn/ssm for each slot (serialized, each into its own x * row), then ONE batched FFN over all M_pad rows. The shared quant stage * (g_act_scale_d etc.) is reused per-slot call since attn/ssm are M=1. */ /* DEBUG: minimal path — only FFN, no attn/ssm, no head. x is already * filled by dequant_embd above. Useful for isolating FFN cost. */ if (std::getenv("QXMX_BATCH_FFN_ONLY")) { for (int L = 0; L < QX_N_LAYER; L++) ffn_forward_b(q, &ffns[L], x, M_pad, g_bvnni_d, g_suma_d, g_act_scale_d, ffn_g, ffn_u, ffn_t); q.wait(); for (int i = 0; i < N; i++) slots[slot_ids[i]].n_tokens++; return; }
int ai[64] = {0}, si[64] = {0}; /* per-slot attn/ssm layer index (max 64) */ for (int L = 0; L < QX_N_LAYER; L++) { bool is_attn = qx_layer_is_attention(L); for (int i = 0; i < N; i++) { slot_state& st = slots[slot_ids[i]]; float* xi = x + (size_t)i * QX_D_MODEL; if (is_attn) { int a = ai[i]; float* rc = rope_c + (size_t)i * (QX_ROT_DIM/2); float* rs = rope_s + (size_t)i * (QX_ROT_DIM/2); attn_forward(q, &attns[a], xi, st.kv_k[a], st.kv_v[a], st.n_tokens, k_ctk, part_buf, proj_buf, attn_q, attn_gate, attn_k, attn_v, attn_out, rc, rs, fd_acc, fd_ml, FD_MAX_SPLIT); ai[i]++; } else { int s = si[i]; ssm_forward(q, &ssms[s], st.ssm_state[s], st.ssm_conv[s], xi, part_buf, proj_buf, alpha_buf, beta_buf, ssm_out_buf, zg_buf); si[i]++; } } phase(t_attnssm); /* Batched FFN: one gemm_gpu_v4 call over all M_pad rows of x. */ ffn_forward_b(q, &ffns[L], x, M_pad, g_bvnni_d, g_suma_d, g_act_scale_d, ffn_g, ffn_u, ffn_t); phase(t_ffn); } /* Per-slot head: norm + lm_head (serialized; 1 ms each is small). Writes * to the engine's logits_batch_buf[i*VOCAB] (USM shared -- GPU-writable). */ for (int i = 0; i < N; i++) { float* xi = x + (size_t)i * QX_D_MODEL; gpu_rmsnorm_quant(q, xi, output_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &output, logits_batch_buf + (size_t)i * QX_N_VOCAB, part_buf); } phase(t_head); if (prof_lvl >= 2) dev_end = q.submit([&](sycl::handler& h) { h.parallel_for(sycl::range<1>(1), [=](sycl::id<1>){}); }); if (!prof) q.wait(); for (int i = 0; i < N; i++) slots[slot_ids[i]].n_tokens++;
auto t1 = std::chrono::steady_clock::now(); t_count++; double total = std::chrono::duration<double,std::milli>(t1-t0).count(); if (prof && (t_count <= 3 || t_count % 20 == 0)) { fprintf(stderr, " [batch decode %d (N=%d): %.1f ms] attn+ssm=%.1f ffn=%.1f head=%.1f", t_count, N, total, t_attnssm, t_ffn, t_head); if (prof_lvl >= 2) { q.wait(); auto s = dev_start.get_profiling_info<sycl::info::event_profiling::command_start>(); auto e = dev_end.get_profiling_info<sycl::info::event_profiling::command_end>(); double dev_ms = (e - s) / 1e6; fprintf(stderr, " | device=%.1f ms dispatch=%.1f ms (%.0f%%)", dev_ms, total - dev_ms, (total - dev_ms) * 100.0 / total); } fprintf(stderr, "\n"); }}
/* Run the block stack (attn/ssm + ffn, NO lm_head) on M token rows of x. * M must be a multiple of 16 (gemm_gpu_v4 constraint). KV cache rows and * ssm recurrent state are updated for positions pos_start..pos_start+M-1. * rope_c/rope_s must hold the [M][ROT_DIM/2] table for those positions. * pass_valid_count < M: only the first pass_valid_count tokens advance the * SSM state and conv state (tail-batched prefill: M=16 padded, pass_valid_count * = the real tail length). Default -1 means "use M" (full-chunk behavior). */void engine::run_chunk_b(slot_state& s, const int* token_ids, int64_t pos_start, int M, int pass_valid_count /*=-1*/) { if (pass_valid_count < 0) pass_valid_count = M; static bool prof = []{ return std::getenv("QXMX_PROFILE") != nullptr; }(); double t_attn = 0, t_ssm = 0, t_ffn = 0; auto phase = [&](double& acc) { if (!prof) return; q.wait(); static auto last = std::chrono::steady_clock::now(); auto t = std::chrono::steady_clock::now(); acc += std::chrono::duration<double, std::milli>(t - last).count(); last = t; }; auto t0 = std::chrono::steady_clock::now(); /* Drain pending GPU writes to x before the host writes the chunk's * embeddings. Same USM-shared host-write ordering fix as forward(). In * the sequential single-chunk case the prior op already waited, but the * scheduler interleaves chunks across slots and a prior chunk's ffn * tmatvec_accum (writes x) could still be in flight. */ q.wait(); for (int t = 0; t < M; t++) dequant_embd(g_model, token_ids[t], x + (size_t)t * QX_D_MODEL); engine_compute_rope_table(pos_start, M, rope_c, rope_s); int ai = 0, si = 0; if (prof) { /* reset phase timer after host embedding work */ } for (int L = 0; L < QX_N_LAYER; L++) { if (qx_layer_is_attention(L)) { attn_forward_b(q, &attns[ai], x, M, s.kv_k[ai], s.kv_v[ai], pos_start, k_ctk, bvnni_d, suma_d, act_scale_d, proj_buf, attn_q, attn_gate, attn_k, attn_v, attn_out, rope_c, rope_s, &fa_ws); ai++; phase(t_attn); } else { ssm_forward_b(q, &ssms[si], s.ssm_state[si], s.ssm_conv[si], x, M, bvnni_d, suma_d, act_scale_d, proj_buf, alpha_buf, beta_buf, ssm_out_buf, zg_buf, /*wyut_ws=*/&wy_ws, /*valid_count=*/pass_valid_count); si++; phase(t_ssm); } ffn_forward_b(q, &ffns[L], x, M, bvnni_d, suma_d, act_scale_d, ffn_g, ffn_u, ffn_t); phase(t_ffn); } if (prof) { auto t1 = std::chrono::steady_clock::now(); double total = std::chrono::duration<double, std::milli>(t1 - t0).count(); static int chunk_count = 0; chunk_count++; if (chunk_count <= 3 || chunk_count % 10 == 0) fprintf(stderr, " [chunk %d M=%d: %.1f ms] attn=%.1f ssm=%.1f ffn=%.1f\n", chunk_count, M, total, t_attn, t_ssm, t_ffn); }}
/* One step of the chunked prefill state machine (M2.5). Drives at most one * _b chunk (or one tail-token forward(), or the final lm_head-on-last-row) * per call so the scheduler can rotate slots at chunk boundaries. i and * last_chunk_m are in-out scheduler-held state; initialize to 0 before the * first call. done=true when `logits` is populated for the last prompt token * (ready for sample()). Mirrors the old prefill() loop exactly, minus the * latent n_tokens++ bug in the no-tail branch (run_chunk_b/forward already * account for all n positions; the lm_head-on-last-row read is not a new * position, so no increment). */void engine::prefill_step(int slot, const int* token_ids, int n, int& i, int& last_chunk_m, bool& done) { done = false; slot_state& st = slots.at(slot); int64_t pos0 = st.n_tokens - i; /* original prefill start position */ const int C = chunk;
/* Run a _b chunk if there are >= 16 tokens left. If this chunk consumes * the last of the prompt (i.e. i+M == n after it), we must compute lm_head * on its final residual row NOW, before returning -- otherwise another * slot's prefill_step could overwrite the engine-wide `x` buffer between * this step and the next, corrupting the logits. (This is the M2.5 * isolation fix: the no-tail lm_head can't be deferred to a later step.) */ auto run_b_chunk = [&](int M) { run_chunk_b(st, token_ids + i, pos0 + i, M); i += M; st.n_tokens += M; last_chunk_m = M; if (i >= n) { /* Last chunk's final token still needs lm_head (no tail forward() * will run it). Read the residual NOW, before `x` is reused. */ float* x_last = x + (size_t)(M - 1) * QX_D_MODEL; gpu_rmsnorm_quant(q, x_last, output_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &output, logits, part_buf); q.wait(); done = true; } };
/* Bulk: full C-token chunks. */ if (i + C <= n) { run_b_chunk(C); return; } /* Mid: a multiple-of-16 chunk for the remaining >= 16 tokens. */ int remaining = n - i; if (remaining >= 16) { run_b_chunk((remaining / 16) * 16); return; } /* Tail: < 16 tokens via forward() (lm_head runs on each; the last leaves * logits populated). */ if (i < n) { forward(slot, token_ids[i]); i++; done = (i >= n); return; } /* n == 0 (empty prompt): nothing to do, but mark done so the scheduler * samples from whatever logits forward() left. In practice the caller * should not submit an empty prompt. */ done = true;}
/* Convenience: run the full prefill to completion (existing single-shot * callers). M2.5 scheduler uses prefill_step() instead. */void engine::prefill(int slot, const int* token_ids, int n) { static bool prof = []{ return std::getenv("QXMX_PROFILE") != nullptr; }(); auto t0 = std::chrono::steady_clock::now(); int i = 0, last_chunk_m = 0; bool done = false; /* Bulk + mid chunks via prefill_step (handles 1-by-1 tail internally for * the M2.5 scheduler path). For the monolithic prefill path, detect the * tail (remaining < 16) and batch it in one M=16 call instead of looping * 1-by-1 through forward() (which costs ~30 ms/tok vs ~4 ms/tok for the * batched call). The 3 padded rows in the batched call have * valid_count=tail_real so the SSM/conv state isn't advanced past the * real last token. See run_chunk_b / ssm_forward_b / gpu_conv1d_step_b. */ while (!done && n - i >= 16) { prefill_step(slot, token_ids, n, i, last_chunk_m, done); } /* Tail: batch the remaining (< 16) tokens in one M=16 call. Pad with the * last real token id (padded rows' output is discarded; valid_count= * tail_real so the SSM/conv state only advances for real tokens). * * pos_tail = n_tokens - tail_real would re-process the bulk chunk's last * few rows (harmless: decode overwrites them), but goes NEGATIVE when the * whole prompt is shorter than 16 tokens (n_tokens=0). Clamp at 0: the * KV rows written past the real tail are don't-care (the first decode * steps overwrite them), so the batched path works for any prompt length. */ if (!done) { slot_state& st = slots.at(slot); int tail_real = n - i; int padded_ids[16]; for (int k = 0; k < tail_real; k++) padded_ids[k] = token_ids[i + k]; int pad_token = token_ids[i + tail_real - 1]; for (int k = tail_real; k < 16; k++) padded_ids[k] = pad_token; int64_t pos_tail = (int64_t)st.n_tokens - tail_real; if (pos_tail < 0) pos_tail = 0; run_chunk_b(st, padded_ids, pos_tail, /*M=*/16, /*valid_count=*/tail_real); st.n_tokens += 16; float* x_last = x + (size_t)(tail_real - 1) * QX_D_MODEL; gpu_rmsnorm_quant(q, x_last, output_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &output, logits, part_buf); q.wait(); st.n_tokens = n; done = true; } if (prof) { auto t1 = std::chrono::steady_clock::now(); double ms = std::chrono::duration<double, std::milli>(t1 - t0).count(); fprintf(stderr, " [prefill %d tokens: %.1f ms, %.0f tok/s]\n", n, ms, n / (ms > 0 ? ms / 1e3 : 1)); }}
int64_t engine::slot_n_tokens(int slot) const { return slots.at(slot).n_tokens;}
/* M2.6 cache-aware prefill. Two phases: * RESTORE: walk blocks [0, n/PC_BLOCK). A block b is cacheable iff there is * an SSM checkpoint AT boundary (b+1) AND the KV block b is cached under the * same hash -- both keyed by the chained hash of ids[0..(b+1)*512). On hit, * DMA host->device into the slot's KV buffer at [b*512, (b+1)*512) and the * slot's ssm_state/conv (one event each, ordered by the in-order queue), * then set n_tokens = (b+1)*512 -- skipping the entire block stack for that * prefix. KV reuse must be CONTIGUOUS from 0: a gap (miss or id-mismatch) * stops restoration, since attention over the delta needs the full prefix KV * present at the fixed offsets. * PREFILL: the remainder [restored_to, n) via the existing prefill_step * loop. After it, snapshot every freshly-prefilled block boundary into the * cache (every-boundary placement). Restored blocks are already cached. * * Snapshot ordering: device->host q.memcpy lands in the in-order queue AFTER * the chunk's GPU kernels and BEFORE the next chunk's kernels, so the KV at * block b's offsets is read before block b+1 overwrites. We q.wait() once at * the end before reading the snapshot host buffers into the cache. * * Returns the number of prefix tokens skipped (restored_to). */ int engine::prefill_cached(int slot, const int* token_ids, int n, prefix_cache* cache, int* computed_out) { if (!cache || n < PC_BLOCK) { prefill(slot, token_ids, n); if (computed_out) *computed_out = n; return 0; } slot_state& st = slots.at(slot); size_t krow = (k_ctk == KCT::fp16) ? (size_t)KROW_BYTES_F16 : (size_t)KROW_BYTES_Q8; size_t vrow = (size_t)VROW_U32 * sizeof(uint32_t); /* Per-layer block sizes (one attn layer's worth of 512 tokens). The total * across n_attn layers is k_blk/v_blk; the cache stores all layers * concatenated, so the per-layer stride into the cache blob is k_blk_layer. */ size_t k_blk_layer = (size_t)PC_BLOCK * krow; size_t v_blk_layer = (size_t)PC_BLOCK * vrow; size_t k_blk = (size_t)n_attn * k_blk_layer; size_t v_blk = (size_t)n_attn * v_blk_layer; size_t s_state = (size_t)QX_SSM_DT_RANK * QX_SSM_STATE * QX_SSM_STATE * sizeof(float); size_t s_conv = (size_t)(QX_SSM_CONV_K-1) * QX_SSM_CONVD * sizeof(float);
int n_full_blocks = n / PC_BLOCK;
/* Snapshot stride: snapshot one GROUP every SNAP_STRIDE blocks (env * QXMX_SNAP_STRIDE, default 1). A group = SNAP_STRIDE consecutive blocks; * at its end boundary we snapshot KV for ALL stride blocks in the group * plus ONE SSM checkpoint (at the group's end). Coarsening cuts the * dominant cost -- the ~157 MB SSM checkpoint DMA + q.wait() per * boundary -- by the stride factor, at the expense of coarser restore * granularity: a future turn restores whole groups only, re-prefilling * up to SNAP_STRIDE-1 extra blocks past the last full group. KV DMA * bytes are unchanged (every block's KV is still cached); only the SSM * checkpoint frequency and q.wait() count drop. */ static int snap_stride = []{ const char* e = std::getenv("QXMX_SNAP_STRIDE"); int v = e ? std::atoi(e) : 1; return v < 1 ? 1 : v; }(); int n_groups = n_full_blocks / snap_stride; /* full groups only */
/* ---- RESTORE (group-by-group) ---- * Walk full groups in order. Group g = blocks [g*stride, (g+1)*stride). * Restore requires: the SSM checkpoint at the group's end boundary * (keyed by the hash of block (g+1)*stride - 1, n_blocks == (g+1)*stride) * AND KV for every block in the group. A gap (any miss / id-mismatch / * boundary mismatch) stops restoration -- contiguity is required. */ int restored_blocks = 0; block_hash parent{0}; std::vector<block_hash> ghash(snap_stride); for (int g = 0; g < n_groups; g++) { block_hash p = parent; for (int s = 0; s < snap_stride; s++) { int b = g * snap_stride + s; ghash[s] = hash_block(p, token_ids + b * PC_BLOCK, PC_BLOCK); p = ghash[s]; } const ssm_checkpoint* cp = cache->ssm_get(ghash[snap_stride - 1]); if (!cp || cp->n_blocks != (g + 1) * snap_stride) break; std::vector<const kv_block*> kvs(snap_stride); bool ok = true; for (int s = 0; s < snap_stride; s++) { int b = g * snap_stride + s; kvs[s] = cache->kv_get(ghash[s], token_ids + b * PC_BLOCK, PC_BLOCK); if (!kvs[s]) { ok = false; break; } } if (!ok) break; /* DMA host->device: KV per block per attn layer, SSM state/conv per * ssm layer. All fire-and-forget into the in-order queue; the next * prefill chunk's kernels are ordered after them. */ for (int s = 0; s < snap_stride; s++) { int b = g * snap_stride + s; for (int la = 0; la < n_attn; la++) { q.memcpy(st.kv_k[la] + (size_t)b * PC_BLOCK * krow, kvs[s]->k.data() + (size_t)la * k_blk_layer, k_blk_layer); q.memcpy(st.kv_v[la] + (size_t)b * PC_BLOCK * VROW_U32, (const uint32_t*)(kvs[s]->v.data() + (size_t)la * v_blk_layer), v_blk_layer); } } for (int ls = 0; ls < n_ssm; ls++) { q.memcpy(st.ssm_state[ls], cp->state.data() + (size_t)ls * s_state, s_state); q.memcpy(st.ssm_conv[ls], cp->conv.data() + (size_t)ls * s_conv, s_conv); } parent = ghash[snap_stride - 1]; restored_blocks = (g + 1) * snap_stride; } int restored_to = restored_blocks * PC_BLOCK; st.n_tokens = restored_to; /* jump (set, not +=) */
/* ---- PREFILL remainder with INLINE snapshot ---- * The SSM state is mutable (updated in-place by run_chunk_b), so a * checkpoint AT a group-end boundary must be captured the MOMENT that * boundary is crossed -- not at the end of the prefill (by then the state * has advanced to n). So we drive prefill_step one call at a time and * snapshot after each call that leaves i on a group-end boundary. * * KV blocks are immutable once written (fixed offset pos*row). At a * group-end boundary we DMA KV for ALL stride blocks of the group plus * the SSM checkpoint, then one q.wait() before the cache copies. */ /* parent hash at the start of the first block to snapshot. */ block_hash snap_parent{0}; for (int b = 0; b < restored_blocks; b++) snap_parent = hash_block(snap_parent, token_ids + b * PC_BLOCK, PC_BLOCK);
/* Snapshot work happens only at group-end boundaries, and only full * groups fire (n_groups = n_full_blocks / snap_stride). If no full group * can complete, skip the snapshot machinery entirely -- no per-call buffer * allocation, no per-boundary hash walk, just a plain prefill of the * remainder. This is the common short-prompt path (n < snap_stride*512) * and the cold-cache long-prompt path when stride is raised to suppress * snapshotting for benchmarking. */ if (restored_to < n && n_groups > 0) { /* Per-call snapshot buffers: KV for one group (stride blocks) + one * SSM checkpoint. Sized to snap_stride; allocated only when at least * one group will fire. (Hoisting to engine members would avoid the * per-call alloc -- ~17 MB at stride 1 -- but the cost is dominated * by the DMA + q.wait(), not the malloc. Deferred.) */ std::vector<uint8_t> k_buf((size_t)snap_stride * k_blk); std::vector<uint8_t> v_buf((size_t)snap_stride * v_blk); std::vector<uint8_t> st_buf((size_t)n_ssm * s_state); std::vector<uint8_t> cv_buf((size_t)n_ssm * s_conv); int i = restored_to, last_chunk_m = 0; bool done = false; int next_snap_block = restored_blocks; /* next block to (maybe) snapshot */ /* Parent hash at the start of the current group (the chain value * BEFORE block next_snap_block's group began). Saved when we enter a * group's first block; consumed at the group's end-boundary snapshot * to re-hash every block in the group from the correct chain root. */ block_hash group_start_parent = snap_parent; while (!done && n - i >= 16) { prefill_step(slot, token_ids, n, i, last_chunk_m, done); /* Walk every block boundary in (prev_i, i]. i advances by a * chunk (>=16) or a tail token; boundaries land on i==B*512. * Advance snap_parent for EVERY boundary (the chain is required * by the next snapshot's hash) but only do the DMA+copy work at * GROUP-end boundaries (block B where B+1 is a multiple of * snap_stride), snapshotting all stride blocks of the group + the * SSM checkpoint at the group end. */ while (next_snap_block < n_full_blocks && i >= (next_snap_block + 1) * PC_BLOCK) { int b = next_snap_block; if (b % snap_stride == 0) /* entering a new group: freeze root */ group_start_parent = snap_parent; const int32_t* bids = token_ids + b * PC_BLOCK; block_hash bh = hash_block(snap_parent, bids, PC_BLOCK); if ((b + 1) % snap_stride == 0) { int g0 = b + 1 - snap_stride; /* first block of group */ /* DMA KV for every block in the group (all still live at * their fixed device offsets) + SSM state/conv, then one * q.wait() before reading the host buffers into the cache. */ for (int s = 0; s < snap_stride; s++) { int gb = g0 + s; for (int la = 0; la < n_attn; la++) { q.memcpy(k_buf.data() + (size_t)s * k_blk + (size_t)la * k_blk_layer, st.kv_k[la] + (size_t)gb * PC_BLOCK * krow, k_blk_layer); q.memcpy(v_buf.data() + (size_t)s * v_blk + (size_t)la * v_blk_layer, st.kv_v[la] + (size_t)gb * PC_BLOCK * VROW_U32, v_blk_layer); } } for (int ls = 0; ls < n_ssm; ls++) { q.memcpy(st_buf.data() + (size_t)ls * s_state, st.ssm_state[ls], s_state); q.memcpy(cv_buf.data() + (size_t)ls * s_conv, st.ssm_conv[ls], s_conv); } q.wait(); /* fill host buffers before the cache copies them */ block_hash gp = group_start_parent; for (int s = 0; s < snap_stride; s++) { int gb = g0 + s; const int32_t* gbids = token_ids + gb * PC_BLOCK; block_hash gbh = hash_block(gp, gbids, PC_BLOCK); cache->kv_put(gbh, gbids, PC_BLOCK, k_buf.data() + (size_t)s * k_blk, v_buf.data() + (size_t)s * v_blk); if (s == snap_stride - 1) cache->ssm_put(gbh, gb + 1, st_buf.data(), cv_buf.data()); gp = gbh; } } snap_parent = bh; /* chain advances even for skipped blocks */ next_snap_block = b + 1; } } if (!done) { /* Tail: batch the remaining tokens in one M=16 call. * valid_count=tail_real so the SSM/conv state isn't advanced * past the real last token. pos_tail clamped at 0 so short * prompts (n_tokens=0) don't underflow (see prefill()). */ int tail_real = n - i; int padded_ids[16]; for (int k = 0; k < tail_real; k++) padded_ids[k] = token_ids[i + k]; int pad_token = token_ids[i + tail_real - 1]; for (int k = tail_real; k < 16; k++) padded_ids[k] = pad_token; int64_t pos_tail = (int64_t)st.n_tokens - tail_real; if (pos_tail < 0) pos_tail = 0; run_chunk_b(st, padded_ids, pos_tail, /*M=*/16, /*valid_count=*/tail_real); st.n_tokens += 16; float* x_last = x + (size_t)(tail_real - 1) * QX_D_MODEL; gpu_rmsnorm_quant(q, x_last, output_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &output, logits, part_buf); q.wait(); st.n_tokens = n; done = true; } } else if (restored_to < n) { /* No full group will complete (n_groups == 0) -> no snapshot fires. * Plain prefill of the remainder. For the tail (remaining < 16), * batch the real tail tokens in one M=16 call instead of looping * 1-by-1 through forward() (which costs ~30 ms/tok vs ~4 ms/tok * for the batched call). The 3 padded rows in the batched call * have valid_count=tail_real so the SSM/conv state isn't advanced * past the real last token. See run_chunk_b / ssm_forward_b / * gpu_conv1d_step_b. */ slot_state& st_tail = st; int i = restored_to, last_chunk_m = 0; bool done = false; while (!done && n - i >= 16) { prefill_step(slot, token_ids, n, i, last_chunk_m, done); } if (!done) { /* Tail: batch the remaining tokens in one M=16 call. * pos_tail clamped at 0 so short prompts don't underflow (see * prefill()). */ int tail_real = n - i; int padded_ids[16]; for (int k = 0; k < tail_real; k++) padded_ids[k] = token_ids[i + k]; int pad_token = token_ids[i + tail_real - 1]; for (int k = tail_real; k < 16; k++) padded_ids[k] = pad_token; int64_t pos_tail = (int64_t)st_tail.n_tokens - tail_real; if (pos_tail < 0) pos_tail = 0; run_chunk_b(st_tail, padded_ids, pos_tail, /*M=*/16, /*valid_count=*/tail_real); st_tail.n_tokens += 16; float* x_last = x + (size_t)(tail_real - 1) * QX_D_MODEL; gpu_rmsnorm_quant(q, x_last, output_norm, QX_D_MODEL, g_bvnni_d, g_suma_d, g_act_scale_d); tmatvec_q(q, &output, logits, part_buf); q.wait(); st_tail.n_tokens = n; done = true; } (void)last_chunk_m; /* silence unused-warning if loop didn't run */ } else if (restored_to == n) { /* Fully restored (n is a multiple of PC_BLOCK and all blocks hit). * The caller still needs logits for the last token; recompute the last * block's prefill (cheap, and the cache already holds its checkpoint). * No new snapshot needed -- block b-1's checkpoint is already cached. */ st.n_tokens = restored_to - PC_BLOCK; int i = restored_to - PC_BLOCK, last_chunk_m = 0; bool done = false; while (!done) prefill_step(slot, token_ids, n, i, last_chunk_m, done); } /* Tokens actually computed on this call: the post-restore prefill (if any) * or the one rechecked block when the whole prompt came from cache. Both * branches leave st.n_tokens == n at completion. */ if (computed_out) { if (restored_to == n) *computed_out = PC_BLOCK; else *computed_out = n - restored_to; } return restored_to;}
/* Explicit instantiations of the templated single-query-token attention * oracle so the M1.4 test (attn_prefill_test) and any other TU that sees the * template decl in qxmx_gpu.h can link without re-defining the kernel. */template void gpu_attn_softmax<KFp16>(sycl::queue&, const float*, const uint8_t*, const uint32_t*, float*, int64_t, float*, float*, int);template void gpu_attn_softmax<KQ8_0>(sycl::queue&, const float*, const uint8_t*, const uint32_t*, float*, int64_t, float*, float*, int);template void gpu_attn_softmax<KFp8 >(sycl::queue&, const float*, const uint8_t*, const uint32_t*, float*, int64_t, float*, float*, int);
/* M1.3: batched per-row rmsnorm+quant (oracle: gpu_rmsnorm_quant). */void gpu_rmsnorm_quant(sycl::queue& q, const float* x, const float* w, int dim, int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d, int n_tok) { constexpr int LWS = 1024; int ngrp = dim / GRP; q.submit([&](sycl::handler& h) { sycl::local_accessor<double,1> lss(sycl::range<1>(LWS), h); sycl::local_accessor<float,1> lamax(sycl::range<1>(LWS), h); h.parallel_for(sycl::nd_range<1>((size_t)n_tok * LWS, LWS), [=](sycl::nd_item<1> it) { int tok = it.get_group(0); int lid = it.get_local_id(0); const float* xr = x + (size_t)tok * dim; double ss = 0; for (int i = lid; i < dim; i += LWS) { float xi = xr[i]; ss += (double)xi * xi; } lss[lid] = ss; sycl::group_barrier(it.get_group()); for (int s = LWS / 2; s > 0; s >>= 1) { if (lid < s) lss[lid] += lss[lid + s]; sycl::group_barrier(it.get_group()); } float inv = 1.0f / sycl::sqrt((float)(lss[0] / dim) + QX_RMS_EPS); sycl::group_barrier(it.get_group()); if (lid < ngrp) { float m = 0; int base = lid * GRP; for (int i = 0; i < GRP; i++) m = sycl::fmax(m, sycl::fabs(xr[base + i] * inv * w[base + i])); lamax[lid] = m; } sycl::group_barrier(it.get_group()); if (lid == 0) { float amax = 0; for (int i = 0; i < ngrp; i++) amax = sycl::fmax(amax, lamax[i]); act_scale_d[tok] = amax / 127.0f; } sycl::group_barrier(it.get_group()); float sc = act_scale_d[tok]; float invq = sc > 0 ? 1.0f / sc : 0.0f; if (lid < ngrp) { int g = lid; int base = g * GRP; int32_t* col = (int32_t*)bvnni_d; int32_t sa = 0; for (int kk = 0; kk < GRP / 4; kk++) { int k4 = g * (GRP / 4) + kk; uint32_t pack = 0; for (int b = 0; b < 4; b++) { float ov = xr[base + kk * 4 + b] * inv * w[base + kk * 4 + b]; int v = (int)sycl::round(ov * invq); v = sycl::max(-128, sycl::min(127, v)); pack |= ((uint32_t)(uint8_t)(int8_t)v) << (8 * b); sa += (int8_t)v; } col[(size_t)k4 * n_tok + tok] = (int32_t)pack; } suma_d[(size_t)g * n_tok + tok] = sa; } }); });}
/* Batched SwiGLU+quant: t = silu(g)*u per row, then pack (oracle: * gpu_swiglu_quant). */void gpu_swiglu_quant(sycl::queue& q, const float* g, const float* u, float* t, int n, int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d, int n_tok) { constexpr int LWS = 1024; int ngrp = n / GRP; q.submit([&](sycl::handler& h) { sycl::local_accessor<float,1> lamax(sycl::range<1>(LWS), h); h.parallel_for(sycl::nd_range<1>((size_t)n_tok * LWS, LWS), [=](sycl::nd_item<1> it) { int tok = it.get_group(0); int lid = it.get_local_id(0); const float* gr = g + (size_t)tok * n; const float* ur = u + (size_t)tok * n; float* tr = t + (size_t)tok * n; for (int i = lid; i < n; i += LWS) { float v = (gr[i] / (1.0f + sycl::exp(-gr[i]))) * ur[i]; tr[i] = v; } sycl::group_barrier(it.get_group()); if (lid < ngrp) { const float* tg = tr + (size_t)lid * GRP; float m = 0; for (int i = 0; i < GRP; i++) m = sycl::fmax(m, sycl::fabs(tg[i])); lamax[lid] = m; } sycl::group_barrier(it.get_group()); if (lid == 0) { float amax = 0; for (int i = 0; i < ngrp; i++) amax = sycl::fmax(amax, lamax[i]); act_scale_d[tok] = amax / 127.0f; } sycl::group_barrier(it.get_group()); float sc = act_scale_d[tok]; float invq = sc > 0 ? 1.0f / sc : 0.0f; if (lid < ngrp) { int grp = lid; const float* tg = tr + (size_t)grp * GRP; int32_t* col = (int32_t*)bvnni_d; int32_t sa = 0; for (int kk = 0; kk < GRP / 4; kk++) { int k4 = grp * (GRP / 4) + kk; uint32_t pack = 0; for (int b = 0; b < 4; b++) { int v = (int)sycl::round(tg[kk * 4 + b] * invq); v = sycl::max(-128, sycl::min(127, v)); pack |= ((uint32_t)(uint8_t)(int8_t)v) << (8 * b); sa += (int8_t)v; } col[(size_t)k4 * n_tok + tok] = (int32_t)pack; } suma_d[(size_t)grp * n_tok + tok] = sa; } }); });}
/* Batched silu_mul+quant: out *= silu(gate) per row, then pack (oracle: * gpu_silu_mul_quant). */void gpu_silu_mul_quant(sycl::queue& q, float* out, const float* gate, int n, int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d, int n_tok) { constexpr int LWS = 1024; int ngrp = n / GRP; q.submit([&](sycl::handler& h) { sycl::local_accessor<float,1> lamax(sycl::range<1>(LWS), h); h.parallel_for(sycl::nd_range<1>((size_t)n_tok * LWS, LWS), [=](sycl::nd_item<1> it) { int tok = it.get_group(0); int lid = it.get_local_id(0); float* outr = out + (size_t)tok * n; const float* gr = gate + (size_t)tok * n; for (int i = lid; i < n; i += LWS) { outr[i] *= gr[i] / (1.0f + sycl::exp(-gr[i])); } sycl::group_barrier(it.get_group()); if (lid < ngrp) { const float* xg = outr + (size_t)lid * GRP; float m = 0; for (int i = 0; i < GRP; i++) m = sycl::fmax(m, sycl::fabs(xg[i])); lamax[lid] = m; } sycl::group_barrier(it.get_group()); if (lid == 0) { float amax = 0; for (int i = 0; i < ngrp; i++) amax = sycl::fmax(amax, lamax[i]); act_scale_d[tok] = amax / 127.0f; } sycl::group_barrier(it.get_group()); float sc = act_scale_d[tok]; float invq = sc > 0 ? 1.0f / sc : 0.0f; if (lid < ngrp) { int g = lid; const float* xg = outr + (size_t)g * GRP; int32_t* col = (int32_t*)bvnni_d; int32_t sa = 0; for (int kk = 0; kk < GRP / 4; kk++) { int k4 = g * (GRP / 4) + kk; uint32_t pack = 0; for (int b = 0; b < 4; b++) { int v = (int)sycl::round(xg[kk * 4 + b] * invq); v = sycl::max(-128, sycl::min(127, v)); pack |= ((uint32_t)(uint8_t)(int8_t)v) << (8 * b); sa += (int8_t)v; } col[(size_t)k4 * n_tok + tok] = (int32_t)pack; } suma_d[(size_t)g * n_tok + tok] = sa; } }); });}
/* Batched sigmoid_mul+quant: out *= sigmoid(gate) per row, then pack * (oracle: gpu_sigmoid_mul_quant). */void gpu_sigmoid_mul_quant(sycl::queue& q, float* out, const float* gate, int n, int8_t* bvnni_d, int32_t* suma_d, float* act_scale_d, int n_tok) { constexpr int LWS = 1024; int ngrp = n / GRP; q.submit([&](sycl::handler& h) { sycl::local_accessor<float,1> lamax(sycl::range<1>(LWS), h); h.parallel_for(sycl::nd_range<1>((size_t)n_tok * LWS, LWS), [=](sycl::nd_item<1> it) { int tok = it.get_group(0); int lid = it.get_local_id(0); float* outr = out + (size_t)tok * n; const float* gr = gate + (size_t)tok * n; for (int i = lid; i < n; i += LWS) { outr[i] *= 1.0f / (1.0f + sycl::exp(-gr[i])); } sycl::group_barrier(it.get_group()); if (lid < ngrp) { const float* xg = outr + (size_t)lid * GRP; float m = 0; for (int i = 0; i < GRP; i++) m = sycl::fmax(m, sycl::fabs(xg[i])); lamax[lid] = m; } sycl::group_barrier(it.get_group()); if (lid == 0) { float amax = 0; for (int i = 0; i < ngrp; i++) amax = sycl::fmax(amax, lamax[i]); act_scale_d[tok] = amax / 127.0f; } sycl::group_barrier(it.get_group()); float sc = act_scale_d[tok]; float invq = sc > 0 ? 1.0f / sc : 0.0f; if (lid < ngrp) { int g = lid; const float* xg = outr + (size_t)g * GRP; int32_t* col = (int32_t*)bvnni_d; int32_t sa = 0; for (int kk = 0; kk < GRP / 4; kk++) { int k4 = g * (GRP / 4) + kk; uint32_t pack = 0; for (int b = 0; b < 4; b++) { int v = (int)sycl::round(xg[kk * 4 + b] * invq); v = sycl::max(-128, sycl::min(127, v)); pack |= ((uint32_t)(uint8_t)(int8_t)v) << (8 * b); sa += (int8_t)v; } col[(size_t)k4 * n_tok + tok] = (int32_t)pack; } suma_d[(size_t)g * n_tok + tok] = sa; } }); });}
/* ========================================================================== * M1.4 — batched KV-cache quant writes + causal FlashAttention prefill. * Each quant kernel is the decode single-row oracle with an added token grid * dimension (one WG per (token, block)); per-(token,block) arithmetic is * identical to qxmx_gpu.cpp's gpu_quant_k_row_* / gpu_quant_v_row. See * qxmx_gpu.h for the cache block-layout constants. * ========================================================================== */
/* fp16 K cache write for n_tok rows (oracle: gpu_fp32_to_fp16 per row). * Bit-exact per element: one half() conversion, same as decode. */void gpu_fp32_to_fp16(sycl::queue& q, const float* src, sycl::half* dst, int n, int n_tok) { q.parallel_for(sycl::range<1>((size_t)n_tok * n), [=](sycl::id<1> i) { int t = (int)(i / n); int e = (int)(i % n); dst[(size_t)t * n + e] = sycl::half(src[(size_t)t * n + e]); });}
/* q8_0 K quant for n_tok rows (oracle: gpu_quant_k_row_q8). */void gpu_quant_k_row_q8(sycl::queue& q, const float* src, uint8_t* dst, int n_tok) { constexpr int LWS = KBLK; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>((size_t)n_tok * NVBLK_K * LWS, LWS), [=](sycl::nd_item<1> it) { int g = it.get_group(0); int tok = g / NVBLK_K; int blk = g % NVBLK_K; int lid = it.get_local_id(0); const float* sb = src + (size_t)tok * QX_KV_DIM + (size_t)blk * KBLK; uint8_t* db = dst + (size_t)tok * KROW_BYTES_Q8 + (size_t)blk * KBLK_BYTES; float v = sb[lid]; float a = sycl::fabs(v); float amax = sycl::reduce_over_group(it.get_group(), a, sycl::maximum<float>{}); float d = amax * (1.0f / 127.0f); float inv = (amax > 0.0f) ? (127.0f / amax) : 0.0f; int code = (int)sycl::round(v * inv); code = sycl::max(-128, sycl::min(127, code)); if (lid == 0) { sycl::half dh = sycl::half(d); uint16_t bits = sycl::bit_cast<uint16_t>(dh); db[0] = (uint8_t)(bits & 0xFF); db[1] = (uint8_t)(bits >> 8); } db[2 + lid] = (uint8_t)(int8_t)code; }); });}
/* fp8 e4m3 K quant for n_tok rows (oracle: gpu_quant_k_row_fp8). */void gpu_quant_k_row_fp8(sycl::queue& q, const float* src, uint8_t* dst, int n_tok) { constexpr int LWS = KBLK; q.submit([&](sycl::handler& h) { h.parallel_for(sycl::nd_range<1>((size_t)n_tok * NVBLK_K * LWS, LWS), [=](sycl::nd_item<1> it) { int g = it.get_group(0); int tok = g / NVBLK_K; int blk = g % NVBLK_K; int lid = it.get_local_id(0); const float* sb = src + (size_t)tok * QX_KV_DIM + (size_t)blk * KBLK; uint8_t* db = dst + (size_t)tok * KROW_BYTES_Q8 + (size_t)blk * KBLK_BYTES; float v = sb[lid]; float a = sycl::fabs(v); float amax = sycl::reduce_over_group(it.get_group(), a, sycl::maximum<float>{}); float scale = amax * (1.0f / 448.0f); float inv = (amax > 0.0f) ? (448.0f / amax) : 0.0f; uint8_t code = fp32_to_e4m3(v * inv); if (lid == 0) { sycl::half sh = sycl::half(scale); uint16_t bits = sycl::bit_cast<uint16_t>(sh); db[0] = (uint8_t)(bits & 0xFF); db[1] = (uint8_t)(bits >> 8); } db[2 + lid] = code; }); });}
/* TurboQuant V quant (WHT + 4-bit Lloyd-Max + norm-corr) for n_tok rows * (oracle: gpu_quant_v_row). Grid = n_tok*NVBLK_V blocks, LWS = 32. */void gpu_quant_v_row(sycl::queue& q, const float* src, uint32_t* dst_u32, int n_tok) { constexpr int LWS = VBLK; /* 32 threads = one block */ constexpr float NRM = 0.17677669529663687f; /* 1/sqrt(32) */ q.submit([&](sycl::handler& h) { sycl::local_accessor<float,1> buf0(sycl::range<1>(LWS), h); sycl::local_accessor<float,1> buf1(sycl::range<1>(LWS), h); h.parallel_for(sycl::nd_range<1>((size_t)n_tok * NVBLK_V * LWS, LWS), [=](sycl::nd_item<1> it) { int g = it.get_group(0); int tok = g / NVBLK_V; int blk = g % NVBLK_V; int lid = it.get_local_id(0); const float* sb = src + (size_t)tok * QX_KV_DIM + (size_t)blk * VBLK; uint32_t* db = dst_u32 + (size_t)tok * VROW_U32 + (size_t)blk * 5; float v_raw = sb[lid]; float orig_ss = v_raw * v_raw; orig_ss = sycl::reduce_over_group(it.get_group(), orig_ss, sycl::plus<float>()); buf0[lid] = v_raw; sycl::group_barrier(it.get_group()); auto* cur = &buf0; auto* nxt = &buf1; for (int bit = 1; bit < LWS; bit <<= 1) { float a = (*cur)[lid]; float b = (*cur)[lid ^ bit]; sycl::group_barrier(it.get_group()); (*nxt)[lid] = (lid & bit) ? (b - a) : (a + b); sycl::group_barrier(it.get_group()); auto* tmp = cur; cur = nxt; nxt = tmp; } float v = (*cur)[lid] * NRM; float a = sycl::fabs(v); float amax = sycl::reduce_over_group(it.get_group(), a, sycl::maximum<float>{}); float scale = amax; float vn = (amax > 0.0f) ? (v / amax) : 0.0f; int code = 0; float best_d = sycl::fabs(vn - V_LM_LUT[0]); for (int i = 1; i < 16; i++) { float d = sycl::fabs(vn - V_LM_LUT[i]); if (d < best_d) { best_d = d; code = i; } } float v_rec_coeff = V_LM_LUT[code] * scale; float rec_ss = v_rec_coeff * v_rec_coeff; rec_ss = sycl::reduce_over_group(it.get_group(), rec_ss, sycl::plus<float>()); float orig_norm = sycl::sqrt(orig_ss); float rec_norm = sycl::sqrt(rec_ss); float norm_corr = (rec_norm > 0.0f) ? (orig_norm / rec_norm) : 1.0f; sycl::half scale_h = sycl::half(scale); sycl::half norm_corr_h = sycl::half(norm_corr); uint32_t parts[4] = {0, 0, 0, 0}; int w = lid >> 3; int off = (lid & 7) * 4; parts[w] |= ((uint32_t)code) << off; uint32_t c0 = sycl::reduce_over_group(it.get_group(), parts[0], sycl::plus<uint32_t>()); uint32_t c1 = sycl::reduce_over_group(it.get_group(), parts[1], sycl::plus<uint32_t>()); uint32_t c2 = sycl::reduce_over_group(it.get_group(), parts[2], sycl::plus<uint32_t>()); uint32_t c3 = sycl::reduce_over_group(it.get_group(), parts[3], sycl::plus<uint32_t>()); if (lid == 0) { uint32_t sc_u = (uint32_t)sycl::bit_cast<uint16_t>(scale_h); uint32_t nc_u = (uint32_t)sycl::bit_cast<uint16_t>(norm_corr_h); db[0] = sc_u | (nc_u << 16); db[1] = c0; db[2] = c1; db[3] = c2; db[4] = c3; } }); });}
} // namespace qx