[CI][AMD][Disagg] Fix Kimi K2.5/K2.6 MXFP4 MLA backends on ROCm nightly for accuracy eval

[avininjamay8]

Purpose

Kimi K2.5/K2.6 MXFP4 disaggregated GSM8K on vllm/vllm-openai-rocm:nightly was failing when MLA went through aiter:

• ROCM_AITER_MLA (and VLLM_ROCM_USE_AITER_MLA=1 on K2.6) — bad prefill / bring-up under load.
• ROCM_AITER_TRITON_MLA — engine did not pass health (same class of failure as full aiter MLA on nightly).
• TRITON_MLA without FLASH_ATTN prefill — health OK, then prefill HTTP 500 from aiter fmha_fwd_bf16_opus_fwd during GSM8K.

Use TRITON MLA for decode and FlashAttention for MLA prefill (-ac.mla_prefill_backend=FLASH_ATTN).

Changes

• models.yaml — Kimi-K2.5-MXFP4: add --attention-backend TRITON_MLA -ac.mla_prefill_backend=FLASH_ATTN to base_flags.
• models.yaml — Kimi-K2.6-MXFP4: drop VLLM_ROCM_USE_AITER_MLA from env; set block-size 16 and the same TRITON + FLASH_ATTN flags (replace ROCM_AITER_MLA / block-size 1).

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 -ac.mla_prefill_backend=FLASH_ATTN \
    --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 -ac.mla_prefill_backend=FLASH_ATTN \
    --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"}}'

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 -ac.mla_prefill_backend=FLASH_ATTN \
    --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 -ac.mla_prefill_backend=FLASH_ATTN \
    --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"}}'

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)

Config	Without fix	With fix
K2.5 1P1D TP8 (proxy)	Fail: ROCM_AITER_TRITON_MLA (init) or TRITON_MLA only (fmha / eval)	0.948
K2.6 1P1D TP8 (proxy)	Fail: ROCM_AITER_MLA / ROCM_AITER_TRITON_MLA (init) or TRITON_MLA only (fmha / eval)	0.948

[lcskrishna]

Why do we need to move prefill backend to Flash attn?

[avininjamay8]

Thats a fair question, actually : TRITON_MLA only fixes decode, without -ac.mla_prefill_backend=FLASH_ATTN, ROCm still uses aiter fmha for MLA prefill, which passes health on nightly but breaks GSM8K : FlashAttention prefill avoids that failure while we keep Triton decode.

[avininjamay8]

@lcskrishna added an alternate solution : #22 . Please review.

[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);