[ROCm][Perf] Fold scalar KV scales in MiniMax-M3 sparse attention

[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 Triton sparse attention currently applies scalar FP8 KV scales to every K and V element loaded inside the attention loop. These scales are loop invariant, so the repeated conversions and multiplications add unnecessary work.

Suggested Fix

For scalar scales (KV_SCALE_MODE == 1), fold the K scale into the attention score scale and apply the V scale once to the normalized FP32 accumulator.

The change covers AMD prefill, common prefill, and common decode. No-scale and per-token/head scale modes remain unchanged through the existing compile-time dispatch.

Test Plan

pytest -q tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py

Compares scalar and per-token/head FP8 KV-cache outputs against a dequantized BF16 reference for both prefill and decode.

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: 4/4 passed (scalar + per-token/head, prefill + decode).

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	315.52	+2.0%
8k/1k	4	TTFT median (ms)	1683	1541	-8.5%
8k/1k	8	Output tok/s	470.73	485.37	+3.1%
8k/1k	8	TTFT median (ms)	2123	2041	-3.9%
128k/1k	4	Output tok/s	93.48	99.85	+6.8%
128k/1k	4	TTFT median (ms)	19229	17552	-8.7%
128k/1k	8	Output tok/s	105.05	113.14	+7.7%
128k/1k	8	TTFT median (ms)	34297	31230	-8.9%

Task	Filter	n-shot	Baseline	This PR
GSM8K (full, 1319)	flexible-extract	8	93.18% ± 0.69	92.95% ± 0.71
GSM8K (full, 1319)	strict-match	8	93.18% ± 0.69	92.87% ± 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.