[ROCm][DSv4.1] Enable DSpark adaptive verification

[tuukkjs]

Purpose

Enable confidence-scheduled adaptive verification for DeepseekV41ForCausalLM on ROCm.

This PR extends the existing validated ROCm flattened-query mechanism to V4.1 and adds adaptive-only graph-support declarations for the two builders. Fixed-width DSpark behavior is unchanged.

AI assistance was used for this PR.

Implementation

• Allow the V4.1 indexer backend to use the ROCm device/CPU query-length mismatch capability.
• Extend ROCm flattened device-query handling to DeepseekV41ForCausalLM.
• Report AttentionCGSupport.ALWAYS for the V4.1 ROCm compressed sparse-MLA and sparse-SWA builders only when adaptive verification is enabled.
• Add focused policy coverage, a fixed-K regression test, and a gfx950 graph-replay test for changed device query boundaries and stable persistent metadata buffers.

Validation stack

• 4x MI355X (gfx950), TP=4, deepseek-ai/DeepSeek-V4.1-Flash; expert parallel enabled for baseline and disabled for the matched mono-off/on (referring to [ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X #60397) comparison
• ROCm nightly vllm/vllm-openai-rocm:nightly-7d0b4e57aac4c4323225b0eee9ae960ce75ffffe (vLLM 0.31.1rc1.dev253+g7d0b4e57a.rocm723)
• nightly source 7d0b4e57aa plus this PR's exact four-file patch
• Matched EP-off accuracy and throughput tests were run with VLLM_ROCM_MONO_DECODE=0 and 1, and the EP-on baseline used mono decode off
• [ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X #60397's mono path does not support EP, so mono's direct effect is measured by the matched EP-off pair

The primary validation matrix below used standard rejection sampling. Adaptive block rejection currently also requires #60945, which fixes an out-of-bounds read when adaptive verification chooses fewer steps than the configured maximum. With that fix, GSM8K accuracy and TP4 throughput matched standard rejection.

vLLM server command and environment

The adaptive arm used the following effective command inside the vllm container (the fixed-K=5 arm changed only enable_adaptive_verification to false; the no-speculation arm omitted --speculative-config; VLLM_ROCM_MONO_DECODE was set to 0 or 1). The EP-on baseline additionally passed --enable-expert-parallel and used VLLM_ROCM_MONO_DECODE=0:

VLLM_ENGINE_READY_TIMEOUT_S=7200 \
VLLM_ROCM_USE_AITER=1 \
VLLM_USE_BREAKABLE_CUDAGRAPH=1 \
VLLM_USE_V2_MODEL_RUNNER=1 \
VLLM_ROCM_MONO_DECODE=<0|1> \
vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --host 0.0.0.0 \
  --port 8000 \
  --served-model-name deepseek-ai/DeepSeek-V4.1-Flash \
  --tokenizer-mode deepseek_v41 \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --distributed-executor-backend mp \
  --max-model-len 8192 \
  --max-num-batched-tokens 4096 \
  --max-num-seqs 256 \
  --gpu-memory-utilization 0.9 \
  --kv-cache-dtype fp8 \
  --block-size 64 \
  --moe-backend aiter \
  --compilation-config '{"mode":3,"cudagraph_mode":"FULL_AND_PIECEWISE","max_cudagraph_capture_size":2048}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","enable_adaptive_verification":true}' \
  --language-model-only \
  --enable-auto-tool-choice \
  --tool-call-parser deepseek_v41 \
  --reasoning-parser deepseek_v41

Test results

Focused tests

• 156 passed on gfx950: the PR's CPU policy tests, the fixed-K regression test, test_dsv41_adaptive_mla_swa_metadata_graph_replay, and the shared adaptive-verification tests.
• Changed-file pre-commit: passed.

Accuracy

GSM8K 5-shot, 1,319 samples per arm, concurrency 16, standard rejection sampling:

EP	Mono decode	Arm	Strict exact match	Flexible exact match
On	Off	No speculation	0.9265 ± 0.0072	0.9272 ± 0.0072
On	Off	Fixed K=5	0.9325 ± 0.0069	0.9325 ± 0.0069
On	Off	Adaptive K<=5	0.9219 ± 0.0074	0.9227 ± 0.0074
Off	Off	No speculation	0.9310 ± 0.0070	0.9303 ± 0.0070
Off	Off	Fixed K=5	0.9249 ± 0.0073	0.9265 ± 0.0072
Off	Off	Adaptive K<=5	0.9189 ± 0.0075	0.9189 ± 0.0075
Off	On	No speculation	0.9204 ± 0.0075	0.9204 ± 0.0075
Off	On	Fixed K=5	0.9287 ± 0.0071	0.9287 ± 0.0071
Off	On	Adaptive K<=5	0.9212 ± 0.0074	0.9219 ± 0.0074

All arms completed every sample. Every within-configuration and same-arm topology difference is within 1.2 combined standard errors.

GSM8K command

lm_eval \
  --model local-completions \
  --model_args 'model=deepseek-ai/DeepSeek-V4.1-Flash,base_url=http://localhost:8000/v1/completions,num_concurrent=16,tokenizer_backend=none,tokenized_requests=False' \
  --tasks gsm8k \
  --num_fewshot 5 \
  --batch_size 1 \
  --seed 42

SPEED-Bench throughput

All 880 qualitative prompts in original order, 512 forced output tokens, a distinct cache salt per concurrency. Each arm and mono mode used one fresh server and one 64-request C32 warmup. Output tokens per second:

Baseline, EP on and mono decode off:

Arm	C4	C16	C32	C64	C128	C256
No speculation	446.5	1,346.4	2,317.4	3,753.6	5,832.9	8,297.5
Fixed K=5	822.3	2,188.0	3,404.2	4,824.2	6,579.9	8,487.3
Adaptive K<=5	806.5	2,265.7	3,588.4	5,331.0	7,448.9	9,475.1

Matched mono comparison, EP off and mono decode off:

Arm	C4	C16	C32	C64	C128	C256
No speculation	463.5	1,338.2	2,299.4	3,713.1	5,688.8	8,017.0
Fixed K=5	821.1	2,154.7	3,345.7	4,760.2	6,527.9	8,031.1
Adaptive K<=5	802.8	2,255.6	3,522.4	5,273.2	7,136.3	9,235.4

Matched mono comparison, EP off and mono decode on:

Arm	C4	C16	C32	C64	C128	C256
No speculation	818.0	2,063.0	2,975.1	3,755.8	5,751.0	8,106.8
Fixed K=5	1,111.4	2,177.5	3,325.4	4,700.0	6,530.2	8,243.4
Adaptive K<=5	1,093.8	2,277.1	3,558.0	5,251.0	7,262.4	9,210.9

All 54 cells completed 880/880 requests with zero failures. In the EP-on baseline, adaptive relative to fixed K=5 was -1.9% at C4 and +3.5% to +13.2% from C16 to C256; it beat no speculation by 27.7% at C128 and 14.2% at C256. In the matched EP-off comparison, adaptive relative to fixed was -2.2% at C4 and +4.7% to +15.0% from C16 to C256 with mono off; with mono on it was -1.6% at C4 and +4.6% to +11.7% from C16 to C256. Adaptive beat no speculation at every measured concurrency in all three configurations.

The mono launches were confirmed active in all three mono-on server logs. Their direct performance effect matched #60397's scope: at C4 mono improved no speculation by 76.5%, fixed K=5 by 35.4%, and adaptive by 36.3%. Above C4 the fixed/adaptive mono-on/off difference was mostly within about 2%, because a full K=5 verify reaches the mono path's 48-row ceiling at eight requests.

SPEED-Bench commands

Each arm used one server and a 64-request C32 random warmup before the measured loop:

vllm bench serve \
  --backend openai-chat \
  --host localhost \
  --port 8000 \
  --endpoint /v1/chat/completions \
  --model deepseek-ai/DeepSeek-V4.1-Flash \
  --tokenizer-mode deepseek_v4 \
  --trust-remote-code \
  --dataset-name random \
  --random-input-len 512 \
  --random-output-len 256 \
  --ignore-eos \
  --temperature 1.0 \
  --request-rate inf \
  --seed 20260925 \
  --num-prompts 64 \
  --max-concurrency 32

DATASET_DIR=<path-to-speed-bench>
ARM=<nospec|fixed_k5|adaptive>

for c in 4 16 32 64 128 256; do
  vllm bench serve \
    --backend openai-chat \
    --host localhost \
    --port 8000 \
    --endpoint /v1/chat/completions \
    --model deepseek-ai/DeepSeek-V4.1-Flash \
    --tokenizer deepseek-ai/DeepSeek-V4.1-Flash \
    --tokenizer-mode deepseek_v4 \
    --trust-remote-code \
    --dataset-name speed_bench \
    --dataset-path "$DATASET_DIR" \
    --speed-bench-dataset-subset qualitative \
    --skip-chat-template \
    --disable-shuffle \
    --no-oversample \
    --temperature 1.0 \
    --ignore-eos \
    --speed-bench-output-len 512 \
    --num-prompts 880 \
    --max-concurrency "$c" \
    --request-rate inf \
    --seed 20260925 \
    --metric-percentiles 95,99 \
    --percentile-metrics ttft,tpot,itl,e2el \
    --extra-body "{\"cache_salt\":\"speedbench-tp4-${ARM}-c${c}\"}" \
    --save-result
done

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

[frida-andersson]

This looks great! One question - since the TP4 validation includes #56831, should that PR land first or be a hard dependency? What happens with TP>1 if this PR lands by itself?

[mergify]

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

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

[tuukkjs]

This looks great! One question - since the TP4 validation includes #56831, should that PR land first or be a hard dependency? What happens with TP>1 if this PR lands by itself?

Yes, #56831 should land first: it is a correctness dependency for adaptive TP>1. If this PR lands alone, adaptive TP=1 works and fixed-K remains unchanged, but adaptive TP>1 is unsafe because ranks can choose different budgets or request boundaries, potentially causing divergent execution or collective hangs.

[ChuanLi1101]

Adaptive TP>1 can run without #56831, and #56831 is still open. Keep this PR blocked until #56831 is merged, then update the branch and rerun the TP tests.

[okorzh-amd]

I believe this is identical code to the function above. create a helper function instead of copy pasting the code

[tuukkjs]

I have kept these overrides explicit because the only shared code is the adaptive-verification condition. Each override has a builder-specific safety rationale and superclass fallback, so a helper would add indirection without centralizing meaningful behavior.

[okorzh-amd]

is the test running in mi355 CI?

[tuukkjs]

It was previously only run manually. I added the focused test_dsv41_adaptive_mla_swa_metadata_graph_replay case to the MI355 DPX MLA Kernels CI job. The local MI355X run passes. The CI result will be available once AMD CI runs.

[okorzh-amd]

what is the plan for V42, V43 and so on?
I think we can have something like this here:

_ROCM_FLATTENED_DEVICE_QUERY_LENS_ARCH_RE = re.compile(r"DeepseekV4\d*ForCausalLM")

[tuukkjs]

I would keep this as an explicit allowlist rather than automatically opt every future V4.x architecture into the capability. This check asserts that the architecture's ROCm metadata and indexer paths safely support device/CPU query-length mismatch and adaptive full-graph replay. A name regex could silently enable an unvalidated future implementation whose backend or metadata layout differs.

New architectures should be added only after they pass the similar focused validation.

[okorzh-amd]

I run this change with DSv41 flash and agentic dataset and it is crushing:

1. Before the fault, adaptive verification was working: acceptance was about 4.5 per step with 8–11 requests running.

00:54:44 Avg prompt tput: 56530.4 toks/s, Avg generation tput: 186.6 toks/s, Reqs Running: 8, Waiting: 56
00:54:44 SpecDecoding metrics: Mean acceptance length: 4.54, ... Per-position acceptance rate: 0.895, 0.747, 0.637, 0.518, 0.431, ...
00:55:04 Avg prompt tput: 32013.2 toks/s, Avg generation tput: 224.2 toks/s, Reqs Running: 9, Waiting: 55
00:55:04 SpecDecoding metrics: Mean acceptance length: 4.57, ... Per-position acceptance rate: 0.890, 0.752, 0.637, 0.530, 0.436, ...

2. The fault, as soon as the batch ramped to 59 running requests:

Warning: Queue error - HSA_STATUS_ERROR_MEMORY_FAULT
Memory access fault by GPU node-8 (Agent handle: 0x192c9d20) on address 0x7eec50800000. Reason: Unknown.
GPU coredump: Pipe handler '/usr/share/apport/apport' not found or not executable.
Falling back to file-based dump. Set HSA_COREDUMP_PATTERN to override.
(RustFrontend) 00:55:14 Avg prompt tput: 32433.0 toks/s, Avg generation tput: 357.7 toks/s, Reqs Running: 59, Waiting: 0, GPU KV cache used: 4.7%, Prefix cache hit rate: 97.3%
(RustFrontend) 00:55:14 SpecDecoding metrics: Mean acceptance length: 3.94, ...
(RustFrontend) 00:55:24 Avg prompt tput: 0.0 toks/s, Avg generation tput: 0.0 toks/s, Reqs Running: 59, Waiting: 0

3. Hang, then the engine died after the 5-minute timeout:

(EngineCore) 00:56:09 [shm_broadcast.py:812] No available shared memory broadcast block found in 60 seconds. ...
(EngineCore) 00:57:09 ... (same, every 60 s) ...
(EngineCore) 00:59:09 ...
(EngineCore) ERROR 01:00:08 [dump_input.py:81] Dumping scheduler stats: SchedulerStats(num_running_reqs=59, num_waiting_reqs=0, ...)
(EngineCore) ERROR 01:00:08 [core.py:1449] EngineCore encountered a fatal error.
Traceback (most recent call last):
  File ".../vllm/v1/executor/multiproc_executor.py", line 435, in get_response
    status, result = mq.dequeue(timeout=dequeue_timeout)
  File ".../vllm/distributed/device_communicators/shm_broadcast.py", line 900, in dequeue
    with self.acquire_read(timeout, indefinite) as buf:
  File ".../vllm/distributed/device_communicators/shm_broadcast.py", line 805, in acquire_read
    self._spin_condition.wait(timeout_ms=read_timeout.timeout_ms())
  File ".../vllm/distributed/device_communicators/shm_broadcast.py", line 758, in timeout_ms
    raise TimeoutError
TimeoutError
The above exception was the direct cause of the following exception:
  File ".../vllm/v1/engine/core.py", line 745, in step_with_batch_queue
    model_output = future.result()
  File ".../vllm/v1/executor/multiproc_executor.py", line 103, in _wait_for_response
    response = self.aggregate(self.get_response())
...
RuntimeError: Engine core process EngineCore died unexpectedly.
(RustFrontend) ERROR 01:00:08 [chat_completions.rs:485] ... chat stream failed error=engine-core error: engine core reported f...

I believe this still need more changes similar to the dependencies in #59824

[tuukkjs]

Please give me the reproducer commands so that I can investigate 🙏 . I haven't seen any errors.

[okorzh-amd]

I'll sent you repro steps in pm.

[tuukkjs]

Thanks. I reproduced the fault only with block rejection. Standard rejection ran cleanly. It appears to be the shared block-rejection out-of-bounds access fixed by #60945, rather than the V4.1 MLA/indexer changes. With that fix, the same TP8 workload ran cleanly at K=5 and K=6.

[tuukkjs]

Adaptive TP>1 can run without #56831, and #56831 is still open. Keep this PR blocked until #56831 is merged, then update the branch and rerun the TP tests.

Update: #56831 is no longer a merge prerequisite for this PR. TP4 focused tests, GSM8K, and multiple full C4–C256 serving sweeps pass without it.

[frida-andersson]

Could you run RULER NIAH at 16K and 64K for fixed K=5 and adaptive? GSM8K looks fine, but the contexts in the tests and the serve runs stay pretty short