Skip to content

[ROCm][Bugfix] Fix GQA=8 crash in mla_decode_fwd_impl by padding nhead to AITER minimum (gfx942/gfx950) - #4

Draft
psakhamo wants to merge 1 commit into
mainfrom
fix/rocm-aiter-gqa8-head-padding
Draft

psakhamo wants to merge 1 commit into
mainfrom
fix/rocm-aiter-gqa8-head-padding

Conversation

@psakhamo

@psakhamo psakhamo commented Jun 6, 2026

Copy link
Copy Markdown
Owner

Purpose

Fix crash in _rocm_aiter_mla_decode_fwd_impl when nhead=8 (GLM-5.1 at TP=8 on ROCm).
Resolves ROCm/aiter#2821 at the vLLM layer as a stopgap until native ASM kernels land.

Root cause

aiter.mla_decode_fwd dispatches to mla_decode_stage1_asm_fwd. For nhead=8 in
non-persistent mode, no matching ASM binary exists — the kernel faults at C-level
with a NULL pointer dereference (Memory access fault on address nil).

Fix

When nhead < 16 and 16 % nhead == 0, pad q from 8→16 heads via repeat_interleave
before calling mla_decode_fwd, then extract real heads from o afterwards.
Mirrors SGLang nsa_backend.py::_forward_aiter (same bug, same fix, see ROCm/aiter#2821).

Test Plan

  • Layer 1 (basic inference): PASS
  • Layer 2 (streaming tool-call finish_reason=tool_calls): PASS
  • SWE-Bench agent run (harbor/swe-agent, sympy task): 46-47 turns completed,
    agent submits patches. No GPU faults during run.
  • Context extended: ~13k → ~20k tokens

Test Result

Before: crash with Memory access fault / SRAM ECC after ~13k tokens.
After: 46-47 agent turns, no GPU faults.

Note: ~20k ceiling remains in ASM binary — native gqa=8 kernel (ROCm/aiter#2821)
needed for full fix.

Related: ROCm/aiter#2821, vllm-project#36855

@psakhamo
psakhamo force-pushed the fix/rocm-aiter-gqa8-head-padding branch from 6a97187 to 24759ee Compare June 6, 2026 20:26
@psakhamo psakhamo changed the title GPUAI-6471 [ROCm][Bugfix] GQA=8 head-padding in _rocm_aiter_mla_decod… [ROCm][Bugfix] Fix GQA=8 crash in mla_decode_fwd_impl by padding nhead to AITER minimum (gfx942/gfx950) Jun 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[MLA] Missing gqa_ratio=8 kernels in mla_decode_stage1_asm_fwd (blocks GLM-5 HiSparse on MI355X TP=8)

1 participant