[AMD][DSV4] feat: enable DSpark with fp8 unified_kv on gfx950

[amd-danli103]

Motivation

Follow-up to #37413. Incremental is only 1 commit, 6 files.

In #37413 we added fp8 two-pool unified_kv on gfx950 (SGLANG_DSV4_UNIFIED_KV_FP8, 640 B/row, ~1.50× KV capacity) and measured it on MTP. We left DSpark as the next SPD path.

--speculative-algorithm DSPARK plus the env builds a second DeepSeekV4TokenToKVPool for the draft worker. That constructor used process-wide is_unified_kv_fp8(), and CUDA-graph capture died on kv.dtype == unified_kv.dtype. Same env, DSpark now comes up. Target decode / prefill / target_verify stay on the two-pool. MTP is unchanged.

Modifications

AMD / HIP unified_kv_triton only. Default off. No change to is_unified_kv_fp8(), DSparkAttention, or the scatter kernel.

• DSpark under the fp8 env. _build_dsv4_kv_pool selects the layout per pool: unified_fp8=is_unified_kv_fp8() and not (is_draft_worker and spec_algorithm.is_dspark()). Target follows the env (640 B two-pool). EAGLE/MTP draft still follows the env.
• Sizing. With DSpark+fp8, _fixed_swa_bytes budgets the target fp8 ring plus the draft SWA ring at 1024 B/row (one layer for shipped drafts), not (T+1)/T * 640. bytes_per_full_token still multiplies (T+1)/T — conservative; DSpark draft has no C4/C128. MTP inflation is the same as [AMD][DSV4] feat: enable fp8 two-pool unified_kv on gfx950 #37413.

What this buys: KV capacity

Same CLI pair, Flash-DSpark, TP=8 DP=8 + dp-attention, gfx950. Differs only in SGLANG_DSV4_UNIFIED_KV_FP8. Same 1.50× as #37413, on DSpark.

arm	unified_fp8	bytes_per_full_token	full_token	ratio
fp8	True	4249.58	53,978,624	1.500×
bf16	False	6373.86	35,985,664	1

Accuracy Tests

GSM8K, all 1319 questions. Flash-DSpark, TP=8 DP=8, SGLANG_RAGGED_VERIFY_MODE=static, radix cache off. Pair differs only in the fp8 env. Capture including draft verify.

arm	n=1319
fp8	0.936
bf16	0.938

Speed Tests and Profiling

Same protocol as #37413: MI355X, Flash-DSpark TP=8 DP=8, radix off, bench_serving random. Pair differs only in SGLANG_DSV4_UNIFIED_KV_FP8. ttt ratio is fp8/bf16. DSpark accept 3.35–4.20, matched across arms.

ISL/OSL: 8k / 1k

conc	bf16 ttt	fp8 ttt	ttt ratio	bf16 TTFT	fp8 TTFT	bf16 ITL	fp8 ITL
8	7942	8513	1.072	1294	920	7.15	6.93
16	12081	11923	0.987	1128	1149	10.08	10.27
64	27140	25444	0.938	4681	4951	11.75	12.53

ISL/OSL: 68k / 200

conc	bf16 ttt	fp8 ttt	ttt ratio	bf16 TTFT	fp8 TTFT	bf16 ITL	fp8 ITL
4	14984	13849	0.924	7087	7808	55.22	59.00
16	73428	69610	0.948	10494	11073	20.97	22.20
48	77576	76542	0.987	24028	24331	90.07	91.38

No new decode kernel. A leftover decode gap is still the aiter reader, same as #37413. The 1.50× is KV capacity.

Reproducing this

MI355X, SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton. fp8 arm adds SGLANG_DSV4_UNIFIED_KV_FP8=1.
1. Unit tests : 160 passed.

HIP_VISIBLE_DEVICES=0 python3 -m pytest \
  test/registered/e2e/dsv4/test_dsv4_unified_fp8_*.py \
  test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py \
  test/registered/unit/model_executor/test_pool_configurator.py \
  test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py \
  -q

2. DSpark server — Flash-DSpark, TP=8 DP=8. This is the CLI the GSM8K and serving numbers were taken on. The bf16 arm is the same command without the last export.

