[ROCm][DSv4] Enable confidence-based adaptive scheduler for DSpark

[larryli2-amd]

Inherits #52362

This work inherits and rebases the ROCm DSpark adaptive-verification implementation from #52362. It is not an independent implementation of the feature.

This branch should only be submitted after explicit coordination with the author of #52362, or after that pull request is closed. If submitted as a replacement, it supersedes #52362 and preserves its original design and attribution.

Summary

Enable confidence-scheduled DSpark adaptive verification for DeepSeek-V4 on ROCm.

The change:

• constructs and loads the AMD DSpark confidence head;
• preserves fail-closed behavior when confidence-head checkpoint weights are absent;
• enables Adaptive Verification CUDA-graph support in the ROCm sparse MLA and sparse SWA metadata builders;
• enables the ROCm flattened device query-length path required by uneven adaptive verification budgets;
• replaces an optimization-removable assertion in compute_confidence() with an explicit RuntimeError;
• adds focused AMD confidence-head and ROCm adaptive-indexer regression tests.

Source Revision

• PR branch: larryli2-amd:enable_adaptive-dspark-rocm
• PR head: 945abd627010c926cbf288f4b374e8b63d7443bf
• Base used when preparing the branch: bbbd0a02c9d1bec965ee818af29a1319baefa69c
• Changed files: 3 runtime files and 2 focused test files
• Total diff: 413 insertions, 6 deletions

Focused Validation

tests/models/test_deepseek_v4_dspark_rocm.py
tests/v1/attention/test_deepseek_v4_rocm_adaptive.py

Result: 10 passed.
The tests cover confidence-head mapping/loading/fail-closed behavior, the explicit missing-head error, ROCm adaptive CUDA-graph support, flattened query-length selection, the single-request uniform path, and stable buffer replay under changed allocations.
Changed-file lint and formatting checks also passed.

Serving Smoke Validation

A live adaptive_k7 ROCm serving smoke test completed successfully:

• /health returned HTTP 200;
• Application startup complete was observed;
• a single request completed with 48 generated tokens;
• no HSA_STATUS_ERROR_MEMORY_FAULT, worker death, queue hang, or traceback was observed;
• the DSpark draft model loaded 97 parameters;
• the ROCm indexer logged use_flattening=True;
• breakable CUDA graphs were enabled.
  Observed speculative-decoding counters during the smoke request:

num_drafts_total:        18
num_draft_tokens_total: 126  # 18 × K=7
num_accepted_tokens:     29

This demonstrates that the DSpark speculative path executed. It is not a throughput, latency, quality, or adaptive-vs-fixed performance measurement.

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:

Image tag:      dspark-av-sep11:52362-ad170197-aiter804648bc
Image ID:       sha256:49067946ae13f4047605ce9edca3d9705903b602f43f8c1fcd0ecdb8927511d5
Base digest:    sha256:72e90adf360ccffb231a098a15eadfb808719c029c6cd43f6386095064552b7e
vLLM base:      2a02f6efe
Original patch: ad170197f40a3e48eddbb92bd195a2e5db684267
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

Each cell used an independent cold server boot. Warmup was:

1. 64 requests at concurrency 32 with output length 256;
2. 64 requests at the cell concurrency with output length 256;
3. 880 measured requests at the cell concurrency.
   A result is valid only when client and server generated-token counts match, no preemptions occur, and no GPU fault is observed.

Resutls

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.

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

After discussion with @tuukkjs , this PR will be merged into #52362 and this PR will be closed.