Conversation
psakhamo
force-pushed
the
fix/rocm-aiter-gqa8-head-padding
branch
from
June 6, 2026 20:26
6a97187 to
24759ee
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
agent submits patches. No GPU faults during run.
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