[MLA] Support 96-head 128-dim reduction

[LiuYinfeng01]

Summary

• add the missing 96 heads × 128 value dimension specialization to MLA_REDUCE_ROUTER
• cover the Kimi K3 PP8 reducer shape in op_tests/test_mla_reduce.py
• explicitly reject non-finite reducer output in addition to checking output/LSE error against the Torch reference

Motivation

Kimi K3 uses 96 MLA query heads with a 128-dimensional value/output head. In a PP8 deployment every pipeline rank still executes attention with all 96 heads, so the phase-2 MLA reducer is called with:

num_heads = 96
head_dim = 128
q_len = 1 (decode)

The reducer already had a 96 × 512 specialization, but MLA_REDUCE_ROUTER did not contain 96 × 128. As a result, a clean JIT build could compile module_mla_reduce from the sources listed in aiter/jit/optCompilerConfig.json, but the resulting module rejected the Kimi shape at dispatch time:

mla_reduce_v1 doesn't support the specified settings: #heads: 96, head dimension: 128

This is an AITER reducer dispatch gap, not a Triton head-count limitation and not a precompiled-binary issue. module_mla_reduce is already correctly sourced from mla_reduce_pybind.cu and reduce.cu in optCompilerConfig.json; no JIT configuration change is needed once the missing template specialization is present.

Validation

Hardware: MI355X (gfx950)

Clean-cache command:

python3 -m pip install tabulate
rm -rf /tmp/aiter-jit-clean
AITER_JIT_DIR=/tmp/aiter-jit-clean AITER_ROCM_ARCH=gfx950 python3 op_tests/test_mla_reduce.py --head_dim 128

Observed clean build:

finish build [module_aiter_core], cost 11.2s
finish build [module_mla_reduce], cost 150.8s

All cases passed against the Torch reference with rtol=2e-2, atol=2e-2, and finite output/LSE checks. The new 96 × 128, q_len=1 coverage passed:

splits=[2]
splits=[3, 2]
splits=[4]
splits=[8, 5, 7]
splits=[33]
splits=[300]
splits=[2, 4, 16, 64]
splits=40x[6]

These cover the simple reducer, massive reducer, ragged/mixed split counts, LDS-spill bucket, and persistent-grid path.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4625 --add-label <label>