perf(qwen35): int8 tensor-core prefill attention for Qwythos

[fansilas]

Summary

Runs the Qwythos hd256 prefill attention on the int8 tensor cores. The merged windowed/tiled prefill attention (#455) removed the O(N²) bandwidth problem; what is left is a compute problem — it evaluates QK^T and PV with scalar FMA plus a 5-shuffle warp reduction per key, which measures ~8 TFLOP/s on sm_120 and is 30.5% of prefill at 32k (nsys, below). This adds a drop-in kernel running the identical mask + online softmax on wmma int8, reusing the pattern the merged int8-MMA flash-decode (fa_split_gqa_mma_i8) already ships.

64k prefill pp 4910.81 → 6937.16 (+41.3%); 32k +39.1%; 4k +26.2%. Decode untouched.
Default on; SPARKINFER_PREFILL_ATTN_MMA=0 restores the #455 path exactly.

Proof of speedup

• Tested on RTX 5090 (sm_120)

Decode tok/s — this PR does not target decode and does not touch the decode path; rows are a no-regression check, not a claimed gain (deltas ≤ 0.15% = run-to-run noise):

	decode tok/s
before (main)	291.85
after (this PR)	291.74

Prefill pp tok/s (Qwythos / Qwen3.5, best context --ctx 65536):

	prefill pp tok/s
before prefill (main)	4910.81
after prefill (this PR)	6937.16

Same binary, same GPU, SPARKINFER_PREFILL_ATTN_MMA=0 vs =1, median of 5 reps via the eval's own sweep path (SPARKINFER_BENCH_SWEEP_CTXS / SPARKINFER_BENCH_SWEEP_REPS, one model load), on an otherwise idle box (verified: 0 compute apps). Prefill pp is launch-sensitive, so single-shot numbers on a loaded box swing ±12%; these are medians.

# RTX 5090 sm_120 · main @ f75deb8 · median of 5 reps · idle box
# SPARKINFER_EVAL_PREFILL=1 SPARKINFER_BENCH_SWEEP_CTXS=4096,32768,65536 SPARKINFER_BENCH_SWEEP_REPS=5 \
#   ./build/runtime/qwen3_gguf_bench /workspace/models35/Qwythos-9B-Claude-Mythos-5-1M-Q4_K_M.gguf 128 sweep
### BEFORE (main: SPARKINFER_PREFILL_ATTN_MMA=0)
ctx=4096    decode tg 293.39   prefill pp 6723.65
ctx=32768   decode tg 292.17   prefill pp 4907.57
ctx=65536   decode tg 291.85   prefill pp 4910.81

### AFTER (this PR, MMA default on)
ctx=4096    decode tg 293.02   prefill pp 8482.92   (+26.2%)
ctx=32768   decode tg 291.73   prefill pp 6825.58   (+39.1%)
ctx=65536   decode tg 291.74   prefill pp 6937.16   (+41.3%)

128k is unchanged: it exceeds SPARKINFER_PREFILL_BATCHED_MAXCTX (65536) and falls back to the token path, which this PR does not touch.

Why the attention, and why tensor cores

nsys (--cuda-graph-trace=node, cuda_gpu_kern_sum) on main @ f75deb8, batched-prefill kernels only:

kernel	4k	32k
pf_gemm (bf16; #422's int8 GEMM is gated to N≤8192)	—	47.2%
win_prefill_windowed (attention, #455)	20.7%	30.5%
deq_q4k+deq_q6k (weight dequant, fixed cost)	34.8%	3.6%
pf_gdn_conv / pf_gdn_scan	9.0% / 8.2%	7.4% / 6.9%

At 32k the windowed attention is 262 ms/layer for ~2.08 TFLOP ⇒ ~8 TFLOP/s. It is not memory-bound — #455 already fixed that — it is scalar-math-bound, and it stages K and V into shared memory as fp32 (2·TK·256·4B = 64 KB), capping it at ~1 block/SM.

This kernel instead:

• feeds K/V to wmma directly from the paged int8 pool — a KV page is exactly 16 tokens and the wmma tile is 16×16, so a page is a fragment with ldm = n_kv_heads*HEAD_DIM. No fp32 KV staging ⇒ smem 64 KB → 31.25 KB, and with REG:80/thread and no spills that is a measured 3 blocks/SM.
• quantizes Q per query row (one scale/row) ⇒ QK^T runs int8×int8→int32; the per-row Q scale, per-token K scale and softmax scale are applied to the int32. * folds the per-token V scale into P and quantizes P per row ⇒ PV also runs int8 on the tensor cores.

The mask (causal + sink/window) and the online-softmax recurrence are identical to #455, and the window is read from the same SPARKINFER_PREFILL_ATTN_WINDOW knob (default 256 blocks), so the scalar prefill, this MMA prefill and the sparse-KV decode (#379) keep one selection. Because the query tile is 16 rows aligned to 16 and a KV page is 16 tokens, n_blk_q = (t+16)/16 is constant across a tile, so the window start is computed once per block and only the causal bound varies per row.

This lever is now exhausted: +41% is essentially the ceiling for removing attention entirely (1/(1−0.305) = +43.9% at 32k), so the remaining prefill time is GEMM, not attention.

Correctness

qwen3_gguf_prefill_check (batched vs token-by-token), 128 continuation positions at 4k — the context the accuracy gate actually measures:

@4k, 128 positions	top-1	KL
main (scalar #455)	0.9141	0.01565
this PR (int8 MMA)	0.9219	0.01367

Slightly better than main on both, and clear of the gate (top-1 ≥ 0.90, KL ≤ 0.20). At only 32 positions the same check reads 29/32 vs 30/32 — that 1-token delta is sampling noise, hence 128.

At ≥8k the batched-vs-token A/B is broken on main, and this PR does not change it. main's own scalar windowed prefill scores top-1 0.5625 / KL 0.2247 at 32k. That is the #393 defect: for Qwythos at seqlen ≥ 8192 the sparse combine never applies sigmoid(q_gate), so the token-path reference is ungated while batched prefill gates via pf_mul_sigmoid. This PR is numerically indistinguishable from main's scalar path there — identical top-1, KL within 0.001:

ctx	main scalar	this PR (MMA)
16k	0.4375 / 0.23494	0.4375 / 0.23593
32k (window off)	0.3438 / 0.24093	0.3438 / 0.24163

So the ≥8k numbers move with #393, not with this PR.

Relation to open PRs

New logic lives entirely in two new files (kernels/csrc/cuda/fused/prefill_attn_mma.cu, kernels/include/sparkinfer/kernels/prefill_attn_mma.h) that no other PR touches. The only shared file is kernels/csrc/cuda/fused/batched_prefill.cu (+7/−0: one include + one guard; every existing line byte-identical). Per eval/copycat_policy.py the guard references open non-draft PRs by a different author sharing a file — this PR shares no file with any open PR, so there is nothing for it to compare. Reporting the rest explicitly rather than leaning on that:

• It supersedes perf(qwen35): fp16 shared-memory tiles for windowed prefill attention (+3–5% pp) #463's niche, and I want that on the record. perf(qwen35): fp16 shared-memory tiles for windowed prefill attention (+3–5% pp) #463 (@inference2026, draft, ex-perf(qwen35): default smem-tiled full prefill attention when window off #461) makes perf(qwen35): windowed prefill attention for Qwythos long context #455's smem-tiled kernel the default when the window is off. This kernel's guard runs first and also covers the window-off case (win_blocks <= 0 ⇒ full causal attention on the tensor cores), so if this lands first, perf(qwen35): fp16 shared-memory tiles for windowed prefill attention (+3–5% pp) #463 no longer changes behaviour in any configuration. I am not claiming perf(qwen35): fp16 shared-memory tiles for windowed prefill attention (+3–5% pp) #463's idea — the overlap is that we improve the same call site — but the practical effect is real and a reviewer should weigh it. perf(qwen35): fp16 shared-memory tiles for windowed prefill attention (+3–5% pp) #463's own table reports +0.3% at 4k (noise); this is +41.3% at 64k, and the window-off path is a debug/A-B knob, not a shipping config. Happy to gate this kernel to win_blocks > 0 and leave the window-off path to perf(qwen35): fp16 shared-memory tiles for windowed prefill attention (+3–5% pp) #463 if maintainers prefer that split.
• Distinct from perf(qwen35): windowed prefill attention for Qwythos long context #455 by technique, not just by lines. perf(qwen35): windowed prefill attention for Qwythos long context #455 (merged) is a memory fix: windowing plus staging K/V into shared memory as fp32. This is a compute fix: int8 tensor cores with K/V read directly from the paged pool and no smem staging at all — the opposite of perf(qwen35): windowed prefill attention for Qwythos long context #455's data movement. No part of their kernel body is reused. What is reused is their window semantics: the same SPARKINFER_PREFILL_ATTN_WINDOW knob and the same rsb window-start formula, which must match or the two paths disagree numerically. Added-line containment against perf(qwen35): windowed prefill attention for Qwythos long context #455 is 0.19 (per-function 0); the shared lines are the drop-in ABI signature, namespace boilerplate, the knob name and that formula. With MMA on, the guard supersedes perf(qwen35): windowed prefill attention for Qwythos long context #455's scalar kernel in the default path — the same relationship perf(qwen35): windowed prefill attention for Qwythos long context #455 has to my merged perf(qwen35): batched prompt prefill for Qwythos (one weight-amortized GEMM pass) #398's naive kernel. perf(qwen35): windowed prefill attention for Qwythos long context #455 stays reachable via SPARKINFER_PREFILL_ATTN_MMA=0.
• Credit for the int8-MMA pattern. The quantize-Q-per-row / int8-wmma / fold-V-scale-into-P pattern is lifted from the merged flash-decode lineage: perf(attn): int8 KV tensor-core GQA flash-decode (+19% decode at 16k) #195 (mine) introduced int8-KV tensor-core GQA flash-decode, and perf(qwythos): GQA-4 int8 MMA flash-decode + tiered KV splits (+4% @64k) #366 (@ai Kernel Engineer) extended it to the GQA-4 / hd256 shape (fa_split_gqa_mma_i8) that this kernel mirrors. New work here is porting it to the prefill shape (M = query rows rather than GQA heads, causal + sink/window masking, per-tile window start). Note this kernel does not copy that reference's hd256 score-stride (it stores the QK tile with ldm=HEAD_DIM but reads it back at row stride 128 — self-consistent only at HEAD_DIM==128); the score buffer here is [BM][GN] with one stride used for both the wmma store and every read.
• Not perf(qwen35): int8 tensor-core prefill GEMM for Qwythos [DRAFT — awaits #398] #422's lane. perf(qwen35): int8 tensor-core prefill GEMM for Qwythos [DRAFT — awaits #398] #422 (merged) is the int8 tensor-core prefill GEMM, gated to N ≤ 8192. This PR does not touch the GEMM — which is exactly why 32k/64k (where the GEMM is still bf16) are the best contexts here, and why the two changes stack rather than compete.
• Not perf(qwen35): weight-amortized Q4_K tile-GEMM prefill for Qwythos above 64k (5.5x pp @128k) #409's lane. perf(qwen35): weight-amortized Q4_K tile-GEMM prefill for Qwythos above 64k (5.5x pp @128k) #409 chunks prompts >64k with a Q4_K dp4a tile GEMM. This PR changes nothing above 64k — 128k still falls back to the token path.

[skyrocket2026]

❌ sparkinfer auto-eval — 0217e1f

metric	value
label	eval:REJECT
Qwen3.5 score	eval-qwen35:REJECT (fail)
Qwen3.6 score	eval-qwen36:none (pass)
Qwen3.5 scored decode (128 ctx)	? tok/s
Qwen3.5 correctness	top-1 0.0% · KL ?
Qwen3.6 vs same-box main	481.85 tok/s → +0.1% (+0.6)
Qwen3.6 scored decode (512 ctx · 512-context)	482.45 tok/s
Qwen3.6 correctness	top-1 95.4% · KL 0.0364
Qwen3.6 128-token no-regression gate	488.84 tok/s vs main 488.74 tok/s · pass
Qwen3.6 512-context no-regression gate	482.45 tok/s vs main 481.85 tok/s · pass
Qwen3.6 4k-context no-regression gate	462.73 tok/s vs main 462.59 tok/s · pass
Qwen3.6 16k-context no-regression gate	450.41 tok/s vs main 450.44 tok/s · pass
Qwen3.6 32k-context no-regression gate	425.56 tok/s vs main 425.37 tok/s · pass
Qwen3.5 optimize	eval:REJECT · ? tok/s · fail
Qwen3.6 optimize	eval:none · 482.45 tok/s · pass
Qwen3.6 optimize — Qwythos-9B (Q4_K_M) guard accuracy	top-1 90.8% · KL 0.0373 · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 128	488.84 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 512	482.45 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 4k	462.73 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 16k	450.41 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 32k	425.56 tok/s · pass

Rejected — primary (Qwythos-9B (Q4_K_M)) produced no verdict (infra error).

RTX 5090 (sm_120) · 128-token decode scored vs same-box main · built from source · correctness vs llama.cpp. Automated — not merged; merge manually after review.

[skyrocket2026]

✅ sparkinfer auto-eval — 0217e1f

metric	value
label	eval:XL
Qwen3.5 score	eval-qwen35:XL (pass)
Qwen3.6 score	eval-qwen36:none (pass)
Qwen3.5 vs same-box main	5509.48 tok/s → +42.1% (+2318.6)
Qwen3.5 scored decode (4096 ctx · 4k-context)	286.77 tok/s
Qwen3.5 scored prefill (65536 ctx · 64k-context)	7828.06 pp tok/s · eval-prefill:XL
Qwen3.5 correctness	top-1 95.7% · KL 0.0296
Qwen3.5 128-token no-regression gate	296.18 tok/s vs main 296.67 tok/s · pass
Qwen3.5 4k-context no-regression gate	286.77 tok/s vs main 286.67 tok/s · pass
Qwen3.5 32k-context no-regression gate	285.44 tok/s vs main 285.45 tok/s · pass
Qwen3.5 64k-context no-regression gate	285.51 tok/s vs main 285.47 tok/s · pass
Qwen3.5 4k prefill no-regression gate	7797.29 pp tok/s vs main 6303.4 pp tok/s · pass
Qwen3.5 32k prefill no-regression gate	7664.45 pp tok/s vs main 5474.93 pp tok/s · pass
Qwen3.5 64k prefill no-regression gate	7828.06 pp tok/s vs main 5509.48 pp tok/s · pass
Qwen3.5 128k prefill no-regression gate	0.0 pp tok/s · pass
Qwen3.6 vs same-box main	462.06 tok/s → +0.2% (+0.7)
Qwen3.6 scored decode (4096 ctx · 4k-context)	462.75 tok/s
Qwen3.6 correctness	top-1 92.2% · KL 0.0372
Qwen3.6 128-token no-regression gate	488.74 tok/s vs main 488.58 tok/s · pass
Qwen3.6 512-context no-regression gate	481.99 tok/s vs main 481.78 tok/s · pass
Qwen3.6 4k-context no-regression gate	462.75 tok/s vs main 462.06 tok/s · pass
Qwen3.6 16k-context no-regression gate	450.43 tok/s vs main 450.24 tok/s · pass
Qwen3.6 32k-context no-regression gate	425.46 tok/s vs main 425.29 tok/s · pass
Qwen3.5 optimize	eval:XL · 7828.06 tok/s · pass
Qwen3.5 optimize — Qwen3.6-35B-A3B guard accuracy	top-1 92.2% · KL 0.0372 · pass
Qwen3.5 optimize — Qwythos-9B (Q4_K_M) 128	296.18 tok/s · pass
Qwen3.5 optimize — Qwythos-9B (Q4_K_M) 4k	286.77 tok/s · pass
Qwen3.5 optimize — Qwythos-9B (Q4_K_M) 32k	285.44 tok/s · pass
Qwen3.5 optimize — Qwythos-9B (Q4_K_M) 64k	285.51 tok/s · pass
Qwen3.6 optimize	eval:none · 462.75 tok/s · pass
Qwen3.6 optimize — Qwythos-9B (Q4_K_M) guard accuracy	top-1 95.7% · KL 0.0296 · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 128	488.74 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 512	481.99 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 4k	462.75 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 16k	450.43 tok/s · pass
Qwen3.6 optimize — Qwen3.6-35B-A3B 32k	425.46 tok/s · pass

Verified speedup over same-box origin/main — 7828.06 tok/s (main was 5509.48 tok/s).

RTX 5090 (sm_120) · 128/512/4k/16k/32k guarded · Qwen3.5 prefill at 4k/32k/64k · scored vs same-box main · built from source · correctness vs llama.cpp. Automated — not merged; merge manually after review.

[skyrocket2026]

✅ Auto-merged as the round's merge-first winner — verified same-box speedup over main, all checks green. Thanks for the contribution!