fix: GLM-5.2 DSA backend crashes on AMD MI355X — .view() on non-contiguous tensors and GlmMoeDsaConfig attribute_map bug

[tanguofu]

Summary

GLM-5.2-FP8 on AMD MI355X (gfx950) fails to start or runs at ~2x slower TPOT due to two bugs in the DSA attention backend and model config initialization path.

Environment

• Hardware: 8× AMD Instinct MI355X (gfx950, 256 CUs each)
• Docker image: lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260623
• Model: GLM-5.2-FP8 (GlmMoeDsaForCausalLM)
• Config: TP8, PP1, MTP steps=2, DSA attention backend, tilelang decode/prefill, FP8 KV cache

Bug 1: .view() on non-contiguous tensors in dsa_backend.py

File: python/sglang/srt/layers/attention/dsa_backend.py

The DSA backend uses .view() to reshape q tensors in 3 code paths:

• forward_extend (line ~1600)
• forward_decode (line ~1781)
• forward_extend_one_shot (line ~2344)

These tensors come from the MLA absorb path (q_b_proj output → .view(-1, num_heads, qk_head_dim)) and are frequently non-contiguous. .view() on a non-contiguous tensor either silently produces wrong results or crashes during CUDA graph capture:

Exception: Capture cuda graph failed: Sizes of tensors must match except in dimension 2. Expected size 1024 but got size 384 for tensor number 1 in the list.

Fix: Replace all 9 .view() calls with .reshape() in these paths. Dimensions are unchanged — only the tensor reshape method changes.

Bug 2: GlmMoeDsaConfig.attribute_map silently overwrites qk_rope_head_dim

File: python/sglang/srt/utils/hf_transformers_patches.py (new patch)

The upstream GlmMoeDsaConfig.attribute_map contains "head_dim": "qk_rope_head_dim", which causes head_dim=192 from config.json to silently overwrite qk_rope_head_dim=64 during config initialization. This results in:

• Wrong fused QKV projection size ([2752, 6144] instead of correct [2624, 6144])
• Weight-loading failures
• qk_head_dim computed as 192 + 192 = 384 instead of 192 + 64 = 256

Fix: Add _patch_glm_moe_dsa_attribute_map() to remove the "head_dim" key from GlmMoeDsaConfig.attribute_map during apply_all().

Why v_head_dim override is NOT needed

An earlier version of this fix proposed overriding v_head_dim from 256 to qk_nope_head_dim (192) in model_config.py. This is incorrect and has been removed.

GLM-5.2's config.json sets v_head_dim=256, which is the correct model definition. The DSA backend uses the attn_mqa (MLA absorb) path, where layer.v_head_dim comes from kv_lora_rank (512) — not from model_config.v_head_dim (256). The model_config.v_head_dim only affects the attn_mha (MHA_ONE_SHOT) path, which DSA+MLA does not use.

Overriding v_head_dim to 192 would break weight loading:

• kv_b_proj expects output 64 × (192 + 256) = 28672, not 64 × (192 + 192) = 24576
• o_proj expects input 64 × 256 = 16384, not 64 × 192 = 12288
• deepseek_weight_loader.py unflatten(0, (-1, 384)) on [28672, 512] → non-integer → crash

Verified on the worker machine (amd-355-worker, 144.202.61.0): with only Bug 1 + Bug 2 fixes (no v_head_dim override), the server starts correctly, weights load successfully, and accuracy is identical to the unpatched baseline.

Performance Impact

Verified on 8× MI355X with identical Docker image and model:

Metric	Unpatched (master)	Patched (this fix)	Improvement
C1/4K TPOT	32.3ms	14.5ms	2.2x
C32/short throughput	333.9 tok/s	765.2 tok/s	2.3x
C32/mid throughput	328.7 tok/s	613.7 tok/s	1.9x
Cold start TTFT	11,263ms	1,390ms	8.1x
AITER "not found tuned config" warnings	27,408	32	856x

Accuracy Verification

45-question test suite (coding/math/knowledge) comparing patched vs unpatched:

Category	Patched	Unpatched	Regression
Coding (HumanEval-style, 15q)	14/15 (93%)*	14/15 (93%)*	None
Math (GSM8K-style, 15q)	13/15 (87%)*	14/15 (93%)*	None
Knowledge (15q)	15/15 (100%)	15/15 (100%)	None

*All "failures" are test-script extraction bugs, not model errors — manual inspection confirms 100% correct answers on both.

Proposed Fix

Branch: fix/glm52-dsa-reshape-no-vhead-override (rebased on latest main)

3 files changed, +115 / -7 lines:

1. python/sglang/srt/layers/attention/dsa_backend.py — .view() → .reshape() (9 replacements across 3 forward paths)
2. python/sglang/srt/utils/hf_transformers_patches.py — new _patch_glm_moe_dsa_attribute_map() function
3. test/registered/unit/utils/test_hf_transformers.py — 3 unit tests for the attribute_map patch

Notes

• The .view() → .reshape() fix is semantically safe: .reshape() returns a view when possible (same as .view()) and falls back to a copy when the tensor is non-contiguous.
• The attribute_map patch is idempotent and no-ops when GlmMoeDsaConfig is not importable.
• v_head_dim is intentionally left as-is from the model config. See "Why v_head_dim override is NOT needed" above.