sglang DSA: enable PD + DP-attention + EAGLE MTP on GLM-5.2 / gfx950 (3 patches, rebased onto main)

[dorado269]

Enables PD disaggregation + DP-attention + EAGLE MTP for GLM-5.2-MXFP4 on gfx950. Without these the combination crashes on the first batch or deadlocks the whole DP group under concurrency.

Supersedes #34, which was branched from cf85272. main has since moved 29 commits ahead and absorbed three of the things that branch carried; this one is rebased onto 8692fb4 with those dropped. See "Relationship to #34" below.

The patches

#	file	what it fixes	shape
1	dsa_indexer_hip_dp_padded_rows.diff	HIP/aiter paged-MQA sizes its output from DP-padded rows while lengths is sized to real rows → Expected lengths.size(0) == B	upstream #32762 (NPU, same bug class): one boolean gates both trim and restore, post-kernel row count asserted before padding is restored
2	dsa_backend_dp_sync_and_page_table_rows.diff	(a) seq_lens.max().item() is a host sync on a branch only some DP ranks take → collectives desynchronize; (b) page table has one row per request, top-k one per token under MTP → assert	ours; no upstream counterpart found
4	draft_cuda_graph_dp_vote.diff	the draft graph/eager choice is made per rank from rank-dependent inputs and diverges on the PD decode leg → deadlock	upstream #32209: the vote rides the MLP-sync all-gather the scheduler already performs — zero extra collectives

Applied at build time by Dockerfile.sglang, defaulting on (APPLY_SGLANG_DSA_PATCHES=1). Build with =0 for a stock engine to A/B.

The apply script verifies every patch reached the bytecode, not just the source: a stale __pycache__ entry silently reverts a patch and has already invalidated one full experiment on this stack.

Validation

Image built from Dockerfile.sglang independently on each of two nodes; nothing patched in the running container. 2 × 8×MI355X (gfx950), GLM-5.2-MXFP4, PD + --dp-size 8 --enable-dp-attention --ep-size 8 + EAGLE MTP with the draft CUDA graph enabled, mooncake RDMA over mlx5 + dma-buf.

Check	Target	Result
Build-time bytecode verification, both nodes	all markers	8/8 + prereq + patch2a
4-prompt correctness probe	4/4, acc_len > 1	4/4, 2.00–3.43
conc=32 × 512 tok	32/32	32/32
conc=128 × 512 tok	128/128	128/128
conc=128 × 512, repeat	—	128/128
Traceback / KVTransferError, either leg	0	0 / 0
DP ranks serving	8	8 every run, 0 retries

The draft CUDA graph was measured in use at 92.0 %, identical on all 8 ranks, on the immediately preceding build of this patch set. That measurement needs an added probe (a different image) so it is not part of the run above — but it is the criterion that matters: forcing the draft path eager passes every functional test while disabling the feature under test, so a green stress result alone cannot distinguish a fix from that workaround. Uniformity across ranks is the property patch 4 exists to produce.

Relationship to #34

#34 carried four patches and two Dockerfiles. Rebasing onto current main applied with zero conflicts, which was misleading — the branch only added files, so git had nothing to compare against changes main made to the same problems in different files. Three overlapped:

• deepseek_nextn_glm52_mtp_bf16.diff — main's patch_glm52_nextn_quark_exclude.py (0d8d0ff) makes the identical edit: same file, same line, same resulting value. Keeping both would have been actively broken, not merely redundant: main's loop runs first, so our context diff would then fail at --fuzz=0 against an already-edited anchor. Dropped; apply_sglang_dsa_patches.sh now asserts it as a prerequisite, because that script is idempotent and a silent skip would surface only at runtime as GLM-5.2 dying at draft weight-load.
• build_mooncake_dmabuf.sh — main's build_mooncake_sglang.sh is a strict superset (same logic line for line, plus a HIP-transport gate check and build-tree cleanup). Dropped.
• Dockerfile.sglang.dmabuf — main's Dockerfile.sglang now compiles the dma-buf branch in and selects it at runtime via MOONCAKE_DISABLE_HIP_DMABUF, which is exactly how these runs drove it (=0). Dropped; the DSA layer moved into Dockerfile.sglang.

Net: 8 files / 1181 lines → 6 files / 1048 lines, one modified file instead of two new Dockerfiles and a duplicated build script.

There is a configuration-only alternative to part of this set

Turning GLM-5.2 MTP IndexShare off avoids the same deadlock without patch 4 or patch 2's page-table half:

--json-model-override-args '{"index_share_for_mtp_iteration":false}'

It works because IndexShare is the source of the divergence: the guard term dsa_topk_indices is None is seeded on the PD decode leg from RDMA-shipped per-request payloads, so it is a function of which requests each rank happens to hold.

patch	substituted?
1	No — independent bug, present regardless
2a DP host-sync	No — a host sync is invisible to the graph/eager decision either mechanism changes
2b page-table rows	Yes, in effect
nextn eh_proj	No — weight-load bug
4	Yes — this is what it targets

Measured with 2b and 4 asserted absent from the bytecode: 4/4, 32/32 ×2, 64/64, accept length 2.98–3.01 (no measurable cost). It needs MTP on the prefill leg too, and that arm ran conc=64, not 128.

Not adopted as the default because it is nearly free only while IndexShare's consumer stays disabled under PD by should_use_dsa_fused_topk. Upstream #31477 removes that limitation; once it lands the override starts costing (~3 % TPOT, reported by AMD's llying — second-hand, not measured by us). Checked with gh on 2026-07-31: open, REVIEW_REQUIRED, unmerged.

Each diff header records where it stands relative to this, so the choice is visible at the point of use.

What this does not establish

• No differential control was run in this validation. Necessity is established in the earlier reproduction kits — patch 4's same-node revert control (0/4, deadlock) among them.
• Patch 2 has no revert-style control — its case was made from runtime state instead: a py-spy dump caught a busy rank blocked inside dsa_backend on .max().item() while idle peers had already advanced into the next collective, and after the fix no rank appears there in a dump again, with PD warmup passing on all 8 ranks. Its second half was forced by experiment — with the max_seqlen_k change alone, the hang persists. A revert control would still be a cheap addition.
• No performance comparison against a DPA-only baseline. Patch 4 adds no collective (that was the point of adopting #32209's placement), but no throughput comparison was run.
• One configuration: context 32768, short prompts, 512-token outputs, --disable-custom-all-reduce (required — the aiter custom all-reduce kernel deadlocks on gfx942/gfx950 during EAGLE verify), MTP on the decode leg only. With prefill MTP off, the rank-split case patch 4 is designed for has not been exercised.

Base image pinning

The base tag stays pinned. These are context diffs applied at --fuzz=0 against sglang 0b3bb0cbe31873994c9f989fddfe2f87ca839fdd; a base bump fails the build at the patch step rather than mis-applying silently. That is intended behaviour. Build with APPLY_SGLANG_DSA_PATCHES=0 if you need a newer base.

Reproduction kits (raw per-request jsonl, both server logs, the build log, and cold-start instructions) live in the infera.yihou.glm5.2.mxfp4 workspace; the one for this run is glm52.mxfp4.spur.mooncake.packup_20260731_main_converged.

🤖 Generated with Claude Code