[ROCm][DSv4] Enable DSpark adaptive verification

[tuukkjs]

Purpose

Enable DSpark confidence-scheduled adaptive verification for DeepSeek-V4 on ROCm. Developed with contributions from @larryli2-amd, who added explicit missing-confidence-head error handling and provided the MI350X performance validation.

This change:

• adds AMD confidence-head construction, inference, checkpoint remapping, and safe fallback when a checkpoint has no confidence head
• enables variable-length FULL graph support in the DeepSeek-V4 ROCm sparse MLA and sparse SWA metadata builders when adaptive verification is enabled, while retaining their existing support classification otherwise
• extends the shared DeepSeek indexer flattening capability to the ROCm DeepSeek-V4 path, deriving adaptive request ownership from device query boundaries
• adds focused confidence-loading and graph-replay regression coverage

Adaptive verification can redistribute a fixed draft-token budget between requests on device, so CPU request boundaries can differ from final device boundaries. The flattened path builds request ownership, context lengths, and block-table rows from device decode lengths. It reuses the shared flattening selection introduced by #52795 and keeps the ROCm capability scoped to the DeepSeek-V4 builder. NVIDIA and unrelated ROCm indexers are unchanged.

The ROCm sparse MLA/SWA declarations are similarly scoped. Their ragged metadata is device-boundary-derived and copied into persistent buffers before FULL graph replay.

AI assistance was used in preparing this PR.

Test plan

• changed-file pre-commit, including Ruff and mypy for Python 3.10–3.13;
• focused AMD confidence-head, ROCm adaptive metadata, native next_n, and adaptive-verification tests;
• complete GSM8K 5-shot evaluation for no speculation, fixed K=5, fixed K=7, and adaptive K<=7;
• SPEED-Bench throughput comparison on 8x MI350X at TP=8 and client concurrency 4/8/16/32/64/128/256.

Test results

Source and focused tests

• rebased PR head ea586873d5210d9e1624c42333be1b02c276c832 onto upstream cd10ed6f9f6b37a8ace9cf380007e66fe12ec0c3;
• changed-file pre-commit passed, including Ruff and mypy for Python 3.10–3.13;
• the two PR-specific test files passed (10 passed) and the shared adaptive-verification suite passed (13 passed).

Performance validation — 8x MI350X

Performance validation contributed by @larryli2-amd used 8x AMD Instinct MI350X at TP=8 with deepseek-ai/DeepSeek-V4-Flash-DSpark.

This PR changes only vLLM; it does not modify or vendor AITER source code. The benchmark environment used a locally rebuilt, pinned AITER wheel installed into the derived image with --no-deps, leaving the base PyTorch and Triton installations unchanged:

• base image: vllm/vllm-openai-rocm@sha256:72e90adf360ccffb231a098a15eadfb808719c029c6cd43f6386095064552b7e;
• AITER version: 0.1.21.post2.dev2+g804648bcb;
• AITER source: 804648bcbf87646ddef4b7add10da9c3b322567c;
• AITER wheel SHA256: 30360d15819a9c8e3ec8e4ae665683abf4c5358d5be1f49a9aa2f7962c0032db;
• PyTorch: 2.12.0+rocm10.0.0.

The pinned wheel was used because the stock AITER post1 environment had a known GPU-fault issue in this setup.

The server used Model Runner V2, breakable CUDA graphs, async scheduling, prefix caching, FP8 KV cache, max_num_seqs=256, max_num_batched_tokens=16384, and FULL_AND_PIECEWISE capture through 2048. Adaptive DSpark used K<=7 and an 8192-token profiling context.

The benchmark used a frozen 880-request SPEED-Bench Throughput-1K request set at temperature 1, top-p 1, and maximum output length 2048. Concurrency ran in descending order from 256 to 4. Each point used an independent cold server boot, two 64-request warmups, matching client/server generated-token counts, zero preemptions, and no GPU faults.

Compared with fixed-length verification, confidence-scheduled DSpark delivered comparable performance at low concurrency and higher throughput at high concurrency. The complete serving command and environment details are preserved in the reproduction comment.

Correctness validation — 8x MI355X

The complete four-arm GSM8K 5-shot validation used TP=8, expert parallel disabled, FP8 KV, probabilistic DSpark, and FULL_AND_PIECEWISE. Pins were upstream 10704541aaf72567fe9d6229b3e3d84d37f2ddba, PR head 3fbcf64cfae8cd9d2180ceed0c92ba025f4dde30, model revision 7872f01b1d1fe23eabc4c98b48bffcef5a386062, and ROCm nightly 73029d42441321b631779db3475031f5ec26dd6c.

All four 1,319-sample evaluations completed:

