Skip to content

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

Open
JohnQinAMD wants to merge 6 commits into
mainfrom
perf/kimi-k3-kda-fb-clean
Open

JohnQinAMD wants to merge 6 commits into
mainfrom
perf/kimi-k3-kda-fb-clean

Conversation

@JohnQinAMD

@JohnQinAMD JohnQinAMD commented Jul 31, 2026

Copy link
Copy Markdown
Owner

Summary

Add a gfx950 Kimi-K3 KDA decode kernel, then incrementally fuse the head-local
128x128 f_b projection into the same workgroup.

The first commit fuses causal convolution, recurrent state update, and gated
RMSNorm. The second removes the low-occupancy f_b GEMV, its global g1
materialization, and the dependency bubble before the state core.

Dispatch and fallback

The public wrapper validates device, shape, dtype, stride, cache layout, batch,
and gate-bound contracts. Unsupported requests are rejected before launch and
remain owned by the vLLM fallback. The kernel is exposed only through the
FlyDSL availability boundary.

Performance

Kimi-K3 TP8 batch-one decode on MI355X:

Chain Time/layer Relative
Unfused control 9.76 us 1.00x
Decode + f_b 7.15 us 1.36x

The accepted endpoint moved from 47.644 to 50.063 tok/s (+5.1%) and reached
19.953 ms TPOT.

Correctness and validation

  • Focused tests cover contract rejection, cache/state mutation, graph replay,
    and fused-vs-unfused output/state comparison.
  • The exact clean AITER source plus the final vLLM perf(flydsl): fuse Kimi-K3 mixed pre-route projections on gfx950 #12-perf(flydsl): select validated Kimi-K3 KDA schedule #14 adapter stack passed
    all 9 focused cases in the immutable MI355X validation image, including both
    canonical convolution-cache layouts.
  • The accepted endpoint stack containing this exact kernel passed frozen
    GSM8K first 100 at 100/100 with zero invalid responses.
  • Ruff check and Ruff format pass all five changed files.
  • git diff --check passes.
  • Both commits carry DCO and AI-assistance trailers.

Fork Black and Ruff CI pass. The PR is ready for review; hardware workflows
were triggered by the ready-for-review transition.

Disclosure

Developed with assistance from OpenAI Codex. The author reviewed the kernel
contracts, numerical path, tests, and benchmark evidence.

Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
Assisted-by: OpenAI Codex
Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
Assisted-by: OpenAI Codex
@github-actions

Copy link
Copy Markdown

🏷️ 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 16 --add-label <label>

Sort public imports and exports, use functools.cache, and annotate intentional FlyDSL stream defaults for the repository-wide lint configuration.

Assisted-by: OpenAI Codex
Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
@JohnQinAMD
JohnQinAMD force-pushed the perf/kimi-k3-kda-fb-clean branch from 37943cc to 762b3e9 Compare August 2, 2026 02:08
op_tests/op_benchmarks/flydsl/ does not exist on main, no CI job references
it, and none of the 84 existing flydsl kernel modules ship a companion
benchmark. The script is kept out of tree; measured numbers go in the PR
description instead.

Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant