Skip to content

[Perf][Kernel][DSv4.1] Swap-AB CuTe DSL decode attention for small TP2/TP4 steps - #62633

Merged
wzhao18 merged 9 commits into
vllm-project:mainfrom
roborluo:dsv41-swapab-decode
Nov 3, 2026
Merged

wzhao18 merged 9 commits into
vllm-project:mainfrom
roborluo:dsv41-swapab-decode

Conversation

@roborluo

@roborluo roborluo commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Overview

Makes DeepSeek-V4.1's per-layer decode attention step 1.4-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 (#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 #57448; it reads the interleaved V4.1 KV records from #61371. Related work that this PR does not duplicate (see Details): #59418, #61615, #57711, #61868.

Claims

  • 1.4-2.15x faster decode attention for 1-16 tokens: 1.8-2.15x with 16 live heads (TP4) and 1.4-1.8x with 32 (TP2), for both NVFP4 and fp8 compressed records (per-layer Q pad + KV insert + attention, table below; the lower ends are B300). End to end at TP2 this is 4.5-7.6% lower median TPOT at concurrency 1-16 (see the A/B comment). 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. On B300 (nightly 76c6f3f) the same harness gives 1.83-2.0x with 16 heads and 1.41-1.64x with 32 heads at 1-16 tokens, for both record formats. 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/models/deepseek_v41/nvidia/test_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 the layer registers at construction is the key a decode step launches, on a compressed cache shaped as the model runner binds it 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 #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. Every field is fixed when the layer is built: the backend has one kernel block size (128), so a compressed page always holds
    128 // compress_ratio slots, whatever the KV cache manager's block size. The layer registers the key in __init__, next to its other JIT warmups, 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 (at least 1.4x on B300).
  • 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 build the model 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 the optional (nightly) 4xB200 step, renamed "(4xB200) DeepSeek-V4.1 NVIDIA Kernels" (models_basic.yaml, key dsv41-nvidia-kernels-4xb200), next to the specialized-decode tests. 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:

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.

roborluo and others added 2 commits October 18, 2026 23:23
…2/TP4 steps

FlashMLA's mega-attention decode kernel takes 64 or 128 query heads. At
TP2 and TP4 a rank holds 32 and 16 live heads of the 64-head model, so
on small decode, verify and draft steps half or three quarters of the
kernel's MMA work goes to padding heads.

Add SwapABDecodeKernel, a CuTe DSL kernel adapted from SGLang's swap-AB
DSv4 decode attention (Apache-2.0), in
vllm/models/deepseek_v41/nvidia/ops/swapab_decode_attn.py. It puts the
KV slots on the tcgen05 MMA's M side and the live heads on N. One CTA
takes a 64-slot tile of a token's SWA list or a 128-slot tile of its
compressed list: it gathers and dequantizes the V4.1 records (fp8 528 B,
NVFP4 288 B) into SMEM, RoPEs Q and computes S = KV Q^T and O = KV^T P,
with P split into bf16 hi and lo parts. A token's tiles form one
thread-block cluster that merges the partials and the attention sink
over DSMEM in a fixed order, applies the inverse RoPE and writes the
MXFP8 output the mega kernel writes, for the live heads' groups only.
It is a single PDL launch with no workspace and no host sync, CUDA-graph
safe and deterministic.

_forward_decode_mega takes the kernel automatically, with no new flag,
when all of these hold: the GPU is SM10x; the layer has a compressed
cache and 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 is built with CUDA 13.2+, as the direct fp8 / fp4 -> bf16x2
conversions need PTX ISA 9.2; VLLM_BATCH_INVARIANT is off; the step has
at most 16 decode tokens; and its SWA list is the window, which DSpark's
non-causal draft steps widen. Every other step, layer and TP size keeps
the mega kernel; prefill and other platforms are unchanged.

The compile key holds the compressed page size, which is final only
once the KV cache is allocated. So kv-source layers register it in
bind_kv_cache, which runs while the runner's JIT warmup registry is
active, and kernel_warmup compiles it before CUDA graph capture.

The output is not bit-identical to the mega kernel's, whose fp32 sums
are ordered differently. Its error against an fp64 reference is held to
at most 1.1x the mega kernel's.

tests/kernels/test_dsv41_swapab_decode.py checks the kernel against an
fp64 reference and the mega kernel on caches written by V4.1's insert
kernels; checks that padding heads' output is untouched and that eager
calls and a CUDA graph replay agree bit for bit; and checks the
dispatch rule and that the warmup key is the launched key.
benchmarks/kernels/benchmark_dsv41_mega_attn.py gains a swap-AB arm.

