sycl: prompt attention scores one work-item per cell (2.0-2.2x the warp kernel, prompt +10% at 2K)

On a B70 the QSA prompt attention is 15% of a 2K prompt and 28% of a 40K one (STRATA_PREFILL_TIMING). The
warp kernel splits each cell's 256 dimensions over a sub-group and reduces them with a 5-step shuffle tree per
query head, 60 shuffles a cell. qsa_prompt_attn_percell gives work-item c cell c's key row (16-byte loads) and
all 12 heads' dot products, reading the query heads from local memory, in 128-cell chunks; the value pass and
the merge are the warp kernel's. Each 16-dimension block is summed in four short chains before it joins the
cell's sum, so its FP32 error stays near the warp kernel's. Prompt path only (prefill.cpp); decode and the draft
layer keep qsa_decode_attn_batch. STRATA_ATTN_PERCELL=0 brings the warp kernel back.

Not bitwise: FP32 in another summation order (as the XMX prompt kernel). qsa_prompt_attn_parity now checks it
too (against FP64 and the warp kernel, int8 / FP16 / Q4_0 pools, 0 failures on a B70: error vs FP64 1.3-1.6x the
warp kernel's, 1.2-2.2x faster) and is registered with ctest; the XMX kernel's refusal of Q4_0 pools no longer
ends the test.

Arc Pro B70, Coder IQ1_M, 256K config, `--pcie-frac 0 --adapt-every 0 STRATA_PREFILL_MMQ=1`, 256 greedy tokens:
2,185-token prompt 1,330 -> 1,481 tok/s (+11%, 2 rounds each). First-token logits against the warp kernel at 8K:
KL 3.9e-10 nats, top-1 equal (upstream's int8 vs FP16 prompt experts on the same prompt: KL 8.4e-10).

sycl/CMakeLists.txt

@@ -356,11 +356,12 @@ if(STRATA_SYCL_PARITY)
   target_link_libraries(conversation_snapshot_test PRIVATE strata_engine)
   add_test(NAME conversation_snapshot_test COMMAND conversation_snapshot_test)
   set_tests_properties(conversation_snapshot_test PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 120)
-  # the prompt-attention parity takes a context length, not --selftest, and needs the tensor-core kernel (not
-  # ported yet: it reports "refused the pools" here). Built, not registered.
+  # the prompt-attention parity takes a context length, not --selftest: the XMX prompt kernel (opt-in in the engine;
+  # it refuses Q4_0 pools) and the prompt path's per-cell kernel, both against the warp kernel and FP64
   strata_resolve(_src src/kernels/qsa_prompt_attn_parity.cpp)
   add_executable(qsa_prompt_attn_parity ${_src})
   target_link_libraries(qsa_prompt_attn_parity PRIVATE strata_kernels)
+  add_test(NAME qsa_prompt_attn_parity COMMAND qsa_prompt_attn_parity 8192 512 2)
   # the hybrid K8V4 KV test (upstream 0.1.27) passes its appends, gathers and attention here and fails only its
   # last step, the same tensor-core prompt kernel. Built, not registered, until that kernel is ported.
   strata_resolve(_src src/kernels/kv_hybrid_parity.cpp)

sycl/src/kernels/cuda/qsa_decode_attn.dp.cpp

@@ -643,6 +643,167 @@ bool pre75_attn() {
 #define STRATA_ATTN_CHUNK(M) attn_chunk_kernel<M>
 #endif
 
+// ---- SYCL port: the prompt path's chunk kernel with the scores computed one work-item per cell (attn_bench).
+// The warp kernel above splits a cell's 256 dimensions over a sub-group and adds them up with a 5-step shuffle
+// tree per query head - 60 shuffles a cell. Here work-item c holds cell c's key row (16-byte loads) and all 12
+// heads' dot products, reading the query heads from local memory, where every work-item reads the same element
+// at once. 128-cell chunks. The values pass is the warp kernel's. FP32, but another summation order than the warp
+// kernel: not bitwise (relative error ~2e-7 in attn_bench); attn_bench: 1.55-1.69x the warp kernel's time.
+constexpr int PC_CHUNK = 128;
+
+// 16 consecutive key values of row `row`, dimensions [d0, d0+16)
+template <int KV_MODE>
+__dpct_inline__ void load16_k(const QsaAttnPools &p, long long row, int d0, float *out) {
+    if constexpr (KV_MODE == 0) {
+        load8_f16(p, false, row, d0, out);
+        load8_f16(p, false, row, d0 + 8, out + 8);
+    } else if constexpr (KV_MODE == 1 || KV_MODE == 3) {
+        const int8_t* codes = p.k_q + row * HD + d0;
+        const float sc = sycl::vec<sycl::half, 1>(sycl::bit_cast<sycl::half, unsigned short>(
+                             p.k_scale[row * (HD / KV_Q8_GROUP) + d0 / KV_Q8_GROUP]))
+                             .convert<float, sycl::rounding_mode::automatic>()[0];
+        const sycl::int4 raw = *reinterpret_cast<const sycl::int4 *>(codes);
+        const int8_t* c = reinterpret_cast<const int8_t*>(&raw);
+#pragma unroll
+        for (int j = 0; j < 16; ++j) out[j] = (float) c[j] * sc;
+    } else {   // q4_0 (rotated): element j of a 32-block in qs[j]'s low nibble, j + 16 in its high one
+        load8_q4(p, false, row, d0, out);
+        load8_q4(p, false, row, d0 + 8, out + 8);
+    }
+}
+
+template <int KV_MODE>
+__dpct_inline__ float load1_v(const QsaAttnPools &p, long long row, int t) {
+    if constexpr (KV_MODE == 0) {
+        return sycl::vec<sycl::half, 1>(sycl::bit_cast<sycl::half, unsigned short>(p.v_pool[row * HD + t]))
+            .convert<float, sycl::rounding_mode::automatic>()[0];
+    } else if constexpr (KV_MODE == 1) {
+        const float sc = sycl::vec<sycl::half, 1>(sycl::bit_cast<sycl::half, unsigned short>(
+                             p.v_scale[row * (HD / KV_Q8_GROUP) + t / KV_Q8_GROUP]))
+                             .convert<float, sycl::rounding_mode::automatic>()[0];
+        return (float) p.v_q[row * HD + t] * sc;
+    } else {
+        constexpr int bytes_per_head = (HD / QK4_0) * sizeof(block_q4_0);
+        const block_q4_0* blk = reinterpret_cast<const block_q4_0*>(p.v_q4 + row * bytes_per_head) + t / QK4_0;
+        const float d = sycl::vec<sycl::half, 1>(sycl::bit_cast<sycl::half, unsigned short>(blk->d))
+                            .convert<float, sycl::rounding_mode::automatic>()[0];
+        const int rem = t % QK4_0;
+        const uint8_t byte = blk->qs[rem < 16 ? rem : rem - 16];
+        return (float) ((rem < 16) ? ((byte & 0x0F) - 8) : ((byte >> 4) - 8)) * d;
+    }
+}
+
+template <int KV_MODE>
+__dpct_inline__ void attn_chunk_percell_kernel(const float *__restrict__ q, QsaAttnPools p,
+                                               const int32_t *__restrict__ ids, const int32_t *__restrict__ step,
+                                               int n_kv_heads, int page_size, float scale, float *__restrict__ part_acc,
+                                               float *__restrict__ part_m, float *__restrict__ part_l, int n_chunks,
+                                               int cap, long long scratch_stride) {
+    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
+    auto sg = item_ct1.get_sub_group();
+    q += (size_t) item_ct1.get_group(0) * (size_t) (n_kv_heads * G) * HD;
+    ids += (size_t) item_ct1.get_group(0) * (size_t) cap;
+    step += (size_t) item_ct1.get_group(0) * kStepCount;
+    part_acc += (size_t) item_ct1.get_group(0) * (size_t) scratch_stride;
+    part_m += (size_t) item_ct1.get_group(0) * (size_t) scratch_stride;
+    part_l += (size_t) item_ct1.get_group(0) * (size_t) scratch_stride;
+    auto &sq = *sycl::ext::oneapi::group_local_memory_for_overwrite<float[G][HD]>(
+        sycl::ext::oneapi::this_work_item::get_work_group<3>());
+    auto &sp = *sycl::ext::oneapi::group_local_memory_for_overwrite<float[G][PC_CHUNK]>(
+        sycl::ext::oneapi::this_work_item::get_work_group<3>());
+    auto &srow = *sycl::ext::oneapi::group_local_memory_for_overwrite<long long[PC_CHUNK]>(
+        sycl::ext::oneapi::this_work_item::get_work_group<3>());
+    const int n_ids = *(step + kStepWidth);
+    const int chunk = item_ct1.get_group(2), kvh = item_ct1.get_group(1);
+    const int t = item_ct1.get_local_id(2);
+    const int lane = (int) sg.get_local_id()[0], warp = (int) sg.get_group_id()[0];
+    const int SGS = (int) sg.get_local_range()[0], NW = THREADS / SGS;
+    const int c0 = chunk * PC_CHUNK;
+    const int n_here = sycl::min(PC_CHUNK, n_ids - c0);
+    const int slot = kvh * n_chunks + chunk;
+    if (n_here <= 0) {
+        if (t < G) { part_m[slot * G + t] = -FLT_MAX; part_l[slot * G + t] = 0.0f; }
+        return;
+    }
+    for (int i = t; i < G * HD; i += THREADS) sq[i / HD][i % HD] = q[(size_t) (kvh * G) * HD + i];
+    for (int c = t; c < PC_CHUNK; c += THREADS) {
+        long long r = -1;
+        if (c < n_here) {
+            const int cell = ids[c0 + c];
+            const long long page = (long long) p.page_table[cell / page_size];
+            if (page >= 0) r = (page * n_kv_heads + kvh) * page_size + (cell % page_size);   // -1: masked (KV streaming)
+        }
+        srow[c] = r;
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+    for (int c = t; c < PC_CHUNK; c += THREADS) {
+        if (c >= n_here || srow[c] < 0) {
+#pragma unroll
+            for (int h = 0; h < G; ++h) sp[h][c] = -FLT_MAX;
+            continue;
+        }
+        float acc_s[G];
+#pragma unroll
+        for (int h = 0; h < G; ++h) acc_s[h] = 0.0f;
+        for (int d = 0; d < HD; d += 16) {
+            float kf[16];
+            load16_k<KV_MODE>(p, srow[c], d, kf);
+#pragma unroll
+            for (int h = 0; h < G; ++h) {   // a 16-dimension block in four short chains, then into the cell's sum:
+                float b0 = 0.0f, b1 = 0.0f, b2 = 0.0f, b3 = 0.0f;   // FP32 error near the warp kernel's tree
+#pragma unroll
+                for (int j = 0; j < 16; j += 4) {
+                    b0 = sycl::fma(kf[j], sq[h][d + j], b0);
+                    b1 = sycl::fma(kf[j + 1], sq[h][d + j + 1], b1);
+                    b2 = sycl::fma(kf[j + 2], sq[h][d + j + 2], b2);
+                    b3 = sycl::fma(kf[j + 3], sq[h][d + j + 3], b3);
+                }
+                acc_s[h] += (b0 + b1) + (b2 + b3);
+            }
+        }
+#pragma unroll
+        for (int h = 0; h < G; ++h) sp[h][c] = acc_s[h] * scale;
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+    for (int h = warp; h < G; h += NW) {
+        float m = -FLT_MAX;
+        for (int c = lane; c < PC_CHUNK; c += SGS) m = sycl::fmax(m, sp[h][c]);
+        m = sycl::reduce_over_group(sg, m, sycl::maximum<float>());
+        float l = 0.0f;
+        for (int c = lane; c < PC_CHUNK; c += SGS) {
+            const float e = (c < n_here && srow[c] >= 0) ? sycl::native::exp(sp[h][c] - m) : 0.0f;
+            sp[h][c] = e;
+            l += e;
+        }
+        l = sycl::reduce_over_group(sg, l, sycl::plus<float>());
+        if (lane == 0) { part_m[slot * G + h] = m; part_l[slot * G + h] = l; }
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+    float acc[G];
+#pragma unroll
+    for (int h = 0; h < G; ++h) acc[h] = 0.0f;
+    for (int c = 0; c < n_here; ++c) {
+        if (srow[c] < 0) continue;   // masked above, weight 0
+        const float v = load1_v<KV_MODE == 3 ? 2 : KV_MODE>(p, srow[c], t);
+#pragma unroll
+        for (int h = 0; h < G; ++h) acc[h] = sycl::fma(sp[h][c], v, acc[h]);
+    }
+#pragma unroll
+    for (int h = 0; h < G; ++h) part_acc[((size_t) slot * G + h) * HD + t] = acc[h];
+}
+
+template <int KV_MODE>
+void launch_percell(dpct::queue_ptr st, const float* q, const QsaAttnPools& pools, const int32_t* ids,
+                    const int32_t* steps, int64_t cap, const QsaShapes& s, float* part_acc, float* part_m, float* part_l,
+                    int n_chunks, long long stride, int64_t n_q, float scale) {
+    const sycl::range<3> grid((size_t) n_q, (size_t) s.n_head_kv, (size_t) n_chunks);
+    st->parallel_for(sycl::nd_range<3>(grid * sycl::range<3>(1, 1, THREADS), sycl::range<3>(1, 1, THREADS)),
+                     [=](sycl::nd_item<3>) [[sycl::reqd_sub_group_size(32)]] {
+                         attn_chunk_percell_kernel<KV_MODE>(q, pools, ids, steps, (int) s.n_head_kv, (int) s.page_size,
+                                                            scale, part_acc, part_m, part_l, n_chunks, (int) cap, stride);
+                     });
+}
+
 }  // namespace
 
 void qsa_decode_attn_batch(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps,
@@ -961,4 +1122,28 @@ void qsa_decode_attn_step(const float* q, const QsaAttnPools& pools, const int32
     const dpct::err0 e = 0;
 }
 
+// SYCL port, prompt path only (prefill.cpp; decode and the draft layer keep qsa_decode_attn_batch): the chunk
+// attention with the per-cell score kernel above, 128-cell chunks in the same scratch as the warp kernel (its stride
+// is the 64-cell layout's, so it fits). FP32, another summation order than the warp kernel: not bitwise.
+void qsa_prompt_attn_percell(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps,
+                             int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream) {
+    if (n_q <= 0) return;
+    const int kv_mode = pools.k_q4 != nullptr ? 2 : (pools.k_q != nullptr && pools.v_q4 != nullptr ? 3
+                        : (pools.k_q != nullptr ? 1 : 0));
+    const long long stride = (long long) qsa_decode_attn_scratch_floats(cap, s);
+    const float scale = 1.0f / sqrtf((float) HD);
+    dpct::queue_ptr st = strata::q_of(stream);
+    const int nc = (int) ((cap + PC_CHUNK - 1) / PC_CHUNK);
+    float* pa = scratch;
+    float* pm = scratch + (size_t) nc * s.n_head * HD;
+    float* pl = pm + (size_t) nc * s.n_head;
+    if (kv_mode == 3) launch_percell<3>(st, q, pools, ids, steps, cap, s, pa, pm, pl, nc, stride, n_q, scale);
+    else if (kv_mode == 2) launch_percell<2>(st, q, pools, ids, steps, cap, s, pa, pm, pl, nc, stride, n_q, scale);
+    else if (kv_mode == 1) launch_percell<1>(st, q, pools, ids, steps, cap, s, pa, pm, pl, nc, stride, n_q, scale);
+    else launch_percell<0>(st, q, pools, ids, steps, cap, s, pa, pm, pl, nc, stride, n_q, scale);
+    st->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, (size_t) n_q, (size_t) s.n_head) * sycl::range<3>(1, 1, HD),
+                                       sycl::range<3>(1, 1, HD)),
+                     [=](sycl::nd_item<3>) { attn_merge_kernel(pa, pm, pl, nc, attn, stride); });
+}
+
 }  // namespace strata::kernels

sycl/src/kernels/qsa_prompt_attn_parity.cpp

@@ -26,6 +26,10 @@
 #include <vector>
 
 namespace k = strata::kernels;
+namespace strata::kernels {   // SYCL port: the prompt path's per-cell attention (qsa_decode_attn.dp.cpp)
+void qsa_prompt_attn_percell(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps,
+                             int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream);
+}
 
 namespace {
 void ck(dpct::err0 e, const char *w) {
@@ -242,17 +246,27 @@ int run(int fmt, int64_t ctx, int64_t nq, int reps) {   // fmt 1 int8, 0 fp16, 2
             k::qsa_decode_attn_batch(d_q + t0 * NH * HD, pl, d_ids + t0 * cap, d_steps + t0 * k::kStepCount, cap, s,
                                      scratch, d_old + t0 * NH * HD, std::min(batch, nq - t0), nullptr);
     };
+    float* d_pc = nullptr;   // SYCL port: qsa_prompt_attn_percell's output
+    ck(DPCT_CHECK_ERROR(d_pc = (float *)sycl::malloc_device(nq * NH * HD * 4, dpct::get_in_order_queue())), "malloc");
+    auto pc_run = [&]() {
+        for (int64_t t0 = 0; t0 < nq; t0 += batch)
+            k::qsa_prompt_attn_percell(d_q + t0 * NH * HD, pl, d_ids + t0 * cap, d_steps + t0 * k::kStepCount, cap, s,
+                                       scratch, d_pc + t0 * NH * HD, std::min(batch, nq - t0), nullptr);
+    };
+    bool refused = false;   // SYCL port: the XMX kernel takes int8 and FP16 pools; the per-cell check runs anyway
     auto new_run = [&]() {
         if (!k::qsa_prompt_attn_batch(d_q, pl, d_ids, d_steps, cap, s, d_new, nq, nullptr)) {
-            std::fprintf(stderr, "qsa_prompt_attn_batch refused the pools\n");
-            std::exit(2);
+            if (!refused) std::fprintf(stderr, "qsa_prompt_attn_batch refused the pools: per-cell check only\n");
+            refused = true;
         }
     };
     old_run();
     new_run();
+    pc_run();
     ck(DPCT_CHECK_ERROR(dpct::get_current_device().queues_wait_and_throw()),
        "run");
-    std::vector<float> o((size_t) (nq * NH * HD)), nw(o.size());
+    std::vector<float> o((size_t) (nq * NH * HD)), nw(o.size()), pcw(o.size());
+    ck(DPCT_CHECK_ERROR(dpct::get_in_order_queue().memcpy(pcw.data(), d_pc, pcw.size() * 4).wait()), "down");
     ck(DPCT_CHECK_ERROR((dpct::get_current_device().queues_wait_and_throw(), dpct::get_in_order_queue())
                             .memcpy(o.data(), d_old, o.size() * 4)
                             .wait()),
@@ -261,13 +275,14 @@ int run(int fmt, int64_t ctx, int64_t nq, int reps) {   // fmt 1 int8, 0 fp16, 2
                             .memcpy(nw.data(), d_new, nw.size() * 4)
                             .wait()),
        "down");
+    if (refused) nw = o;   // SYCL port: nothing of the XMX kernel to compare
     {   // a hash of the new kernel's output bits (bitwise A/B of two builds / switches)
         uint64_t hh = 1469598103934665603ull;
         for (float x : nw) { uint32_t u; std::memcpy(&u, &x, 4); hh = (hh ^ u) * 1099511628211ull; }
         std::printf("NEWHASH %016llx\n", (unsigned long long) hh);
     }
     // 1. FP64 reference on a sample of queries
-    double err_old = 0, err_new = 0, ref_scale = 0;
+    double err_old = 0, err_new = 0, err_pc = 0, ref_scale = 0;
     for (int64_t i = 0; i < nq; i += std::max<int64_t>(1, nq / 16)) {
         const int64_t w = steps[i * k::kStepCount + k::kStepWidth];
         const int32_t* sel = ids.data() + i * cap;
@@ -301,13 +316,15 @@ int run(int fmt, int64_t ctx, int64_t nq, int reps) {   // fmt 1 int8, 0 fp16, 2
                 ref_scale = std::max(ref_scale, std::fabs(r));
                 err_old = std::max(err_old, std::fabs(o[at] - r));
                 err_new = std::max(err_new, std::fabs(nw[at] - r));
+                err_pc = std::max(err_pc, std::fabs(pcw[at] - r));
             }
         }
     }
     // 2. new vs old everywhere
-    double diff = 0, scale = 0;
+    double diff = 0, scale = 0, diff_pc = 0;
     for (size_t i = 0; i < o.size(); ++i) {
         diff = std::max(diff, (double) std::fabs(o[i] - nw[i]));
+        diff_pc = std::max(diff_pc, (double) std::fabs(o[i] - pcw[i]));
         scale = std::max(scale, (double) std::fabs(o[i]));
     }
     // 3. speed
@@ -331,7 +348,20 @@ int run(int fmt, int64_t ctx, int64_t nq, int reps) {   // fmt 1 int8, 0 fp16, 2
         (e1->get_profiling_info<sycl::info::event_profiling::command_end>() -
          e0->get_profiling_info<sycl::info::event_profiling::command_start>()) /
         1000000.0f;
-    const bool ok1 = err_new <= std::max(4.0 * err_old, 1e-6 * ref_scale);
+    float ms_pc = 0;
+    dpct::sync_barrier(e0);
+    for (int r = 0; r < reps; ++r) pc_run();
+    dpct::sync_barrier(e1);
+    ck(DPCT_CHECK_ERROR(e1->wait_and_throw()), "time");
+    ms_pc = (e1->get_profiling_info<sycl::info::event_profiling::command_end>() -
+             e0->get_profiling_info<sycl::info::event_profiling::command_start>()) / 1000000.0f;
+    // SYCL port: the per-cell kernel - FP32 like the warp kernel, another summation order
+    const bool ok_pc = err_pc <= std::max(4.0 * err_old, 1e-6 * ref_scale) && diff_pc <= 1e-4 * scale;
+    std::printf("%s %s per-cell: vs FP64 %.3g (warp kernel %.3g), vs the warp kernel %.2g of scale; %.3f -> %.3f ms per chunk (%.2fx)\n",
+                ok_pc ? "PASS" : "FAIL", fmt == 2 ? "q4_0" : fmt == 1 ? "int8" : "fp16", err_pc, err_old, diff_pc / scale,
+                ms_old / reps, ms_pc / reps, ms_old / ms_pc);
+    sycl::free((void *)d_pc, dpct::get_in_order_queue());
+    const bool ok1 = err_new <= std::max(4.0 * err_old, 1e-6 * ref_scale) && ok_pc;
     const bool ok2 = diff <= 1e-4 * scale;
     std::printf("%s %s ctx %lld, %lld queries: vs FP64 old %.3g new %.3g (output scale %.3g); new vs old %.3g (%.2g of "
                 "scale); %.3f -> %.3f ms per chunk (%.2fx)\n",

sycl/src/prefill/prefill.cpp

@@ -40,6 +40,10 @@
 #include "strata/prefill/moe_mmq.hpp"
 #include "strata/core/peer_experts.hpp"
 #include "strata/prefill/kernels.hpp"
+namespace strata::kernels {   // SYCL port, kept out of the shared headers: qsa_decode_attn.dp.cpp
+void qsa_prompt_attn_percell(const float* q, const QsaAttnPools& pools, const int32_t* ids, const int32_t* steps,
+                             int64_t cap, const QsaShapes& s, float* scratch, float* attn, int64_t n_q, void* stream);
+}
 
 #include <algorithm>
 #include <set>
@@ -3432,6 +3436,15 @@ bool Prefill::run_impl(const int64_t *tokens, int64_t n, int64_t pos0,
                                                                             m.attn, T, m.cs))
                         for (int64_t t0 = 0; t0 < T; t0 += m.attn_batch) {
                             const int64_t nb = std::min(m.attn_batch, T - t0);
+                            // SYCL port: the scores one work-item per cell (qsa_prompt_attn_percell, 128-cell chunks;
+                            // FP32, not bitwise with the warp kernel). STRATA_ATTN_PERCELL=0: the warp kernel.
+                            static const bool percell = [] { const char* v = std::getenv("STRATA_ATTN_PERCELL"); return !v || v[0] != '0'; }();
+                            if (percell) {
+                                strata::kernels::qsa_prompt_attn_percell(m.q + t0 * ZV, pools, m.sel_ids + t0 * m.cap,
+                                                                         m.steps_dev + t0 * strata::kernels::kStepCount, m.cap,
+                                                                         s, m.attn_scratch, m.attn + t0 * ZV, nb, m.cs);
+                                continue;
+                            }
                             strata::kernels::qsa_decode_attn_batch(m.q + t0 * ZV, pools, m.sel_ids + t0 * m.cap,
                                                                    m.steps_dev + t0 * strata::kernels::kStepCount, m.cap,
                                                                    s, m.attn_scratch, m.attn + t0 * ZV, nb, m.cs);