[Perf] Deepseek V41 Attention MegaKernel

[zyongye]

Purpose

Integrate FlashMLA PR #221's fused norm + RoPE + sparse-attention + inverse-RoPE + FP8-cast kernels into DeepSeek V4.1 attention and add the nvfp4_ds_mla KV cache (V4.1 fp8 528 B SWA + V4.1 fp4 288 B compressed)..

• flashmla.cmake pins vllm-project/FlashMLA@c112cc1 (the fork's sync of Add kernels for DeepSeek v4.1 deepseek-ai/FlashMLA#221 behind the stable torch.ops._flashmla_C ABI): same build structure as before, source list mirroring the fork's setup.py, both Python interface files vendored, include/cccl added for CUDA 13 host compiles. vllm/v1/attention/ops/flashmla.py adds thin wrappers for the fused prefill/decode ops.
• New DeepseekV4FlashMLAFusedAttention (vllm/models/deepseek_v4_1/nvidia/flashmla_fused.py), selected by --attention-config '{"dsv4_fused_attention": true}'. It permutes wq_b rows and wo_a columns at load into the kernel's chunk-interleaved layouts, feeds the kernel's FP8 output straight into the DeepGEMM wo_a einsum, and keeps the split-KV path as a fallback below dsv4_fused_decode_min_tokens. The base layer gains _alloc_attn_out / _prepare_q_and_insert_kv / _finish_o_proj hooks so platform subclasses can change the eager-region output shape. Every decode token carries its own metadata row (positions, SWA indices, compressed-cache indices), so the kernel runs on the flattened token batch at any concurrency.
• nvfp4_ds_mla KV layout table, SlidingWindowMLASpec plumbing, Triton insert/gather kernels for both V4.1 formats, SM100 gating, DSpark context insert.
• Bug fix surfaced by the evals: combine_topk_swa_indices only wrote the valid prefix of each row and relied on the caller's -1 pre-fill, but both prefill paths reuse an out= workspace. flash_mla_sparse_fwd caps reads at topk_length; the upstream fused prefill reads the whole row and propagates NaN from stale indices into uninitialized gathered rows, which turned 72 % of gsm8k generations into 4096 tokens of NaN-argmax output. The kernel now writes the sentinels itself.
• DeepseekV41ForConditionalGeneration joins the breakable-CUDA-graph default list; positions_int32 added to the sparse SWA metadata.

Performance

All numbers: GB200, TP4, DeepSeek V4.1 (ckpt20260903), fp8_ds_mla unless stated, same branch for fused and unfused.

Kernel level (benchmarks/kernels/benchmark_dsv41_fused_attention.py, h_q = 64, topk 128 SWA + 512 compressed, µs per call). The fused kernel replaces the split-KV attention kernel plus the separate inverse-RoPE + FP8 quant kernels; the unfused columns time the attention kernel alone and the whole chain.

decode s_q	fused	split-KV attention only	split-KV attention + inv-RoPE + quant
1 (CUDA graph)	22.7	22.7	43.1
64 (CUDA graph)	24.7	24.7	59.7
256 (eager)	50.1	48.2	400
1024 (eager)	155	161	469

prefill s_q (topk 640, eager)	fused	sparse_fwd + inv-RoPE + quant
184	30	375
2123	163	790

Eager numbers include ~50 µs of stable-ABI boxed dispatch per op call (both fused and the fork's sparse_decode_fwd pay it); under CUDA graphs the kernels alone are 22.7 µs at s_q = 1.

Serving, bs1 (vllm bench serve, 20 random prompts, concurrency 1, no speculative decoding, --max-model-len 40960), TPOT mean / p50 ms:

variant	8k in / 1k out	32k in / 1k out	TTFT p50 8k / 32k
unfused	6.20 / 6.14	6.24 / 6.17	230 / 591 ms
fused	6.46 / 6.45	6.49 / 6.49	227 / 552 ms

At bs1 the fused layer is ~4 % slower per token even though its kernel is faster: the attention runs in the eager region of the breakable CUDA graph, and the fused path adds the Q layout kernel, the op call with fresh output allocations and the DeepGEMM wo_a einsum to that region, whereas the base layer's o-projection is inside the captured graph. Follow-up: persistent out_fp8 / out_sf buffers and the einsum in _finish_o_proj. dsv4_fused_attention therefore defaults to off.

Serving, high concurrency (gsm8k, 1319 five-shot prompts sent at concurrency 1024, ~1000 decode tokens per step, 4096-token cap): unfused 67 s, fused 73 s, fused + nvfp4_ds_mla 58 s wall time for the whole eval.

KV capacity: nvfp4_ds_mla raises the KV token budget at 8k context from 2.95 M to 4.08 M tokens (1.38x).

Accuracy

gsm8k, 1319 problems, 5-shot, greedy:

variant	flexible	strict
unfused (dsv4_fused_attention: false)	0.889	0.889
fused, fp8_ds_mla	0.900	0.899
fused, nvfp4_ds_mla	0.905	0.905

Before the sentinel fix the fused run scored 0.229 (949 / 1319 generations hit the 4096-token cap with empty text). The long-context accuracy of the fp4 compressed path (gpqa, 32k task) is not yet measured.

Test Plan

.venv/bin/python -m pytest tests/v1/attention/test_deepseek_v4_swa_visible.py tests/kernels/attention/test_flashmla_sparse.py \
  tests/kernels/attention/test_flashmla.py tests/kernels/test_compressor_kv_cache.py \
  tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py tests/v1/attention/test_dspark_noncausal_sparse_mla.py -k "not builder"
pre-commit run --files <changed files>   # ruff, mypy, typos all pass
vigil -c recipe/dsv41/sra_vigil_tp4_gsm8k_fused.yaml            # + _fused_nvfp4 / _worktree_unfused variants
vigil -c recipe/dsv41/sra_vigil_tp4_8k1k_32k1k_bs1_fused.yaml   # + _worktree_unfused variant
.venv/bin/python -m benchmarks.kernels.benchmark_dsv41_fused_attention --decode-s-q 1 64 256 1024 [--cudagraph]

124 existing SWA / FlashMLA / compressor / DSpark tests pass on the new FlashMLA build. Development-time kernel tests (layout permutations, Q padding kernel, V4.1 fp8/fp4 insert + gather against a torch port of FlashMLA's reference quantizer, fused vs split-KV decode/prefill equivalence, combined-index sentinels) were run on GB200 and passed but are not part of this PR.

AI assistance: this PR was developed with Claude Code (Claude Fable 5.1); every change was reviewed and the tests and evals above were run by the submitter.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[mergify]

Documentation preview: https://vllm--56344.org.readthedocs.build/en/56344/

[Harry-Chen]

Kernel level (benchmarks/kernels/benchmark_dsv41_fused_attention.py, h_q = 64, topk 128 SWA + 512 compressed, µs per call). The fused kernel replaces the split-KV attention kernel plus the separate inverse-RoPE + FP8 quant kernels; the unfused columns time the attention kernel alone and the whole chain.

decode s_q fused split-KV attention only split-KV attention + inv-RoPE + quant
1 (CUDA graph) 22.7 22.7 43.1
64 (CUDA graph) 24.7 24.7 59.7
256 (eager) 50.1 48.2 400
1024 (eager) 155 161 469
prefill s_q (topk 640, eager) fused sparse_fwd + inv-RoPE + quant
184 30 375
2123 163 790
Eager numbers include ~50 µs of stable-ABI boxed dispatch per op call (both fused and the fork's sparse_decode_fwd pay it); under CUDA graphs the kernels alone are 22.7 µs at s_q = 1.

Stable ABI dispatch should not be so expensive: #26946 (comment). And if it is 50us per call, then 50.1us does not make sense. Any misunderstanding here?

[zyongye]

Kernel level (benchmarks/kernels/benchmark_dsv41_fused_attention.py, h_q = 64, topk 128 SWA + 512 compressed, µs per call). The fused kernel replaces the split-KV attention kernel plus the separate inverse-RoPE + FP8 quant kernels; the unfused columns time the attention kernel alone and the whole chain.
decode s_q fused split-KV attention only split-KV attention + inv-RoPE + quant
1 (CUDA graph) 22.7 22.7 43.1
64 (CUDA graph) 24.7 24.7 59.7
256 (eager) 50.1 48.2 400
1024 (eager) 155 161 469
prefill s_q (topk 640, eager) fused sparse_fwd + inv-RoPE + quant
184 30 375
2123 163 790
Eager numbers include ~50 µs of stable-ABI boxed dispatch per op call (both fused and the fork's sparse_decode_fwd pay it); under CUDA graphs the kernels alone are 22.7 µs at s_q = 1.

Stable ABI dispatch should not be so expensive: #26946 (comment). And if it is 50us per call, then 50.1us does not make sense. Any misunderstanding here?

Yea I was originally use the upstream flashmla branch for initial integration and later switch to our branch after that's done. So will update the number here.

[Fangzhou-Ai]

Hi @zyongye nice work! Can you also add mega kernel support to ROCm path? cc @ChuanLi1101

[LucasWilkinson]

should we rethink how we do model specific attention configs? this feels like it could get out of hand very fast

[zyongye]

yea lemme work on this.