[ROCm][Bugfix] Use BF16 MLA prefill for short prompts

[LiuYinfeng01]

Summary

Route ROCm MLA prefill requests shorter than 1024 tokens through the BF16 FlashAttention implementation. Keep the optimized AITER FP8 MLA prefill path for long, context-free requests.

Problem

AiterMLAImpl.forward_mha() previously selected AITER FP8 prefill for every request without chunked context, including short greedy-decoding prompts. On Kimi-K3 PP8, the GSM8K-100 5-shot prompts were 603–698 tokens. FP8 prefill changed borderline logits and produced run-dependent wrong answers; rerunning an individual failed prompt could return the correct answer.

The 1024-token threshold covers the observed short-prompt range while leaving long-context prefill on AITER FP8. The controlled A/B below reports the measured short-request cost rather than assuming it is negligible.

Scope

This is only the short-prompt MLA prefill fix. It does not change FlyDSL MoE stage-2. The remaining MoE accuracy work is tracked separately in ROCm/AITER#4502.

Validation environment

• 8 × AMD Instinct MI355X (gfx950)
• Kimi-K3, TP=1, PP=8
• vLLM base: beca88e59ea75a7aa1af72a5ae50188fa91d4e3d
• AITER: 026eac044851de98f05d1fd6d90a10fd8be8fe84
• AITER MLA reducer support for the 96-head × 128-dim Kimi shape
• max_model_len=102400, max_num_batched_tokens=4096, max_num_seqs=1
• TRITON_MLA decode and automatic MLA prefill selection
• AITER A8W4 SiTUv2 MoE enabled; no CK stage-2 fallback

Relevant environment:

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_MLA=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1
export AITER_SITUV2_A8W4=1
export AITER_BF16_FP8_MOE_BOUND=0
export VLLM_USE_BREAKABLE_CUDAGRAPH=1

Serve configuration:

vllm serve /data/Kimi-K3/Kimi-K3   --trust-remote-code   --host 127.0.0.1 --port 8010   --tensor-parallel-size 1   --pipeline-parallel-size 8   --max-model-len 102400   --max-num-batched-tokens 4096   --max-num-seqs 1   --gpu-memory-utilization 0.70   --attention-backend TRITON_MLA

Accuracy reproduction

Use vLLM's in-tree GSM8K evaluator against the running endpoint:

python - <<'PY'
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k

print(evaluate_gsm8k(
    num_questions=100,
    num_shots=5,
    max_tokens=256,
    model="/data/Kimi-K3/Kimi-K3",
    host="http://127.0.0.1",
    port=8010,
    temperature=0.0,
    seed=42,
    request_timeout_seconds=1200,
    max_concurrency=1,
))
PY

Observed on the same image and launch configuration:

• Before this change: 98/100, 0 invalid responses.
• With this change: 99/100, 0 invalid responses.
• The observed query lengths were 603–698, so all evaluated prompts use BF16 prefill after this change.
• This PR does not claim 100/100 or full determinism; the remaining run-to-run mismatch is associated with the separate FlyDSL MoE path.

Performance validation

• For a 100K request, chunked prefill uses 4096-token chunks. Those chunks remain above the 1024 threshold and continue to use AITER FP8 prefill.
• Measured steady-state 100K TTFT on this configuration was 6.40–6.45 seconds.
• Short-request cost is reported in the controlled A/B section; no isolated prefill microbenchmark is claimed.

Checks

• python -m py_compile vllm/v1/attention/backends/mla/rocm_aiter_mla.py
• Kimi-K3 PP8 GSM8K-100, 5-shot, greedy, seed 42
• Kimi-K3 PP8 100K-context steady-state TTFT

Controlled A/B

The two images used the same hardware, model, vLLM base, AITER commit, 96×128 reducer, environment variables, launch parameters, evaluator, prompts, seed, and concurrency. The only image difference was:

# A: original source from vLLM beca88e59
# if has_context:

# B: FROM the exact A image, then replace only this source file
COPY rocm_aiter_mla.py \
  /opt/vllm-main/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
# if prefill_metadata.max_query_len < 1024 or has_context:

Run the serve command and in-tree evaluator above once with A and once with B after server readiness.

A — original FP8 short prefill:

• Accuracy: 98/100
• Invalid responses: 0/100
• Total evaluation latency: 416.478 s
• Output tokens: 9,699
• Output throughput: 23.288 tok/s

B — BF16 prefill below 1024 tokens:

• Accuracy: 99/100
• Invalid responses: 0/100
• Total evaluation latency: 444.953 s
• Output tokens: 9,750
• Output throughput: 21.912 tok/s

Delta (B versus A):

• Accuracy: +1 percentage point
• Invalid responses: unchanged
• Total evaluation latency: +6.84%
• Output tokens: +0.53%
• Output throughput: -5.91%

The end-to-end latency includes autoregressive generation and differing output-token counts; it is not an isolated prefill latency measurement. Because the remaining FlyDSL MoE path is nondeterministic, this single-run A/B demonstrates the observed improvement but does not establish a deterministic accuracy guarantee. A repeated-run confidence interval is still required before claiming a stable expected accuracy gain.

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