export SGLANG_DEFAULT_THINKING=1
export SGLANG_DSV4_REASONING_EFFORT=high
export SGLANG_USE_ROCM700A=0
export SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton
export AITER_BF16_FP8_MOE_BOUND=0
export SGLANG_OPT_USE_AITER_BATCHED_GEMM=1
export SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS=1
export SGLANG_TIMEOUT_KEEP_ALIVE=900
export SGLANG_RAGGED_VERIFY_MODE=static
export SGLANG_DSV4_FP4_EXPERTS=true
export SGLANG_DSV4_UNIFIED_KV_FP8=1      # bf16 arm omits this (or sets 0)

python3 -m sglang.launch_server \
  --model-path <DeepSeek-V4-Flash-DSpark> \
  --served-model-name deepseek-ai/DeepSeek-V4-Flash-DSpark \
  --host 0.0.0.0 --port 8891 --trust-remote-code \
  --tensor-parallel-size 8 --dp-size 8 \
  --enable-dp-attention --enable-dp-lm-head \
  --attention-backend dsv4 --page-size 256 \
  --swa-full-tokens-ratio 0.10 --kv-cache-dtype fp8_e4m3 \
  --disable-shared-experts-fusion \
  --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 \
  --chunked-prefill-size 8192 --mem-fraction-static 0.89 \
  --max-running-requests 64 --cuda-graph-max-bs-decode 64 \
  --speculative-algorithm DSPARK --speculative-dspark-block-size 5 \
  --disable-radix-cache --watchdog-timeout 3600

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 #35463746142
Latest PR Test (Extra): ❌ Run #35463746035
Latest PR Test (AMD ROCm 10): ❌ Run #35463746148

[amd-danli103]

hi @1am9trash , could you pls help to add a run-CI tag for this pr which supports DSpark + fp8 unified_kv? Sincerely thanks!

[HaiShaw]

/tag-and-rerun-ci

[HaiShaw]

@ch-wan @hnyls2002 @ispobock @1am9trash please help review.

[HaiShaw]

LGTM

[HaiShaw]

@amd-danli103 @1am9trash
https://github.com/sgl-project/sglang/pull/38901/conflicts

[amd-danli103]

Re-validated on MI355X (gfx950, 8 GPUs, ROCm 7.2) after merging main at #40205. That fix is required here: without it every speculative arm dies in AssertionError: No indexer pool for compression ratio 4 before the server binds.

Model is DeepSeek-V4-Pro-0813, which bundles the DSpark draft, so no --speculative-draft-model-path. Feature gate is SGLANG_DSV4_UNIFIED_KV_FP8=1 with SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton.

Common base for all four arms:

python3 -m sglang.launch_server \
  --model-path <DeepSeek-V4-Pro-0813> --trust-remote-code \
  --tensor-parallel-size 8 --attention-backend dsv4 \
  --page-size 256 --swa-full-tokens-ratio 0.10 --kv-cache-dtype fp8_e4m3 \
  --disable-shared-experts-fusion --chunked-prefill-size 8192 \
  --mem-fraction-static 0.89 --max-running-requests 64 \
  --cuda-graph-max-bs-decode 64 \
  --speculative-algorithm DSPARK --speculative-dspark-block-size 5

Arms differ only in the tail: DP adds --dp-size 8 --enable-dp-attention --enable-dp-lm-head; hicache adds --enable-hierarchical-cache --hicache-ratio 2 --hicache-mem-layout page_first_direct --hicache-io-backend direct --hicache-write-policy write_through; the non-hicache arms run --disable-radix-cache.

GSM8K, all 1319 questions, 5-shot, parallel 64:

arm	dp	hicache	accuracy	invalid	latency	output tok/s
fp8	8	off	0.9348	0.000	133.5 s	862
fp8	-	off	0.9469	0.000	71.0 s	1627
fp8 + hicache	8	on	0.9378	0.000	134.5 s	854
fp8 + hicache	-	on	0.9393	0.000	60.2 s	1922

All four came up, answered a completion request, and finished the sweep with zero invalid outputs. For reference, the no-speculation baseline measured while validating #40205 (same model and server flags, speculation off) was 0.9340.

Unit tests on the same tree, 135 passed plus 32 subtests.

[HaiShaw]

https://github.com/sgl-project/sglang/actions/runs/35463746035/job/106032320395?pr=38901 is fixed in #40354

[HaiShaw]

https://github.com/sgl-project/sglang/actions/runs/35463746035/job/106032320423?pr=38901 is being addressed in #40418

[HaiShaw]

PR Test Base, and other PR Test Extra look okay.