[ROCm][Perf] Fuse Kimi-K3 MLA output gate

[JohnQinAMD]

Purpose

Wire the AITER Kimi-K3 MLA output-gate fusion into the AMD model path while
preserving the generic implementation.

MLAModules gains an optional output-gate callback. Its default remains the
existing projection, sigmoid, and multiply expression; only the Kimi-K3 AMD
adapter installs the AITER callback. Unsupported AITER installations, devices,
shapes, dtypes, or layouts use the existing expression. NVIDIA dispatch and
generic MLA split policy are unchanged.

Depends on ROCm/aiter#4497.

Test plan

Tested on 8x MI355X (gfx950) with the public Kimi-K3 image and
moonshotai/Kimi-K3@9f62e4e9. The AITER #4497 kernel head is d2f4c61f; the
command below pins the vLLM base and PR revisions.

Fetch and verify the vLLM source:

git clone https://github.com/vllm-project/vllm.git vllm
git -C vllm fetch origin refs/pull/50664/head:pr-50664
test "$(git -C vllm rev-parse pr-50664)" = \
  d367e2727440ee838c2035883e6a1a0b6c8755a7
git -C vllm checkout --detach 6c91de36897932ba9b5adb11992235a2a789e009
git -C vllm diff 6c91de36897932ba9b5adb11992235a2a789e009..pr-50664 \
  --binary | git -C vllm apply -
git -C vllm diff --check

Use complete verified vLLM/AITER trees as overlays on the public image,
including AITER csrc/ and hsa/; do not copy individual files. Verify
import vllm, vllm._C, aiter before testing.

python -m pytest -q \
  ../aiter/op_tests/flydsl_tests/test_kimi_k3_mla_gate_epilogue.py \
  tests/models/kimi_k3/test_amd_mla_gate.py

Serve both source-isolated arms with identical flags:

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FP4BMM=1 \
AITER_SITUV2_A8W4=1 AITER_BF16_FP8_MOE_BOUND=0 \
VLLM_USE_BREAKABLE_CUDAGRAPH=0 \
vllm serve /model --served-model-name moonshotai/Kimi-K3 \
  --tensor-parallel-size 8 --trust-remote-code --moe-backend auto \
  --gpu-memory-utilization 0.95 --max-num-seqs 128 \
  --max-num-batched-tokens 4096 --max-model-len 1048576 \
  --enable-prefix-caching --kv-cache-dtype fp8 --reasoning-parser kimi_k3

After one 8K/128 warmup, run three 8K/1K batch-one trials with seeds 1--3,
temperature zero, and --ignore-eos. Run full lm-eval==0.4.12 GSM8K on both
arms: 1,319 questions, 5-shot, greedy completion, 2,048 generated tokens,
concurrency 128, and seed 42.

Test results

• Combined focused suite: 11/11 passed.
• Changed-file pre-commit, git diff --check, and DCO passed.

Measurement	Control	Candidate	Change
AITER gate microbenchmark	12.7591 us	5.4204 us	2.3539x
TP8 decode throughput	36.0500 tok/s/GPU	36.3890 tok/s/GPU	+0.940%
TPOT	27.7393 ms	27.4808 ms	-0.2585 ms

Endpoint values are medians of three fixed-seed 8K/1K trials after warmup.

Full GSM8K: control 1277/1319, candidate 1273/1319, zero invalid or
transport failures; 8 wins/12 losses (p=0.5034). This paired run found no
statistically detectable accuracy difference.

Limits

ROCm/aiter#4497 owns the kernel; this PR owns only the generic callback seam and Kimi-K3
AMD adapter. Output-buffer ownership is intentionally outside this change.
Operator tests establish the fused BF16 numerical contract; GSM8K is the
model-quality gate.

Tool assistance

OpenAI Codex assisted with implementation, tests, benchmarking, and drafting
this description.

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

[AndreasKaratzas]

Nobody is going to easily review a FULLY vibe coded PR.

[jamesETsmith]

Thanks for this @JohnQinAMD, I was about to implement the fusion myself until I saw you'd already done it.

The tests in the PR also only test the mock gate and not the actual code path. I think it would make sense to consturct KimiMLAAttenntion and ensure that your changes are working. I think if you removed output_gate=kimi_k3_mla_output_gate, in vllm/models/kimi_k3/amd/linear.py all your tests would still pass.

[jamesETsmith]

If I'm reading it correctly, I don't think the signature of _get_aiter_mla_gate in this patch matches what the function can return (either None or (kimi_k3_mla_gate, supports_kimi_k3_mla_gate))

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @JohnQinAMD.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork