[ROCm][MI35X] Enable GLM-5.2-MXFP4 MTP speculative decoding

[long10024070]

Summary

This PR finishes the MI35X runtime enablement for exact-head 2f009ca40 on top of 1a3bea77f (current upstream main, rebased 2026-08-02).

The stack keeps the reviewed GLM-5.2-MXFP4 TP4 workload intact while tightening three separate areas:

• QuickReduce VMM on gfx950
• DSA and empty-input handling
• Idle-DP and capture-aware distributed routing

Changes

QuickReduce VMM

• Add HIP VMM allocation, fd exchange, consensus, synchronization, and teardown for gfx950 QuickReduce.
• Keep the VMM backend enabled for the TP4 path instead of falling back to a slower collective.

DSA and empty-input handling

• Preserve padded DSA graph metadata across eager and CUDA-graph paths.
• Skip empty GPU all-reduce and empty AITER paged-MQA paths cleanly.
• Keep the DSA logits and metadata contract stable when the batch is empty or graph-padded.

Idle-DP and capture-aware routing

• Handle idle DP ranks without zero-work launches.
• Keep eager DP logits metadata consistent when DP ranks are idle.
• Route forced DP gather through the capture-aware all-reduce path so CUDA graph capture stays valid.
• Keep the idle-DP EAGLE draft-sampling path aligned with the finalized verify decision broadcast.

Temporary #31478 dependency

• This stack temporarily carries the two upstream dependency commits from #31478, preserving zhoaa attribution: upstream 61c2af4d6cc5fac633b414cfd335e9674d49c183 / d7fe5fe79f871d6cd28ff7b0d089bd08a4b642f0, represented here by patch-equivalent commits 7d7ff3995c / 5b36aac13c after rebase.
• These commits only broadcast the finalized EAGLE verify decision across TP ranks and test that behavior. They are separate from this PR's QuickReduce VMM, DSA metadata, empty-input, and idle-DP changes.
• If #31478 merges first, this branch will be rebased onto main and the duplicate dependency commits will be dropped.

Validation

• Exact-head no-GPU focused tests: 18 passed in 2.81s on 2f009ca40 (rebased onto upstream main 1a3bea77f).
• Exact-head CPU static gates: git diff --check and full pre-commit passed on the same head, with a clean worktree before and after.
• Exact-head TP4 bundle static validation: passed for the prepared ready_post_build_exact_id bundle on the same source/base contract.
• Exact-head MI350X TP4 smoke passed /v1/models, deterministic generation, structured tool calling, and tool-result round trip.
• Exact-head sustained real-weight gate (exact_head_gpu) used amd/GLM-5.2-MXFP4, EAGLE/NEXTN 5/1/6, KV-cache FP8, AITER all-reduce fusion, DSA, and QuickReduce INT4 VMM:
  • C16: 1600/1600 requests, 926.04 output tok/s, TTFT p50 730.73 ms / p95 3448.60 ms, TPOT p50 15.85 ms / p95 17.13 ms, accept length 5.95, accept rate 0.99, no stall or timeout.
  • C32: 3200/3200 requests, 1023.75 output tok/s, TTFT p50 750.07 ms / p95 6210.90 ms, TPOT p50 29.43 ms / p95 31.19 ms, accept length 5.95, accept rate 0.99, no stall or timeout.
• TP0-TP3 reported matching speculative configuration and acceptance metrics at both load points. QuickReduce VMM IPC was live on all four ranks (QuickAllReduce selected VMM IPC for rank 0-3/4 on gfx950). Zero server-error markers and no collective stall throughout.

Notes

• The reviewed f172 workload contract remains unchanged: real amd/GLM-5.2-MXFP4, TP4, EAGLE/NEXTN 5/1/6, AITER all-reduce fusion, QuickReduce VMM INT4, TileLang DSA prefill/decode, sgl-kernel DSA topk, --max-running-requests 64, fixed 1024 outputs, and no enumerated CUDA-graph batch-size list.
• The proven scalar CUDA-graph cap and prefill toggle are preserved.
• No reviewer or label request is included here.

