vllm - 💡(How to fix) Fix [Bug][ROCm] GLM-5 MXFP4 sparse MLA decode crash on MI355x [1 comments, 2 participants]

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 #36855 - Fix for head repeat in sparse MLA (addresses the gqa:8 crash)
• https://github.com/ROCm/aiter/issues/726 - Upstream feature request for native nhead=8 support
• https://github.com/ROCm/aiter/issues/2563 - Upstream issue for ZeroDivisionError in paged MQA logits
• Issue #34553 - GLM-5 FP8 on H200 OOM in sparse_attn_indexer (CUDA, different issue)
• Issue #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 #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)