Signed-off-by: bofengl <bofengl@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Resolve the import conflict with vllm-project#60118 in flash_mla_mega_attn.py: keep
main's `cdiv` import next to this PR's `has_cutedsl` import.

Signed-off-by: bofengl <bofengl@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added deepseek Related to DeepSeek models performance Performance-related issues DSv4.1 Related to DeepSeek-V4.1 models labels Oct 20, 2026
@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

@LopezCastroRoberto LopezCastroRoberto left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you share some e2e numbers (TPOT at TP2/TP4) and accuracy evals?

@roborluo

Copy link
Copy Markdown
Contributor Author

Will do. Planned data:

  • TPOT: A/B (base vs this PR, identical flags) at TP2 (fp8 KV) and TP4 (NVFP4 and fp8 KV), at concurrencies where decode steps stay inside the 16-token gate (c1–c16 with spec decode off; with DSpark each step carries concurrency × (drafts + 1) verify tokens, so the cap drops accordingly), plus one point above it where both arms run the mega kernel — that one should be parity by construction. The per-layer 1.6–2.15× won't translate 1:1 to TPOT (attention is one slice of the decode step), so I'll report the measured deltas either way.
  • Accuracy: GSM8K on both arms at TP2 and TP4, with concurrency capped so the swap-AB kernel is actually on the decode path (at full concurrency few steps are ≤16 tokens, so the standard eval barely exercises it). Since the output is intentionally not bit-identical to the mega kernel (different fp32 reduction order), this is the right check on top of the fp64-reference unit tests.

@roborluo

roborluo commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor Author

Correction (11/03): the "TP2 fp8" rows below used NVFP4 compressed records (--kv-cache-dtype fp8 resolves to nvfp4_ds_mla on this layer). fp8_ds_mla numbers, repeats and spread on current main are in #62633 (comment).

TPOT and accuracy, as promised.

Setup. B300 (8-GPU node), vLLM nightly image at commit 00b6c4c; arm A is that commit stock, arm B is this PR's runtime hunks applied on top (they apply cleanly). Per TP config the two arms ran back to back on the same node. Server: vllm serve with --tokenizer-mode deepseek_v41, --attention-config '{"backend":"FLASHMLA_MEGA_ATTN_DSV41","indexer_kv_dtype":"mxfp4","indexer_sparse_logits":true}', --engram-config '{"cpu_offload":true}', --max-model-len 16384 --max-num-batched-tokens 8192 --no-enable-prefix-caching; TP2 with --kv-cache-dtype fp8, TP4 with --kv-cache-dtype nvfp4_ds_mla. Load: vllm bench serve --dataset-name random --random-input-len 1024 --random-output-len 512, no spec decode, so decode steps carry exactly concurrency tokens; c64 is the above-gate parity point.

Median TPOT (ms), stock → PR:

Concurrency TP2 fp8 Δ TP4 NVFP4 Δ
1 4.89 → 4.46 −8.8% 3.91 → 3.71 −5.1%
2 5.47 → 4.98 −9.0% 4.28 → 4.04 −5.6%
4 5.96 → 5.46 −8.4% 4.66 → 4.39 −5.8%
8 7.08 → 6.58 −7.1% 5.35 → 5.13 −4.1%
16 8.96 → 8.56 −4.5% 6.67 → 6.53 −2.1%
64 (above gate) 16.58 → 16.85 +1.6% 12.91 → 13.05 +1.1%

Inside the gate the win shrinks as concurrency grows, as expected (attention is a shrinking share of the step). At c64 both arms run only the mega kernel; the +1.1–1.6% there reads as run noise — two reruns of the PR arm's TP2 c64 on another node gave 16.81/16.83 ms with clean means, and I'm happy to pin it down further if you'd like. E2E TPOT gains are, as expected, much smaller than the 1.6–2.15x per-layer attention-step numbers in the description.

GSM8K (lm-eval local-completions, 5-shot strict-match, all 1319 questions, num_concurrent=8 so decode steps stay inside the gate):

stock this PR
TP2 fp8 0.9310 ± 0.0070 0.9227 ± 0.0074
TP4 NVFP4 0.9287 ± 0.0071 0.9249 ± 0.0073

Both deltas are within one standard error and within this model's run-to-run spread on GSM8K (e.g. the main spread reported on #63018 was 0.9249–0.9356); the engine is not run-deterministic at these settings, and the PR's fp64-reference unit tests bound the kernel-level error.

@@ -206,6 +250,25 @@ def packed_kv_cache_dtype(self) -> CacheDType:
# it is what an unspecific --kv-cache-dtype resolves to here.
return "nvfp4_ds_mla"

def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