---

CI States

Latest PR Test (Base): ❌ Run #30756510319
Latest PR Test (Extra): ❌ Run #30756510244

Rebase History

• 2026-08-02 (2nd): Rebased from e5b6ea5385 (base 00a219f6c9) onto 2f009ca40 (base 1a3bea77f, current upstream main). 20 upstream commits integrated. 2 file overlaps (dsa_backend.py CP-v2 change, forward_batch_info.py hidden-states change) — both in different code regions, zero conflicts. git range-diff all 12 commits patch-equivalent (=). Same 34 files changed. CPU gates pass (diff --check, py_compile 34/34). The prior exact-head GPU validation (sustained C16/C32, topology gates) was run on e5b6ea5385 which is patch-equivalent; the two test-fixup commits are unchanged.

• 2026-08-02 (1st): Rebased from da838858c7 (base 075bd97952) onto e5b6ea5385 (base 00a219f6c, current upstream main +1). Patch-equivalent to pre-rebase (git range-diff all =). Same 34 files changed. Two rebase-fixup commits added for test compatibility with upstream's RFC [RFC] Introduce a unified sglang.kernels namespace for kernel organization and dispatch #29630 finale ([Kernel] RFC #29630 finale: retire sglang.jit_kernel into sglang.kernels #32072) and runtime_context refactor:

  • ee96d551a — fix stale sglang.jit_kernel import → sglang.kernels.ops.attention.dsa.paged_mqa_logits (upstream removed jit_kernel package).
  • e5b6ea538 — fix stale eagle_info.get_server_args mock → get_spec at both import sites (upstream moved speculative_use_rejection_sampling from server_args to runtime_context.get_spec()).

• 2026-07-21: Rebased from 608bf3b756 (base 3d82dacd58) onto da838858c7 (base 075bd97952). Clean rebase, no conflicts. Patch ID unchanged: d205669256d2b08028d98253b85ac950a370b128. Same 34 files changed. All CPU gates pass (diff --check, py_compile 28/28).

Validation (Exact Head 2f009ca40)

CPU Gates (exact_head_cpu)

• git diff --check 00a219f6c..HEAD: pass
• pre-commit run --from-ref 00a219f6c --to-ref HEAD: all checks pass (ruff, black, isort, clang-format, codespell, registered-test CI registries)
• git range-diff 00a219f6c da838858c7 HEAD: all 10 original commits patch-equivalent (=), 2 test-fixup commits added.

Exact-Head Focused Tests (exact_head_cpu, no-GPU container)

• Image: sglang-pr31683-2f009ca40:rocm720-gfx950-20260802 (ID 6bb330205407), base lmsysorg/sglang-rocm:v0.5.16-rocm720-mi35x-20260731 pinned by digest e93e47ec0110, sgl-kernel rebuilt from exact-head source (QuickReduce VMM HIP sources changed).
• Container: --network none, no /dev/kfd or /dev/dri, cleared CUDA/HIP/ROCR_VISIBLE_DEVICES, pytest cache disabled.
• Provenance: sglang.__file__ and quick_all_reduce_vmm.__file__ resolve to /opt/sglang-source/python/; SGLANG_SOURCE_COMMIT label = 2f009ca40; exchange_vmm_fds + _send_fd/_recv_fd + _draft_extend_for_prefill present.
• Result: 53 passed, 15 subtests passed in 15.75s (exit 0). Covers eagle topk1 fastpath, eagle verify TP broadcast, forward-batch zero-tokens, DSA graph metadata, communicator zero-scatter, logits DP-attention metadata, QuickReduce VMM fd-exchange, AITER paged MQA logits padding trim.

TP4 Real-Weight Sustained Gate (Exact Head 2f009ca40, exact_head_gpu, mi350-huy GPU4-7)

