perf(flydsl): fuse Kimi-K3 MLA output gate on gfx950

[JohnQinAMD]

Purpose

Add a fixed-shape gfx950 FlyDSL Kimi-K3 MLA decode epilogue. It fuses the BF16
[1,7168] x [1536,7168] projection, BF16 materialization, sigmoid, and
elementwise multiply with the [1,1536] attention output. Accumulation is
FP32; the returned boundary remains BF16.

The wrapper fails closed outside the exact device, shape, dtype, and stride
contract.

Test plan

• Validate primary dispatch, output reuse and aliasing, changed-input graph
  replay, and invalid contracts against the production PyTorch boundary.
• Benchmark both complete boundaries with 13 rotating weights totaling
  286.26 MB.

IMAGE='vllm/vllm-openai-rocm:kimi-k3@sha256:5aa7e626ff73672f5ca7aae46754570488c23d33ca1ac90756a1d2d1a3fe099b'
docker run --rm --ipc=host --shm-size=16g \
  --device=/dev/kfd --device=/dev/dri --group-add=video --group-add=render \
  -e HIP_VISIBLE_DEVICES=0 -e AITER_JIT_DIR=/tmp/aiter-jit-mla-gate \
  --entrypoint bash "$IMAGE" -lc '
git clone -q https://github.com/ROCm/aiter.git /tmp/aiter && cd /tmp/aiter
git fetch -q origin pull/4497/head && git checkout -q --detach FETCH_HEAD
test "$(git rev-parse HEAD)" = d2f4c61f5ee17e6fe8436b8cf9e32748152745f3
git submodule update -q --init --recursive
AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=gfx950 python3 setup.py develop
python3 -m pytest -q op_tests/flydsl_tests/test_kimi_k3_mla_gate_epilogue.py
'

Test results

MI355X/gfx950, TP8 batch-one decode, BF16 input/weight/attention/output; 21
alternating paired trials:

Complete boundary	p50	Relative
Projection + sigmoid gate	12.7591 us	1.000x
Fused epilogue	5.4204 us	2.3539x

• Focused suite: 8/8 passed.
• Output relative RMSE: 0 against the production BF16 boundary.
• Ruff, formatting, git diff --check, and DCO checks passed.

Overlap and limits

vllm-project/vllm#50664 owns serving dispatch. No model-level accuracy change
is expected because the production BF16 materialization is preserved; no
endpoint gain is attributed to this primitive alone.

Tool assistance

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

[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 4497 --add-label <label>

[Copilot]

🟡 Not ready to approve

The new FlyDSL kernel imports vendored-removed FlyDSL modules and the tests need tighter gating and a reference that matches the intended BF16 rounding boundaries.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Pull request overview

Adds a gfx950-only BF16 FlyDSL kernel and Python wrapper for the Kimi-K3 full-MLA output gate (gate projection + stable sigmoid + attention multiply), plus targeted unit tests to validate dispatch, graph replay, and out= aliasing behavior.

Changes:

• Introduces a fixed-shape gfx950 BF16 FlyDSL kernel implementing projection + sigmoid + multiply in one launch.
• Adds a narrow Python dispatch/launcher (supports_* + kimi_k3_mla_gate) that raises NotImplementedError outside the supported envelope.
• Adds focused FlyDSL tests for dispatch, accuracy, out= reuse/aliasing, CUDA graph replay, and rejection paths.

File summaries

File	Description
op_tests/flydsl_tests/test_kimi_k3_mla_gate_epilogue.py	Adds gfx950-focused tests for the new MLA gate specialization (dispatch/accuracy/graph/out=).
aiter/ops/flydsl/kimi_k3_mla_gate.py	Adds the narrow dispatch predicate and launcher wrapper for the fused MLA gate epilogue kernel.
aiter/ops/flydsl/kernels/kimi_k3_mla_gate_epilogue_gfx950.py	Implements the gfx950 BF16 FlyDSL kernel for the fused gate projection + sigmoid + attention multiply.

Review details

Suppressed comments (2)

op_tests/flydsl_tests/test_kimi_k3_mla_gate_epilogue.py:70

• Same FP32-only reference issue as the primary accuracy test: the kernel rounds projection and sigmoid to BF16 before multiplying, but the reference rounds only once at the end, so it may not catch BF16 boundary mismatches.

    reference = (
        torch.sigmoid(hidden.float() @ weight.float().T) * attention_input.float()
    ).to(torch.bfloat16)

op_tests/flydsl_tests/test_kimi_k3_mla_gate_epilogue.py:98

• Same FP32-only reference issue in the graph replay test: reference should mirror the kernel's BF16 rounding boundaries (projection->bf16, sigmoid->bf16) to make the test able to falsify BF16 rounding regressions.

    reference = (
        torch.sigmoid(hidden.float() @ weight.float().T) * attention.float()
    ).to(torch.bfloat16)

• Files reviewed: 3/3 changed files
• Comments generated: 3
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.