[ZhengGong-amd]

Overview

On ROCm, every DeepSeek-V4 / V4.1 sparse-MLA decode build() packs its ragged indices into scratch tensors and then copies them into the FULL decode graph buffers, and the ROCm SWA builders cannot refresh those buffers after a draft step; this PR packs straight into the graph buffers and adds the in-place draft refresh, so the ROCm SWA builders set supports_draft_decode_metadata_update = True.

build_ragged_indices_from_dense_out() runs the count → scan → compact kernels from #58058 into caller-owned buffers, writing counts into indptr[1:] and scanning them in place; build_ragged_indices_from_dense() becomes an allocating wrapper, so the -1 filtering has one implementation and prefill output is unchanged. build_ragged_into_graph_buffers() and refresh_decode_swa_ragged() sit beside it in rocm_aiter_mla_sparse.py and are shared by the V4 C128A/SWA and V4.1 SWA builders; each ROCm SWA update_draft_decode_metadata() is super() plus one call. The gfx950 width narrowing is unchanged. The draft refresh comes from #52628 (Tres Popp, co-author), closed in favour of this PR.

When every draft attention group can update in place, the MRV2 draft speculators record the update in their FULL draft graph instead of rebuilding the draft metadata eagerly on every draft step; since #60692 that includes DSpark, whose draft layers are SWA-only on V4 and V4.1.

Claims

• The packing runs 7 times per decode step on V4.1 and 3 on V4. Each call drops from 5 kernels, 3 allocations and 2 hipMemcpyAsync to 3 kernels, saving 14.2–14.6 µs (27 %) at every shape on MI355X.
• End to end on V4.1-Flash the effect is marginal: −1.5 % per verify step with DSpark and −0.4 % TPOT without spec at concurrency 1, within run-to-run spread at concurrency 4 and 64.
• All of it comes from the packing. The in-graph draft update removes main's eager draft rebuild (on 948 of 950 draft steps under DSpark) and its Attention backend(s) DEEPSEEK_SPARSE_SWA cannot update draft decode metadata in place; falling back to rebuilding attention metadata eagerly. log, but gives no measurable speedup: with the flag forced off, DSpark is within 0.5 % of the PR at every concurrency.
• GSM8K and DSpark acceptance agree with main as closely as a repeated main run does.

Validation

Unit, MI355X (gfx950), main 5c4b94beec, main's and the PR's tests/ against each tree: the PR's tests pass on the PR; main's tests fail on the PR only where they import the removed _copy_ragged_to_graph_buffers; forcing the flag off fails only test_rocm_swa_builder_enables_fused_draft_decode; a no-op refresh_decode_swa_ragged() fails the four draft-update cases.

Micro, MI355X, main 5c4b94beec, the production helpers at the e2e configs' V4.1 SWA shapes (1–384 rows × 128), 5 % -1 sentinels, median of runs 2–4, same indptr and indices on both arms: main 53.5–53.9 µs per call, PR 39.0–39.4 µs, −14.2 to −14.6 µs (27 %) at every shape.

e2e, V4.1-Flash, TP=2, MI355X, main 5c4b94beec, --kv-cache-dtype fp8, VLLM_ROCM_USE_AITER=1, DSpark as {"method": "dspark"}. vllm bench serve, random 1031 in / 509 out, 32 × concurrency prompts, one server on the host at a time, fresh server per run, median (range) of runs 2–4. "Flag off" is this PR with supports_draft_decode_metadata_update = False. With DSpark the metric is ms per verify step, mean TPOT × mean acceptance length: on random prompts acceptance moves between 1.64 and 1.94 from run to run, and raw TPOT and throughput follow it.

concurrency	metric	main	flag off	PR	PR vs main
1	mean TPOT, ms	6.76 (6.75–6.78)	–	6.73 (6.65–6.74)	−0.4 %
4	median TPOT, ms	8.34 (8.33–8.37)	–	8.32 (8.31–8.33)	−0.2 %
1, DSpark	ms / verify step	9.64 (9.60–9.67)	9.51 (9.46–9.55)	9.49 (9.43–9.57)	−1.5 %
4, DSpark	ms / verify step	13.92 (13.75–13.92)	13.83 (13.75–13.85)	13.89 (13.84–13.92)	−0.2 %
64, DSpark	ms / verify step	50.78 (50.53–50.97)	50.80 (50.47–50.82)	50.79 (50.76–50.82)	0.0 %

GSM8K 5-shot strict-match, greedy, seed=1234, same setup:

V4.1-Flash	main	main (repeat)	flag off	PR
DSpark	0.9303	0.9348	0.9265	0.9272
DSpark mean acceptance length	3.448	3.451	3.442	3.431
no spec	0.9219	0.9219	–	0.9227