Mode	Flexible exact match	Strict exact match
No speculation	0.9431 +/- 0.0064	0.9431 +/- 0.0064
Fixed K=5	0.9416 +/- 0.0065	0.9409 +/- 0.0065
Fixed K=7	0.9462 +/- 0.0062	0.9462 +/- 0.0062
Adaptive K<=7	0.9401 +/- 0.0065	0.9401 +/- 0.0065

The strict-score range was 0.0061, within one reported standard error.

Current ROCm qualification

On current ROCm main/nightly, DeepSeek-V4 defaults to Model Runner V1, while DSpark requires Model Runner V2. Set VLLM_USE_V2_MODEL_RUNNER=1 when launching this feature.

ROCm breakable CUDA graphs were disabled by default for every architecture in #53155. DeepSeek V4 adaptive throughput regresses without that execution mode, so the validated launch configuration also sets VLLM_USE_BREAKABLE_CUDAGRAPH=1.

[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

[larryli2-amd]

Hi @tuukkjs
I made some changes to the code based on your PR, and it does indeed implement a confidence-based adaptive scheduler for DSpark on ROCm. The results are below:


The results are expected: Compared with the non-adaptive scheduler, DSpark with the confidence-based adaptive scheduler delivers comparable performance at low concurrency and significant speedup at high concurrency.

It is my PR:
#56942. If you have time, please merge my PR into yours. If you’re too busy, I’ll merge your PR into mine. Either way, I’ll find someone from the vLLM team to review our PRs so we can get this feature merged into vLLM as soon as possible, since a lot of people are waiting for it.

[larryli2-amd]

Reproduction Environment

Hardware

• 8 × AMD Instinct MI350X
• Tensor parallel size: 8
• Data parallel size: 1

Model

• Model: deepseek-ai/DeepSeek-V4-Flash-DSpark
• Local validation path: /models/DeepSeek-V4-Flash-DSpark
• Architecture: DeepseekV4ForCausalLM
• Quantization: FP4 expert weights with FP8 mixed precision
• Checkpoint size: 166,878,536,440 bytes across 48 safetensor shards
• config.json SHA256: 6c8f3d2d3b48707541b88f32f22ef3f0f8a6b57d8523281e2b8d3cdb0ae9a023
• model.safetensors.index.json SHA256: 98efab455cf08dfbbbaaba6f570e1bf10bf927d2b4c3c453a59c2f6f0e3be92b

Container Provenance

The performance validation used a locally derived, pinned image:

vllm docker image: docker pull vllm/vllm-openai-rocm@sha256:72e90adf360ccffb231a098a15eadfb808719c029c6cd43f6386095064552b7e
AITer:          0.1.21.post2.dev2+g804648bcb
AITer commit:   804648bcbf87646ddef4b7add10da9c3b322567c
AITer wheel:    SHA256 30360d15819a9c8e3ec8e4ae665683abf4c5358d5be1f49a9aa2f7962c0032db
PyTorch:        2.12.0+rocm10.0.0

The exact image tag is a locally derived artifact rather than a public registry image. A byte-identical reproduction requires access to that image or reconstruction from the pinned base digest and the exact AITer wheel above.

Serving Configuration

export VLLM_ROCM_USE_AITER=1
export VLLM_USE_V2_MODEL_RUNNER=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=1
export VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN=8192
export AITER_USE_CK_MOE_SORTING=0
export AITER_USE_FLYDSL_MOE_SORTING=0
export VLLM_ROCM_AITER_MOE_DISPATCH_POLICY=0
vllm serve /models/DeepSeek-V4-Flash-DSpark \
  --served-model-name dsv4flash \
  --tensor-parallel-size 8 \
  --data-parallel-size 1 \
  --async-scheduling \
  --enable-prefix-caching \
  --max-model-len 32768 \
  --max-num-seqs 256 \
  --max-num-batched-tokens 16384 \
  --distributed-executor-backend mp \
  --gpu-memory-utilization 0.8 \
  --kv-cache-dtype fp8 \
  --trust-remote-code \
  --moe-backend aiter \
  --tokenizer-mode deepseek_v4 \
  --reasoning-parser deepseek_v4 \
  --compilation-config '{"mode":3,"cudagraph_mode":"FULL_AND_PIECEWISE","max_cudagraph_capture_size":2048}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":7,"enable_adaptive_verification":true}'

Throughput Methodology

The benchmark used a frozen 880-request SPEED-Bench Throughput-1K request set:

Temperature: 1.0
Top-p: 1.0
Maximum output length: 2048
Concurrency order: 256, 128, 64, 32, 16, 8, 4

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

[larryli2-amd]

Hi @dllehr-amd
Please review this PR. Many Thanks.