perf(flydsl): fuse Kimi-K3 B1 latent MoE tail

[JohnQinAMD]

Purpose

Add a fixed-shape gfx950 FlyDSL Kimi-K3 TP8 latent-MoE tail. One kernel
fuses RMSNorm of [1,3584], the BF16 [7168,3584] up-projection, BF16
materialization, and the shared-expert add into [1,7168]. Norm and dot
products accumulate in FP32.

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

Test plan

• Compare the fused boundary with the production PyTorch chain, including
  graph replay with changed input, output reuse, input immutability,
  unsupported contracts, and non-contiguous input.
• Benchmark the complete unfused and fused boundaries with eight rotating
  411-MB weights to limit cache reuse.

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-latent-tail-bf16 \
  --entrypoint bash "$IMAGE" -lc '
git clone -q https://github.com/ROCm/aiter.git /tmp/aiter && cd /tmp/aiter
git fetch -q origin pull/4496/head && git checkout -q --detach FETCH_HEAD
test "$(git rev-parse HEAD)" = ae0e0e12f16c9a818cf384e8ca427c789ca853cf
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_latent_moe_tail.py
'

Test results

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

Complete boundary	p50	Relative
RMSNorm + projection + add	13.6019 us	1.000x
Fused tail	9.9237 us	1.3706x

• Focused suite: 6/6 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#50657 owns serving dispatch and model-level evidence; no
endpoint gain is attributed to this primitive alone. This optimization is
limited to the exact Kimi-K3 BF16 gfx950 boundary above.

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 4496 --add-label <label>

[Copilot]

Pull request overview

Adds a narrow, gfx950-only FlyDSL implementation of the Kimi-K3 batch-1 latent-MoE “tail” (routed RMSNorm → up-projection GEMV → shared add), plus validation that it matches a reference and is safely gated behind a strict support predicate.

Changes:

• Introduces latent_moe_tail() / supports_latent_moe_tail() Python API with strict shape/dtype/layout/device gating and optional out= reuse.
• Adds a fixed-shape gfx950 FlyDSL kernel module implementing RMSNorm + GEMV + add in one launch.
• Adds dedicated FlyDSL tests for numerical matching, gating behavior, contiguity rejection, and graph-capture output reuse.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
op_tests/flydsl_tests/test_latent_moe_tail.py	New targeted tests for correctness, support gating, and CUDAGraph/out reuse behavior.
aiter/ops/flydsl/latent_moe_tail.py	New Python wrapper and support predicate for the fused latent-MoE tail specialization.
aiter/ops/flydsl/kernels/latent_moe_tail_gfx950.py	New fixed-shape gfx950 FlyDSL kernel implementing RMSNorm + BF16 projection + add.
aiter/ops/flydsl/init.py	Exposes the new latent_moe_tail public symbols when FlyDSL is available.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.