[AMD] [GLM-5.3-Flash Day 0] Enable zero-RoPE MHA prefill on ROCm

[Jacob0226]

Summary

GLM-5.3-Flash's 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 tensor while the helper still allocates a qk_head_dim-wide k and appends it. So, this PR returns k_nope directly.

k_pe	attention_backend	today	with this PR
None	aiter	len(k_rope.shape) → AttributeError	early return k_nope
None	other	k[..., self.qk_nope_head_dim:] = k_pe → TypeError	early return k_nope
zero-width	aiter	`tl.arange(0, rope_dim)' → Triton crashes	early return k_nope
zero-width	other	k[..., self.qk_nope_head_dim:] = k_pe → TypeError	early return k_nope

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

Unit base 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.

End to end, since no single PR here can serve the model on its own: main plus the eight open Day-0 PRs — #39338, #39339, #39340, #39341, #38545, #38546, #38547 and #39317 — which all merge cleanly onto it. rocm/sgl-dev:v0.5.19-rocm720-mi35x-20260909, MI355X TP4, decode CUDA graphs on, 1,319 GSM8K examples, one scoring pass each. Only the checkpoint differs between the two rows:

Checkpoint	GSM8K
zai-org/GLM-5.3-Flash, block FP8	97.27%
amd/GLM-5.3-Flash-Quark-MXFP4	96.51%

Accuracy: not attributable to this PR in isolation, and there is no standalone baseline — GLM-5.3-Flash does not start on ROCm without the rest of the stack.

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
pre-commit	validate registered test CI registries reports the same 582 pre-existing taxonomy violations as pristine main, i.e. none added

---

CI States

Latest PR Test (Base): ❌ Run #34937435623
Latest PR Test (Extra): ❌ Run #34937435289
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.