[okorzh-amd]

Overview

With --monokernel on gfx950 (MI355X), DeepSeek-V4.1 decode batches of up to 32 rows now build their SWA metadata with swa_decode_meta's single Triton launch (#60731). Before, they took vLLM's 8 launches per SWA group. #60731 enabled the kernel only on gfx942, but the kernel has no arch-specific code. This PR converges with #63994: it takes #63994's gate, swa_decode_meta_enabled() (on_gfx942() or on_gfx950()), and adds the gate test, the MI355 CI step and the GPU validation below.

Claims

• 49 fewer launches per decode step at c1: 114 to 65. That is 7 SWA groups (4 target, 3 DSpark draft), each going from 8 launches to 1. Counted in a torch-profiler trace of each arm.
• Decode step (ITL p50): c1 -100 us (-1.75%), c2 -80 us (-1.16%).
  • Setup: DeepSeek-V4.1-Flash, MI355X TP4, --monokernel, DSpark k=6 (synthetic AL 3.87, AV off), 32k in / 1k out, on main nightly 5d35ce3.
  • Runs: 2 rounds with the arms' GPU halves swapped. The per-round deltas are -1.85% and -1.66% at c1, and -1.70% and -0.61% at c2.
• Same metadata, same accuracy.
  • On MI355X, the kernel test matches vLLM's ops.
  • With AV off, GSM8K matches base at 4 parallel requests, where every decode-only step takes the new path (6 runs per arm), and at 8 parallel (10 runs per arm).
  • Acceptance length (AL) on first passes and on passes with a per-request cache_salt is 3.577 in both arms.
  • With AV on (1 server per arm), GSM8K and AL match as well.
• Unchanged: gfx942, runs without --monokernel, other arches, and prefill, mixed or >32-row batches. The gate test covers arch x flag; the row and prefill guards are [ROCm][Perf][DSv4.1] Mono decode on MI300X and MI325X (gfx942) #60731's and untouched. There is no new flag.

Validation

Tests (MI355X, gfx950)

Test	This PR	Before (PR base)
tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py: kernel vs vLLM's ops, 300 random batches per case	-k swa_decode_meta: 2 passed; whole file: 3 passed	(kernel unchanged)
tests/models/test_dsv41_mono_decode.py, with the new gate test (arch x --monokernel)	82 passed	81 passed, 1 failed: [gfx950-True] assert [] == [False]
Extra, not in the PR: the kernel test with the target's 7 rows per request (k=6), the k=6 draft block, and 4 seeds x 500 batches	8 passed: 4,000 batches; target rows 1-4 and 7/14/21/28, draft rows 6/12/18/24 at width 192

• CI:
  • The kernel test (-k swa_decode_meta) is added to :amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics. Until now CI ran it only in the MI300 (gfx942) kernels step.
  • The gate test runs in the existing (MI355) Basic Models (Other) mirror. It stubs launch and the platform, so no CI job runs build() with the real kernel on gfx950. The trace and the GSM8K runs below cover that path outside CI.
• Other checks:
  • On MI355X, the gate test kills 6 of 6 gate mutants: on_mi3xx() in place of the arch check, gfx942 only, any ROCm arch, no --monokernel check, no cache_clear(), and build() back on enabled().
  • pre-commit passes on the 5 files at the first commit, and ruff passes on the second.
  • The PR applies cleanly to current main.

Performance. DeepSeek-V4.1-Flash on MI355X, one TP4 server per half of the node. The servers run --monokernel with DSpark k=6, synthetic AL 3.87 and AV off. The workload is random 32k in / 1k out, --ignore-eos, 32 prompts per run, 3 runs per point on the same server. The decode step is ITL p50.

ITL p50 (ms), each cell base / PR	Round 1 (base on GPUs 0-3)	Round 2 (base on GPUs 4-7)	Mean base	Mean PR	Delta
c1	5.687 / 5.581	5.681 / 5.587	5.684	5.584	-100 us (-1.75%)
c2	6.917 / 6.800	6.865 / 6.824	6.891	6.812	-80 us (-1.16%)

• Noise:
  • The 3 runs on one server span up to 0.3% at c1 and up to 1.0% at c2 (max-min).
  • The same arm moves 0.1% between GPU halves at c1 and up to 0.75% at c2.
• Image: vllm/vllm-openai-rocm:nightly-rocm100-5d35ce3934a7804b298559d2d99eec33792e1eb9. That is the PR base, which already has [Bugfix][ROCm][DSv4.1] Await krdy before K2's combine reads klen in index-source launches #63576, plus 23 commits.
  • The PR arm is the same image with the PR's 3 vllm files copied in. The end-to-end runs used the first commit. The second only swaps on_mi3xx() for on_gfx942() or on_gfx950(), which give the same value on gfx942 and gfx950 (_ON_MI3XX is exactly those two arches). The tests above ran on the second commit.
  • Those files are byte-identical between the PR base and 5d35ce3.
  • AV-on perf was not measured.

