[AMD] Fix deepseek-v4 mtp accept length issue

[1am9trash]

Co-authored-by: @bingxche
Waiting CI checking by Bingxche help: https://github.com/sgl-project/sglang/actions/runs/27679583497/job/81863337423

Motivation

On DeepSeek-V4 with MTP enabled, the unified_kv_triton attention backend reaches an average accept length of only ~2.17, while the triton backend reaches ~3.04. Accuracy is unaffected, but the lower accept length significantly hurts decode throughput.

Per-position acceptance rate shows the draft chain collapsing from token 2 onward on unified_kv_triton.

backend	token 0	token 1	token 2	token 3	accept length
unified_kv_triton (before)	1.00	0.894	0.223	0.054	2.17
triton	1.00	0.999	0.710	0.333	3.04

Modifications

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

• The SWA ring write target (swa_loc) is precomputed once per forward. This path never accounted for multi-step draft decode: across draft steps the positions advance, but the cached swa_loc does not, so every draft step writes its KV to the same ring slot. Later steps overwrite the KV that earlier draft tokens produced, breaking the draft chain at depth >= 2 and collapsing accept length.
• In get_unified_swa_loc, detect multi-step draft decode and recompute the ring slot from the live per-step positions. It only triggers on draft layers, so the overhead is negligible.

test/registered/amd/test_deepseek_v4_pro_fp4_mtp.py

• Also add a DeepSeek-V4-Pro FP4 + MTP accuracy/accept-length CI test to the nightly MI355 ROCm 7.2 suite.

Accuracy Tests

Local run new CI task command:

SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton python test/registered/amd/test_deepseek_v4_pro_fp4_mtp.py

Both accept length and accuracy check passed.

• TestDeepseekV4ProFp4MTP.test_a_gsm8k

  metrics={'accuracy': np.float64(0.9537528430629265), 'invalid': np.float64(0.0), 'latency': 91.38883687090129, 'output_throughput': 1356.5551794376104}
  avg_spec_accept_length=2.93023201257419
  

• TestDeepseekV4ProFp4MTP.test_b_bs_1_speed

  +-------------+--------+------------+-----------------+
  | Latency (s) | Tokens | Acc Length | Speed (token/s) |
  +-------------+--------+------------+-----------------+
  |    7.949    |  1030  |   3.012    |     129.57      |
  +-------------+--------+------------+-----------------+
  acc_length=3.01 speed=129.57
  

Speed Tests and Profiling

conc16, 8k1k:

metric	before PR	after PR
Accept length	2.18	3.08
Median ITL (ms)	17.36	11.51
Total token throughput (tok/s)	6355.88	7324.37

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #27682217928
Latest PR Test (Extra): ❌ Run #27682217663

[gemini-code-assist]

Code Review

This pull request modifies the DeepSeek-V4 attention backend (deepseek_v4_backend_hip_radix.py) to bypass cached swa_loc and recompute it from live per-step positions during multi-step draft decoding, preventing ring slot reuse issues. It also introduces a new integration test suite (test_deepseek_v4_pro_fp4_mtp.py) for validating DeepSeek-V4-Pro FP4 with MTP on AMD MI35x GPUs. There are no review comments, and I have no feedback to provide.

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.

[bingxche]

testing: https://github.com/sgl-project/sglang/actions/runs/27679583497/job/81863337423

[1am9trash]

/tag-and-rerun-ci