[None][perf] Kimi K3: overlap KDA/MLA attention decode onto the aux stream

[hyukn]

@coderabbitai summary

Description

The Kimi K3 attention layers ran 100% single-stream. KimiK3DecoderLayer already creates one side CUDA stream per model, but it was passed only to KimiK3MoERuntime — so the identical sm100_fp8_fp4_gemm_1d1d_impl kernel was 97% hidden inside MoE and 94% exposed inside KDA.

Measured on the pre-change build (16x GB300 TEP16, ISL 8192 / OSL 1024, c=1):

family	segments/step	device-busy	two-stream residency
KDA	69	4.667 ms	0.0000 ms
MLA	24	2.161 ms	0.0000 ms

This routes that existing stream into KimiKDARuntime and KimiK3MLAAttention and forks three data-independent decode op pairs via maybe_execute_in_parallel. Pure scheduling — no dtype, precision, kv-cache or backend change. Python only, no rebuild.

Changes

Three forks, each on a pair that reads disjoint inputs and writes disjoint outputs:

1. KDA in-projection ∥ cache staging (modeling_kimi_linear.py, _forward_decode). The in-projection GEMV chain reads only x2d; the conv-window gather/repack plus recurrent-state gather read only the cache pools and slot_indices. Both feed the KDA decode kernel, so the staging can hide behind the GEMV instead of following it.
2. MLA Q chain ∥ KV chain (kimi_k3_mla_attention.py, _forward_decode_fused). q_a_proj → q_a_layernorm → q_b_proj and kv_a_proj_with_mqa → kv_a_layernorm → cat both read only hidden_states. DeepSeek MLA fuses q_a/kv_a into one GEMM and can therefore overlap only the two layernorms; K3 keeps them separate, so the whole chains overlap.
3. MLA absorb-bmm ∥ mla_rope_generation. The bmm writes only fused_q[..., :kv_lora_rank]; the rope kernel writes only the rope tail, the paged cache and the scheduler buffers. Gated off under an FP8 KV cache — the extra CTAs the grid gains for cache_type == FP8 read the latent slice back to quantize it, i.e. they consume the bmm's output, so overlapping there would race. Same gate the DeepSeek MLA decode path applies in mla.py.

Supporting details:

• The forks live inside forward() so the captured CUDA graph records them, matching how KimiK3MoERuntime already does it. All use disable_on_compile=True; with no stream supplied or multi-stream off, both legs run in order on the default stream with an identical result.
• One stream still serves everything. Each fork joins before the next one starts, and the attention forks join before the MoE block reuses the stream.
• record_stream on every tensor that crosses streams (cs, the gathered state, latent_cache, and pre-fork on fused_q/cu_*_seqlens/fmha_scheduler_counter) so the caching allocator cannot recycle storage mid-use.

Performance

16x GB300 TEP16, ISL 8192 / OSL 1024, median TPOT. Reproduced across both a rack change and a rebuild:

observation	job	rack	build	c=1 TPOT	delta
1	2720192	d157	2719879	—	-3.720%
2	2720932	d127	2720688	13.2106 → 12.6789 ms	-4.025%

The two agree to 0.316%, against a measured c=1 noise band of 2.63%.

A third observation and the full concurrency curve were taken afterwards on the same recipe (job 2721639), confirming the gain holds across the sweep and regresses nothing:

concurrency	before	after	delta	noise band
1	13.5070 ms	12.6934 ms	-6.02%	2.63%
2	15.9358 ms	16.1229 ms	+1.17%	5.08% (within noise)
4	17.5480 ms	16.8362 ms	-4.06%	4.01%
8	22.7114 ms	21.7538 ms	-4.22%	3.63%

c=8 is the only genuinely independent corroborator: stream_interval=10 makes median ITL exactly 10x TPOT, so ITL, tps/user and output_throughput at c=1 are all restatements of the same metric. Per-point noise bands are from three zero-change runs on the unmodified build.

Post-change nsys on the same recipe confirms the forks execute on device rather than merely being issued:

family	two-stream residency before	after	share of busy time
KDA	0.0000 ms/step	0.7016 ms/step	16.1%
MLA	0.0000 ms/step	0.2829 ms/step	14.3%

with the aux stream appearing in both families' stream sets.

Test Coverage

Draft — two verification gaps remain before this is merge-ready.

Done:

• The KDA fork (Bump onnx from 1.12.0 to 1.13.0 #1) and the MLA Q/KV fork (Add static libraries for batch manager #2) were checked bitwise-identical against their serial equivalents under real CUDA-graph capture on GB300.
• Fork Update TRT-LLM code #3 is argued disjoint from source: mlaKernels.cu writes qkv_output[dst_q_idx] write-only, and both head_dim_idx and c_k = 512 are multiples of ELTS_PER_VEC = 8, so no vector store can straddle the bmm's slice.
• End-to-end serving at the canonical decode point plus the full concurrency sweep, reported above.

Outstanding:

• Fork Update TRT-LLM code #3 (absorb-bmm ∥ mla_rope_generation) has a source-level disjointness argument but no numerical probe of its own.
• No model-level greedy-token comparison has been run.

Both should land before this leaves draft.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.