[None][fix] Kimi K3: eager CUDA-graph buffer allocation and prebuilt fused-verify constants - #17421
[None][fix] Kimi K3: eager CUDA-graph buffer allocation and prebuilt fused-verify constants#17421brnguyen2 wants to merge 2 commits into
Conversation
…fused-verify constants Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughKimi KDA decode staging now uses fixed pool-slot capacity. Fused-verification convolution weights are built before decode execution and required from a prebuilt cache. Parity runtime setup now mirrors production weight preparation. ChangesKimi KDA capture-safe runtime
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant WeightLoader
participant KDAInitializer
participant ConvWeightCache
participant DecodeKernel
WeightLoader->>ConvWeightCache: Build convolution constants
KDAInitializer->>ConvWeightCache: Require prebuilt weights
DecodeKernel->>ConvWeightCache: Retrieve prebuilt weights
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/models/modeling_kimi_linear.py (1)
1351-1376: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winEnforce the slot-capacity invariant before the first allocation.
Line 1365 uses
max(conv_pool.shape[0], B). If the first decode hasB > conv_pool.shape[0], this allocates beyond pool capacity and bypasses the required hard failure. Allocate exactlyconv_pool.shape[0]rows and assertB <= conv_pool.shape[0]before allocation.Proposed fix
+ assert B <= conv_pool.shape[0], ( + f"KDA decode batch has {B} rows but the conv pool has " + f"{conv_pool.shape[0]} slots" + ) buf = self._cs_dense if buf is None: ... buf = torch.empty( - 3, max(conv_pool.shape[0], B), d, W - 1, dtype=torch.bfloat16, device=x2d.device + 3, conv_pool.shape[0], d, W - 1, dtype=torch.bfloat16, device=x2d.device )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py` around lines 1351 - 1376, Update the _cs_dense initialization path to assert B <= conv_pool.shape[0] before allocating, enforcing the slot-capacity invariant on the first decode. Allocate the staging buffer with exactly conv_pool.shape[0] rows instead of max(conv_pool.shape[0], B), while preserving the existing capture-safe fallback and post-allocation capacity assertion.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py`:
- Around line 1629-1640: Update the _get_mtp_conv_weights method signature to
include the return type annotation -> Tuple[torch.Tensor, ...], preserving its
existing behavior and error handling.
---
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py`:
- Around line 1351-1376: Update the _cs_dense initialization path to assert B <=
conv_pool.shape[0] before allocating, enforcing the slot-capacity invariant on
the first decode. Allocate the staging buffer with exactly conv_pool.shape[0]
rows instead of max(conv_pool.shape[0], B), while preserving the existing
capture-safe fallback and post-allocation capacity assertion.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6e46eb69-4366-4b5f-a439-913d9f4494fd
📒 Files selected for processing (2)
tensorrt_llm/_torch/models/modeling_kimi_linear.pytests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py
|
PR_Github #64626 [ run ] triggered by Bot. Commit: |
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #64635 [ run ] triggered by Bot. Commit: |
|
PR_Github #64626 [ run ] completed with state |
|
PR_Github #64635 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64658 [ run ] triggered by Bot. Commit: |
|
Re CodeRabbit's outside-diff note on the |
|
PR_Github #64658 [ run ] completed with state |
|
/bot run |
|
PR_Github #64807 [ run ] triggered by Bot. Commit: |
|
PR_Github #64807 [ run ] completed with state
|
2 participants
Summary
Follow-up to #17269, hardening two CUDA-graph-unsafe lazy-allocation patterns in the Kimi K3 KDA runtime (reviewer feedback on the decode fast path):
_cs_densesilent-realloc hazard. The decode kernel's persistent conv-window staging buffer was lazily allocated with abuf.shape[1] < Bregrow branch. Previously captured CUDA graphs hold the old pointer, so if the "first allocation is max-size" invariant ever broke, a regrow would leave those graphs writing into freed memory. The buffer is now allocated exactly once (sized to the conv pool's slot count, which bounds the decode batch:slot_indicesare distinct pool rows and this is the plain one-token-per-request path), and a hard assert replaces the regrow so a broken invariant fails loudly instead of corrupting memory. The capture-time fallback to the reference path is kept: a first decode under capture must not allocate.Lazy fused-verify conv-weight compute.
_get_mtp_conv_weights()computed the fp32 conv constants at first use, which under CUDA graph capture would bake capture-pool pointers into the cached tuple (previously guarded by a capture-time-only RuntimeError). The constants are now always prebuilt at weight-load time:load_weights()builds them unconditionally for every KDA layer — closing a gap where FP8 KDA weight read with the fused decode glue disabled ran neither finalize variant yet could still reach the fused verify path — and the getter raises unconditionally if they are missing.Behavior-neutral by construction: no numeric changes, only allocation timing and failure mode.
Changes
tensorrt_llm/_torch/models/modeling_kimi_linear.py_forward_decode: allocate_cs_denseonce; hard assert instead of silent realloc._build_mtp_conv_weights(); called from_build_decode_kernel_constants()and unconditionally per KDA layer inload_weights()._get_mtp_conv_weights(): unconditional RuntimeError when the constants were not prebuilt.tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py: the test-built runtime now prebuilds the constants after weight init, mirroring production.Validation
Validated on a Blackwell (SM103) node:
tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py,test_kimi_kda_verify_parity.py,test_kda_mtp_decode_cute_parity.pyandtests/unittest/_torch/modules/kimi_kda/: 50 passed, 1 pre-existing unrelated skip._cs_denseunder capture): PASS.max_draft_len=2, exercises the fused-verify constants end to end): PASS.Dev Engineer Review
_cs_denseat conv-pool slot capacity._get_mtp_conv_weights()fail when required constants are unavailable.QA Engineer Review
_make_runtimeintests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py.tests/integration/test_lists/entry was provided.