Per sample, the PR matches main's verdict on 1265 of 1319 with DSpark (repeat 1259) and 1250 without (repeat 1257). V4 was last measured on main 72e7874fa6, before #60692: at concurrency 61, TP=1 and MTP × 3, the arms' ranges overlapped and GSM8K matched main as closely as a repeat.

Details

With --monokernel (off by default), V4.1 decode batches of up to 32 rows get their SWA metadata from swa_decode_meta, on gfx942 since #60731 and on gfx950 since #63885, which also writes into the graph buffers; this PR covers the default configuration, larger V4.1 batches, and V4. #57860 fuses indptr into a different V4.1 pack kernel and composes with this.

---

Found by Hyperloom; reviewed and measured by hand.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ZhengGong-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[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.

[tpopp]

Thank you for carrying this on. I think effects would be better seen with lower concurrency benchmarks in this case.

[tpopp]

I'm not really caught up on how the new model runner system should be structured. Do you know where common code could be shared? I think the _build_ragged_into_graph_buffers and update_draft_decode_metadata are exactly the same in deepseek_v4 and deepseek_v41.

[ZhengGong-amd]

Done in ec746e9. _build_ragged_into_graph_buffers is now build_ragged_into_graph_buffers() in vllm/v1/attention/ops/rocm_aiter_mla_sparse.py, next to build_ragged_indices_from_dense_out(), which both modules already import. The draft-step body lives there as refresh_decode_swa_ragged(), so each ROCm SWA builder's update_draft_decode_metadata() is now super() plus one call. The draft-update unit test now runs on both the V4 and V4.1 builders, and its cases fail when refresh_decode_swa_ragged() is made a no-op.

[tpopp]

Am I mistaken that, based on the runs, setting this flag to true and update_draft_decode_metadata are dead code because 4.1 uses Dspark rather than MTP?

[ZhengGong-amd]

When you wrote this, yes: V4.1 runs DSpark, and the flag only mattered for the MTP fused-AR draft path. #60692 (merged Oct 12) changed that. DSparkSpeculator subclasses DFlashSpeculator, which now checks supports_draft_decode_metadata_update on every draft attention group. If all of them support it, it records update_draft_decode_metadata() inside the FULL draft CUDA graph and skips the eager _build_uniform_attn_metadata() on replay.

The DSpark draft layers are SWA-only on both models: V4.1 treats layers past compress_ratios as ratio 0, and V4-Flash's entry at index 43 is 0. So their only builder is the ROCm SWA builder this PR enables, and on today's main the flag switches that path on for ROCm rather than being dead code.

That also raises the bar on validation. None of the Oct 1 numbers exercised this path, because they predate #60692. I'll rerun with DSpark before asking for another look.

[ZhengGong-amd]

It is reached: since #60692 DSpark records update_draft_decode_metadata in its draft graph, so this removes the eager draft rebuild (948 of 950 draft steps on main) and the fallback log. It gives no measurable speedup on V4.1, though (flag off is within 0.5 % per verify step at concurrency 1, 4 and 64); the gain is from the build() packing.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ZhengGong-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ZhengGong-amd]

Thanks for carrying #52628 over to this one. On lower concurrency: agreed. The build() saving is fixed per step (−14 to −15 µs per call, 3 calls/step on V4 and 7 on V4.1), so at concurrency 61 it disappears into a ~30 ms TPOT.

The next round will cover:

• concurrency 1 and 5, with and without DSpark;
• a null arm (this PR with the flag forced off), so the build() packing and the DSpark draft-update effect show up separately;
• draft acceptance next to GSM8K. A stale ragged buffer on the draft side would lower acceptance long before it shows up in accuracy, because verification keeps the target's greedy output either way.

The shared-helper refactor is in ec746e9, and the DSpark answer is in the inline thread.

[ZhengGong-amd]

Reran on main 5c4b94beec (V4.1-Flash, TP=2, MI355X) at concurrency 1 and 4 with and without DSpark, 64 with DSpark, plus this PR with supports_draft_decode_metadata_update = False. Tables are in the description.

• Micro: the decode build() packing is 27 % faster (53.5–53.9 → 39.0–39.4 µs per call, 7 calls per V4.1 step).
• End to end the gain is marginal: −0.4 % TPOT without spec and −1.5 % per verify step with DSpark at concurrency 1; within noise at 4 and 64.
• The draft update itself gives no measurable speedup (flag off is within 0.5 %); it only removes the eager draft rebuild and the fallback log.
• GSM8K and DSpark acceptance match main.

Could you take another look?