Summary

The investigation and this report was generated with the help of Claude Code.

The ROCM_AITER_MLA_SPARSE attention backend — the only backend available for glm_moe_dsa sparse-MLA models on ROCm — silently produces corrupt output when the prefill prompt exceeds approximately 20,000 tokens. This is a distinct bug from the decode-path regression reported in #39303 (which was fixed by PR #39509 reverting aiter to v0.1.10.post3). The new bug persists on the fixed nightly and is in the prefill path, not the decode path.

Environment

• Docker image: docker.io/rocm/vllm-dev:nightly (image id 38936227491d)
• vLLM version: 0.19.1rc1.dev296+gbcc2306ce
• aiter version: 0.1.10.post3 (the post-#39509 fix)
• Model: zai-org/GLM-5.1-FP8 (754B MoE, FP8 E4M3, glm_moe_dsa architecture)
• Hardware: 4× AMD Instinct MI355X OAM (gfx950:sramecc+:xnack-), 288 GiB HBM3e each
• TP: 4 (required — aiter ASM MLA kernels only exist for gqa=16, see #39303)
• Node: mi355-gpu-37
• Date: 2026-04-16

Patches applied

Two patches are required on this nightly for the model to start at all:

1. pa_mqa_logits.py ZeroDivisionError clamp — GLM at TP=4 has heads_per_rank=16 < ChunkQ=64, causing heads // ChunkQ == 0 → TileQCount == 0 → ZeroDivisionError during CUDA graph capture. Fix: if heads < ChunkQ: ChunkQ = heads; TileQCount = max(1, ...). This is needed on every nightly tested and is orthogonal to this bug.

2. rocm_aiter_mla_sparse.py workspace + skip_kv_gather patch — see Root Cause section below. This patch improves the threshold from ~20K to ~25K tokens but does not fully fix the issue.

Reproduction

Minimal reproducer

"""
Reproducer: ROCM_AITER_MLA_SPARSE prefill corruption at large prompt sizes.
Run against a vLLM server serving GLM-5.1-FP8 on MI355X (gfx950).
"""
import requests, hashlib

VLLM_URL = "http://localhost:8005/v1/chat/completions"

def make_diverse_prompt(n_chars):
    """Generate non-repetitive text to avoid tokenizer compression."""
    parts = []
    for i in range(n_chars // 100 + 1):
        h = hashlib.md5(str(i).encode()).hexdigest()
        parts.append(
            f"Function process_{h[:8]}(data_{i}: List[Dict], "
            f"config_{h[8:16]}: Optional[str] = None) -> Tuple[int, str]: "
            f"Validates input parameter {h[16:24]} against schema version "
            f"{i % 50}.{i % 10}, applies transformation rule "
            f"#{i*7 % 1000} from module {h[24:32]}, then returns "
            f"status code {200 + i % 5} with message. "
        )
    return " ".join(parts)[:n_chars]

for target_chars in [20000, 40000, 60000, 80000, 100000]:
    prompt = make_diverse_prompt(target_chars)
    resp = requests.post(VLLM_URL, json={
        "model": "glm-5.1-fp8",
        "messages": [
            {"role": "system", "content": prompt},
            {"role": "user", "content": "What is 2+2? Answer in one sentence."},
        ],
        "max_tokens": 1000,
        "temperature": 0.7,
    }, timeout=300)
    r = resp.json()
    content = r["choices"][0]["message"].get("content") or ""
    reasoning = r["choices"][0]["message"].get("reasoning") or ""
    text = reasoning + content
    words = text.split()
    unique = len(set(w.lower() for w in words)) if words else 0
    ratio = unique / len(words) if words else 0
    pt = r["usage"]["prompt_tokens"]
    ct = r["usage"]["completion_tokens"]
    status = "OK" if ratio > 0.15 else "GARBAGE"
    print(f"prompt={pt:6d}  comp={ct:4d}  unique_ratio={ratio:.3f}  {status}")
    print(f"  first 80 chars: {text[:80]}")
    print()

Results

Without rocm_aiter_mla_sparse.py patch (upstream nightly as-is + pa_mqa_logits clamp only):

Garbage output at 23K: AFC} 8^PU^P}^U^P>^P> ^P> ^P> ^P> ^P> ^P> ^P> ...

With rocm_aiter_mla_sparse.py workspace + skip_kv_gather patch:

Threshold improved from ~20K to ~25K, but still fails for prompts > ~25K tokens.

Control: small prompt + long generation (same server, same patches):

6,000 tokens of coherent computing history essay. Sliding-window unique-ratio never drops below 0.74 across the entire output. The decode path is fine — the corruption is entirely in the prefill path.

What was ruled out

Root cause analysis

The upstream ROCm path vs CUDA reference

The file vllm/v1/attention/ops/rocm_aiter_mla_sparse.py is the ROCm-specific implementation of the sparse-MLA indexer. The CUDA reference lives in vllm/model_executor/layers/sparse_attn_indexer.py. Comparing the two reveals two bugs in the ROCm path, both related to the prefill sub-chunk handling:

Bug 1: Missing skip_kv_gather check (causes garbage on sub-chunk reuse)

The CUDA path (lines 113-130 of sparse_attn_indexer.py):

workspace_manager = current_workspace_manager()
k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
    ((total_seq_lens, head_dim), fp8_dtype),
    ((total_seq_lens, 4), torch.uint8),
)
for chunk in prefill_metadata.chunk_specs:
    k_fp8 = k_fp8_full[: chunk.total_seq_lens]
    k_scale = k_scale_full[: chunk.total_seq_lens]
    if not chunk.skip_kv_gather:          # <-- KEY: skip on sub-chunks
        ops.cp_gather_indexer_k_quant_cache(...)

