[AMD] Cherry-pick AMD fixes into release/v0.5.15

[michaelzhang-ai]

Summary

Cherry-picks the finalized AMD-related PR list from main into release/v0.5.15. 9 PRs included. Applied in chronological (landing) order, each with -x provenance. #30237 and #30333 were verified already present in release/v0.5.15.

PR	Title	Source commit (on main)	Apply
#30313	[AMD] Cap DSV4 Flash max_total_num_tokens	dabd4cfcfd349eb8c3c3a15210fd74978335a812	clean
#30302	[AMD][MORI-EP] Skip LocalExpertCount kernel in decode graph when not recording	9ddea8d9efb3c16bfa35df07681de5c05e5eb041	clean
#30374	[AMD] Fix DeepSeekV4 server cutlass error	40a68521c9c325cf2757c05e7a476ad4e54f8038	clean
#29275	Fix gfx95 bpreshuffle FP8 activation scale layout	8d2b66fd9071f29434f5f0a149be1f6829907c2f	clean
#30265	[AMD] Fix GLM-5.2 MTP Quark excludes	07ef650ef7b066f8bab81d531acb1edc8231902d	conflict resolved
#30557	[AMD] Fix AITER custom all-gather CUDA-graph capture crash under torch_memory_saver	bd7e54d7379e437cf5f027382d6ca214e046626b	clean
#29479	[AMD] fix dsv4 indexer dtype dispatch on gfx950	336b64ecce300a3cefc615d84b6780e22f83a89c	clean
#30339	[AMD] Fix stale SWA ring buffer on radix prefix reuse for DeepSeek-V4 (unified_kv)	462b6171bd80902f68a0056c41bf95e0cec91400	clean
#29417	[AMD] Enable unified-KV HiCache on DeepSeek-V4	8d0fd341507710d628bf3e05d88ae87253970b78	clean

Conflict resolutions (please review)

• [AMD] Fix GLM-5.2 MTP Quark excludes #30265: kept release's get_flags import, added WeightsMapper (verified present in release models/utils.py), and added only GlmMoeDsaForCausalLMNextN to the arch lists (the upstream Longcat entries were pre-existing context on main, absent in release, so not introduced).

Not included

• [AMD] Enable JIT staged HiCache write-back and fix CPU-index crash #28534 ([AMD] Enable JIT staged HiCache write-back and fix CPU-index crash): dropped. It was authored on top of [mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py #30249 ("[mem_cache][6/N] move MHA host-pool into pool_host/mha.py", not in release), so its diff and its added test reference the post-move path mem_cache/pool_host/mha.py, which doesn't exist in release/v0.5.15. Rather than carry it without [mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py #30249, it is deferred; can be revisited together with [mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py #30249.
• Enable RDNA3/4 (gfx1100/gfx1201) for ROCm kernels #30415 (Enable RDNA3/4 gfx1100/gfx1201): dropped from the finalized list.
• [diffusion] Enable breakable CUDA graph (BCG) for diffusion DiTs #27436 ([diffusion] breakable CUDA graph for DiTs, 33c3dfd7): deferred — depends on [diffusion] Fix Z-Image accuracy #29742 ("fix z-Image accuracy", not in release); release's patchify_and_embed() returns a 5-tuple while [diffusion] Enable breakable CUDA graph (BCG) for diffusion DiTs #27436 expects the 7->8-tuple form, so it can't be cherry-picked cleanly on its own.

Test plan

• AMD CI on release/v0.5.15 (pr-test-amd) green
• Sanity-check [AMD] Fix GLM-5.2 MTP Quark excludes #30265 on ROCm hardware
• Follow-ups: [AMD] Enable JIT staged HiCache write-back and fix CPU-index crash #28534 (with [mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py #30249), and [diffusion] Enable breakable CUDA graph (BCG) for diffusion DiTs #27436 (with [diffusion] Fix Z-Image accuracy #29742)

Note: blocked/pending #24651 per the release coordination thread.

---

CI States

