[CI][ROCm][Disagg] Fix Kimi K2.5/K2.6 MXFP4 MLA backends on ROCm nightly for accuracy eval (disable aiter OPUS FMHA)

[avininjamay8]

Purpose

Kimi K2.5/K2.6 MXFP4 disaggregated GSM8K on vllm/vllm-openai-rocm:nightly was failing on several MLA backends:

• ROCM_AITER_MLA (and VLLM_ROCM_USE_AITER_MLA=1 on K2.6) — prefill / bring-up failures under load.
• ROCM_AITER_TRITON_MLA — health gate failed at engine init.
• TRITON_MLA only — health passed, but GSM8K then failed at MLA prefill with HTTP 500 and fmha_fwd_bf16_opus_fwd TypeError.

With --attention-backend TRITON_MLA, vLLM uses Triton MLA for decode only. On ROCm, MLA prefill is selected separately and defaults to ROCM_AITER_FA, which calls amd-aiter's fused multi-head attention (FMHA) kernels. On nightly (aiter 0.1.19), that prefill path picks the OPUS FMHA kernel (fmha_fwd_bf16_opus_*), which is broken and triggers the TypeError above.

We therefore use TRITON MLA for decode, keep ROCM_AITER_FA for MLA prefill, and set AITER_DISABLE_FMHA_OPUS=1 so aiter skips OPUS and uses its older fmha_v3 / Composable Kernel (CK) FMHA implementation instead.

For K2.6, this stack also needs --block-size 16 — with --block-size 1, nightly failed health with mla_gluon[bh16bn128] requires batch_size=1, got 512.

Changes

• Kimi-K2.5-MXFP4 — add AITER_DISABLE_FMHA_OPUS: "1" to env; add --attention-backend TRITON_MLA to base_flags.
• Kimi-K2.6-MXFP4 — remove VLLM_ROCM_USE_AITER_MLA from env; add AITER_DISABLE_FMHA_OPUS: "1"; TRITON MLA decode and --block-size 16 (keep existing K2.6 fp8 KV and mm-encoder-tp-mode data flags).

Test Plan

MoRIIO 1P1D TP8 disaggregated serving, ROUTER_TYPE=proxy, RUN_AFTER_HEALTH=accuracy, image vllm/vllm-openai-rocm:nightly.

Kimi-K2.5-MXFP4 — TP8 1P1D

# PREFILL — TP8
vllm serve /data/models2/Kimi-K2.5-MXFP4 \
    --host $PREFILL_IP --port 8100 --tensor-parallel-size 8 \
    --trust-remote-code \
    --attention-backend TRITON_MLA \
    --gpu-memory-utilization 0.85 \
    --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer",
      "kv_connector_extra_config":{"proxy_ip":"$PREFILL_IP","proxy_ping_port":"36367",
      "http_port":"8100","handshake_port":"6301","notify_port":"61005"}}'

# DECODE — TP8
vllm serve /data/models2/Kimi-K2.5-MXFP4 \
    --host $DECODE_IP --port 8200 --tensor-parallel-size 8 \
    --trust-remote-code \
    --attention-backend TRITON_MLA \
    --gpu-memory-utilization 0.85 \
    --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_consumer",
      "kv_connector_extra_config":{"proxy_ip":"$PREFILL_IP","proxy_ping_port":"36367",
      "http_port":"8200","handshake_port":"7301","notify_port":"62005"}}'

Set AITER_DISABLE_FMHA_OPUS=1 in the prefill/decode environment before launch.

Kimi-K2.6-MXFP4 — TP8 1P1D

# PREFILL — TP8
vllm serve /data/models2/Kimi-K2.6-MXFP4 \
    --host $PREFILL_IP --port 8100 --tensor-parallel-size 8 \
    --trust-remote-code --kv-cache-dtype fp8 --mm-encoder-tp-mode data --block-size 16 \
    --attention-backend TRITON_MLA \
    --gpu-memory-utilization 0.9 \
    --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer",
      "kv_connector_extra_config":{"proxy_ip":"$PREFILL_IP","proxy_ping_port":"36367",
      "http_port":"8100","handshake_port":"6301","notify_port":"61005"}}'

# DECODE — TP8
vllm serve /data/models2/Kimi-K2.6-MXFP4 \
    --host $DECODE_IP --port 8200 --tensor-parallel-size 8 \
    --trust-remote-code --kv-cache-dtype fp8 --mm-encoder-tp-mode data --block-size 16 \
    --attention-backend TRITON_MLA \
    --gpu-memory-utilization 0.9 \
    --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_consumer",
      "kv_connector_extra_config":{"proxy_ip":"$PREFILL_IP","proxy_ping_port":"36367",
      "http_port":"8200","handshake_port":"7301","notify_port":"62005"}}'

Set AITER_DISABLE_FMHA_OPUS=1 in the prefill/decode environment before launch.

GSM8K evaluation

lm_eval --model local-completions \
    --tasks gsm8k \
    --model_args "model=/data/models2/Kimi-K2.5-MXFP4,base_url=http://127.0.0.1:10001/v1/completions,num_concurrent=16,max_retries=3,tokenized_requests=False,trust_remote_code=True,timeout=7200" \
    --limit 250

Use the K2.6 model path the same way.

Test Result

Accuracy results (GSM8K, flexible-extract, limit 250)

Verified on vllm/vllm-openai-rocm:nightly (0.26.1rc1.dev602+g65b7662d3, aiter 0.1.19, Aug 11 2026).

Config	Without fix	With fix
K2.5 1P1D TP8 (proxy)	Fail: ROCM_AITER_TRITON_MLA (init) or TRITON_MLA only (aiter OPUS FMHA / eval)	0.956
K2.6 1P1D TP8 (proxy)	Fail: ROCM_AITER_MLA / ROCM_AITER_TRITON_MLA (init) or TRITON_MLA only (aiter OPUS FMHA / eval)	0.936

[avininjamay8]

@lcskrishna requesting a review as alternate to #21

[avininjamay8]

Note: an earlier draft included --disable-custom-all-reduce as a workaround for vLLM vllm-project#50999 vs AITER custom all-reduce. That is no longer needed after vLLM vllm-project#51635 (merged Aug 10, 2026);