[tianxiaojiang4]

Motivation

PR #31477 added a decode-local remap so a GLM-5.2 MTP IndexShare seed can enter the allocator-local fused TopK domain under PD disaggregation. The remap is gated on is_cuda():

def should_remap_pd_dsa_seed_to_local_slots() -> bool:
    return (
        is_cuda()
        and envs.SGLANG_DSA_FUSE_TOPK.get()
        ...

On ROCm that gate is False, so the remap never applies, pd_index_share_seed stays True, and should_use_dsa_fused_topk() returns False. The #31477 fix is inert on gfx942/gfx950 and PD decode runs unfused.

@kpham-sgl confirmed in #31477 that the intent was to exclude NPU, not ROCm:

I meant to gate off NPU only but ROCm paths are untested. Can you submit a PR with tested results?

This is that PR.

Modifications

One predicate, in should_remap_pd_dsa_seed_to_local_slots():

-        is_cuda()
+        (is_cuda() or is_hip())

NPU still declines, since is_cuda() and is_hip() are both False there. The remap itself is untouched.

Happy to switch this to is_cuda_alike(), which is the same predicate (lru_cached is_cuda() or is_hip()) and has ~24 existing call sites. I kept the expansion because it is the exact form the validation below ran, and because the unit test patches the two predicates individually.

Two things I checked before proposing this, since "it compiles on ROCm" is not the same as "it is correct on ROCm":

• The remap reads req_to_token, which is page_size=1 token slots on every backend, so the slot domain is not platform-dependent.
• dsa_drop_wide_page_table already excludes HIP independently (is_cuda() and not _is_hip), so turning fusion on here cannot drop a page table the ROCm indexer still reads.

Also note SGLANG_OPT_USE_TOPK_V2 is force-disabled on ROCm for DSA models, so the fused path exercised here is fast_topk_transform_fused, not _topk_transform_v2_paged.

Tests

test_pd_decode_fused_topk_remaps_wire_positions_to_local_slots is parameterised over CUDA / ROCm / neither, asserting the third still declines the seed and passes the wire positions through unremapped. It stays a CPU unit test under register_cpu_ci — the platform predicates are patched, so no AMD hardware is needed in CI.

Correctness on MI300X (gfx942, ROCm 7.2)

Ran against a real ROCm build, which is the arm NVIDIA CI cannot cover:

• Gate admits ROCm; with is_hip forced False (pre-patch behaviour) it declines.
• SGLANG_DSA_FUSE_TOPK=0, hisparse, and the prefill role each still veto.
• On-device remap produces the correct allocator-local slots; an out-of-range position and an unallocated slot each fail the seed closed rather than consuming a stale slot.
• The upstream wire unit test passes: 25 passed, 10 subtests.
• test_dsa_transform_index.py: 7 passed.
• test_dsa_indexer.py and attention/unittests/dsa/test_dsa.py fail on ROCm — byte-identically with the gate reverted, so not attributable to this change. I ran a stock-gate control arm specifically to attribute rather than assume; they are NVIDIA-only paths (flashinfer, sgl-kernel topk-v2, flashmla_kv, fa3, tilelang).

Accuracy

GLM-5.2-MXFP4, 8×MI355X (gfx950, ROCm 7.2.0), 1P1D at TP2, PD + EAGLE/MTP. Every comparison is paired on verified-identical sample sets (McNemar), with full coverage and zero router circuit-breaker events on both arms.

path	benchmark	n	patched	stock	diff	paired 95% CI	McNemar p	discordant b/c
dense	GSM8K	1319	0.9295	0.9287	+0.08 pp	[−1.1, +1.3] pp	1.000	34/33
sparse	LongBench-v2 short	180	0.6111	0.6000	+1.11 pp	[−5.4, +7.6] pp	0.868	19/17

Needle-in-a-haystack at ~188k prompt tokens: 3/3 depths retrieved on both arms.

No regression on either path; the nominal direction favours the patch on both. Stated honestly, this excludes regressions worse than ~1.1 pp dense and ~5.4 pp sparse. It does not prove exact equality — n=180 is the entire short split, and resolving a 5 pp sparse effect would need ~1500 per arm, which this benchmark cannot supply at any sampling.

Two notes on why the two paths are reported separately rather than averaged:

• GSM8K prompts are ~100 tokens, below GLM-5.2's index_topk = 2048 dense-attention threshold, so DSA runs dense and the top-k shortlist is never consulted. GSM8K is a gross-breakage guard; it cannot validate a shortlist change.
• LongBench-v2 short is ~24k tokens median, well above the threshold, so it is the sparse-path test — and it is the metric [Spec][PD] Enable fused TopK for GLM-5.2 MTP IndexShare #31477's own table reports.

The LongBench absolute is not comparable to #31477's 60.94% → 62.50%: different hardware, prompt template and sample selection, and at n=64 the binomial SE alone is ~6 pp. The paired within-run delta is the comparable quantity.

Speed

Same topology, 300 s at concurrency 4, acceptance simulation off:

metric (p50)	stock	patched	ratio
TPOT	23.16 ms	7.94 ms	2.92×
ITL	23.16 ms	7.94 ms	2.92×
TTFT	2125.5 ms	2112.3 ms	1.01× (flat)
MTP accept length	3.12	3.12	—
request errors	0	0	—

TTFT flat and acceptance unchanged are both expected: this is a decode-side change, and unchanged acceptance is the load-bearing correctness signal, since a wrong remap would hand the draft loop bad KV slots.

Per-worker evidence that the one line changed exactly the intended worker:

worker	role	"Disabling fused DSA top-k" stock → patched
w0	prefill	2 → 2
w1	decode	2 → 0

The prefill warning persisting in both arms is correct by design — the gate requires disaggregation_mode == "decode", so prefill draft-extend never fuses, on CUDA either. This signature reproduced across three independent job pairs.

Decode step time is derived as TPOT × accept_len (~72 ms → ~24.8 ms) rather than measured: step_time_dict is a mean over decode_log_interval steps taken from the async scheduler loop's inter-log gap, so percentiles over it are percentiles of averages. Where it is usable it corroborates — 26.87 ms against 24.8 ms derived.

Caveats

• RDMA is unusable in the ROCm serving images on this cluster (vendored libionic targets a newer ionic_rdma ABI than the deployed driver), so PD KV transfer falls back to XGMI. This is single-node 1P1D, where both workers share a host and KV never crosses the fabric, so it does not affect these numbers — but TTFT on this cluster is transfer-dominated in general.
• The accuracy runs used DISAGG_TRANSFER_BACKEND=mori; the speed pair used mooncake. Measured small where compared (unpatched c8: 33.85 vs 34.14 tok/s/user, 0.9%), unquantified with the patch active.

---

CI States

Latest PR Test (Base): ❌ Run #33095503271
Latest PR Test (Extra): ❌ Run #33095502916
Latest PR Test (AMD ROCm 7.2): ❌ Run #33095503597

[tianxiaojiang4]

@kpham-sgl could you take a look at this one when you get a chance? I do not have permissions to add you as a reviewer directly.

[tianxiaojiang4]

Something good to know: further digging into the 3x performance, what the DSA-topk fusion helped most was not kernel launching elimination. It's the two index-transform kernels (transform_index_page_table_decode_kernel and transform_index_page_table_prefill_kernel)that translates index to kvcache address I believe got jit recompiled during runtime due to growing context length.

[kpham-sgl]

/rerun-test test/registered/unit/disaggregation/test_disaggregation_wire.py

[github-actions]

Results for /rerun-test test/registered/unit/disaggregation/test_disaggregation_wire.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/disaggregation/test_disaggregation_wire.py

[kpham-sgl]

LGTM