[Bug][ROCm] GLM-5 MXFP4 sparse MLA decode crash on MI355x

[ChuanLi1101]

Summary

GLM-5 (GlmMoeDsaForCausalLM) with MXFP4 quantization crashes during the decode phase on ROCm MI355x (gfx950) with 8 GPUs. The model loads and prefills successfully, but decode consistently fails with either ZeroDivisionError or Memory access fault.

Root Cause

GLM-5 has 64 attention heads and 32 sparse indexer heads (index_n_heads=32). At TP=8, the MLA decode kernel receives 8 heads per GPU, but AITER's sparse MLA kernels require num_heads >= 16:

1. MLA decode kernel: mla_decode_stage1_asm_fwd only supports gqa >= 16. With TP=8, gqa=8 triggers RuntimeError: get_heuristic_kernel_mla: cannot get heuristic kernel! gqa:8
2. FP8 paged MQA logits (indexer): deepgemm_fp8_paged_mqa_logits_stage1 computes TileQCount = heads // ChunkQ with default ChunkQ=64. When heads < 64, TileQCount=0 causes ZeroDivisionError at SplitKV = (max(1, TotalCuCount // TileQCount) + 4) // 5 * 5 * WavePerEU

Even with TP=4 (16 heads per GPU, satisfying the >= 16 requirement for the MLA kernel), a Memory access fault persists during decode, suggesting additional issues in the sparse attention indexer's forward_hip path.

Environment

• GPU: 8x AMD MI355X (gfx950)
• ROCm: 7.2.1
• vLLM: main branch (latest)
• Model: GLM-5-MXFP4 (zai-org/GLM-5-MXFP4 or equivalent Quark checkpoint)
• Config: num_attention_heads=64, index_n_heads=32, kv_lora_rank=512, n_routed_experts=256

Reproduction

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_LINEAR=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_ROCM_USE_AITER_MLA=1

vllm serve /path/to/GLM-5-MXFP4 \
    --tensor-parallel-size 8 \
    --block-size 1 \
    --gpu-memory-utilization 0.90 \
    --enforce-eager

Then send any chat completion request - prefill succeeds but decode crashes.

Error Traces

ZeroDivisionError (TP=8, indexer path)

File "aiter/ops/triton/attention/pa_mqa_logits.py", line 198, in deepgemm_fp8_paged_mqa_logits_stage1
    SplitKV = (max(1, TotalCuCount // TileQCount) + 4) // 5 * 5 * WavePerEU
ZeroDivisionError: integer division or modulo by zero

Memory access fault (TP=4 or TP=8, during decode)

Memory access fault by GPU node-X (Agent handle: ...) on address ...

Related

• PR [ROCm] Fix AITER sparse MLA crash for num_heads < 16 (e.g. GLM-5 TP=8) #36855 - Fix for head repeat in sparse MLA (addresses the gqa:8 crash)
• [Feature]: decode MLA kernel support nhead=8 ROCm/aiter#726 - Upstream feature request for native nhead=8 support
• ZeroDivisionError in deepgemm_fp8_paged_mqa_logits_stage1 with 8 attention heads (GLM-5 TP=8) ROCm/aiter#2563 - Upstream issue for ZeroDivisionError in paged MQA logits
• Issue [Bug]: GLM-5 FP8 on H200 CUDA OOM in sparse_attn_indexer at High Concurrency #34553 - GLM-5 FP8 on H200 OOM in sparse_attn_indexer (CUDA, different issue)
• Issue [Bug]: vllm serve quantized GLM-5 failed #36220 - GLM-5 quantized serving failure (AWQ/INT4, different quantization)

Proposed Fix

Two complementary approaches (both needed):

1. MLA decode kernel: Head repeat padding from 8->16 (temporary workaround in PR [ROCm] Fix AITER sparse MLA crash for num_heads < 16 (e.g. GLM-5 TP=8) #36855) until AITER supports nhead < 16 natively
2. Indexer MQA logits: Fall back to PyTorch reference implementation when heads < 16 (implemented in fix/rocm-glm5-mxfp4-optimizations branch)

[github-actions]

CC @hongxiayang @tjtanaa @vllmellm for ROCm-related issue

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[github-actions]

This issue has been automatically closed due to inactivity. Please feel free to reopen if you feel it is still relevant. Thank you!