Launch counts. These come from a torch-profiler trace of each arm: rank 0, c1, 133 decode steps each. The SWA-build kernel counts are the same on every step. The total, fill and copy counts below are the per-step mode; a few steps add the same extra launches in both arms. The PR is 49 launches lower on every step.

Per decode step	base	PR
All launches (eager kernels plus graph launches): 115 of 133 steps; the others are 117/133 for base, 68/84 for PR	114	65
_compute_swa_indices_and_lens_kernel (target) and the draft's non-causal index kernel	4 + 3	0
torch.ge (compare_scalar_kernel<long>) / FillFunctor<int>	7 / 9	0 / 2
_count_valid_row_entries_kernel / rocprim scan / _compact_dense_row_to_ragged_kernel	7 / 7 / 7	0 / 0 / 0
__amd_rocclr_copyBuffer	24	10
_decode_swa_meta_kernel	0	7
SWA metadata launches. The fill and copy rows also hold 2 fills and 10 copies per step outside the SWA build, in both arms.	56	7
_token_request_mapping_kernel (control)	2	2
GPU time per step (median, profiler on). Base: the target index, compare, count, scan and compact kernels; the draft index kernels, fills and copies are not counted. PR: _decode_swa_meta_kernel.	136.5 us	36.8 us

Accuracy and AL (AV off unless noted). GSM8K-1319 is 8-shot at temperature 0, over /v1/completions. AL is real block-rejection acceptance, read from /metrics.

