Repository navigation
[Performance] CuTe DSL MLA FP8-KV decode kernel slower at long KV after #19338 #20510
Description
Summary
Since #19338 (9257aae, "Upgrade Blackwell cuteDSL MLA kernel for packed q heads and compact input"), the CuTe DSL
FP8-KV MLA decode kernel (mla_decode_fp8.py, used by trtllm::cute_dsl_mla_decode_fp8_blackwell) is 10–30% slower at
long KV for a small number of heads per rank, whenever the split-KV grid fills the GPU. The output is unchanged
(bit-identical for the same split plan). The BF16 kernel (mla_decode_fp16.py) has no such regression, and the
packed-q target of #19338 (H 96, q 8) does get faster.
Setup
- Blackwell (SM10x) GPU, CUDA graphs.
- Before: bcdc2d5 (its
mla_decode_fp8.pyis byte-identical to 9257aae~1). After: e88d19c (current main). - MLA decode geometry: latent 512 + rope 64, page size 32, FP8 E4M3 KV / Q, BF16 output, 12 heads per rank (TP8 of a
96-head model like the one in [None][feat] Upgrade Blackwell cuteDSL MLA kernel for packed q heads and compact input #19338), q = 1 and q = 8 (MTP verify), KV 128K and 1M tokens, batch 1–16.
Repro 1: kernel only, the module's own benchmark run()
Same arguments on both builds: q = 1 (no folding / packing involved), uniform split_kv so that B * split_kv is about
one wave of 2-CTA clusters, CUTEDSL_BENCH_CUDA_GRAPH=1, skip_ref_check.
fp8_standalone.py
import inspect, os
import cutlass
from tensorrt_llm._torch.cute_dsl_kernels.blackwell.attention.mla import mla_decode_fp8 as k
os.environ.setdefault("CUTEDSL_BENCH_CUDA_GRAPH", "1")
params = inspect.signature(k.run).parameters
for L in (131072, 1048576):
for B, split in [(1, 32), (2, 32), (4, 25), (8, 12), (16, 6)]:
kw = dict(batch_size=B, seq_len_q=1, seq_len_k=L, num_heads=12, latent_dim=512, rope_dim=64,
in_dtype=cutlass.Float8E4M3FN, out_dtype=cutlass.BFloat16, acc_dtype=cutlass.Float32,
lse_dtype=cutlass.Float32, mma_qk_tiler_mn=(128, 128), mma_pv_tiler_mn=(128, 256),
split_kv=split, is_persistent=False, is_var_seq=False, is_var_split_kv=False, page_size=32,
softmax_scale=0.0722, output_scale=1.0, skip_correction_threshold=0.0, tolerance=1e-2,
warmup_iterations=5, iterations=30, skip_ref_check=True, use_cold_l2=False)
# arguments added after #19338
for n, d in dict(is_var_q=False, query_lens=None, emit_softmax_stats=False, kv_bounds=None).items():
if n in params:
kw[n] = d
print(f"B{B} L{L} split{split}: {k.run(**kw):.1f} us")H 12, q 1 (mean of 2 runs, µs):
| KV | B (split) | before | after | |
|---|---|---|---|---|
| 128K | 1 (32) | 41.2 | 39.5 | −4% |
| 128K | 2 (32) | 43.7 | 43.8 | 0% |
| 128K | 4 (25) | 54.2 | 59.1 | +9% |
| 128K | 8 (12) | 86.4 | 97.0 | +12% |
| 128K | 16 (6) | 151.7 | 170.8 | +13% |
| 1M | 1 (32) | 199.1 | 200.8 | +1% |
| 1M | 2 (32) | 201.4 | 204.9 | +2% |
| 1M | 4 (25) | 286.5 | 334.7 | +17% |
| 1M | 8 (12) | 555.8 | 591.7 | +6% |
| 1M | 16 (6) | 1095.5 | 1342.8 | +23% |
With 32–64 clusters (B 1–2, split 32) the two kernels are on par; with about one full wave the new kernel is slower.
Repro 2: same runner, same plan, only the kernel swapped
On current main, through CuteDSLNVMlaDecodeBlackwellRunner with a per-request split plan (block_split_kvs from
cute_dsl_mla's planner, split cap _MAX_SPLIT_KV raised to 128 so that one request can fill the GPU), we replaced
only the FP8 split-KV kernel module by the 9257aae~1 version (adapted only to read softmax_scale / output_scale
from the device tensors and to clamp the plan to the grid, as the new kernel does). Our build carries a few small local
changes on both sides (identical in both columns); the Nsight Compute table below is kernel-only. For the same plan
the output o is bit-identical between the two kernels; op time (µs, CUDA graph):
| H 12, q 8 | plan (splits / request) | new kernel (main) | previous kernel |
|---|---|---|---|
| B1 128K | 94 | 32.0 | 32.6 |
| B2 128K | 49 | 50.9 | 44.9 |
| B4 128K | 25 | 84.6 | 63.6 |
| B8 128K | 12 | 146.1 | 106.7 |
| B16 128K | 6 | 250.0 | 192.4 |
| B1 1M | 100 | 127.6 | 99.1 |
| B2 1M | 50 | 241.0 | 176.4 |
| B4 1M | 25 | 461.8 | 341.2 |
| B8 1M | 12 | 881.5 | 735.0 |
| B16 1M | 6 | 1755.2 | 1461.3 |
q 1 gives the same picture (B1 1M 124.4 vs 93.3, B4 128K 76.5 vs 55.8, B16 128K 241.4 vs 165.3, B16 1M 1741.5 vs
1277.9). For comparison, H 96 q 8 (the #19338 target, 6 packed tiles instead of 8) is faster with the new kernel
(B1 1M 629 → 500 µs, B4 128K 349 → 267 µs), and the BF16 kernel is unchanged (B1 1M 193 vs 192 µs).
Analysis (Nsight Compute, split-KV kernel only, H 12 q 8, same plan)
| B1 1M before | B1 1M after | B4 128K before | B4 128K after | |
|---|---|---|---|---|
| duration (µs) | 95.9 | 125.8 | 56.0 | 78.9 |
| threads / CTA | 384 | 512 | 384 | 512 |
| registers / thread | 168 | 128 | 168 | 128 |
| tensor pipe active (% of active cycles) | 76.3 | 52.5 | 69.9 | 44.0 |
| SM throughput (%) | 70.7 | 51.5 | 62.8 | 42.8 |
| DRAM read (MB) | 607 | 850 | 304 | 356 |
| L2 hit / miss sectors (M) | 21.3 / 19.2 | 12.8 / 27.1 | ||
| TMA load bytes (MB) | 1148 | 1148 | ||
| instructions (M) | 28.0 | 31.8 | 15.0 | 20.0 |
| barrier stall samples | 918 | 22,700 | 661 | 11,631 |
| long-scoreboard stall samples | 52k | 74k |
- The new compute pattern (16 warps: a second softmax warpgroup synchronized with the first through named barriers,
QK and PV MMA in separate warps, registers 192/256/48 → 152/152/56) shows ~25× more barrier stall samples and the
tensor pipe busy drops from 76% to 52%. - The kernel issues the same TMA bytes, but DRAM reads grow by 40% (607 → 850 MB for 604 MB of KV) and L2 hits drop
from 21.3M to 12.8M sectors. In MLA the V operand re-reads the latent part of the KV pages already loaded for K; it
looks like the previous pipeline still finds those lines in L2 while the new one more often does not when the whole
GPU is streaming KV, which is consistent with the gap disappearing at 32–64 clusters (Repro 1, B 1–2). We have not
verified this beyond the counters above. - At H 12 the q 8 rows (96) fit one 128-row tile in both kernels (folding before, packing after), so the packed-q
change cannot reduce work for this shape; only the new compute pattern is exercised.
Questions
- Is the slowdown expected for few heads per rank (H·q ≤ 128) at long KV?
- Would it be acceptable to keep the previous FP8 pipeline for shapes where packing gives no fewer M tiles than folding
(e.g. H 12 q 1 / q 8), or is there a tuning of the new pipeline (stages, warpgroup ordering, K/V load order for L2
reuse of the V re-read) that recovers it?
We are happy to run more configurations or profiles.
Activity
lucifer1004 commentedon Oct 19, 2026
Follow-up with a workaround we validated and the trade-off it carries, to help decide how to fix this upstream.
Where the time goes. #19338 moved the FP8 kernel from 12 warps / 384 threads per CTA to 16 warps / 512 threads (a second softmax warpgroup alternating with the first via named barriers; MMA split into separate QK and PV warps) and changed the register split for softmax / correction / other from 192 / 256 / 48 to 152 / 152 / 56. Tiles (128x128 QK, 128x256 PV), pipeline stages and dequant placement are unchanged. On the split-KV kernel with the same plan (B1 1M / B4 128K): duration 96 → 126 / 56 → 79 µs, tensor pipe 76 → 52% / 70 → 44%, barrier-stall samples 918 → 22.7k / 661 → 11.6k, DRAM read 607 → 850 MB (+40%; KV itself is 604 MB) with identical TMA bytes, L2 hit sectors 21.3M → 12.8M. Our reading (not separately verified): V re-reads the latent that K already loaded; the old pipeline still hits it in L2, the new one evicts more of it once the whole GPU is busy. With H 96 / q 8 (the #19338 target) the new kernel is faster (6 vs 8 M tiles; B1 1M 629 → 500 µs), but with H 12 both kernels use one M tile, so packing does not help.
Workaround we tried: keep the pre-#19338 kernel as a dedicated path, dispatched on static shape. The old kernel is kept as a separate module and selected by the runner when KV is FP8 and the old kernel needs no more M tiles per request than the packed kernel (otherwise, and whenever
can_implementrejects, fall back to the current kernel). It reuses the main tactic, workspace, device-side scale (element 0, no host read, CUDA-graph safe) and the per-request split plan (clamped to [1, split_kv]), plus askip_empty_splitsearly-exit for CTAs whose split slot is empty. Results on the same setup as in the description (single GPU, CUDA graphs), H 12, page 32, q 8, FP8, µs, op / op + fused v_b-gate merge:Output is bit-identical to the pre-#19338 kernel on all 22 q8 cases with the same effective split count (FP8 and BF16) and on 6/6 q1 cases, and bit-identical to the main kernel under the same plan. q 1 (draft shape) behaves the same way (e.g. B4 128K 55.8 vs main 77 µs). H 96 / q 8 falls back to the main kernel and is unchanged.
The catch: skewed batches. The dedicated path uses a "one wave of clusters / tiles per request" grid, like the old code, so a skewed batch (one 1M request plus short ones) gets the old 257 µs instead of main's 140 µs. Using the main tactic's split count plus
skip_empty_splitsinstead gives 116 µs on the skewed batch (fastest of the three) but adds empty CTAs on uniform batches (B4 128K +8%, B16 32K +13%, B8 128K +7%). We cannot pick per batch on the host because KV lengths are not known there under CUDA graphs.Suggested directions, in our order of preference:
Separately, at short KV the main planner's minimum of 4 tiles per split gives B1 2K 5 splits / 19.2 µs vs 17 splits / 16.7 µs before; a minimum of 1 tile restores 16.6 µs without changing long KV (FP8 only measured).
saiyambharara commentedon Oct 22, 2026
@lucifer1004 I'd like to take this one if you haven't started, and I'd test one thing before restructuring anything.
Comparing
mla_decode_fp8.pywith 9257aae~1, the stage counts and the separate K and V loader warps are unchanged. What changed is that the old single MMA warp issued QK(i) then PV(i-1) in program order, while the QK warp (8) and PV warp (11) now run independent loops that only meet through the pipelines. So K can run further ahead of V, and the V re-read of the latent can land after L2 has moved on. That would fit your unchanged TMA bytes with +40% DRAM reads, and the gap closing at 32 to 64 clusters. It wouldn't explain the barrier stalls, so I'd expect only part of the gap.Proposal: first bound how far the QK warp can run ahead of the PV warp (back to the old one-tile skew) and re-measure L2 hit sectors and duration on your B4 128K and B1 1M shapes. That tests the L2 reading you flagged as unverified. Depending on what it recovers, the 12-warp variant from your option 1 comes next, and your option 2 PR can still land first.
I don't have SM10x hardware. If a shared B200 or a CI runner is possible, or you could run my branch on your side, that would help a lot, but I'm happy to rent otherwise. For context, I've landed a tcgen05 backend fix in Triton (triton-lang/triton#11973) and CubeCL kernels in burn (tracel-ai/burn#6045), and I have #20838 open here. First numbers within a week of getting a GPU, and a draft PR once the output matches the pre-#19338 kernel bit for bit.
Does that match what you saw, and is anyone else on this?
saiyambharara commentedon Oct 26, 2026
@pengbowang-nv quick check before I put time into this: is the CuTe DSL MLA decode kernel staying the long-term path, or is decode moving to PrimTS? I'd rather not rework a kernel that's being replaced.
pengbowang-nv commentedon Oct 26, 2026
Hi @saiyambharara we'll be gradually enabling PrimTS versions of kernel, but we would like the perf to be verified first and make sure prim-ts enablement will bring better or equal perf. Meanwhile I haven't been able to repro this perf gap on my side. Will do more test before I can reach a conclusion. Thanks!