[Perf][Kimi-K3] Cut MLA decode concat/cache epilogue latency

[zyongye]

Purpose

Reduce the latency of the fused Kimi-K3 MLA decode q-concat + latent cache insert kernel, which sits between the absorbed-q BMM and the decode FMHA on every full-attention layer.

Three changes, all inside fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu:

1. Warps-per-row split, chosen from the token count. At decode sizes the kernel had one warp per (token, head+1) row — 9 warps at concurrency 1 — with each lane serially moving 2-3 chunks of a 576-wide row, leaving it memory-latency-bound. Batches up to 64 tokens now split each row across 3 warps (exactly one 8-element load→convert→store per lane); larger calls keep one warp per row.
2. Dependency-wait restructure. Only ql_nope/q_pe come from the producing GEMM chain the kernel is PDL-launched behind. slot_mapping, the scales, the rope table, and the cache-row inputs are older data, so they are now read before cudaGridDependencySynchronize, and the cache-slot warps skip the wait entirely — the post-wait critical path is a single round trip on the query warps.
3. Early dependent-launch trigger in the split regime. The query warps fire cudaTriggerProgrammaticLaunchCompletion right after their dependency wait, so the consuming FMHA launches and runs its prologue while the last round trip is in flight. Cache warps never trigger (a no-wait warp triggering would release the consumer before the producer's data arrived) and complete via exit; the one-warp-per-row fallback keeps the end-of-kernel trigger.

Per-element math and element mapping are unchanged, so outputs are bitwise identical; prefill and the larger-batch decode fallback behavior are unchanged.

Not duplicating open work: #53526 moves the trigger for the num_tokens == 1 case on the same kernel; this PR covers that case as a consequence of the split-regime trigger (no C=1 specialization) and additionally restructures the row parallelism and the dependency wait.

Test Plan

• pytest tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py on B300 (SM103), including a new test_decode_epilogue_row_split_boundary covering both sides of the 64-token dispatch.
• A/B serving run on Kimi-K3 TP8, single 8×B300 node: 8192-in/1024-out, concurrency 1, fp8 KV cache, no speculative decoding; identical tree except this diff; Nsight Systems traces on both sides.

Test Result

• Kernel tests: 40 passed.
• Decode epilogue kernel time (trace mean): 2.95 → 2.06 us (−30%).
• Full-attention layer pre-attention window (residual-completion → FMHA start, nsys median over 24 layers × 145 stable steps): 17.89 → 16.80 us/layer.
• Stable decode step median: 8.7788 → 8.7399 ms; bench median ITL: 8.7746 → 8.7394 ms (−0.40%).
• No model-output change expected or observed (bitwise-identical kernel semantics); greedy spot checks unchanged.

AI assistance was used for this PR (Claude Code); every changed line was reviewed and the tests/benchmarks above were run by the submitter.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Performance

  • Improved decode processing for small batches of up to 64 tokens.
  • Maintained the existing processing path for larger batches.

• Bug Fixes

  • Improved fused decode epilogue behavior when writing query and cache data across multiple warps.

• Tests

  • Added coverage for the 64-token boundary and larger 65-token batches to verify consistent results.

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

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #86813 for commit 33458afeb50f.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #86929 for commit 5a6c86d017e7.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Team

Run ID: b6656665-e755-41bd-8575-a7b184980d39

📥 Commits

Reviewing files that changed from the base of the PR and between 0e3ac49 and 5a6c86d.

📒 Files selected for processing (2)

• csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
• tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Walkthrough

Walkthrough

The fused Kimi-K3 MLA decode kernels now split rows across three warps for batches of up to 64 tokens. Wrappers select the split mode for bf16 and fp8 paths. A boundary test covers 64 and 65 tokens.

Changes

Kimi-K3 MLA decode optimization

Layer / File(s)	Summary
Generalize latent row writes csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu	writeLatent576 now accepts an explicit lane and lane stride. Existing prefill and other decode callers pass a 32-lane stride.
Split decode rows csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu	The decode kernel maps split warps to row slots and writes query and cache data with split-aware lanes. Grid synchronization runs only in the head branch.
Dispatch and boundary validation csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu, tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py	The bf16 and fp8 wrappers select SPLIT=3 for up to 64 tokens and SPLIT=1 otherwise. The test checks both 64- and 65-token cases.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: 🟡 Moderate · up to 5a6c8

For decode calls of up to 64 tokens, downstream attention may be signaled before all query fragments are ready, which could cause incorrect outputs or incomplete KV-cache state. This bounded synchronization risk should be fixed or explicitly accepted before merging.

Sequence Diagram(s)

sequenceDiagram
  participant DecodeWrapper
  participant DecodeKernel
  participant LatentWriter
  DecodeWrapper->>DecodeKernel: Launch SPLIT=3 or SPLIT=1
  DecodeKernel->>DecodeKernel: Map warp to row slot and split part
  DecodeKernel->>LatentWriter: Write query or cache row with lane_stride

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 9 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: reducing latency in the Kimi-K3 MLA decode concatenation and cache epilogue kernel.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #86973 for commit 015c2252c0ba.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87010 for commit ed1457de2e22.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87117 for commit 36f6a601a6e9.