The upstream ROCm path (lines ~553-570):

for chunk in prefill_metadata.chunk_specs:
    k_fp8 = torch.empty(                  # <-- BUG: fresh buffer each iteration
        [chunk.total_seq_lens, head_dim],
        device=k.device, dtype=fp8_dtype,
    )
    k_scale = torch.empty(...)            # <-- BUG: fresh buffer, uninitialized
    ops.cp_gather_indexer_k_quant_cache(   # <-- BUG: always gathers (no skip check)
        kv_cache, k_fp8, k_scale,
        chunk.block_table, chunk.cu_seq_lens,
    )

When split_indexer_prefill_chunks produces sub-chunks with skip_kv_gather=True (because the KV data was already gathered by a prior sub-chunk), the ROCm path:

1. Allocates a fresh torch.empty buffer (uninitialized memory)
2. Always re-gathers instead of skipping
3. The fresh buffer does NOT contain the previously gathered data

For the first sub-chunk this accidentally works (it gathers fresh). For subsequent sub-chunks marked skip_kv_gather=True, the CUDA path reuses the shared workspace (which still holds valid data from the first sub-chunk), but the ROCm path allocates garbage-filled memory and re-gathers into a differently-sized buffer.

Bug 2: Missing workspace reservation during profiling

The CUDA path reserves workspace during the profiling run:

if not isinstance(attn_metadata, dict):
    current_workspace_manager().get_simultaneous(
        ((total_seq_lens, head_dim), torch.float8_e4m3fn),
        ((total_seq_lens, 4), torch.uint8),
    )

The upstream ROCm path has no equivalent reservation, so the workspace manager's peak-memory accounting during profiling does not include the indexer buffers. At runtime, when the workspace allocates for real, it may compete with the KV cache for the same memory, potentially causing silent corruption at large sizes.

Why the patch helps but doesn't fully fix it

The rocm_aiter_mla_sparse_ops_patched.py patch (authored 2026-04-08) addresses both bugs above by mirroring the CUDA reference. This pushes the correctness threshold from ~20K to ~25K prompt tokens. However, the underlying rocm_fp8_mqa_logits kernel appears to have its own size-dependent correctness issue beyond ~20-25K total sequence length in the logits computation, which is NOT present in the CUDA path (which uses a different kernel).

The remaining corruption past ~25K is likely in one of:

• aiter.ops.triton.attention.fp8_mqa_logits.rocm_fp8_mqa_logits — the FP8 MQA logits kernel called per sub-chunk
• The radix-topk implementation used on ROCm (vs persistent_topk on CUDA)
• An indexer metadata builder issue with total_seq_lens exceeding some internal buffer

Impact

• GLM-5.1-FP8 is unusable for applications with prompts > ~25K tokens on MI355X. This includes all agentic coding assistants (system prompt + tools + history routinely exceeds 30K tokens), RAG pipelines, and long-document QA.
• ROCM_AITER_MLA_SPARSE is the only backend for sparse-MLA models — there is no TRITON_MLA fallback. The backend selector at platforms/rocm.py:362 hardcodes return [AttentionBackendEnum.ROCM_AITER_MLA_SPARSE] when use_sparse=True.
• Short-prompt use cases (< 20K tokens) and decode-only workloads are unaffected.

Suggested fixes

1. Immediate (P0): Port the two CUDA-path fixes into rocm_aiter_mla_sparse.py:

   • Add skip_kv_gather check mirroring sparse_attn_indexer.py:123
   • Add workspace reservation during profiling mirroring sparse_attn_indexer.py:59
   • Use persistent shared workspace buffers instead of per-iteration torch.empty

   Diff available: compare the upstream rocm_aiter_mla_sparse.py against vllm/model_executor/layers/sparse_attn_indexer.py lines 55-130.

2. Follow-up: Investigate the residual corruption past ~25K tokens after applying the workspace patch. This may be in rocm_fp8_mqa_logits itself or in the ROCm-specific topk path.

3. Defense in depth: Add a TRITON_MLA fallback path for use_sparse=True models, even if slower, so there is a correct reference to test against. Currently the only backend is ROCM_AITER_MLA_SPARSE with no escape hatch.

Additional context

• PR #39509 (revert aiter to v0.1.10.post3) fixed the decode regression from #39303. That fix is confirmed working — decode coherence past the 2048 boundary is verified on this nightly. This new bug is orthogonal.
• The pa_mqa_logits.py ZeroDivisionError clamp (heads < ChunkQ at TP=4) is still required on all nightlies for GLM models and should be upstreamed separately.
• The topk_fixed.py patch (moe_fused_gate 256-expert limit for GLM's 256 experts) is needed on aiter v0.1.12 but NOT on v0.1.10.post3 (different dispatch logic).