#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