[Perf][Kimi-K3] Use the one-shot custom all-gather for TP logits

[zyongye]

Purpose

Kimi-K3's vocab-parallel logits all-gather went through the NCCL ring collective, which costs ~21us per decode step on an NVLink-connected TP8 group (latency-bound: 7 hops for a ~40KB-per-rank shard). This adds a KimiK3LogitsProcessor that routes the gather through the one-shot custom all-gather (MNNVL Lamport when available), which does the same rank-order concat in a single exchange, with a transparent fallback to the base collective whenever the custom kernel is unavailable or declines the input (size, dtype, registration). Gather semantics are identical, including the padded-vocab slice afterwards.

Not duplicating open work: searched open PRs; #48572 warms the spec-sized logits all-gather at init and #54433 changes DSpark top-k — neither changes the collective used here.

Test Plan

• pytest tests/models/kimi_k3/ unit suites on B300, run on a development tree containing this change.
• A/B Kimi-K3 TP8 serving (single 8x B300 node, 8192-in/1024-out, concurrency 1, fp8 KV cache) with Nsight Systems traces; combined with the pure-decode gather skip (submitted separately), since the two were measured in one run.

Test Result

• Logits all-gather: 20.2 -> 10.8 us per decode step (trace mean); the mnnvl_lamport_all_gather kernel launches early enough to overlap the LM-head GEMM tail.
• Combined with the separate decode-gather skip: stable-step median 8.7972 -> 8.7747 ms, bench median ITL 8.791 -> 8.770 ms (-0.25%).
• Greedy output spot checks unchanged (the collective is bitwise-equivalent).

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

🤖 Generated with Claude Code

Summary by CodeRabbit

• Improvements
  • Improved logits gathering for Kimi K3 models across tensor-parallel configurations.
  • Added fallback handling when the optimized gathering path is unavailable, preserving standard logits processing behavior.

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

[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: e121fdfb-6eae-4fcf-b390-31e011f2a300

📥 Commits

Reviewing files that changed from the base of the PR and between e3e1241 and eb9bed1.

📒 Files selected for processing (1)

• vllm/models/kimi_k3/nvidia/model.py

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

---

📝 Walkthrough

Walkthrough

Changes

Kimi K3 tensor-parallel logits processing

Layer / File(s)	Summary
Custom logits processor and model wiring vllm/models/kimi_k3/nvidia/model.py	The model accesses the tensor-parallel communicator. KimiK3LogitsProcessor performs custom all-gather and reshaping for supported 2D logits, then falls back to LogitsProcessor when required. KimiLinearForCausalLM uses the new processor.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to eb9be

The change accelerates Kimi K3 tensor-parallel logits gathering while preserving the existing behavior when the optimized path is unavailable; no actionable merge-blocking risk remains after normal checks and review.

Suggested reviewers: gau-nernst, zjy0516, gcanlin

Sequence Diagram(s)

sequenceDiagram
  participant KimiLinearForCausalLM
  participant KimiK3LogitsProcessor
  participant TPGroup
  participant LogitsProcessor
  KimiLinearForCausalLM->>KimiK3LogitsProcessor: process logits
  KimiK3LogitsProcessor->>TPGroup: custom all-gather for 2D logits
  TPGroup-->>KimiK3LogitsProcessor: gathered logits
  KimiK3LogitsProcessor->>KimiK3LogitsProcessor: reshape multi-token results
  KimiK3LogitsProcessor->>LogitsProcessor: fallback gather when unsupported or unavailable

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 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: using a one-shot custom all-gather for Kimi-K3 tensor-parallel logits.
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.}