[ykamiset]

Overview

MiniMax-M3 LBHNC interleaves K and V inside each block, so the V cache view has fewer rows than K. AITER pa_decode_gluon requires those leading dimensions to match and fails on the sparse decode path.

Claims

• _run_gluon_decode now passes one page count to pa_decode_gluon by trimming K to V's length.
• The dropped K rows are the last block's V pages. No block-table page id addresses them (max id is 127 * 16 + 7 = 2039).
• The trim stays inside _run_gluon_decode, so _sides_are_packed and minimax_m3_sparse_block_page_stride still see the original shapes and keep a block page stride of 16.

Validation

Reproduced on MiniMax-M3 with the AITER sparse paged-attention decode path. Before the change, K and V enter the kernel as:

k_cache  [2048, 1, 8,  16, 16]
v_cache  [2040, 1, 1, 128, 16]

File "vllm/models/minimax_m3/amd/sparse_attention_msa.py", line 223, in forward
    minimax_m3_sparse_attn_prefill_aiter(
File "vllm/models/minimax_m3/amd/ops/sparse_pa.py", line 509, in minimax_m3_sparse_attn_prefill_aiter
    _run_gluon_decode(
File "vllm/models/minimax_m3/amd/ops/sparse_pa.py", line 590, in _run_gluon_decode
    pa_decode_gluon(
File "aiter/ops/triton/gluon/pa_decode_gluon.py", line 5393, in pa_decode_gluon
    assert key_cache.shape[0] == value_cache.shape[0], (
AssertionError: key/value caches must have the same number of blocks, but got key=torch.Size([890768, 1, 8, 16, 16]), value=torch.Size([890760, 1, 1, 128, 16])

pa_decode_gluon asserts the two leading dimensions match, so this fails even though every page id addresses both tensors.

Server Command to run

export GPTOSS_SWIGLU_MXFP4_BF16_BOUND=1
export HIP_VISIBLE_DEVICES=0,1,2,3
export VLLM_ENGINE_READY_TIMEOUT_S=3600
export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=3600
export VLLM_USE_BREAKABLE_CUDAGRAPH=0
export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4
export VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16=0
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB=256
export VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1
export VLLM_CUSTOM_SCOPES_FOR_PROFILING=1
export ROCPROFILER_QUEUE_INTERPOSITION=0
 
mkdir -p /app/traces
export KINETO_CONFIG=/app/libkineto.conf
 
vllm serve /data/inference/model/MiniMax-M3-MXFP4-AttnFP8 \
--served-model-name MiniMaxAI/MiniMax-M3 \
--port 8190 \
--tensor-parallel-size 4 \
--trust-remote-code \
--block-size 128 \
--enable-chunked-prefill \
--language-model-only \
--no-enable-prefix-caching \
--max-num-batched-tokens 65536 \
--default-chat-template-kwargs '{"thinking_mode":"enabled"}' \
--max-num-seqs 128 \
--kv-cache-dtype fp8 \
--attention-backend ROCM_AITER_UNIFIED_ATTN \
--attention-config '{"indexer_kv_dtype": "fp8"}' \
--moe-backend aiter \
--reasoning-parser minimax_m3 \
--compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [1,2,3,4,5,7,8,9,10,11,12,13,14,15,16,17,18,20,22,24,26,28,30,32,34,36,48,64,80,96,112,128,144,160,176,192,208,224,240,256,272,288,304,320,336,352,368,384,400,416,432,448,464,480,496,512]}'

Accuracy test passes after this fix:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.956	±	0.0056
		strict-match	5	exact_match	↑	0.956	±	0.0056

Details

LBHNC places a token's V pages a fixed half-block after its K pages. The gluon decode kernel carries one page id and applies it to both caches, so V's view is rebased by that half-block and is shorter than K. Using V's length for both the cache views and the FP8 scale arguments satisfies the kernel without changing the stride logic that reads the untrimmed shapes.

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.

[Fangzhou-Ai]

FYI: #62712 hits this on CI now that its MiniMax model tests run under LBHNC + AITER pa_decode_gluon (equal-block assert). Pointed that PR at this draft as the fix; avoided opening a duplicate.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ykamiset.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[benenzhu]

We hit the same assertion with MiniMax-M3.1-Flash on MI355X: the released vllm/vllm-openai-rocm:minimax-m31-flash image ships aiter 0.1.26, so fp8 KV + VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1 (needed by the ROCm mono decode) dies at startup with key=[8192, 1, 8, 16, 16], value=[8184, 1, 1, 128, 16]. This change fixes it there (server starts with mono decode on; direct decode calls with 64-1024 requests x 4 query tokens run). Backported to the M3.1 branch in #63270.