GSM8K-1319	base	PR	Delta
4 parallel: every decode-only step takes the new path. 6 runs per arm on 4 servers.	0.9143	0.9151	+0.0008
8 parallel: steps with 5 or more requests keep vLLM's build. 10 runs per arm on 4 servers.	0.9146	0.9123	-0.0023 (SE 0.0022). Per-server means: -0.006, below the 0.0077 binomial SE of a single run.
Mono forced off, so the aiter attention reads the kernel's ragged views. 4 parallel, 2 runs per arm, 1 server each.	0.9132	0.9143	+0.0011
Mono forced off, 8 parallel (mostly vLLM's views). 4 runs per arm, 1 server each.	0.9166	0.9177	+0.0011
AV on, 8 parallel: 2 runs per arm, 1 server each	0.9143	0.9166	+0.0023
AV on, 4 parallel: 1 run per arm	0.9227	0.9158	-0.0068 (per-run sd is about 0.005)
AL with AV off, on a first pass (only the shared 8-shot prefix cached) or with a per-request cache_salt (nothing cached). 8 runs per arm: 6 at 8 parallel, 2 at 4 parallel.	3.577	3.577	0
AL with AV on, first pass	3.373	3.383	+0.010

• Run-to-run variation:
  • Any two AV-off runs with mono on, at the same parallelism, on the same server or not, differ on 53-88 of the 1319 questions (53-92 including the mono-off runs).
  • Greedy decode on this stack is not run-to-run deterministic on either arm. The first token can differ even with a fresh cache_salt, so a text-identity check was not possible.
• The mono-off rows:
  • With --monokernel, mono consumes only the dense indices and lengths.
  • The 4-parallel mono-off row is therefore the end-to-end check of the ragged indices, indptr and is_valid_token.

Sync check. One TP4 server per arm ran with VLLM_GPU_SYNC_CHECK=error, which polices the worker's execute_model and sample_tokens after warmup. Each served c1 and c2 (8 requests each, 512 out), 300 GSM8K questions at 4 parallel requests, and 4 greedy probes. Neither arm reported a sync, logged an error or failed a request.

Reproduction commands

IMG=vllm/vllm-openai-rocm:nightly-rocm100-5d35ce3934a7804b298559d2d99eec33792e1eb9
# PR arm: the same image with this PR's 3 files copied over
# /usr/local/lib/python3.12/dist-packages/vllm/models/deepseek_v41/amd/{ops/__init__.py,ops/swa_decode_meta.py,rocm.py}
# Mono-off arms: in .../vllm/models/deepseek_v41/amd/mono_decode.py, in _decode_step, change
# `if not 1 <= rows <= _max_rows():` to `if True:`, so mono declines every step.

# Server, 4x MI355X, inside the image
VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=True \
VLLM_USE_BREAKABLE_CUDAGRAPH=1 VLLM_USE_RUST_FRONTEND=1 AITER_TRITON_LOG_LEVEL=ERROR OMP_NUM_THREADS=1 \
VLLM_ENGINE_READY_TIMEOUT_S=3600 \
vllm serve deepseek-ai/DeepSeek-V4.1-Flash --tensor-parallel-size 4 --monokernel --moe-backend aiter \
  --max-model-len 1048576 --block-size 128 --max-num-batched-tokens 16384 --max-num-seqs 256 \
  --max-cudagraph-capture-size 1024 --gpu-memory-utilization 0.9 --language-model-only \
  --tokenizer-mode deepseek_v41 --tool-call-parser deepseek_v41 --reasoning-parser deepseek_v41 --enable-auto-tool-choice \
  --attention-config '{"indexer_kv_dtype":"mxfp4","indexer_sparse_logits":true}' --engram-config '{"cpu_offload":false}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":6,"draft_sample_method":"probabilistic","rejection_sample_method":"synthetic","synthetic_acceptance_length":3.87,"enable_adaptive_verification":false}'
# Accuracy servers: "rejection_sample_method":"block", without synthetic_acceptance_length.
# AV-on round: "enable_adaptive_verification":true.

# Perf: warm up (c2, 4 prompts, 256 out, seed 7), then run each point 3 times:
vllm bench serve --backend openai --model deepseek-ai/DeepSeek-V4.1-Flash --tokenizer-mode deepseek_v41 \
  --dataset-name random --random-input-len 32768 --random-output-len 1024 --ignore-eos \
  --max-concurrency 1 --num-prompts 32 --seed 101 --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,90,99
# c2: --max-concurrency 2 --seed 102

# GSM8K: sglang few_shot_gsm8k prompt and scoring, 8 shots, 1319 questions, temperature 0, max 512 tokens,
# 8 or 4 parallel requests.
# AL = 1 + delta(vllm:spec_decode_num_accepted_tokens_total) / delta(vllm:spec_decode_num_drafts_total)
# Uncached passes: a unique "cache_salt" per request.

# Tests on MI355X
pytest tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py -k swa_decode_meta
pytest tests/models/test_dsv41_mono_decode.py
# Fail-before: keep the new test, restore the pre-PR vllm files, re-run.
git checkout $(git merge-base HEAD origin/main) -- vllm/models/deepseek_v41/amd/ops/__init__.py \
  vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py vllm/models/deepseek_v41/amd/rocm.py
pytest tests/models/test_dsv41_mono_decode.py   # 81 passed, 1 failed ([gfx950-True])

Details

Change. ops.swa_decode_meta_enabled(), the gate from #63994, returns true with --monokernel on ROCm when on_gfx942() or on_gfx950() holds. Only swa_decode_meta.build uses it. The other gfx942 helpers keep ops.enabled(). On gfx942 the two gates agree.

Limits.

• The kernel takes at most 32 rows; gfx950 mono takes up to 48.
  • With AV off at k=6, steps with 5 or more requests keep vLLM's build: the target at 35 or more rows, and the draft, which is padded to the 36-row graph. gfx942 behaves the same.
  • With AV on, the cutoff is 32 rows, not a request count.
• AV-on --monokernel deployments now take this path. The AV-on round served without errors and matched base on accuracy and AL.

Related PRs.

• [ROCm][Perf][DSv4.1] One-launch decode SWA metadata on gfx950 #63994 (cpersson-amd) made the same change, and the two are converged here. This PR takes [ROCm][Perf][DSv4.1] One-launch decode SWA metadata on gfx950 #63994's swa_decode_meta_enabled() gate, with Carl Persson as co-author of that commit. It adds the gate test, the MI355 CI step ([ROCm][Perf][DSv4.1] One-launch decode SWA metadata on gfx950 #63994 used the optional test-amd.yaml step) and the GPU validation. [ROCm][Perf][DSv4.1] One-launch decode SWA metadata on gfx950 #63994 can be closed.
• [ROCm][DSV4][Perf] Pack sparse-MLA ragged decode metadata into graph buffers #58352: it does not conflict textually, but the gains overlap. If it lands first, about 20 of the 49 launches remain (inferred, not measured).
• [Core] Shared abstractions for specialized decode kernels #61311 touches 4 of this PR's 5 files: models_basic.yaml, ops/__init__.py, rocm.py and test_dsv41_mono_decode.py.
  • It renames KernelConfig.monokernel and record_monokernel to specialized_decode, and replaces test_dsv41_mono_decode.py with test_dsv41_specialized_decode.py.
  • Whichever lands second rebases: rename _monokernel in swa_decode_meta_enabled() and move the gate test.
• [CI][ROCm] Check AMD monokernels stage by stage against float64 references #63486 (draft) moves the DSv4.1 mono CI and test_dsv41_mono_decode.py. If it lands first, I drop the models_basic.yaml hunk and move the gate test.
• [ROCm][Perf][DSv4.1] Run DSpark's fused Markov sampling on gfx950 with --monokernel #63891 (draft) runs fused Markov sampling on gfx950 with --monokernel. It adds its own mi3xx_enabled() on on_mi3xx(), and the same review applies to it. Both PRs touch the same models_basic.yaml step and the ops/__init__.py docstring; the second to land merges those hunks.

Why the metadata is the same, and who reads it

• What the kernel writes. _decode_swa_meta_kernel is integer Triton with no arch intrinsics.
  • For rows [0, n) it writes is_valid_token and the dense window indices and lengths. It sets the lengths from row n to the end of the buffer to 0. It writes indptr[:n+1] and the ragged indices up to indptr[n].
  • These are the same values that torch.ge, the length fill, _compute_swa_indices_and_lens_kernel (or ComputeDSparkNoncausalSWAIndicesKernel for the draft block), build_ragged_indices_from_dense and _copy_ragged_to_graph_buffers produce.
  • build() returns views of the same persistent buffers, with the same lengths: ragged [:max(n*width, 1)] and indptr[:n+1].
• Who reads it.
  • Mono: with --monokernel, every target and DSpark-draft layer runs mono on decode-only steps of 48 rows or fewer that replay a FULL graph or run eagerly. Mono reads only the dense indices and lengths, slot_mapping and token_to_req.
  • aiter sparse MLA: on steps that mono declines, such as PIECEWISE steps, it (VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA) reads the ragged indices between indptr[t] and indptr[t+1]. It sizes its splits from the ragged view's numel(), which is n*width on both paths.
  • Top-k build: it reads is_valid_token on compressed layers.
  • Draft and target builders do not share buffers.
• The only bytes that differ are the ragged entries past indptr[n]. vLLM's build copies them from a torch.empty tensor; the kernel leaves the previous contents. No reader reads past indptr.
• [ROCm][DSv4/4.1] Fuse inverse RoPE and MXFP8 quant into the aiter sparse MLA and sparse MLA gating changes #60154 is on main but not in the measured image. It fuses inverse RoPE and the MXFP8 quant into the aiter sparse-MLA output, and does not change the index views that call reads. So this argument holds on current main.

An AL caveat for anyone reproducing (affects main, not this PR)

On 4 of the 10 AV-off servers, AL fell from about 3.58 to 3.42-3.45 on GSM8K passes that repeated the same prompts. Those passes had about 96% prefix-cache hits, against about 89% on a first pass. The 4 servers were 2 base and 2 PR; one of the PR servers had mono forced off.

• The two dropped servers that ran later passes stayed low.
• On the one dropped server that also ran passes with a per-request cache_salt, those passes stayed at 3.58.
• Accuracy did not move.
• With AV on, the PR server's repeat pass also fell (3.383 to 3.279); base's did not (3.373 to 3.364).

Compare AL on first or salted passes only. I will look into the drop separately.

---

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.

[frida-andersson]

This overlaps with #63994, please converge to one PR. The piece this draft has that #63994 does not is the CPU gate test

[frida-andersson]

Don't gate this on mi3xx_enabled() / on_mi3xx(). Check the CDNA version or the gfx arch with get_cdna_version() or on_gfx942() / on_gfx950(). #63994 already does the second in swa_decode_meta_enabled()

[okorzh-amd]

@frida-andersson Done in 768a354. swa_decode_meta.build now checks swa_decode_meta_enabled() from #63994, which is --monokernel and on_gfx942() or on_gfx950(). The gate test fakes on_gfx950 instead of on_mi3xx, so on MI355X it also fails if the gate goes back to on_mi3xx().

[claude]

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.

[okorzh-amd]

@frida-andersson Converged with #63994 here, as proposed. The second commit takes #63994's swa_decode_meta_enabled() gate, with @cpersson-amd l as co-author. This PR adds the CPU gate test, the MI355 CI step and the MI355X GPU validation, now in the description:

• c1 -100 us (-1.75%)
• 49 fewer launches per decode step
• GSM8K and acceptance length unchanged

[dllehr-amd]

can we check for cdna 3 or 4 instead? that way we won't need to add additional imports as we go forward

[okorzh-amd]

updated. thanks

[dllehr-amd]

This also looks good, thanks @okorzh-amd

[github-actions]

✅ @okorzh-amd, 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.

[mergify]

Hi @okorzh-amd, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[AndreasKaratzas]

Test def looks good, but i m trying to make sure that the UT is also looking good. Let's not skip the NITs all the time . I can merge it rn but @okorzh-amd could you first lmk any comments on the lines i highlighted?

[AndreasKaratzas]

this looks a bit suboptimal, if you choose cdna def then why not parametrize it like that ? Also what are we testing on gfx90a here?

[okorzh-amd]

@AndreasKaratzas Agreed, thanks. Done in 67aec3e and d331b82

[okorzh-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98447 for commit c8a836b628d6.

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