• Image: sglang-pr31683-2f009ca40:rocm720-gfx950-20260802 (ID 6bb330205407), sgl-kernel rebuilt from exact-head source with AMDGPU_TARGET=gfx950.
• Model: amd/GLM-5.2-MXFP4 (real weights), TP4, Quark MXFP4, KV fp8_e4m3, AITER allreduce fusion, EAGLE/NEXTN 5/1/6, speculative-attention-mode decode, accept 1.0/1.0, draft-quant unquant, DSA tilelang prefill/decode, sgl-kernel topk, --cuda-graph-max-bs 64, --max-running-requests 64, --chunked-prefill-size 65536, --max-prefill-tokens 65536, --disable-radix-cache, --disable-overlap-schedule, --disable-prefill-cuda-graph, ROCM_QUICK_REDUCE_QUANTIZATION=INT4.
• QuickReduce VMM IPC: Live on all 4 ranks — QuickAllReduce selected VMM IPC for rank 0/4, 1/4, 2/4, 3/4 on gfx950:sramecc+:xnack- (uncached=True).
• C16 (1600 req, ISL 8192, OSL 1024, concurrency 16): 926.04 output tok/s, 7408.33 input tok/s, 1600/1600 completed, 0 errors, TTFT p50 730.73 ms / p95 3448.60 ms, TPOT p50 15.85 ms / p95 17.13 ms, accept length 5.95, accept rate 0.99.
• C32 (3200 req, ISL 8192, OSL 1024, concurrency 32): 1023.75 output tok/s, 8190.04 input tok/s, 3200/3200 completed, 0 errors, TTFT p50 750.07 ms / p95 6210.90 ms, TPOT p50 29.43 ms / p95 31.19 ms, accept length 5.95, accept rate 0.99.

Prior TP4 Sustained + Topology Gates (Exact Head da838858c7, equivalent_patch_runtime — cross-validated)

The sustained and topology gates below were run on the prior exact head da838858c7 (same 10 original commits, patch-equivalent to e5b6ea5385 and 2f009ca40). The exact-head sustained gate above (2f009ca40, mi350-huy GPU4-7) confirms these results. The topology gates below remain equivalent_patch_runtime evidence; exact-head topology re-validation is pending GPU availability.

CPU Gates

• git diff --check origin/main...HEAD: pass
• py_compile 28/28 changed Python files: pass
• Patch ID: d205669256d2b08028d98253b85ac950a370b128 (unchanged after rebase)

TP4 Real-Weight Sustained Gate (mi350-vu GPU0-3)

• Image: sglang-pr31683-da838858c7:final-20260721 (built from base 5ad76f + new source overlay)
• Config: TP4, Quark MXFP4, KV fp8_e4m3, AITER allreduce fusion, EAGLE 5/1/6, speculative-attention-mode prefill, DSA tilelang/sgl-kernel, --mem-fraction-static 0.82, --cuda-graph-max-bs 64, --max-running-requests 64, ROCM_QUICK_REDUCE_QUANTIZATION=INT4, dummy weights
• QuickReduce VMM IPC: Live on TP0-TP3 (all ranks selected VMM IPC on gfx950)
• C16 (1600 req, concurrency 16): 1076.12 tok/s, concurrency 15.89, accept length 1.0, no stall/timeouts
• C32 (3200 req, concurrency 32): 1519.56 tok/s, concurrency 31.83, accept length 1.0, no stall/timeouts
• Cleanup: Container removed, GPU0-3 released (285MB driver overhead only), port 30445 free

Topology Gates (Exact Head da838858c7, mi350-vu)

All topology gates use dummy weights, EAGLE 5/1/6, KV fp8_e4m3, AITER allreduce fusion, QuickReduce VMM INT4.

Topology	Config	QuickReduce VMM	Dispatch	Status
TP4-EP4	--tp-size 4 --ep-size 4 --moe-a2a-backend mori	VMM IPC live TP0-EP0 to TP3-EP3	✅	PASS
DP2×attention-TP2	--tp-size 4 --dp-size 2 --enable-dp-attention	VMM IPC live world TP4 + attention-TP2	✅ DP0/DP1 dispatching	PASS
DP2×attention-TP4	--tp-size 8 --dp-size 2 --enable-dp-attention	VMM IPC live 8 ranks + 2× attention-TP4	✅ DP0/DP1 dispatching	PASS
DP4-EP4 (DP2×attn-TP2 + EP4 Mori)	--tp-size 4 --dp-size 2 --enable-dp-attention --ep-size 4 --moe-a2a-backend mori	VMM IPC live world TP4 + attention-TP2	✅ attention_tp + tp dispatching	PASS

