[ROCm] Add DeepSeek-V4 DSpark speculative decoding support for AMD GPU(MI350X / MI355X, gfx950)

[haic0]

Summary

This PR is now a small follow-up on top of the DSpark implementation that already landed in main via #30261. It keeps the ROCm/AMD DeepSeek V4 DSpark fixes that were still missing after rebasing the branch onto current main.

• Fix DeepSeek V4 HIP target-verify metadata to use the active DSpark draft token count instead of the static configured speculative token count.
• Prepare DSpark KV-injection prefill position metadata before the target prefill forward so the correct cache locations and hidden states are used.
• Stage DSpark prefill extend_lens / prefix_lens through pinned CPU tensors before non-blocking transfer to the AMD GPU device, avoiding extra synchronization in the prefill path.

The earlier large DSpark implementation commit was dropped from this PR because the equivalent newer implementation is already present on main.

Test plan

• python3 -m compileall -q python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
• Docker environment used for validation: sglang-rocm-v0514 with AMD MI350X/ROCm.
• Verified DeepSeek V4 Flash DSpark server startup and /health readiness.
• Ran ShareGPT DSpark benchmark successfully for 10 requests and compared against baseline; DSpark showed about 1.57x output throughput improvement in the validated run.

Reproduce

DSpark Enabled

unset SGLANG_HACK_FLASHMLA_BACKEND
export SGLANG_DEFAULT_THINKING=1
export SGLANG_DSV4_REASONING_EFFORT=max
export SGLANG_USE_ROCM700A=0
export AITER_BF16_FP8_MOE_BOUND=0

sglang serve \
  --trust-remote-code \
  --model-path deepseek-ai/DeepSeek-V4-Flash-DSpark \
  --speculative-algorithm DSPARK \
  --tp 8 \
  --attention-backend dsv4 \
  --page-size 256 \
  --mem-fraction-static 0.8 \
  --disable-radix-cache \
  --swa-full-tokens-ratio 0.05 \
  --speculative-moe-runner-backend triton \
  --disable-shared-experts-fusion \
  --kv-cache-dtype fp8_e4m3 \
  --cuda-graph-backend-decode=disabled \
  --max-running-requests 128 \
  --host 0.0.0.0 \
  --port 8888

Baseline

export SGLANG_DEFAULT_THINKING=1
export SGLANG_DSV4_REASONING_EFFORT=max
export SGLANG_USE_ROCM700A=0
export SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton
export AITER_BF16_FP8_MOE_BOUND=0

sglang serve \
  --trust-remote-code \
  --model-path deepseek-ai/DeepSeek-V4-Flash \
  --tp 8 \
  --attention-backend dsv4 \
  --page-size 256 \
  --mem-fraction-static 0.80 \
  --swa-full-tokens-ratio 0.05 \
  --disable-shared-experts-fusion \
  --kv-cache-dtype fp8_e4m3 \
  --cuda-graph-backend-decode=disabled \
  --max-running-requests 128 \
  --host 0.0.0.0 \
  --port 8888

Benchmark Command

python3 -m sglang.bench_serving \
  --backend sglang \
  --base-url http://0.0.0.0:8888 \
  --dataset-name sharegpt \
  --dataset-path /tmp/ShareGPT_V3_unfiltered_cleaned_split.json \
  --num-prompts 10 \
  --request-rate inf \
  --max-concurrency 8 \
  --sharegpt-output-len 128 \
  --warmup-requests 5 \
  --disable-ignore-eos

DSpark Results

Successful requests: 10
Benchmark duration: 19.84 s
Request throughput: 0.50 req/s
Output token throughput: 64.53 tok/s
Total token throughput: 218.64 tok/s
Mean TTFT: 1559.49 ms
Mean TPOT: 70.84 ms
Accept length: 2.02

Baseline Results

Successful requests: 10
Benchmark duration: 31.21 s
Request throughput: 0.32 req/s
Output token throughput: 41.01 tok/s
Total token throughput: 138.94 tok/s
Mean TTFT: 2118.59 ms
Mean TPOT: 107.23 ms

---

CI States

Latest PR Test (Base): ❌ Run #29302344043
Latest PR Test (Extra): ❌ Run #29302343884

[gemini-code-assist]

Code Review

This pull request introduces 'DSpark', a new speculative decoding implementation for DeepSeek-V4 models, including support for ragged verification, online SPS profiling, and various performance-optimized kernels. My review identified two critical performance and safety issues: a synchronous device-to-host transfer bottleneck in the attention backend and an out-of-bounds memory access risk in the Triton kernel for ROCm/AMD environments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!