[ROCm][Perf] Split MiniMax-M3 prefill index-score K loop

[akii96]

Note: Validation required cherry-picking two fixes on top of the nightly image:

• PR #51585 — Fix CUDA graph capture failure (zeroed CPU query offsets broke mixed-attention metadata builders)
• PR #51632 — Fix Triton fused shared expert alignment (incorrect expert count caused silent accuracy loss)

Base image: vllm/vllm-openai-rocm:nightly-3ee2df30337a301164c46ae444b76ee67e71c106

Purpose

MiniMax-M3 prefill computes index scores by scanning every visible 128-token index-K block for each query block.

The current gfx942 kernel uses 64 query tokens per program and processes the full K range in one loop. At long context lengths, this repeatedly loads index-K data and leaves insufficient parallel work for the GPU.

Suggested Fix

For gfx942 prefill with at least 128 query tokens and an 8K sequence length:

• Increase the query block size from 64 to 128.
• Split the state-free K loop across multiple programs.
• Select a power-of-two split count using the query workload, sequence length, and available compute units.
• Cap the split count at 16.

Each split writes a separate range of score columns, so no atomics or reduction pass are required. Other architectures and smaller shapes retain the existing 64-token, single-split configuration.

Test Plan

pytest -q tests/kernels/attention/test_minimax_m3.py \
  -k prefill_index_topk_correctness

The test covers the existing short-context path and the gfx942 long-context split path. It compares every score written by the kernel, including causal masking, and verifies the final top-k block sets.

Serve command

export HIP_VISIBLE_DEVICES=0,1,2,3
export VLLM_ROCM_USE_AITER=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=0
export VLLM_ROCM_USE_AITER_MOE=0
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1

vllm serve EmbeddedLLM/MiniMax-M3-FP8-dynamic \
  --served-model-name minimax-m3 \
  --tensor-parallel-size 4 \
  --block-size 128 \
  --max-model-len 262144 \
  --gpu-memory-utilization 0.92 \
  --enable-chunked-prefill \
  --max-num-batched-tokens 32768 \
  --max-num-seqs 32 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --tool-call-parser minimax_m3 \
  --enable-auto-tool-choice \
  --reasoning-parser minimax_m3 \
  --kv-cache-dtype fp8 \
  --language-model-only \
  --port 8000 \
  --attention-backend ROCM_AITER_UNIFIED_ATTN

Test Result

Kernel accuracy: 2/2 passed (short-context fallback + long-context split-K).

End-to-end serving on 4x MI325X (gfx942), TP4, EmbeddedLLM/MiniMax-M3-FP8-dynamic, FP8 KV cache:

ISL/OSL	Concurrency	Metric	Baseline	This PR	Change
8k/1k	4	Output tok/s	309.18	310.03	+0.3%
8k/1k	8	Output tok/s	470.73	469.60	-0.2%
128k/1k	4	Output tok/s	93.48	95.24	+1.9%
128k/1k	4	TTFT median (ms)	19229	18766	-2.4%
128k/1k	8	Output tok/s	105.05	106.97	+1.8%
128k/1k	8	TTFT median (ms)	34297	33430	-2.5%
256k/1k	2	Output tok/s	44.19	45.69	+3.4%
256k/1k	2	TTFT median (ms)	25539	24367	-4.6%
256k/1k	4	Output tok/s	49.73	51.58	+3.7%
256k/1k	4	TTFT median (ms)	42422	40513	-4.5%

Long-context prefill improvement scales with sequence length.

Short-context workloads are unaffected (selector falls back to unsplit path). Long-context prefill shows a consistent TTFT improvement.

Task	Filter	n-shot	Baseline	This PR
GSM8K (full, 1319)	flexible-extract	8	93.18% ± 0.69	92.87% ± 0.71
GSM8K (full, 1319)	strict-match	8	93.18% ± 0.69	92.80% ± 0.71

[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.