[roborluo]

Overview

Makes DeepSeek-V4.1's per-layer decode attention step 1.6-2.15x faster for decode steps of up to 16 tokens at TP2 / TP4, by adding a swap-AB CuTe DSL decode attention kernel to the FlashMLA mega-attention layer (vllm-project#56935) on SM10x. Steps of up to 16
tokens with 16 or 32 live heads per rank (TP4 / TP2 of the 64-head model) now take it instead of the mega kernel, which pads the heads to 64.
Part of vllm-project#57448; it reads the interleaved V4.1 KV records from vllm-project#61371. Related work that this PR does not duplicate (see Details): vllm-project#59418, vllm-project#61615, vllm-project#57711, vllm-project#61868.

Claims

• 1.6-2.15x faster decode attention for 1-16 tokens: 1.9-2.15x with 16 live heads (TP4) and 1.6-1.8x with 32 (TP2), for both NVFP4 and fp8 compressed records (per-layer Q pad + KV insert + attention, table below). The mega kernel's FP8 decode maps
  query heads to the MMA's M dimension, which is at least 64, so at TP2 and TP4 half or three quarters of its MMA work goes to padding heads; swap-AB puts the KV slots on M and only the live heads on N.
• Accuracy: the output's error against an fp64 reference is at most 1.1x the mega kernel's, and the output is within 2% of the mega kernel's, over 64 test cases.
• Automatic and contained. There is no new flag, environment variable or config option. Every other decode step, layer and TP size keeps the mega kernel; prefill and other platforms are unchanged. The new kernel is CUDA-graph safe and deterministic,
  and adds no host sync.

Validation

Performance. This PR adds a swap-AB arm to the existing benchmarks/kernels/benchmark_dsv41_mega_attn.py. Each arm's per-step ops (Q pad + KV insert + attention) are timed as a CUDA graph replay, the same harness that produced the mega vs split-KV
table in flash_mla_mega_attn.py.

for heads in 16 32; do
  for rec in 288 528; do
    .venv/bin/python benchmarks/kernels/benchmark_dsv41_mega_attn.py \
        --local-heads $heads --extra-bytes $rec --decode-s-q 1 2 4 8 16 32 64
  done
done

Speedup = mega step time / swap-AB step time (median of two runs, which differ by at most 9% per cell):

Live heads	Compressed record	1 token	2	4	8	16	32	64
16 (TP4)	NVFP4 (288 B)	2.15x	2.00x	2.06x	1.94x	2.04x	1.25x	0.94x
16 (TP4)	fp8 (528 B)	2.13x	1.96x	2.02x	1.90x	1.99x	1.22x	0.91x
32 (TP2)	fp8 (528 B)	1.76x	1.66x	1.71x	1.62x	1.71x	0.98x	0.73x

At 32 tokens swap-AB still wins with 16 heads (1.22-1.25x) but only breaks even with 32 (0.98-1.02x), and at 64 it loses with both, so _SWAPAB_MAX_TOKENS = 16 is the largest step size where it wins for both head counts. As in the existing harness,
every replay reuses the same buffers, so the KV reads hit a warm L2 for both kernels; timed alone with the L2 flushed before every call, the attention kernel is 1.6-2.3x faster than the mega kernel over 48 shapes with 1-16 tokens (geomean 1.94x).

Correctness: unit tests. The kernel tests skip unless the GPU is SM10x and the CuTe DSL is built with CUDA 13.2+.

.venv/bin/python -m pytest -v tests/kernels/test_dsv41_swapab_decode.py
.venv/bin/python -m pytest -v "tests/kernels/test_compressor_kv_cache.py::test_v41_cache_writers_match_flashmla_decode"

Test (cases)	Checks	Result
test_matches_reference (64)	16/32 heads x 1/2/5/16 tokens x NVFP4/fp8 compressed records x 64/128-slot pages x partial/full lengths. Output is finite; Frobenius and max-abs error vs an fp64 reference are <= 1.1x the mega kernel's; within 2% of the
mega output	64 passed
test_untouched_padding (2)	Padding heads' output groups keep their bytes	2 passed
test_cuda_graph_and_determinism (12)	Repeated eager calls and a CUDA graph replay agree bit for bit, also after the inputs change in place	12 passed
test_dispatch_rule (9, CPU)	Which kernel each dispatch condition selects	9 passed
test_warmup_key_matches_launch (4, CPU)	The key bind_kv_cache registers is the key a decode step launches	4 passed
test_v41_cache_writers_match_flashmla_decode (3, existing)	The V4.1 cache writers and FlashMLA agree on the interleaved records the kernel reads	3 passed

Details

Motivation. Mega attention is the default DeepSeek V4.1 decode path on SM10x. Its FP8 decode kernel takes 64 or 128 query heads, so at TP2 and TP4 half or three quarters of its MMA work goes to padding heads. The mega vs split-KV note in
flash_mla_mega_attn.py already records that this padded-head cost dominates small decode steps. Small steps are what low-concurrency decode and spec-decode verify and draft steps run.

Why a separate kernel, not a FlashMLA change. FlashMLA's decode maps query heads to the MMA's M dimension, which is at least 64. Serving 16 or 32 heads without padding needs the opposite mapping, with KV slots on M and heads on N. That is a
different kernel, not a smaller tile. In FlashMLA it would be CUDA C++ shipped through a vllm-project/FlashMLA release and a vLLM bump, as vllm-project#61371 was. The CuTe DSL kernel instead compiles from Python at JIT warmup, needs no build change, and lives next
to its only caller in the DSv4.1 model package. It is adapted from SGLang's swap-AB DSv4 decode attention (Apache-2.0, attribution in the file header).

Kernel design

• MMA shape. SwapABDecodeKernel (vllm/models/deepseek_v41/nvidia/ops/swapab_decode_attn.py) puts the KV slots on the tcgen05 MMA's M side (128 rows) and the live heads on N (16 or 32).
• One CTA per (token, KV tile). Tiles are 64 slots of the SWA list or 128 slots of the compressed list. Each CTA:
  • gathers its slots' V4.1 records in place (fp8 528 B, NVFP4 288 B, each token's scales after its data);
  • dequantizes them exactly to bf16 in SMEM;
  • RoPEs Q in fp32 from the mega kernel's chunk-interleaved, head-padded layout;
  • computes S = KV Q^T and O = KV^T P, with P split into bf16 hi and lo parts.
• One cluster per token. A token's tiles form one thread-block cluster: 6 CTAs for V4.1-Flash's 128-slot window plus 512 compressed slots. The cluster merges the partials and the attention sink over DSMEM in a fixed peer order. It then applies the
  inverse RoPE and writes MXFP8 with the mega kernel's rounding rule into the o-fused buffers wo_a reads, for the live heads' groups only.
• One launch. The kernel uses PDL and has no workspace, second kernel, allocation or host sync.
• JIT warmup. The compile key is (live heads, slot-list widths, compressed record size, SWA and compressed page sizes); the token count is the grid's y dimension. The compressed page is final only once the KV cache is allocated. So kv-source layers
  register the key in bind_kv_cache, which both model runners call while the JIT warmup registry is active, and kernel_warmup compiles it before CUDA graph capture. V4.1-Flash needs two keys, one per compressed page size.

Dispatch. DeepseekV4MegaAttnAttention._forward_decode_mega takes the swap-AB kernel only when all of these hold:

• the GPU is SM10x;
• the layer has a compressed cache;
• the layer has 16 or 32 live heads;
• a token's tiles fit one portable cluster: ceil(window / 64) + ceil(index_topk / 128) <= 8;
• the CuTe DSL imports and was built with CUDA 13.2+, which the direct fp8 / fp4 to bf16x2 conversions need (PTX ISA 9.2);
• VLLM_BATCH_INVARIANT is off;
• the step has at most 16 decode tokens (_SWAPAB_MAX_TOKENS). Under CUDA graphs this is the padded count, so each captured graph holds exactly one of the two kernels;
• the step's SWA list is exactly the window. DSpark's non-causal draft steps widen it, so they keep the mega kernel.

The first six conditions are fixed when the layer is built. The last two are host-side shape checks on each step. If any condition fails, the step runs the mega kernel exactly as before.
Complexity. The change adds about 1,260 lines of kernel code (about 400 of them are 23 small inline-PTX wrappers) and about 100 lines in the layer. Nothing changes in the model runner, the scheduler, FlashMLA or shared attention code. At runtime the
kernel module is imported only by SM10x layers with a compressed cache, 16 or 32 heads and an importable CuTe DSL. Most of the maintenance cost is the inline PTX and the SMEM / TMEM layout math. Both are documented next to the code, and the tests pin
them to an fp64 reference and to the mega kernel.

