perf(decode): n>=4 tiled QKV + o_proj GEMV for NVFP4 batched decode (extract of #214 onto main)#266
perf(decode): n>=4 tiled QKV + o_proj GEMV for NVFP4 batched decode (extract of #214 onto main)#266rsafier wants to merge 1 commit into
Conversation
[BETA] Automated PR review — SeedSource AI pipelineAdvisory only, not a merge gate. AI-generated (internal risk-triage, Atlas serve gate, GLM-5.2 review), attributed per Atlas convention. Triage
Serve gate
GLM-5.2 code reviewCorrectness concerns:
Risk: Medium. Buffer sizing is the critical unverified assumption. Quality: Tiling logic is clean; the k==1 remainder fallback duplicating the sequential path is acceptable but could call into a shared helper to avoid drift. Pipeline: build+serve on GB10, GLM-5.2 via NVIDIA NIM, internal risk-triage classifier. External diff treated as untrusted input. Public repo. |
rsafier
force-pushed
the
feat/qkv-batched-tiled-on-main
branch
from
4c29f45 to
dd1e987
Compare
Extends the proven n=2/n=3 batched NVFP4 GEMV paths to n>=4 concurrent decode by tiling: greedy 3s then a 2-or-1 remainder. Each tile reads the q/k/v/o projection weights ONCE for its 2-3 tokens, so n=8 costs 3 weight-matrix reads instead of 8 — amortizing the attention-projection bandwidth that otherwise caps decode scaling at C>=4. - multi_seq/qkv.rs: add `ms_qkv_batch_tiled` + dispatch branch (n>=4 && all-nvfp4). Byte-identical scratch layout to `ms_qkv_batch3` (same ssm_qkvz/attn_output scratch, same q_proj_bytes/kv_bytes/ per_seq_qkv offsets, stream-serialized), so it inherits that path's correctness; k==1 remainder falls back to the per-seq projection. - multi_seq/attn.rs: matching n>=4 o_proj tiling (attn_out/o_out are contiguous, so each tile is just an offset into the same batch3/batch2 kernels the n==2/3 branches use). Focused extraction of the still-unabsorbed batched-decode feature from loader/MMQ rework). No behavior change for n<4 or non-NVFP4 weights. fmt + clippy -D warnings + spark-model lib tests (79) all pass. GPU numeric A/B (tiled vs per-seq, C4/C8 + needle) recommended before merge.
rsafier
force-pushed
the
feat/qkv-batched-tiled-on-main
branch
from
dd1e987 to
2020785
Compare
GPU A/B — tiled-GEMV (this PR) vs wide-GEMM (#305), on current
|
| Concurrency | wide-GEMM (#305) per-seq | tiled (this PR) per-seq | speedup | wide agg | tiled agg |
|---|---|---|---|---|---|
| C=1 | 91.2 | 90.6 | 1.00× (n<4 → shared path) | 89.2 | 88.6 |
| C=4 | 15.8 | 26.7 | 1.69× | 59.9 | 98.6 |
| C=8 | 9.9 | 15.1 | 1.53× | 78.3 | 119.0 |
| C=16 | 5.9 | 8.8 | 1.49× | 91.2 | 134.7 |
Result: the tiled batch3/batch2 GEMV beats the wide GEMM by ~1.5–1.7× at every batched size — opposite to the "one weight read wins" intuition. The wide GEMM's M-tile padding (94% waste at M=4) outweighs its single-read advantage; note wide's C=4 aggregate (59.9) is actually below its C=1 (89.2) — it degrades under moderate batch, which the tiling fixes. Even at C=16 (≈ #305's γ=16 target) tiled wins.
Resolution in this sync
This branch is now rebased onto current main with the n≥4 conflict resolved toward tiled — i.e. it supersedes wide_verify_gemm for that branch (the function is removed; the n==2/3 batch kernels are untouched). Justified by the table above.
Caveats / remaining before merge
- Single model / prompt / greedy / one run. Directionally very strong (and since this is a MoE+SSM model where attention is only part of decode, the raw attention-kernel win is likely larger than the end-to-end 1.5×), but a multi-model sweep would firm up the exact factor.
- Numeric identity: the outstanding tiled-vs-per-seq output-identity check (this PR's own correctness argument) is still the must-do — temp-0 output diverges from the wide kernel (expected: different accumulation order; both coherent/correct), so it must be pinned against the
n==1per-seq reference, not against wide. - DFlash interaction:
wide_verify_gemmwas introduced by feat(dflash): unified context commit (commit_ctx), env-gated (stacked on #304) #305 for the DFlash γ=16 verify step, which flows through this samen≥4branch. Removing it routes DFlash verify through the tiled path too. That's functionally fine (tiled handles any n≥4) and faster here, but a--dflash-on run should confirm no verify-path regression before this lands, since feat(dflash): unified context commit (commit_ctx), env-gated (stacked on #304) #305's kernel was tuned for that specific call site.
DFlash verify-path check — no regression from removing
|
| binary | verify step | decode | accept | crash |
|---|---|---|---|---|
| wide-GEMM (#305) | ~87.0 ms | 58.5 tok/s | 0–13% | none |
| tiled (this PR) | 86.8 ms | 59.8 tok/s | 0–7% | none |
- Verify-step timing identical within noise (~87 ms). feat(dflash): unified context commit (commit_ctx), env-gated (stacked on #304) #305's wide GEMM shows no advantage even in its own DFlash-verify context — the verify step is dominated by drafter-propose + SSM/MoE forward, so the o_proj/QKV kernel washes out there (unlike concurrent batched decode, where tiled is 1.5×).
- Accept rate is kernel-independent (identical math, ~5–13% — that low rate is the known 35B-A3B/drafter divergence, unrelated to this PR).
- Both coherent, no illegal-address / crash through the tiled verify batch.
Conclusion: removing wide_verify_gemm and using tiled everywhere is safe — faster for concurrent decode, equal for DFlash verify, no downside. This PR is synced onto current main with that resolution. Remaining pre-merge item is the tiled-vs-n==1-per-seq numeric-identity check.
High-acceptance verify test — tiled still winsThe DFlash run above had low accept (~5%, adaptive suspended), so the verify path barely ran. To stress a productive n≥4 verify, ran MTP forced past its throughput gate (
Tiled is ~1.18× faster end-to-end even with the verify running every step at high acceptance, and accept is at parity (the 1.72↔1.93 spread is temp-0 numeric drift between the two accumulation kernels over a single run, not a systematic edge). Both coherent, no crash. Verify-path summary across regimes:
Net: tiled is faster or equal in every regime tested — concurrent decode, low-accept DFlash verify, and high-accept MTP verify. Removing |
Second model (Holo-3.1) — the win is model-dependentRan the same n≥4 concurrent-decode sweep on
On Holo the tiled advantage is marginal (1.0–1.10×) — not the 1.49–1.69× seen on Ornith. The cause is in the wide-GEMM baseline: on Ornith (qwen3_6_moe) wide degraded at moderate batch (C=4 aggregate below C=1) and tiling recovered it; on Holo (qwen3_5_moe) wide already scales cleanly (C=4 agg 104.6 vs C=1 76.6), leaving little for tiling to gain. Revised takeawayThe tiled kernel win is model/shape-dependent, not universal: large on Ornith, near-parity on Holo, and never a regression across everything tested (concurrent decode on 2 models, low-accept DFlash verify, high-accept MTP verify). So tiled is a safe drop-in — it helps some models substantially and costs nothing elsewhere — but the honest headline is "up to ~1.7×, ≥ parity everywhere," not a flat 1.5×. Whether that justifies carrying the tiled path over #305's wide GEMM is a maintainer call; the data says there's no downside to doing so. |
Tiebreaker: the win is the BF16→FP4 requant path, NOT model/shapeRan the A/B on
Full heretic sweep: C=4 1.07×, C=8 1.06×, C=16 1.01× — i.e. near-parity, like Holo, not Ornith's 1.5–1.7×. Conclusion: the earlier "model/shape-dependent" framing was wrong. Same shape, opposite result. The discriminator is the attention weight source:
Implication for the PR: the tiled kernel is only a meaningful win for runtime-requantized-BF16 attention checkpoints (a niche: e.g. the AEON "Ultimate" BF16 repacks). For the mainstream native-NVFP4/FP8 releases it's parity. That narrows the value proposition considerably — and suggests the alternative fix worth weighing is building the transposed NVFP4 copy during the BF16→FP4 requant (so |
Follow-up: keeping Ornith's attention in native BF16 doesn't help eitherSince the Ornith win traced to the BF16→FP4 requant (not shape), the obvious question was: just don't requant — drop Ornith per-seq, all attention paths (Qwen3.6-35B-A3B, C = concurrent decode):
Native BF16 lands at exactly 15.8 — same as the pathological wide path. BF16 attention has no fast batched kernel at n≥4 either; it falls to the per-token seq-loop (n full-weight GEMVs at 2 bytes/weight), which is just as slow as the transpose-less wide-GEMM and forfeits the FP4 bandwidth savings. So dropping the requant is a net loss, not a fix. Complete picture for the n≥4 attention path
Bottom line for merge: tiled is ≥ parity everywhere tested (2 native-quant models, DFlash verify, MTP verify) and the only fast option for BF16-repack checkpoints at concurrency. It's a safe drop-in with a real (if niche) benefit. The equally-valid alternative — emit the transposed NVFP4 copy during the BF16→FP4 requant so |
2 participants
Focused extraction of the still-unabsorbed batched-decode feature from #214 onto post-#229
main.Why a focused extraction (not a rebase of #214)
#214's base predates #229. A mechanical
rebase --onto mainproduces ~39K deletions because #229 added the MMQ/q4k-vendor/NVFP4 FFN stack that #214 doesn't have — the loader/kernel surface diverged too far to replay safely. Comparing #214's own.rsagainst currentmainshows that main already absorbed the n=2/n=3 batched QKV paths (ms_qkv_batch2/ms_qkv_batch3live inmulti_seq/qkv.rs). The one genuinely-novel, still-absent piece is the n≥4 tiling that extends those proven kernels to higher decode concurrency. This PR ports exactly that, adapted to main's currentmulti_seqstructure.What it does
For n≥4 concurrent decode with NVFP4 q/k/v/o weights, tile the proven
batch3/batch2GEMV kernels — greedy 3s, then a 2-or-1 remainder. Each tile reads the projection weight once for its 2–3 tokens, so n=8 costs 3 weight-matrix reads instead of 8. This amortizes the attention-projection bandwidth that otherwise makes decode collapse back to one weight read per token at C≥4.multi_seq/qkv.rs: newms_qkv_batch_tiled+ an >= 4 && <all nvfp4>dispatch branch. Scratch layout is byte-identical toms_qkv_batch3(samessm_qkvz()/attn_output()scratch, sameq_proj_bytes/kv_bytes/per_seq_qkvoffsets, all stream-serialized), so it inherits that path's correctness. Thek==1remainder falls back to the per-seq projection.multi_seq/attn.rs: matching n≥4o_projtiling (attn_out/o_outare contiguous, so each tile is just an offset into the samebatch3/batch2kernels the n==2/3 branches already use).No behavior change for n<4 or non-NVFP4 weights (existing branches untouched; the new branch sits between the n==2 case and the seq-loop fallback).
Correctness argument
Correct by construction: identical kernels + offsets to the already-shipped
ms_qkv_batch3, which this repo already trusts for n=3. The scratch buffers are already sized for k≤3 (main'sms_qkv_batch3usesattn_output().offset(3*kv_bytes)); the tiled path never exceeds a k=3 tile. Cross-tile scratch reuse is safe because each tile'scopy_d2d_asyncintoqkv_bufis ordered before the next tile's GEMV on the same stream.Validation
cargo fmt -p spark-model --checkclean.cargo clippy -p spark-model -- -D warningsclean.cargo test -p spark-model --lib→ 79 passed, 0 failed.#[ignore]), so this is a manual pre-merge step.Supersedes the batched-decode portion of #214 (the rest was absorbed by #229). Signed commit.