Latest PR Test (Base): ❌ Run #29067430598
Latest PR Test (Extra): ❌ Run #29067430426

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[HaiShaw]

@michaelzhang-ai remove #30415

[HaiShaw]

remove #30415

[michaelzhang-ai]

@HaiShaw done — #30415 has been removed. I rebased it out and force-pushed, so the PR now contains only the 10 finalized commits (head 7302618):

#30313, #30302, #30374, #29275, #30265, #28534, #30557, #29479, #30339, #29417

#30415 (RDNA3/4 gfx1100/gfx1201) is no longer included, per the finalized list. Note #27436 is also deferred here — it depends on #29742 ("fix z-Image accuracy", not in release/v0.5.15): release's patchify_and_embed() returns a 5-tuple while #27436 expects the 7→8-tuple form, so it can't be cherry-picked cleanly on its own. Happy to bring it in as a #29742 → #27436 pair as a follow-up.

PTAL when you have a moment — thanks!

[michaelzhang-ai]

Detail on why #27436 is deferred

#27436 (33c3dfd7) was written on top of #29742 "[diffusion] fix: fix z-Image accuracy" (fa185ed8), which is not in release/v0.5.15. Concretely:

• Release's ZImageTransformer2DModel.patchify_and_embed() returns a 5-tuple, and forward() unpacks 5 values.
• [diffusion] Fix Z-Image accuracy #29742 rewrote that path (571 lines in zimage.py) to return a 7-tuple (adding all_image_attn_lens/all_cap_attn_lens) and added the build_varlen_mask_meta_from_lengths/build_varlen_mask_meta_from_ranges helpers in layer.py.
• [diffusion] Enable breakable CUDA graph (BCG) for diffusion DiTs #27436 builds on that 7-tuple and extends it to 8 (adding caption_valid_mask).

So cherry-picking #27436 onto release produces a tuple-arity mismatch in the z-image forward path. The layer.py conflict I could resolve cleanly (only DynamicVarlenMaskMeta is genuinely #27436's — the two build_varlen_mask_meta_* helpers it "conflicts" on are actually #29742 context with no callers in #27436), but the zimage.py conflicts can't be resolved faithfully without #29742.

#29742 is large — 1067 insertions across 11 files: a new Triton kernel zimage_native_norm.py, encoder changes, perf baselines, and tests. (#29777 "Support SP for Krea-2" touches layer.py too but is separate/independent — Krea-2 — and not a strict prerequisite.)

Options for #27436:

1. Bring it in as a [diffusion] Fix Z-Image accuracy #29742 → [diffusion] Enable breakable CUDA graph (BCG) for diffusion DiTs #27436 pair (faithful, but pulls in the large z-image accuracy change).
2. Adapt [diffusion] Enable breakable CUDA graph (BCG) for diffusion DiTs #27436 to release's 5-tuple form (not faithful; risks the z-image accuracy/BCG path).
3. Leave it deferred for a later diffusion-focused batch.

@HaiShaw

[michaelzhang-ai]

Inline notes on the two conflict resolutions in this cherry-pick (for reviewer traceability).

[michaelzhang-ai]

