[AMD] DeepSeek-V4.1-Flash on MI350X (gfx950)

[kevin-mii]

Enable DeepSeek-V4.1 on AMD MI350X (gfx950), based on dsv4.1 at 099615b20e.

Changes

• HIP JIT kernels and Triton implementations for mHC, FP4 indexing, routing, quantization, and gfx950 dense GEMMs.
• Fuse query RoPE into padded attention output on HIP decode and gfx950 large prefill, preserving AMD BF16 rounding.
• HIP attention metadata, sorted top-k selection, V4.1 FP8/FP4 KV stores, low-ratio decode/verification, and multistream preparation.
• AMD model and AITER MoE dispatch, fused sorting, and shared-expert reduction.
• Select attention graph variants from the DP-group maximum sequence length so ranks replay matching graphs.
• Fix the shared-expert fusion guard for ROCm DeepSeek-V3 and apply the review-requested ROCm source-pointer casts.

Scope and external dependency

Changes outside python/sglang and test are limited to the ROCm Dockerfile, the AITER synchronization patch, and the DSV4.1 EP4 tuning CSV. The pinned AITER stage-1 kernel can leave LDS-DMA loads in flight across a barrier. The patch drains those loads and fences their scheduling; dropping it would remove a correctness fix.

The PR no longer changes global Docker environment defaults, pre-commit configuration, or documentation. The model-specific CSV was restored after measuring the performance effect of removing it. The existing AITER pin and cherry-picks come from the target branch. No workflow changes are introduced relative to the target.

