[micah-wil]

Overview

Skip AITER's tuned GEMM (tgemm.mm) on gfx950 when M·N ≥ 2^31 (Also skip AITER tuned GEMM when A or B spans >= 4 GiB) and fall back to F.linear (hipBLASLt). This works around an int32 output-offset overflow in AITER's gfx950 FlyDSL A16W16 kernel (fix pending in ROCm/aiter#6853). Since AITER v0.1.26, that overflow crashes amd/MiniMax-M3-MXFP4 at startup with --mm-encoder-tp-mode data.

Claims

• Fixes the Memory access fault ... Write access to a read-only page crash during the encoder-cache profile run of MiniMax-M3-MXFP4 on MI355X (TP4) with --mm-encoder-tp-mode data.
• No performance cost: only GEMMs with ≥ 2^31 output elements are rerouted, and hipBLASLt matches the FlyDSL kernel's throughput at those sizes. Normal LLM prefill/decode GEMMs never reach this size.

Validation

Environment: MI355X, vllm/vllm-openai-rocm:nightly, AITER 0.1.26.

Kernel-level repro (tgemm.mm vs F.linear, N=6144, K=1280, bf16, with bias; this matches the tuned minimax_m3_mxfp4_bf16_tuned_gemm.csv rows):

M	M·N vs 2^31	tgemm.mm (AITER)	F.linear (hipBLASLt)
340,000	below	1221 TFLOP/s	1206 TFLOP/s
360,000	above	Memory access fault	n/a
768,000	above	Memory access fault	1229 TFLOP/s

End-to-end with this patch:

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 VLLM_ROCM_USE_AITER_MOE=1 \
vllm serve amd/MiniMax-M3-MXFP4 -tp 4 --trust-remote-code --block-size 128 \
  --attention-backend TRITON_ATTN --moe-backend aiter --kv-cache-dtype fp8 \
  --mm-encoder-tp-mode data --mm-encoder-attn-backend ROCM_AITER_FA \
  --tool-call-parser minimax_m3 --enable-auto-tool-choice --reasoning-parser minimax_m3 

Before: crashes at startup with the memory access fault. After: startup completes (178 GiB KV cache on TP0), and three chat requests with a 2016×2016 image return 200 with no faults. Confirmed there is no accuracy change.

No unit test is added: reproducing the bug needs gfx950 hardware plus AITER's model-specific tuned CSV, and a mocked test would only restate the condition.

Details

Root cause. With --mm-encoder-tp-mode data, each rank runs the ViT unsharded, so the profile run issues an M=768000, N=6144, K=1280 GEMM with bias. In AITER v0.1.26, getPaddedM maps it to the 8192,6144,1280 FlyDSL row in minimax_m3_mxfp4_bf16_tuned_gemm.csv. That kernel computes global_row * n + global_col in int32, which wraps once M·N ≥ 2^31 and writes out of bounds. In TP mode N is sharded 4 ways, so M·N stays under the limit.

Why route around it here. rocm_unquantized_gemm_impl is the only vLLM call site of tgemm.mm. The guard is limited to gfx950 and to the exact overflow condition, so all other shapes keep AITER's tuned dispatch. Skipping FlyDSL for this call also avoids the FlyDSL 0.4.0 from_dlpack leak described in ROCm/aiter#6853, which pinned about 10.6 GiB after the ViT profile run.

The guard can be removed once vLLM's minimum AITER includes ROCm/aiter#6853.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[micah-wil]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98431 for commit 025611841ad1.

[micah-wil]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98439 for commit e44957c7fa6f.

[AndreasKaratzas]

/ci retry

[github-actions]

✅ Triggered Buildkite CI #98540 for commit bd01b0b27762, running 1 failed step(s) from Buildkite CI #98439.

[mergify]

Hi @micah-wil, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[micah-wil]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98553 for commit dec78774a147.