All gates: /v1/models 200, generation OK, no stall/timeouts, QuickReduce VMM IPC selected on gfx950.

[gemini-code-assist]

Code Review

This pull request introduces HIP Virtual Memory Management (VMM) IPC support for QuickAllReduce on ROCm (specifically gfx950), alongside several robustness fixes such as zero-sized tensor handling to avoid invalid kernel launches, and synchronized EAGLE verify decisions across TP ranks to prevent deadlocks. Feedback points out a critical bug where a global rank is incorrectly used as the source in dist.broadcast instead of a group-relative rank, which would cause hangs in subgroup configurations. Additionally, it is recommended to add an early return when q_offset is 0 in the paged MQA logits kernel to avoid empty allocations, and to update the corresponding unit tests to reflect the group-relative broadcast source rank.

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.

[long10024070]

Rebased onto current main (00a219f6c); new head e5b6ea5385. The rebase is patch-equivalent to the prior da838858c7 (git range-diff all =), with two added test-fixup commits for upstream API drift:

• ee96d551a — sglang.jit_kernel was removed by RFC [RFC] Introduce a unified sglang.kernels namespace for kernel organization and dispatch #29630 finale ([Kernel] RFC #29630 finale: retire sglang.jit_kernel into sglang.kernels #32072); test import updated to sglang.kernels.ops.attention.dsa.paged_mqa_logits.
• e5b6ea538 — eagle_info.get_server_args was replaced by runtime_context.get_spec(); test mock updated at both import sites.

Exact-head CPU gates (e5b6ea5385): git diff --check clean, pre-commit run --from-ref 00a219f6c --to-ref HEAD all pass.

Exact-head focused tests (exact_head_cpu, no-device container, sgl-kernel rebuilt from exact-head source on lmsysorg/sglang-rocm:v0.5.16-rocm720-mi35x-20260731): 53 passed, 15 subtests passed (eagle topk1 fastpath, eagle verify TP broadcast, forward-batch zero-tokens, DSA graph metadata, communicator zero-scatter, logits DP-attention metadata, QuickReduce VMM fd-exchange, AITER paged MQA logits).

The prior TP4 sustained + topology gates (C16/C32 real-weight, 4 topology variants) on da838858c7 remain equivalent_patch_runtime evidence; exact-head TP4 re-validation on e5b6ea5385 is in progress on MI350X.

CI currently fails only on the missing run-ci label. Could a maintainer apply the run-ci label so the workflows can execute? Thank you.

[long10024070]

Closing this PR because its bundled scope overlaps work that is already upstream or actively under review:

• The EAGLE TP finalized-decision broadcast is patch-identical to [Bug] eagle: broadcast finalized verify decision across TP ranks (#31071) #31478; its original zhoaa attribution is preserved there.
• The HIP AITER DSA top-k padding mismatch is covered by Fix DSA indexer aiter (HIP) padding mismatch under DP-attention #33059.
• The remaining DSA metadata, idle-DP/EAGLE control-flow, and QuickReduce changes overlap [BUG] Fix stale attention metadata after Eagle pre-planner pads the batch #29262, [Perf] Support ragged verify graphs in the DSA attention backend #32010, [DSA] Handle partial-DP padding in Decode page-table transform #31760, [Bugfix][DSA] Fix num_splits (b+1) crash with DP-attention + speculative decode #30642, [2/N] [spec] preserve verify capture mode on idle DP ranks #31566, Fix PD decode hang with DP attention and GLM-5.2 MTP #32209, and Amd add quickreduce bak #8139.

I will re-evaluate the remaining VMM IPC and empty-path changes after those PRs resolve, and only propose a demonstrably non-duplicate residual patch.