[hyukjlee]

Purpose

An opt-in specialized decode for Kimi-K3 on MI355X (gfx950), TP8, following RFC #60904. It uses the same specialized_decode shape as #61867 (Kimi-K3, NVIDIA) and #60889 (GLM-5.2, ROCm), in layer mode. It is enabled with additional_config={"decode_backend": "latency_specialized"} and has no other options.

A pure spec-verify decode step of at most 8 rows runs each MoE decoder layer as persistent FlyDSL MonoKernels. The TP all-reduces run in-kernel over peer memory:

• Whole layers (68 KDA layers): K1 (AttnRes, in_proj, f_b, spec conv, KDA recurrence, gated RMSNorm), then K2.
• FFN layers (24 MLA layers): vLLM's AttnRes and MLA up to o_proj's input, then K2.
• K2: o_proj + all-reduce + MLP AttnRes, then the latent MoE (router, top-16 MXFP4 experts, shared experts, latent all-reduce + RMSNorm, up_proj, final all-reduce).

The kernels read vLLM's loaded weights and KDA states in place. Every other step (prefill, mixed, wider and PIECEWISE steps) runs vLLM's layers. With decode_backend=auto (the default) nothing changes.

Layout: no common code

vllm/models/kimi_k3/amd/specialized_decode/
    README.md        supported deployments, refusals, steps, fail-stop
    hooks.py         SpecializedDecodeMixin: the model-side hooks
    model_state.py   KimiK3ModelState: builds the executor after load, binds KV caches
    runtime/
        __init__.py  DecodeConfig, specialized_decode_config(), unsupported()
        decode.py    SpecializedDecode: plan() -> Step, Step.layer(i), prepare_kv(), after_step()
        guards.py    weight / KV cache refusals, fail-stop
        layers.py    layer kernels: vLLM tensors -> launches
        runner.py    per-rank launch state;  peer.py  peer memory, TP agreement
    mono/            kernel source (K1, K2, MoE stages, device primitives)
tests/models/kimi_k3/amd_specialized_decode/

The existing Kimi-K3 files only call into the folder (26 added lines):

• linear.py: specialized_decode_config() in __init__ before any module allocates weights, plan_specialized_decode() in forward, (step.layer(i) or layer)(...) in the layer loop, and after_specialized_decode_step() in compute_logits.
• model.py: get_model_state_cls.
• mla.py: skip_o_proj, which returns o_proj's input for the MLA layers' K2.

Outside vllm/models/kimi_k3/ the only change is one Buildkite step.

Refusal and steps

• Refusal. An unsupported deployment fails at startup with one ValueError that lists every reason. Refusals run in three stages: the config in __init__, the loaded weights when the model state builds the executor, and the KV caches at every bind. Every TP rank agrees on each verdict. It is default-deny for PP/DP/EP/EPLB/PCP, DBO, LoRA, KV connectors, routed-expert capture, sleep mode, CPU offload, weight updates, adaptive or dynamic speculative decoding, ReplaySSM/RecoverSSM, the V1 runner, STOCK_TORCH_COMPILE, other model shapes, and other GPUs or partitioned GPUs.
• Steps. plan() runs once per step. It picks the kernels for pure spec-verify steps of whole uniform requests at a built width, i.e. multiples of k+1 up to 8 rows.
• CUDA graphs. vLLM's capture sizes are not changed. Under FULL graphs the plan is baked in at capture. A replay with fewer requests than its graph pads whole requests (null state slot, one accepted token), and the kernels serve those like real ones.

Fail-stop

Every in-kernel wait gives up after 10 s. The rank records it in its sticky health words, so a stalled peer fails every rank instead of hanging them. In compute_logits each rank queues a non-blocking copy of its FAILED word. Every 16 steps the TP ranks agree on what has landed over the TP CPU group, and then every rank raises at the same step. This keeps all ranks out of the next collective.

Detection takes up to about 16 steps. A zero-delay check through a core hook can follow in a separate PR.

Not a duplicate