don't think this is the best place to do the warmup

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, moved in 0300f89. The key is now registered in __init__ (_register_swapab_warmup), next to the layer's other register_warmup calls, and the bind_kv_cache override is gone.

I had deferred it because I took the compressed page to be known only once the cache was allocated. It is fixed earlier: FlashMLAMegaAttnBackend supports one kernel block size (128), and the compressed cache is always viewed in kernel blocks of 128 // compress_ratio slots (get_num_kernel_states), whatever the manager's block size. So every key field comes from construction-time state.

  • test_warmup_key_matches_launch now checks the construction-time key against the key a decode step launches on a compressed cache shaped by MLAAttentionSpec.get_num_kernel_states, for both record formats and compress ratios 1/2.
  • B300, TP4 NVFP4 server: each rank compiled exactly the two expected keys (extra_page 128 and 64). There were no swap-AB compiles once requests were being served (c1/c4/c16).

roborluo and others added 3 commits November 2, 2026 13:49
Every field of the compile key is fixed when the layer is built: the
compressed cache is read in the mega backend's 128-token kernel blocks,
so its page is 128 // compress_ratio slots whatever the manager block
size is. Registering it next to the layer's other JIT warmups removes
the bind_kv_cache override.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: bofengl <bofengl@nvidia.com>
Signed-off-by: bofengl <bofengl@nvidia.com>
Move tests/kernels/test_dsv41_swapab_decode.py to
tests/models/deepseek_v41/nvidia/test_swapab_decode.py and run it in a
single-GPU B200 step. test_warmup_key_matches_launch now checks the key
registered at construction against the key a decode step launches on a
compressed cache shaped by MLAAttentionSpec.get_num_kernel_states.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: bofengl <bofengl@nvidia.com>
@roborluo
roborluo requested a review from a team November 2, 2026 22:05
@mergify mergify Bot added the ci/build label Nov 2, 2026
…1971

vllm-project#61971 moved the mega attention layer to int64 RoPE positions and
dropped positions_int32, so the swap-AB call merged in from main named
a variable that no longer exists. The kernel now reads int64 positions
like the mega kernel (it only uses them as an int64 row offset), and the
benchmark and tests pass int64 positions.

The CPU tests build the layer with object.__new__; initialize its
nn.Module bookkeeping so monkeypatch can undo instance attributes on
current torch.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: bofengl <bofengl@nvidia.com>
Comment thread .buildkite/test_areas/models_basic.yaml Outdated
- pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_numerics.py
- pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_layers.py

- label: ":nvidia: (B200) DeepSeek-V4.1 Swap-AB Decode"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how about merging this with nvidia: (4xB200) DeepSeek-V4.1 Specialized Decode and make it Deepseek v4.1 nvidia kernels test? What do you think?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, done in 0490667. The swap-AB tests now run in that step, renamed :nvidia: (4xB200) DeepSeek-V4.1 NVIDIA Kernels (key dsv41-nvidia-kernels-4xb200). Its dependencies now also cover vllm/model_executor/warmup/ and the swap-AB test file, and the timeout went from 40 to 50 minutes. One trade-off: the step is optional, so the swap-AB tests now run nightly rather than on every PR that touches these files.

@wzhao18 wzhao18 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm. left a nit suggestion.

@wzhao18 wzhao18 added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 3, 2026
@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

✅ @roborluo, CI is now available for this PR.

  • /ci run starts upstream CI; /amd-ci run starts AMD CI only.
  • Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
  • /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
  • /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
  • /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

@roborluo

roborluo commented Nov 3, 2026

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99430 for commit f90b642e67f0.

ℹ️ This PR is 28 commits behind upstream main. Merged upstream main into this branch as f90b642e67f0 before starting CI. Pull the branch before pushing more changes.

@roborluo roborluo changed the title [Perf][Kernel][DSv4.1] Swap-AB CuTe DSL decode attention for small TP2/TP4 steps [Perf][Kernel][DSv4.1] Swap-AB CuTe DSL decode attention for small TP2/TP4 steps Nov 3, 2026
Fold the single-GPU swap-AB step into the 4xB200 specialized-decode step
and rename it to "DeepSeek-V4.1 NVIDIA Kernels", per review.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: bofengl <bofengl@nvidia.com>
@roborluo

roborluo commented Nov 3, 2026

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99459 for commit fbb376c69222.

ℹ️ This PR is 11 commits behind upstream main. Merged upstream main into this branch as fbb376c69222 before starting CI. Pull the branch before pushing more changes.

@wzhao18
wzhao18 enabled auto-merge (squash) November 3, 2026 15:57
@wzhao18
wzhao18 disabled auto-merge November 3, 2026 15:57
@wzhao18 wzhao18 added ready ONLY add when PR is ready to merge/full CI is needed and removed ready ONLY add when PR is ready to merge/full CI is needed labels Nov 3, 2026
@wzhao18
wzhao18 enabled auto-merge (squash) November 3, 2026 16:27
@wzhao18

wzhao18 commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

@roborluo Not sure why it is not merging. Maybe ask in pr-merge-request channel in vllm to ask for help merging.

@roborluo

roborluo commented Nov 3, 2026

Copy link
Copy Markdown
Contributor Author

Correction to the 10/28 numbers, and a re-run on current main.

  1. The 10/28 "TP2 fp8" rows were NVFP4 records. On this layer --kv-cache-dtype fp8 resolves to the packed nvfp4_ds_mla format; fp8 compressed records need --kv-cache-dtype fp8_ds_mla. So the TP2 TPOT column and the TP2 GSM8K row of that comment measured NVFP4 records. The TP4 column was NVFP4 as labeled.
  2. Merge break, fixed in 176da21. After [Perf][DSv4.1] Bump FlashMLA to pass int64 positions to MegaAttn #61971, the merged branch still passed the removed positions_int32 to the swap-AB kernel. The kernel input, layer call, benchmark and tests now take the int64 positions. Everything below is on that head.

Re-run. B300, vLLM nightly 76c6f3f2 (includes #61971); stock vs stock plus this PR's vllm/ diff; TP2; flags as on 10/28 except --kv-cache-dtype. One job on one node ran the arms back to back: fp8_ds_mla stock → PR, nvfp4_ds_mla stock → PR, then that sequence again. No failed requests.

Median TPOT (ms), stock → PR:

Concurrency fp8_ds_mla rep 1 fp8_ds_mla rep 2 Δ nvfp4_ds_mla rep 1 nvfp4_ds_mla rep 2 Δ
1 4.96 → 4.58 4.95 → 4.58 −7.6% 5.01 → 4.55 5.00 → 4.82 −6.4%
2 5.42 → 5.03 5.41 → 5.05 −6.9% 5.46 → 4.99 5.44 → 5.32 −5.4%
4 5.93 → 5.55 5.92 → 5.54 −6.4% 6.01 → 5.56 6.02 → 5.56 −7.6%
8 7.03 → 6.60 7.07 → 6.61 −6.3% 7.11 → 6.66 7.12 → 6.65 −6.5%
16 8.89 → 8.40 8.87 → 8.55 −4.6% 8.98 → 8.56 8.96 → 8.53 −4.7%

Δ is the change in the summed TPOT of both repeats.

  • Spread. Stock repeats agree within 0.04 ms. PR repeats agree within 0.15 ms, except NVFP4 rep 2 at c1/c2 (4.82/5.32 vs 4.55/4.99). In that window per-request TPOT was uniform (p99 equal to the median), and the same server measured 5.56 ms again at c4, so I've kept it in the average rather than dropping it. An earlier fp8_ds_mla pair on this image gave c1 4.87 → 4.51 and c16 8.89 → 8.41.
  • Same gain for both record formats. This matches the kernel microbenchmark, where swap-AB saves about the same time per layer for 288 B and 528 B records.
  • Dispatch check. A measurement-only logging probe in every server confirmed that each decode step of 1-16 tokens ran swap-AB on the compressed layers in the PR arm, for both formats. Only the pure-SWA layers (compress_ratio == 0) kept the mega kernel, by design.
  • Above the gate. At c64 both arms run the mega kernel. One earlier pair on this image differed by +0.7% (fp8_ds_mla) and +4.1% (NVFP4); this is run-to-run noise at that point, not an effect of the PR.
  • Older image. One earlier fp8_ds_mla pair on the older 00b6c4c5 image, run without the dispatch probe, showed a smaller gain (−1.6% at c1, parity at c16). I haven't reproduced that on the current image, so I'm not drawing on it.

GSM8K, fp8_ds_mla at TP2 (same settings as 10/28: 5-shot strict-match, all 1319 questions, num_concurrent=8 so decode steps stay inside the gate): stock 0.9295 → PR 0.9325.

Kernel level on B300. The description's benchmark gives 1.83-2.0x with 16 heads and 1.41-1.64x with 32 heads at 1-16 tokens, for both record formats. That is below the 1.6-1.8x the description claimed for TP2, so I've updated the description to give the B300 range.

@wzhao18
wzhao18 disabled auto-merge November 3, 2026 18:24
@wzhao18
wzhao18 enabled auto-merge (squash) November 3, 2026 18:24

@khluu khluu left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

approve for buildkite changes

@wzhao18
wzhao18 merged commit 4b01cba into vllm-project:main Nov 3, 2026
133 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci/build deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models performance Performance-related issues ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants