[gcanlin]

Purpose

fused_norm_rope_attn_rope_cast_decode schedules one cluster per query token, so a small decode batch leaves most of the GPU idle. On one B200 (h_q=64, topk 128+512, CUDA graphs) the kernel costs the same 26.6 us at s_q = 1, 6 and 96 — 96x the work for free. The kernel header explains why there is no split-KV today:

Split-KV is not supported (the fused RoPE + FP8-quant epilogue cannot be combined across splits)

This PR adds an opt-in fp32 partial-output mode so the epilogue can move to a combine kernel, and with it a KV split.

What changes

• mega_num_splits extends the job space from (query token) to (query token, KV split). Each job walks only its own KV block range; the grid grows accordingly. A split with no blocks of its own re-runs the last block with contributes=false so the pipeline and its barriers stay uniform, and its result is discarded.
• When mega_o_accum / mega_lse_accum are supplied, store_o writes an fp32 partial normalised by that split's l_i plus a log2-space LSE, and skips the FP8 cast. This is exactly the protocol smxx/decode/combine already consumes — including leaving the attention sink to the combine.
• The inverse RoPE stays in the main kernel: it is linear and every split shares the query position, so it commutes with the cross-split weighted sum.
• Unsplit behaviour is unchanged. Without the partial buffers the fused FP8 epilogue runs exactly as before.

Measured

Kernel alone, one B200, h_q=64, topk 128+512, CUDA graphs, us:

s_q	unsplit	split2	split4	split8	split16
1	26.6	20.5	16.4	14.4	14.3
6	26.7	22.5	18.4	16.4	14.4
16	28.7	22.5	18.5	18.4	26.4
96	30.7	42.9	47.1	68.4	117.6
384	77.5	110.9	158.9	242.4	434.4

End to end through vLLM (DeepSeek-V4.1-Flash, TP4 + EP + MegaMoE + SP, FLASHMLA_MEGA_ATTN_DSV41, vllm bench serve, random 4096/256, 3 seeds per lane, cold server per configuration), with a combine kernel on the caller side and splitting gated to <= 8 decode tokens:

lane	metric	unsplit	split-KV	delta
C1	out tok/s	120.12 / 121.67 / 121.72	130.69 / 132.55 / 132.55	+8.9%
C1	TPOT ms	7.77 / 7.78 / 7.78	7.11 / 7.10 / 7.10	-8.7%
C16	out tok/s	1301.84 / 1304.29 / 1299.16	1299.77 / 1290.00 / 1290.72	-0.6% (noise)

Callers must gate this tightly. The partials cost num_splits * s_q * h_q * d_v * 4 bytes per layer per step: 1.05 MB/layer at one decode token with 8 splits, but 16.8 MB/layer at 16 tokens. Splitting at 16 tokens regressed that lane by ~35% end to end even though the kernel microbenchmark above still shows a win there — the microbenchmark times one layer in isolation and cannot see the aggregate HBM pressure.

Test plan

pytest tests/test_fused_norm_rope_attn_rope_cast_splitkv.py -q   # 20 passed

The new test's oracle is the kernel itself at mega_num_splits=1, which takes the same fp32-partial path over the whole KV range — that isolates the split logic from the FP8 epilogue, so a real regression appears far below quantisation noise. It covers s_q in {1, 6, 16} x splits in {2, 4, 8} x sink {off, on}, plus a degenerate case with more splits than KV blocks, plus a guard that the unsplit path still produces FP8.

Against the unsplit FP8 path: LSE max abs error 9.5e-07, and the FP8 codes are 99.6% bit-identical with non-cancelling KV values. With randn KV they are 91-94% identical; that is not a logic error but fp32 re-association showing through — |O| is ~20x smaller there because the V terms cancel, while the absolute error stays ~6e-4, and the LSE (which sums every token's contribution) stays exact at every split count.

tests/test_fused_norm_rope_attn_rope_cast.py needs tile_kernels / tilelang / kernelkit, which I could not install in this environment; the new test deliberately depends only on torch so it runs wherever the extension builds.

Limitations / follow-ups

• h_q == 64 (CLUSTER_SIZE == 1) only. The API rejects the partial buffers otherwise.
• No combine kernel in this repo yet. The combine used for the numbers above lives in the caller. The natural follow-up is porting it to smxx/decode/combine with PDL, which should also widen the end-to-end gain — that combine is currently launch-bound and does not overlap the main kernel's tail the way the existing one does.
• Prefill is untouched.
• The attention-sink path is validated synthetically (split vs unsplit), not against a model eval.

Duplicate-work check

Searched open PRs and issues on this repo for split-KV / partial-output work on the fused mega decode kernel and found none. The non-fused sm100/decode/sparse/head64 path already ships split and no-split instantiations; this brings the same capability to the fused kernel and reuses its partial/LSE protocol rather than inventing a new one.

AI assistance

This change was developed with AI assistance (Claude Code). The submitter has reviewed every changed line.

🤖 Generated with Claude Code