[kevin-mii]

Motivation

Kimi-K3.1 on MI35x is fastest when each step size runs the launch built for it. #102 covers verify steps of 2 .. 48 requests with the wide layer launch, and leaves a step of one request (6 rows with DFLASH at 6 draft tokens) on SGLang's path. This PR serves those steps on Kimi-K3's per-step launches, which K3.1 shares (the shapes are the same; K3.1 adds the DSA layers, which keep SGLang's attention up to o_proj).

It carries Kimi-K3's gfx950 launches from #13 (@billishyahao), with the work on top of them that the K3.1 bring-up measured, and supersedes #13.

Modifications

• kimi_k3: Kimi-K3's launches at TP8, for decode and verify steps of up to 8 rows (from [Feat] [AMD] kimi_k3: Kimi-K3 layer monokernel on gfx950 #13, plus):
  • K1 on the KDA layers: in_proj, the conv window, the recurrence and the gated norm. It reads the committed state in place, or the ReplaySSM ring under --enable-linear-replayssm-spec; the conv window is staged and handed back in one Triton launch each.
  • mla_pre, one launch a MLA (DSA) layer: the AttnRes seam, fused_qkv_a, the gates, the q / kv norms, q_b, the indexer projections and the q absorb (q_nope @ w_kc). SGLang takes its outputs through _mono_pre / q_nope_out.
  • The layer launch (o_proj to the MoE output): the seam, the router, top-16 by candidate rank with the route table from a bitmask, the latent MoE (latent quantized once by its producers), the shared expert, and the residual stream written in place.
  • Sentinel all-reduces on o_proj, the latent and the final reduction; pre-AR loads; tagged x rows (SGLANG_AMD_MONO_K3_K2_OPT).
• kimi_k31/spec.py: steps of up to 8 rows run the kimi_k3 launches; DSpark and DFLASH verify steps above 8 rows keep [Feat] [MI355X] kimi_k31: Wide-row layer launch for Kimi-K3.1 verify steps #102's wide launch. SGLANG_AMD_MONO_K31_WIDE=0 serves the narrow steps only, without ReplaySSM. The wide launch's requirements (ReplaySSM, 6 or 8 draft tokens, a speculative algorithm) are config refusals that name the switch.
• muse_spark/draft.py: the DFlash2 draft spec names its hidden size in requires.hf_config. Kimi-K3.1's DFLASH draft is also DFlash2DraftModel with model type qwen3, so under --enable-amd-monokernel it picked up the Muse Spark spec and the load failed ('DFlashMLAAttention' object has no attribute 'qkv_proj'). SGLang keeps a draft's forward when the spec names another model ([Feat] [AMD] model: Add Kimi-K3.1 MI35x support sglang#48427).
• common: the sentinel-word helpers (sync.py), wave and CTA prefix scans (ops.py).
• Tests: test/kernels/kimi_k3 (K1, K2 and mla_pre against torch references, check_k2, time_k2), test/unit/test_kimi_k3_spec.py, test_kimi_k31_spec.py (narrow / wide routing), and the forward drift pins re-read against the current SGLang.
• .github/CODEOWNERS, README.

Tested against SGLang sgl-project/sglang#48427 (kimi-k3.1-day-0-amd, which carries the _mono_pre / q_nope_out hooks) plus its binding fix.

Accuracy Tests

Kimi-K3.1-Pro, DFLASH (K31-DFlash2-MLA5-...-step55000, 6 draft tokens), TP8 on 8x MI350X, --enable-linear-replayssm-spec, on sgl-project/sglang#48427 plus its binding fix:

Check	Monokernel off	This PR
GSM8K 200, 16 threads (narrow and wide steps)	99.0%	99.0%
GPU faults	0	0
Plugin unit tests		same 11 failures as main (other models' contracts against this SGLang); every kimi_k3 / kimi_k31 test passes

Speed Tests and Profiling

Step time of one request (cc1, 10K-token prompts, 1024 out, real acceptance, median over 2-3 rounds), same SGLang tree and session:

Build	cc1 step
Monokernel off	19.06 ms
This PR (1ce5a34, then the final commit)	13.54 / 13.47 ms
The same launches before the merge onto main	13.43 ms

The gap to the pre-merge build is within run-to-run spread (13.42 .. 13.51 ms across earlier runs of that build). Verify steps of 2 .. 48 requests run #102's wide launch unchanged and were not re-measured here.

Not covered: DCP8 on the narrow path (the MLA front is declined under DCP; K1 and the layer launch were measured at DCP1 only).

Checklist

• pre-commit run --all-files passes.
• Tests under test/ cover the change.
• The SGLang commit or PR this was tested against is named above.
• Accuracy and speed results are attached.
• A new model adds one line to .github/CODEOWNERS that names its owner.

🤖 Generated with Claude Code

[kevin-mii]

Review pass before merge (an independent read of the diff, then fixes in this commit):

Fixed here:

• DCP8 on narrow steps: under DCP, SGLang's ROCm MLA forward replicates q_b and computes every head's q itself (q_replicate_active), while mla_pre hands over this rank's 12. Under DCP the MLA front now declines, and SGLang's projection front runs.
• The wide launch's merged-in_proj check now covers every KDA layer it runs, not only the first.
• The K3.1 drift test no longer duplicates the two outer-forward pins, which now live in the Kimi-K3 test.
• Forward drift digests: re-pinned after reading the diff against SGLang. The new branches (MegaMoE, prefill-CP shards, the CP causal-LM forward, the attention's tuple result) are not taken by a routed step: the spec refuses an a2a backend and attention or prefill CP, and forward_pre_o_proj returns a tensor.
• Comments: commit hashes, timing notes and references to files that are not in the repo are dropped.
• time_k2.py defaults to the serving form (fold, route 2, rtab 2, K2_OPT).
• README: a Kimi-K3.1 row, its recipe and its layout.

Left as they are, for follow-ups:

• wide.wide_widths ([Feat] [MI355X] kimi_k31: Wide-row layer launch for Kimi-K3.1 verify steps #102) declares 312 .. 384 rows, which have no L=6 bucket. These are only reachable above 48 requests or with CUDA graphs off.
• Kimi-K3's _backbone still collects the DFLASH taps in AuxHiddenStateList, not the packed buffer (same logits).
• The K3 A/B switches (SGLANG_AMD_MONO_K3_*) stay, as other models' switches do.
• Kimi-K3.1 now needs Kimi-K3's environment too (SGLANG_AITER_K3_OPT=1, AITER_SITUV2_A8W4=1), as the README recipe shows.