[AMD] [GLM-5.3-Flash Day 0] Handle the zero-width RoPE tail in the ROCm MHA k concat

[Raiden-Makoto]

Summary

Replacement for #37530, which GitHub closed when its deleted support base was merged into main. The original author and review history remain credited there.

GLM-5.3-Flash's 11 full-attention layers are norope (qk_rope_head_dim == 0), so k_pe reaches _concat_and_cast_mha_k_rocm as None or zero-width while the helper still allocates a qk_head_dim-wide k and appends it. #31531 split this helper out of the shared CUDA path on 2026-08-08, when no norope model existed and neither side carried a guard; 0b9c384 on this branch then added one to the CUDA twin _concat_and_cast_mha_k only. This restores parity.

k_pe	current_attention_backend	today	with this PR
None	aiter	concat_and_cast_mha_k_triton handed a None rope tensor	early return
None	other	k[..., qk_nope_head_dim:] = k_pe raises	early return
zero-width	aiter	Triton concat driven over a zero-width tail	early return
zero-width	other	correct, but one new_empty + one k_nope-sized copy per call	early return

Scope

forward_mha_rocm.py only. CUDA and every other backend are untouched. ROCm models with qk_rope_head_dim > 0 — DeepSeek V3/R1, the roped GLM-5.x MLA layers — never reach the guard and keep current behaviour. No new environment variable, no architecture gate, no dispatch change. k_pe widens to torch.Tensor | None to match the CUDA signature.

On MI355X this is the default short-prefill path, not an edge case:

Dispatch step	Resolves to
handle_attention_dsa	reads backend.use_mha
HybridLinearAttnBackend.use_mha	property, forwards to the DSA backend — the hybrid wrapper does not block this
DeepseekSparseAttnBackend.set_dsa_prefill_impl	True on _IS_GFX95, listed beside SM90/SM100, when max_kv_len <= SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD (2048)
resolve_rocm_forward_method	MHA_ONE_SHOT → MHA_ONE_SHOT_ROCM
forward_normal_one_shot_rocm_prepare	_concat_and_cast_mha_k_rocm(k_nope, k_pe)
no dtype cast	CUDA casts because its concat branch can allocate k in the KV-pool dtype for fa3 on an FP8 KV cache; the ROCm branch is k_nope.new_empty(*k_shape) with no override, so k_nope.contiguous() is dtype-identical to what it would have built

Test plan

Combined main-based integration (RM/glm53-day0-main-integration tested at 9d0d62b744; current integration head 4e2f84ff3c; TP8/TP4 validation predates the above-cap guard, whose below-cap path is unchanged, pinned AITER d9e5ef7ce0, fresh caches): TP8 GSM8K 97.50% in 87.0s with 0.00% errors; TP4 GSM8K 96.97% in 320.2s with 0.00% errors. Decode CUDA graphs were enabled in both runs.

Recovery base: main @ 30e7a3072d.
Recovery validation on MI355X: scoped pre-commit passed; 7 unit tests passed.

rocm/sgl-dev:v0.5.18-rocm720-mi35x-20260901, the image family pr-test-amd-rocm720.yml resolves to, on xinyuan/glm-5.3-flash-support @ 515e865189 with no other PRs applied. Both builds run the same test file; the guard is the only variable. cd test && python3 registered/unit/models/test_nope_mha_k_cast.py.

Build	Result
Baseline	3 failed, 4 passed — both None cases and zero-width on aiter fail
This PR	7 passed — 3 existing CUDA cases, 4 new ROCm

Zero-width on the non-aiter backend passes either way, per the Summary table's last row; it is kept so both backends are pinned to one behaviour rather than one of them resting on that accident.

Check	Detail
Registration	register_cpu_ci(suite="base-a-test-cpu"), unchanged — the four cases are mocked and GPU-free, and this change removes a kernel call rather than adding one, so there is no gfx950 work for an AMD-suite test to cover
pre-commit	all hooks pass on both files, including validate registered test CI registries
Accuracy	not measured — GLM-5.3-Flash cannot yet be served on ROCm from this branch; GSM8K lands with the DSA / k-pool enablement to follow

---

CI States

Latest PR Test (Base): ⏳ Run #34768725511
Latest PR Test (Extra): ❌ Run #34768725378
Latest PR Test (AMD ROCm 10): ❌ Run #34768725524

[Raiden-Makoto]

Closing this replacement because @Jacob0226 will reopen and maintain the original PR, #37530. Please continue review and CI there.

[Jacob0226]

Thanks @Raiden-Makoto. Reopening as #39338 rather than #37530: #37530's base branch xinyuan/glm-5.3-flash-support has been deleted, so GitHub will not let it reopen. #39338 is off main and its changed files are byte-identical to this PR's head.