Conflict resolution (#30265): the 3-way merge widened this hunk with upstream LongcatFlash* context lines that are not in release/v0.5.15. I added only what #30265 genuinely introduces — GlmMoeDsaForCausalLMNextN — here (and at the two other arch-list sites, ~L542 draft-model remap and ~L745), and did not pull in the Longcat context.

Verify (where the conflict was) — compare the two diffs (jumps to this file):

• Upstream [AMD] Fix GLM-5.2 MTP Quark excludes #30265 on main: 07ef650#diff-be7f3f1b02373a39149a134fd553cc61243cbe11d997dc72e319c6123dba4618
• As applied here (on release): 8501969#diff-be7f3f1b02373a39149a134fd553cc61243cbe11d997dc72e319c6123dba4618

[michaelzhang-ai]

Conflict resolution (#30265): import-line clash. Kept release's from sglang.srt.runtime_context import get_flags, get_parallel (get_flags is used in the file body) and merged in only WeightsMapper on this line — it's needed by the new GlmMoeDsaForCausalLMNextN class and is verified present in release models/utils.py. Dropped #30265's get_server_args/get_stream import additions since the cherry-picked code doesn't reference them.

Verify (where the conflict was) — compare the two diffs (jumps to this file):

• Upstream [AMD] Fix GLM-5.2 MTP Quark excludes #30265 on main: 07ef650#diff-0ccc8d832ece5a51e8d7a0a62b9ba5c056033186e11a0e758f3c11bd5a751106
• As applied here (on release): 8501969#diff-0ccc8d832ece5a51e8d7a0a62b9ba5c056033186e11a0e758f3c11bd5a751106

[michaelzhang-ai]

Upstream (main) vs release (resolved) — conflict-resolution compare

Update: #28534 has been dropped from this PR (it depended on the #30249 host-pool move, not in release). The only remaining conflict resolution is #30265, shown below.

For human verification, here are the exact upstream-vs-release deltas for the conflicted region(s), with direct links to the exact lines on both sides.

1. #30265 — python/sglang/srt/configs/model_config.py · is_deepseek_dsa() arch tuple

Inline delta (upstream main → release):

 "GlmMoeDsaForCausalLM",
 "GlmMoeDsaForCausalLMNextN",
-"LongcatFlashForCausalLM",
-"LongcatFlashForCausalLMNextN",
 )

Only GlmMoeDsaForCausalLMNextN (the #30265 addition) is kept. The two LongcatFlash* lines are pre-existing main context absent from release/v0.5.15, so they are not introduced. (Same treatment at the other two arch-list sites.)

Exact code:

• Upstream (main): model_config.py L112-L115
• Release (resolved): model_config.py L112-L113

2. #30265 — python/sglang/srt/models/glm4_moe.py · import block

Inline delta (upstream main → release):

 from sglang.srt.models.utils import WeightsMapper, apply_qk_norm
-from sglang.srt.runtime_context import (
-    get_parallel,
-    get_server_args,
-    get_stream,
-)
+from sglang.srt.runtime_context import get_flags, get_parallel

WeightsMapper (needed by the new GlmMoeDsaForCausalLMNextN class) is taken from upstream; get_flags is preserved because release's file body uses it (disable_shared_experts_fusion, enable_dp_lm_head); get_server_args/get_stream are dropped since the cherry-picked code doesn't reference them.

Exact code:

• Upstream (main): glm4_moe.py L85-L90
• Release (resolved): glm4_moe.py L85-L86

3. #28534 — MHA JIT gate (modify/delete) — DROPPED

#28534 was removed from this PR. It was authored on top of #30249 ([mem_cache][6/N] move MHA host-pool into pool_host/mha.py), which is not in release, so its diff and its added test referenced mem_cache/pool_host/mha.py (nonexistent in release). Deferred to be revisited with #30249.

---

Native GitHub views for cross-checking:

• Split left/right of release-base vs resolved: the Files changed tab.

[HaiShaw]

@michaelzhang-ai let's remove #28534 too

[michaelzhang-ai]

Update: #28534 ([AMD] Enable JIT staged HiCache write-back and fix CPU-index crash) has been dropped from this cherry-pick. The PR now contains 9 commits (head 4bc3c9c).

Why: #28534 was authored on top of #30249 ([mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py), which is not in release/v0.5.15. On main the MHA host-pool lives in mem_cache/pool_host/mha.py (created by #30249); on release it still lives in mem_cache/memory_pool_host.py. So #28534's diff — and its added test test/registered/jit/test_hicache_page_first_write_back.py (which does from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost) — reference a path that doesn't exist in release. Rather than carry #28534 in an adapted form without its #30249 base, we're deferring it.

Follow-up: revisit #28534 together with #30249 if the ROCm JIT staged HiCache write-back + CPU-index crash fix are wanted in this release line.

No other included commit depends on #28534 (verified: #30557/#29479/#30339/#29417 don't touch any of its files), so the drop is clean.

[HaiShaw]

LGTM

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!