• [Perf] Kimi K3 Specialized Decode: CuTeDSL Optimized GB300 (1.5-2x) #61867 is Kimi-K3 on NVIDIA (CuTeDSL, SM100); this is the ROCm/FlyDSL counterpart under kimi_k3/amd/.
• [ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X #60397, [ROCm] MonoKernel for GLM 5.2 #60889, [ROCm][Perf][DSv4] Standalone mono decode layer for DeepSeek-V4-Pro on MI355X #61482 and [ROCm][Perf][DSv4.1] Mono decode on MI300X and MI325X (gfx942) #60731 target other models (DeepSeek-V4.x, GLM-5.2), with different attention, residual and MoE.
• [Core] Shared abstractions for specialized decode kernels #61311 proposes shared abstractions. This PR keeps its runtime per model, as [Perf] Kimi K3 Specialized Decode: CuTeDSL Optimized GB300 (1.5-2x) #61867 and [ROCm] MonoKernel for GLM 5.2 #60889 do, and can move onto shared helpers once those land.

Test Plan

pytest -v -s tests/models/kimi_k3          # 8x MI355X; includes amd_specialized_decode/
pre-commit run --files <changed files>     # incl. mypy

• CPU: the switch (auto / unknown value), widths, layer pattern, every refusal (one at a time, then all at once), PIECEWISE warning, cache refusal, step eligibility, plan() routing (PIECEWISE, too wide, prefill, PP input, unbound or re-bound KV cache), weight-view refusals.
• GPU: K1 against vLLM's KDA op chain; K2 at TP8 against o_proj + all-reduce + AttnRes + KimiMoE, with 300 bit-identical graph replays; a stalled rank finishes every launch and fails the step; warm-up compiles without launching.
• Engine: an 8-layer dummy K3 at TP8 stalls rank 3 inside a served decode step. A request fails, every rank raises at the same step, and the server exits.

Kimi-K3, TP8, DCP8, DSpark MTP, fp8 KV, --stream-interval 10, MI355X, FULL CUDA graphs, Model Runner V2. "Off" means the same image, tree and flags with decode_backend=auto.

Test Result

This revision ran on 8x MI355X with vllm/vllm-openai-rocm:nightly-rocm100-1b2a6663, which is 20 commits behind this branch's base with no kimi_k3, mamba, gdn_attn or model-state changes in between. This branch's vllm/models/kimi_k3/ was mounted over the image.

• pytest tests/models/kimi_k3: 420 passed, 0 failed (89 skipped). This covers the new suite and every existing Kimi-K3 test.
• ISA: all 17 production K1/K2 builds are byte-identical to the previous revision (df3693c), so the kernel results below carry over.
• Startup refusal: TP4 + V1 runner + no speculative decoding gives one ValueError with all three reasons. decode_backend="fast" is rejected.
• 8k/1k, c1, k=7, 3 seeds: TPOT p50 5.44 → 4.31 ms (−20.8%), p90 5.48 → 4.35 ms, output 170.2 → 210.3 tok/s (+23.6%); the server logs Kimi-K3 specialized decode: 92 layers, widths [8], 30 builds.
• GSM8K, 5-shot, 1319 q:
  • c1, k=7, on: 0.9651 ± 0.0051 (an earlier run of the same cell on this revision: 0.9538 ± 0.0058).
  • Padded replays (k=1, c3, vLLM's default capture sizes, so 3-request steps replay the 4-request graph): off 0.9621 vs on 0.9613 (±0.0053).

Accuracy, previous revision (same kernels, ISA byte-identical)

The kernel path serves steps of at most 8 rows, so these configs fuse every decode step (conc x (k+1) <= 8):

eval	config	off	on
GSM8K 5-shot	c4, k=1	0.9644	0.9629 / 0.9568
GSM8K 5-shot	c2, k=3	0.9522	0.9568
GSM8K 5-shot	c1, k=7	0.9598	0.9629
GPQA-Diamond (198 q, CoT, 0-shot)	c4, k=1	0.9242	0.9394
AIME24 + AIME25 (60 q), avg@4, T=1.0	c4, k=1	0.7958	0.8083
NIAH single_1 / single_2 / multikey_1 @ 8k	c4, k=1	1.00 / 1.00 / 1.00	1.00 / 1.00 / 1.00
NIAH single_1 / multikey_2 / multikey_3 @ 32k / 128k / 256k	c1, k=7		1.000 (all)

Paired GSM8K (1319 q, c4/k1): 18 right only with the kernels on, 27 right only with them off, McNemar exact p = 0.23.

Performance, previous revision

vllm bench serve, random, 1024 output tokens, ignore-eos, 2 seeds. TPOT p50 in ms:

conc, k	widths	8k off	8k on	64k off	64k on
1, 7	8	5.47	4.35 (-20.5%)	5.67	4.47 (-21.2%)
2, 3	4, 8	21.58	16.79 (-22.2%)*
4, 1	8	26.55	22.18 (-16.5%)*
8, 1	(16 rows: vLLM's path)	32.60	32.58

* Synthetic acceptance length 1.0, so TPOT is the decode step time at that width. c1/k7 uses the CI's synthetic acceptance length 3.84.

Per-layer time (graph-replayed, profiler): K1 36.7 us (vs about 54 us plus gaps for vLLM's ops), K2 104.8 us (vs about 178 us).

AI assistance (Claude) was used to develop this change.

🤖 Generated with Claude Code

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

[Fangzhou-Ai]

Great work! Please use the main switch VLLM_ROCM_MONO_DECODE from #60397 to gate the mono path.

[Fangzhou-Ai]

can we follow the #60397 with a more organized way for all these kernel files? Putting everything under a common folder make it harder to identify a higher abstract for all mono kernels in the future.

[Fangzhou-Ai]

can we use VLLM_ROCM_MONO_DECODE as #60397 ?

[hyukjlee]

sure will unify the env settings

[Fangzhou-Ai]

gms8k at conc1 is not sufficient to justify this PR, please use the maximal conc you can take for mono kernel to test. On top of that please do SBS GPQA-Diamond, AIME and NIAH evaluation as well.

[hyukjlee]

@Fangzhou-Ai thank you for your valuable comments. I have updated the followings:

• env config settings and the overall layout to match PR#60397
• added CI to guard the precision
• tested all accuracy tests mentioned and updated the results in PR description

could you plz check if the updates are aligned with your expectation?

[seungrokj]

@hyukjlee @Fangzhou-Ai

Ran a eval suite in parallel. (C=1)

Kimi-K3 mono decode (VLLM_ROCM_MONO_K3=1): C=1 accuracy on MI355X

Setup

• Image vllm/vllm-openai-rocm:nightly-rocm100-81198e97ba7eee2a22540caaa756b7fdddcb4d93, with vllm/models/kimi_k3/amd/ and the envs.py hunk from this PR at e2e886c37. git apply did not apply cleanly to the image's older base, so I copied the directory over instead.
• MI355X, TP8. Serving flags from InferenceX recipe kimik3-fp4-mi355x-vllm-agentic-mtp: fp8 KV cache, DSpark MTP with num_speculative_tokens=7 and block rejection, max-num-seqs 2.
• VLLM_ROCM_MONO_DECODE=1 and VLLM_USE_BREAKABLE_CUDAGRAPH=1. py-spy on the workers shows the decode step running in _forward_mono_kda and _mono_k2.forward.
• Concurrency 1: one request at a time, at most 8 rows per decode step, so every decode step stays within the mono limit.

Results (C=1, mono on)

Eval	Protocol	Score	Mean accept len
GSM8K	5-shot, 1319 q, temp 0	strict 0.9697 ± 0.0047 / flexible 0.9689	4.72
RULER NIAH single_1	40 samples/len, greedy, thinking on	32K 1.000 · 128K 1.000 · 256K 1.000	2.36
RULER NIAH multikey_2	40 samples/len, greedy, thinking on	32K 1.000 · 128K 1.000 · 256K 1.000	2.61
RULER NIAH multikey_3	40 samples/len, greedy, thinking on	32K 1.000 · 128K 1.000 · 256K 1.000	2.62
GPQA-Diamond, Moonshot settings	198 q × 2 orders, temp 1.0, top_p 0.95, 128K, effort max	mean 0.9040 over 1 seeds; min 0.9040 / max 0.9040	2.62
GPQA-Diamond, Moonshot settings (re-extracted)	same samples, case-sensitive / bold / boxed answer extraction	mean 0.9268 over 1 seeds; empty responses 0/396
AIME25	30 q, boxed, temp 0.6, top_p 0.95, 64K, effort max (K3 tokenizer default)	mean 0.9333 over 4 seeds (σ 0.0000); min 0.9333 / max 0.9333	3.16

GPQA-Diamond with Moonshot's settings (effort max, temperature 1.0, top_p 0.95; no token limit is stated, so I used 128K) scores 92.7 ± 1.7 after re-grading. Moonshot's official number is 93.5, so the gap is −0.8, about half a standard error. No answer came back empty.

Moonshot publishes no official Kimi-K3 numbers for AIME, GSM8K or RULER/NIAH. For reference, Kimi K2.5 reports 96.1 on AIME25 (96K budget, 32 runs averaged). Of the 8 AIME25 misses (2 per seed), problem 23 got 150 instead of 149 on all four seeds. Problems 21 and 28 were each wrong once with a boxed but incorrect answer. Problems 8 and 13 each came back empty once, likely because the reasoning ran out of the 64K budget.

[mergify]

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

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

[seungrokj]

@hyukjlee plz follow #60904

[hyukjlee]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96326 for commit 0b5df303a4f0.

ℹ️ This PR is 58 commits behind upstream main. Merged upstream main into this branch as 0b5df303a4f0 before starting CI. Pull the branch before pushing more changes.

[Fangzhou-Ai]

LGTM