[MiniMax-M3] optimize sparse vLLM metadata for improved performance of vllm-atom

[lirui927]

Motivation

Improve MiniMax-M3 sparse attention performance in the vLLM plugin path by reducing per-step metadata overhead and enabling CUDA Graph capture for the uniform single-token decode path.

Technical Details

• Marked MinimaxM3SparseAttentionMetadataBuilder as UNIFORM_SINGLE_TOKEN_DECODE instead of disabling CUDA Graph support entirely.
• Added a decode-only metadata fast path for max_query_len == 1, avoiding decode/prefill splitting and CPU-side max reductions for plain decode batches.
• Reused preallocated qo_indptr for prefill metadata and replaced exact CPU-side max calculations with existing common metadata upper bounds.
• Moved MiniMax-M3 sparse block-size validation to builder initialization.
• Reused page-16 KV cache views between sparse KV/index insertion and sparse attention execution to avoid duplicate view setup.
• support index cache of vLLM-ATOM

Test Plan

• Launch MiniMax-M3 vLLM service with ATOM adapters enabled with index cache.
• Run GSM8K accuracy validation:
  • MIXFP8
  • MIXFP4
  • MIXFP8 + kv_fp8
  • MIXFP4 + kv_fp8
• Performance Test

MODEL=/shared/data/amd_int/models/MiniMax-M3-MXFP8
export AITER_QUICK_REDUCE_QUANTIZATION=INT4
vllm serve "$MODEL" \
    --dtype auto \
    --load-format auto \
    --host localhost \
    --port 8001 \
    --tensor-parallel-size 4 \
    --gpu-memory-utilization 0.85 \
    --max-model-len 32768 \
    --max-num-batched-tokens 32768 \
    --block-size 128 \
    --no-async-scheduling \
    --kv-cache-dtype fp8 \
    --no-enable-prefix-caching \
    --language-model-only \
    --no-trust-remote-code \
    --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' \
    --additional-config '{"online_quant_config": {"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head", "model.embed_tokens", "vision_tower", "multi_modal_projector", "patch_merge_mlp", "*block_sparse_moe"]}}' \
    --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
    2>&1 | tee log_m3_mxfp8_vllm.log

MODEL=/shared/data/amd_int/models/MiniMax-M3-MXFP4
export AITER_QUICK_REDUCE_QUANTIZATION=INT4
vllm serve "$MODEL" \
    --dtype auto \
    --load-format auto\
    --host localhost \
    --port 8001 \
    --tensor-parallel-size 4 \
    --gpu-memory-utilization 0.85 \
    --max-model-len 32768 \
    --max-num-batched-tokens 32768 \
    --block-size 128 \
    --kv-cache-dtype fp8 \
    --no-enable-prefix-caching \
    --language-model-only \
    --no-trust-remote-code \
    --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' \
    --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
    2>&1 | tee log_m3_mxfp4_vllm_0625.log

Test Result

GSM8K accuracy

MIXFP4

MIXFP4-kv_fp8

MIXFP8

MIXFP8-kv_fp8

Performance

MIXFP4-8192-1024-16

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.