AMD runtime switches remain in srt/environ.py: SGLANG_OPT_HIP_FUSE_SINKHORN_INTO_NORM, SGLANG_OPT_HIP_FUSED_MOE_REDUCE_ADD, and SGLANG_OPT_HIP_FUSED_DECODE_GLUE default on; SGLANG_OPT_HIP_ATTN_KV_SPLITS=4 (0 uses AITER's cost model).

Latest optimization port and rebase

Rebased all 16 existing commits onto 099615b20e; git range-diff reports each patch unchanged. This incorporates the shared deterministic-router FP32 output and graph token-count fixes.

On MI350X with ROCm 7.2.4:

• RoPE, model dispatch, and decode/prefill graph helpers: 33 passed, 10 subtests. RoPE checks exact BF16 output, untouched input/padding, changed-position graph replay, and prefill output buffers larger than 2 GiB.
• AMD regression directory: 105 passed, 4 skipped, 86 subtests, using the same environment below.
• Deterministic router: FP32 output preserves 16 close scores both with and without batch-invariant mode.
• RoPE+copy versus fused store, graph-replay microbenchmarks: decode 1.95–2.67x; prefill 4096 tokens 28.85 -> 23.59 us and 65536 tokens 516.31 -> 467.83 us. These are operator timings, not model throughput.
• RoPE tests also pass independently on all eight MI350X GPUs: 5 tests and 8 subtests per GPU. This checks each GPU, not collective communication.
• Changed-file pre-commit checks pass.
• Eight-rank DSpark collective regression currently aborts inside RCCL graph capture. It also aborts on pre-rebase 217338f4b5; a temporary test using the production PyNccl capture context also aborts. No test or production workaround is committed. The earlier successful eight-rank result below has not been reproduced in this run.

The other optimization gaps remain separate follow-ups: Engram prefetch, routed-MoE quantization overlap, mHC statistics scheduling, CP dense indexing, and broader bounded replay support. Full-model validation still needs weights and a launch/benchmark configuration.

Earlier validation (before the latest rebase)

On eight MI350X VF GPUs with PyTorch 2.11.0+rocm7.2 and ROCm 7.2.4:

• Targeted indexer, FP4 rounding, KV-store, mHC, quantization, MoE, metadata, and graph-policy suites passed.
• Eight-rank DSpark sharded selection: 36 passed and 36 NVLink-only skips per rank using the upstream test configuration. Experimental PyNccl capture configurations aborted inside RCCL; those test edits were reverted.
• Pipelined mHC/quantization tests required TRITON_HIP_USE_ASYNC_COPY=0; the default produced HIP error 401 before and after rebasing. No production workaround is added here.
• Cleanup AMD suite: 105 passed, 4 skipped, 86 subtests with TRITON_HIP_USE_ASYNC_COPY=0 SGLANG_USE_ROCM700A=1 ROCM_QUICK_REDUCE_QUANTIZATION=INT8.
• Cleanup model/graph regressions: 34 passed and 17 subtests with the restored ROCm defaults. An initial combined run aborted in AITER mHC; the separate rerun passed.
• Pre-commit passes for all PR-changed files, including test registration with GITHUB_BASE_REF=dsv4.1. The all-files run hit an unrelated Rust Clippy file-descriptor limit.
• Moved the 15 new MI35x regression files into test/registered/amd/dsv41/ to satisfy the target branch taxonomy without changing their CI runner assignments.
• Original and cleaned AITER patches apply to the pinned source and produce identical Python ASTs.

Not executed: full-model inference, accuracy, throughput, and context-parallel inference after the rebase (no weights or benchmark configuration available); a fresh Docker build; CUDA and ROCm 10 validation. The context-parallel fallback has request-mapping coverage only. The local AITER installation does not have the stage-1 patch applied, so the kernel regressions do not validate that race fix.

CSV restoration measurements

Using current AITER 0138f88b2 with the synchronization fix on both paths, five gfx950 graph-replay microbenchmarks measured 1.06-1.29x operator speedup from the CSV. At one token, default 48.37 us -> tuned 37.44 us; at 4096 tokens, 1411.56 us -> 1248.79 us. All 12 rows survive normal AITER config merging and do not collide with existing keys.

The upstream tuning draft includes the benchmark and full methodology. This is synthetic operator timing, not full-model accuracy or throughput validation. Tuned/default maximum absolute differences were 0.0273-0.0410; the numerical tolerances still need independent validation. The synchronization tests pass both patched and unpatched current main, so they exercise the path without reproducing the reported intermittent race.

A second comparison on the installed SGLang-pinned AITER 4ad998328 also favored the CSV: M=1: 43.05 -> 37.17 us (1.16x), M=32: 291.22 -> 262.54 us (1.11x), M=128: 343.24 -> 316.97 us (1.08x), M=512: 419.93 -> 399.23 us (1.05x), M=4096: 1359.34 -> 1239.05 us (1.10x). This installation does not include the synchronization patch; the same installation was used for both sides. These are a8w4 measurements with AITER_BF16_FP8_MOE_BOUND=0, not the default small-batch bf16 route.

Historical results before rebase and cleanup

The original PR reported the following results on four MI350X GPUs. These have not been reproduced on the current revision, and depended on the earlier deployment setup, including global Docker defaults that are no longer changed by this PR.

DSpark, blog protocol (same method as the DeepSeek-V4.1-Flash kernel post: its prompt.json of 4096 random ids at seed 42, 1024 output tokens with ignore_eos, SGLANG_SIMULATE_ACC_LEN=5.5 / match-expected, SGLANG_RAGGED_VERIFY_MODE=static, block size 5, its benchmark.py with 6 runs per launch after a discarded warm-up, throughput from the first to the last streamed event):

BS=1, random 4k/1k	DSpark off, EP4	DSpark on, EP4	DSpark on, TP4 MoE
output tokens/s (median)	147.11	548.15 (2 launches, 12 runs, 538.75-553.43)	555.74 (1 launch, 549.88-559.98)
measured accept length	-	5.505	5.505

For reference the post reports 223.50 / 853.49 / 873.63 on 4x GB300 with the same protocol.

High-throughput recipe (--disable-radix-cache, real text): decode bs1 / bs8 / bs32 127 / 851 / 2382 -> 149 / 937 / 2694 tok/s; gsm8k 5-shot (1319) 0.898-0.916 across gated trees, in the 0.90 band. Prefill first-token logprob spread 0 across batch sizes 1-32 at 64/1024/2048 tokens; greedy decode identical across runs and page layouts (decode top-k is ordered by position, not KV slot).

Follow-ups

Upstream drafts: AITER synchronization fix #5561 and EP4 tuning #5562. Remove the downstream patch/CSV when the AITER pin includes the respective fixes. Deployment defaults remain a separate concern. Encoder SWA bounded replay and Engram KV prefetch remain CUDA-only. The previously reported DP prefill row asymmetry and eager idle-batch Engram assertion remain outside this cleanup.

---

CI States

Latest PR Test (Base): ❌ Run #35039152924
Latest PR Test (Extra): ❌ Run #35039152740
Latest PR Test (AMD ROCm 10): ❌ Run #35039152911