perf(flydsl): fuse Kimi-K3 KDA decode and f_b projection

[JohnQinAMD]

Purpose

Add gfx950 FlyDSL specializations for Kimi-K3 KDA decode and for the
incremental boundary that also computes the head-local 128 x 128 f_b
projection. The latter removes a low-occupancy GEMV, global g1
materialization, and its dependency bubble.

The wrapper owns the exact device, shape, dtype, stride, cache, state-index,
and gate contract and fails closed otherwise. The two kernels keep separate
static ABIs and schedules.

Test plan

• Compare output plus convolution/recurrent-state mutation with a PyTorch
  oracle at batch 1, 8, and 16, including graph replay and invalid contracts.
• Benchmark separate f_b + decode against the complete fused boundary.

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

Test results

MI355X/gfx950, TP8 batch-one decode, BF16 output/cache/gates, FP32 recurrent
state; 21 alternating paired trials:

Complete boundary	p50	Relative
Separate f_b + decode	8.1055 us	1.000x
Fused f_b + decode	6.6118 us	1.2259x

• Focused suite: 12/12 passed.
• Output and recurrent-state relative RMSE: 0 after one matched update.
• Ruff, formatting, git diff --check, and DCO checks passed.

Overlap and limits

vllm-project/vllm#50654 overlaps the decode-core ownership. Reviewers should
choose one core implementation; this PR's distinct incremental fused-f_b
boundary remains independently useful. No endpoint gain is attributed here.

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

[Copilot]

Pull request overview

Adds a new gfx950-only FlyDSL decode specialization for Kimi-K3 KDA, and an optional fused variant that folds the head-local 128×128 f_b projection into the same workgroup to remove an intermediate global materialization and kernel launch.

Changes:

• Introduces public Python wrappers for flydsl_kimi_k3_kda_decode and flydsl_kimi_k3_kda_decode_with_f_b, plus a is_flydsl_kimi_k3_kda_decode_supported capability predicate.
• Adds two new FlyDSL kernels implementing (1) KDA decode and (2) fused f_b projection + KDA decode on gfx950.
• Adds focused correctness tests comparing fused vs. reference behavior and validating contract rejection.

Reviewed changes

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

Show a summary per file

File	Description
op_tests/flydsl_tests/test_kimi_k3_kda_decode.py	Adds correctness/contract tests for the new decode and fused-f_b decode paths.
aiter/ops/flydsl/kimi_k3_kda_decode.py	New high-level API wrapper with device/shape/dtype/stride validation and kernel launch.
aiter/ops/flydsl/kernels/kimi_k3_kda_decode.py	New FlyDSL kernel implementing fused conv + recurrent state update + RMSNorm/gate for decode.
aiter/ops/flydsl/kernels/kimi_k3_kda_decode_fb.py	New FlyDSL kernel that fuses the f_b projection into the decode workgroup.
aiter/ops/flydsl/init.py	Exposes the new public symbols behind the FlyDSL availability gate.

Suppressed comments (1)

aiter/ops/flydsl/kimi_k3_kda_decode.py:376

• batch = f_a.shape[0] if f_a.ndim == 2 else -1 leads to the same misleading failure mode as the decode-only wrapper: rank!=2 inputs trip the "non-empty batch dimension" check rather than reporting the expected rank/shape for f_a. Adding an explicit f_a.ndim check makes the error actionable.

    if not f_a.is_cuda:
        raise ValueError("`f_a` must be a CUDA tensor.")
    device = f_a.device
    batch = f_a.shape[0] if f_a.ndim == 2 else -1
    _check_same_device((("f_b_weight", f_b_weight),), device)

---

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