Tradeoffs and limitations.

• Cost grows with tokens. The kernel runs one cluster of up to 8 CTAs per token, and its 150+ KiB of SMEM means one CTA per SM. Hence the 16-token cutoff, which is the same for TP2 / TP4 and both record formats: in the table above swap-AB is at
  least 1.6x faster at every step size up to 16 tokens.
• Not bit-identical. The fp32 sums are ordered differently from the mega kernel's, so the output depends on which kernel the step size selects. Batch-invariant runs therefore keep the mega kernel at every size.
• Unused Q padding. The fused KV insert still zero-pads Q to 64 heads, which swap-AB does not read. This keeps one Q contract for both kernels.
• Toolchain. The kernel needs a CuTe DSL built with CUDA 13.2+. With an older build the layer keeps the mega kernel and the kernel tests skip.
• Startup. V4.1-Flash adds two JIT compiles per worker, at warmup.

Compatibility.

• Scope. Only DeepseekV4MegaAttnAttention changes. The FlashMLA sparse layer, FlashInfer sparse MLA, and the SM90, SM120 and ROCm paths are untouched. TP1 (64 heads) and TP8 (8 heads) keep the mega kernel.
• Model runners. Both V1 and V2 bind the KV cache under the JIT warmup registry, so both register the key.
• CUDA graphs, batch invariance, spec decode. See Dispatch above.
• Unchanged. Prefill, the KV cache layout and checkpoints.

CI. The new tests run in "(B200) Miscellaneous Kernels". The kernel tests need SM10x, the FlashMLA mega ops and a CuTe DSL built with CUDA 13.2+, and skip otherwise; with vLLM's pinned nvidia-cutlass-dsl[cu13]==4.8.0 they run. The optional "(B200)
LM Eval Spec Decode" step runs the V4.1-Flash NVFP4 DSpark TP4 GSM8K config. It runs at full concurrency, where few steps are small enough for swap-AB.

Not a duplicate. On 2026-10-18 I searched GitHub (open PRs plus PRs since 2026-08-15) for "swap-AB", "swapab", "small-head decode", "heads padded 64 decode", "mega attention DSv4.1" and "V4.1 decode attention". No PR adds an SM10x small-head decode
kernel for the mega-attention path. Related work:

• [Perf][DSv4.1] Small-head sparse decode on SM90 vllm-project/vllm#59418 (open, small-head sparse decode on SM90) is a Triton kernel for Hopper. It targets 8 or 16 heads in the plain FlashMLA sparse layer, which reads the 584 B fp8_ds_mla records used off SM10x. This PR targets SM10x tcgen05 and 16 or 32 heads
  in the mega-attention layer, which reads the V4.1 records and fuses the inverse RoPE and MXFP8 output.
• [Perf] Enable small-head FlashMLA sparse prefill for DS V4.1 on SM90 vllm-project/vllm#61615 (open) is small-head FlashMLA sparse prefill on SM90. This PR changes decode only, on SM10x.
• [Perf][DSv4.1] Opt-in split-KV decode for the mega attention backend vllm-project/vllm#57711 (closed by its author on 2026-09-25) was an opt-in split-KV mode inside the mega decode kernel, behind the VLLM_DSV41_MEGA_SPLITKV env flags and blocked on Add split-KV to the fused mega decode kernel vllm-project/FlashMLA#25. That kernel still pads the heads to 64. This PR leaves
  FlashMLA alone, removes the padded-head work and adds no flags.
• [Perf][DSv4.1] DeepSeek-V4.1-Flash specialized decode on GB300 vllm-project/vllm#61868 (open) is an opt-in GB300 TP4 DSpark executor whose sparse-MLA decode uses FlashInfer's trtllm_batch_decode_sparse_mla_dsv4. That is a separate opt-in path; this PR changes the default mega-attention path.

AI assistance. This PR was prepared with AI assistance (Claude Code, Claude Opus 5.5). It drafted the port of the kernel onto main, the tests, the benchmark arm, the /pr-checklist pass and this description. The submitter reviewed every changed
line and ran the validation above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.