[None][feat] Add Kimi K3 KDA prefill/MTP decode CuTe DSL kernels and fused attention-residual kernel - #17225
[None][feat] Add Kimi K3 KDA prefill/MTP decode CuTe DSL kernels and fused attention-residual kernel#17225brnguyen2 wants to merge 11 commits into
Conversation
Wave-A pre-flight (plan item A4), based directly on main. Cherry of feat/kimi_k3 b420908 plus its later sm_100-family pinning fix. All-new files plus CMake registrations. Same content as k3/14813-preflight-a4-attn-res-kernel minus dependency-branch context in the shared CMake files. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Add Blackwell CuTe DSL kernels for KDA chunked prefill (fused K1/K2/K3/K4 stages, standalone K4 persistent variant, and the A_kk inverse solve) and for KDA MTP speculative decode with replay-cache semantics, plus the custom-op wrappers (kda_prefill, KDA MTP ops) that JIT-compile and dispatch them. Prefill supports varlen batches (including empty and small-tail batches) with shape-independent JIT caching and OOB guards on the beta/state tiles; the op layer validates cu_seqlens/chunk_indices dtype consistency to avoid mixed-dtype capture corruption. Co-authored-by: Tao Li <tali@nvidia.com> Co-authored-by: Pengbo Wang <sephw@nvidia.com> Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>
Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>
Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #63559 [ run ] triggered by Bot. Commit: |
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe PR adds a Blackwell Kimi K3 attention-residual CUDA operator, CuTe DSL KDA prefill kernels, and a multi-token speculative decode operator. It also adds build integration, validation, caching, state handling, custom operator registration, fake implementations, and PTX patch tests. ChangesKimi K3 acceleration
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant PyTorch
participant kda_prefill
participant KdaPrefillRunner
participant fused_kernel123
participant k4_persistent_kernel
PyTorch->>kda_prefill: pass KDA prefill tensors
kda_prefill->>KdaPrefillRunner: forward with runtime options
KdaPrefillRunner->>fused_kernel123: execute fused K1-K3 path
KdaPrefillRunner->>k4_persistent_kernel: execute persistent K4 path
k4_persistent_kernel-->>PyTorch: return output and final state
sequenceDiagram
participant PyTorch
participant kda_mtp_decode
participant kda_decode_mtp_kernel
participant ReplayCache
participant RecurrentState
PyTorch->>kda_mtp_decode: pass tokens, metadata, state, and caches
kda_mtp_decode->>kda_decode_mtp_kernel: validate and launch
kda_decode_mtp_kernel->>ReplayCache: replay accepted-token values
kda_decode_mtp_kernel->>RecurrentState: update and commit state
kda_decode_mtp_kernel-->>PyTorch: return output
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2⚔️ Resolve merge conflicts 💡
🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 18
🧹 Nitpick comments (12)
cpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txt (1)
30-31: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueConfirm that
--use_fast_mathis required for this kernel.The kernel computes RMSNorm statistics and an online softmax.
--use_fast_mathrelaxes division and reciprocal accuracy and flushes denormals. If the parity tests pass only inside a tolerance, record the measured tolerance in a comment. If the flag gives no measurable speedup, prefer removing it and keeping the explicitrsqrtf/exp2fcalls, which already provide the fast paths.🤖 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 `@cpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txt` around lines 30 - 31, Validate whether --use_fast_math in the kimi_k3_attn_res_src CUDA compile options is necessary by comparing parity accuracy and measurable performance with and without it. If it provides no speedup, remove the option; otherwise retain it and document the measured parity tolerance in a nearby comment, preserving the explicit rsqrtf and exp2f fast paths.cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu (1)
48-137: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused helpers and constants.
ATTN_RES_BLOCK,ATTN_RES_WARPS,warp_reduce_sum,block_reduce_sum,v_addr,tcgen05_after_thread_sync,umma_arrive_noelect, andfloat2_subhave no reference in the rest of the file. The guidelines prohibit dead code and unused defined functions. Delete them, or keep only the ones a follow-up kernel needs and note that in a comment.#!/bin/bash # Confirm that the helpers have no other reference in the repository. for sym in ATTN_RES_BLOCK ATTN_RES_WARPS warp_reduce_sum block_reduce_sum v_addr \ tcgen05_after_thread_sync umma_arrive_noelect float2_sub; do echo "== $sym" rg -nP "\b${sym}\b" cpp | head -20 doneAs per coding guidelines: "do not leave defined functions unused" and "avoid dead code".
🤖 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 `@cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu` around lines 48 - 137, Remove the unused constants ATTN_RES_BLOCK and ATTN_RES_WARPS and the unreferenced helpers warp_reduce_sum, block_reduce_sum, v_addr, tcgen05_after_thread_sync, umma_arrive_noelect, and float2_sub from this file. Preserve float2_mul, float2_fma, and float2_add if they are used by the remaining implementation.Source: Coding guidelines
cpp/tensorrt_llm/thop/attnResOp.cpp (1)
129-140: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a fake or Meta implementation before using
trtllm::attn_res_fwdin compile or export paths.
No fake/meta registration or Python caller exists. Such a caller fails during fake-tensor dispatch because only the CUDA implementation is registered.🤖 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 `@cpp/tensorrt_llm/thop/attnResOp.cpp` around lines 129 - 140, Register a Fake/Meta implementation for the trtllm::attn_res_fwd operator alongside its schema and CUDA implementation, defining output metadata for all four returned tensors from the input tensors and parameters. Ensure compile and export fake-tensor dispatch can resolve attn_res_fwd without invoking the CUDA kernel, and add the corresponding Python caller or registration path if required by the existing operator integration.tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.py (1)
253-451: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused Neumann diagonal inversion.
_invert_diag_neumannhas no caller in this file.akk_inv_kernelinverts the diagonal blocks with_invert_diag_forward_fp32(Lines 798-801), and the module docstring states that the diagonal blocks use FP32 forward substitution to avoid BF16 Neumann cancellation. The 200 lines of Neumann code therefore document a rejected approach as live code.Delete the function, or move the rationale into a comment next to
_invert_diag_forward_fp32.🤖 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/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.py` around lines 253 - 451, Remove the unused `_invert_diag_neumann` function and its entire implementation. Keep `akk_inv_kernel` using `_invert_diag_forward_fp32` for diagonal inversion, and preserve the existing rationale by retaining or adding a concise comment near `_invert_diag_forward_fp32` that diagonal blocks use FP32 forward substitution instead of BF16 Neumann inversion.tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py (1)
445-449: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused launcher parameters and cache slots.
Several declared items are never read:
_launch_k4_persistentacceptsS_inanduse_fast_sync; the body uses neither._launch_fused_k123_invacceptscute_wrappers; the body uses only the explicitakk_in_view/akk_out_viewarguments._get_buffersseeds_k123_fns,_akk_inv_fnincute_wrappers; the K123 and akk_inv functions are cached in the module-level_fused_k123_cacheand_akk_inv_cacheinstead.Delete the unused parameters and cache slots, or wire the per-buffer function caches as the comment at Lines 445-446 describes.
Also applies to: 473-485, 616-617
🤖 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/custom_ops/cute_dsl_kimi_k3_custom_ops.py` around lines 445 - 449, Remove the unused S_in and use_fast_sync parameters from _launch_k4_persistent, and remove cute_wrappers from _launch_fused_k123_inv and its call sites. In _get_buffers, delete the unused _k123_fns and _akk_inv_fn cache slots while retaining _k4_fn if it is still read; rely on the module-level _fused_k123_cache and _akk_inv_cache for those functions.tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py (1)
81-92: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the debug switches from the shipped kernel.
DEBUG_K4_LEVELandDEBUG_SKIP_MMA5are module-level constants that disable the whole K4 phase (Line 1807) or MMA5 (Line 1911). A wrong value produces silently wrong output with no runtime signal. Delete both flags and the branches they guard, or move them behind an environment variable that logs when it is active.The comment on Line 92 also states 28 warps.
THREADS // WARP_SIZEis 20.I can prepare the patch that drops both flags and the dead branches. Tell me if you want it.
🤖 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/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py` around lines 81 - 92, Remove the module-level DEBUG_K4_LEVEL and DEBUG_SKIP_MMA5 switches and all branches they control in the K4 and MMA5 execution paths, preserving the normal full computation without debug bypasses. Correct the NUM_WARPS comment to reflect THREADS // WARP_SIZE evaluating to 20.tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py (1)
77-82: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDelete the unused Akk inversion helpers and register constants.
k4_persistent_kernelperforms no Akk inversion;akk_inverse.pyowns that stage. The following symbols have no reader in this file:
mma_tf32_m16n8k8,inv_internal_barrier,_invert_diag,_matmul_AB_inv,_chain_mma_B_inv,_chain_mma_A_inv,_store_neg_C_inv,_shuffle_C_to_B_inv,_store_C_temp_inv,_load_C_temp_inv.- The constants
SB,AKK_PAD,AKK_STRIDE,TEMP_COLS,NUM_TEMPS.NUM_REGS_WG0,NUM_REGS_WG1,NUM_REGS_WG2,MAX_REGS; the kernel makes nosetmaxregistercall.- The
nv_a_slkernel parameter, whose only use is commented out at Line 536.Removing them cuts about 300 lines and prevents drift against the real inversion kernel.
Also applies to: 87-378, 423-426
🤖 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/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py` around lines 77 - 82, Remove the unused Akk inversion helpers, including mma_tf32_m16n8k8, inv_internal_barrier, and the listed inversion functions, plus the unused register and temporary constants. Remove NUM_REGS_WG0, NUM_REGS_WG1, NUM_REGS_WG2, MAX_REGS, and the nv_a_sl parameter from k4_persistent_kernel and its call sites; also delete related commented-out usage. Preserve only symbols required by the kernel, since Akk inversion is handled by akk_inverse.py.tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py (2)
281-293: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a static assertion for the hardcoded
KERNEL_WIDTH == 4assumption.
KERNEL_WIDTHis aConstexpr[int]parameter, but several stages unroll exactly three history taps plus one current tap instead of deriving the count fromKERNEL_WIDTH. This block does it for Q, Lines 353-364 for K, Lines 382-388 for the K state shift, and Lines 478-481 for the V weights. Other stages, such as Line 297, do loop overKERNEL_WIDTH - 1.Today only
cute_dsl_kimi_k3_kda_mtp_ops.pyLine 451 prevents a wrong result, by rejectingW != 4in the host wrapper. The kernel itself would silently compute an incorrect convolution for any other width.Add an explicit compile-time check near the top of the kernel so the constraint fails loudly at the kernel boundary.
♻️ Proposed guard near Line 142
vec_size = TILE_K // 32 + # Several stages below unroll the conv taps for W == 4 explicitly. + assert cutlass.const_expr(KERNEL_WIDTH == 4), "kda_decode_mtp_kernel requires KERNEL_WIDTH == 4" num_v_tiles = V // TILE_V🤖 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/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py` around lines 281 - 293, Add a compile-time static assertion near the top of the kernel that requires KERNEL_WIDTH to equal 4, using the existing KERNEL_WIDTH constexpr parameter and assertion mechanism. Place the guard at the kernel boundary before the unrolled Q, K, state-shift, and V-weight logic so unsupported widths fail during kernel compilation.
471-473: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueWarp 7 is idle during the V convolution stage.
_v_idx = tidx - 96maps threads 96 through 255 onto the V lanes. The guard_v_idx < VwithV == 128admits onlytidxin[96, 224), so warps 3 through 6 cover all 128 lanes and warp 7 does no work in this stage. The block has 256 threads, so 32 threads idle here.This is safe because no barrier appears inside either branch of the
warp_idx < 3split. Please confirm the mapping is intentional. If it is, add a short comment, becauseNUM_THREADS = 256andV = 128make the 96-thread offset non-obvious.🤖 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/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py` around lines 471 - 473, Confirm that leaving warp 7 idle in the V convolution stage is intentional, then add a concise explanatory comment beside the _v_idx = tidx - 96 mapping describing the 96-thread offset and resulting inactive threads for NUM_THREADS = 256 and V = 128.tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py (3)
320-328: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winReduce duplicate DLPack construction on the decode hot path.
_layout_keybuilds a full DLPack-imported tensor for each of about 20 tensors only to readdynamic_shapes_maskanddynamic_strides_mask. The launch block at Line 610 then builds the same arguments again. This doubles host-side work in a per-decode-step path.Build each argument once, keep it in a local, and derive both the cache key and the launch argument from that single object.
🤖 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/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py` around lines 320 - 328, Update the layout-key and decode launch flow around _layout_key so each tensor’s DLPack argument is constructed once and retained for reuse. Derive dynamic_shapes_mask, dynamic_strides_mask, and the cache-key components from that local argument, then pass the same argument to the launch block instead of rebuilding it at Line 610; preserve the existing dynamic_layout selection and key behavior.
178-208: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate this helper and prefer built-in generics.
_require_stride_layoutdeclares 24 parameters with no type annotations and no return annotation. The coding guidelines require annotating every function and usingNonefor non-returning functions. Annotate the tensor parameters astorch.Tensorand the dimension parameters asint, and add-> None.Line 56 also imports
OptionalandTuple. The guidelines prefer built-in generic types and|, and the project targets Python 3.10+. Usefloat | Noneandtuple[...]instead.As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions, ... prefer built-in generic types and|".🤖 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/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py` around lines 178 - 208, Annotate _require_stride_layout so all data/state/cache arguments are torch.Tensor, all shape/count parameters are int, and the function returns None. Replace Optional-based annotations with the Python 3.10 form using |, and replace Tuple annotations with built-in tuple[...] types; remove the now-unused typing imports.Source: Coding guidelines
64-81: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueRemove the redundant
IS_CUTLASS_DSL_AVAILABLEguards. The module-levelImportErrormakes the guards at lines 81 and 642 unreachable when the DSL is unavailable. Keep the explicit import-time failure and remove the duplicate guards.🤖 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/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py` around lines 64 - 81, Remove the redundant IS_CUTLASS_DSL_AVAILABLE conditional guards surrounding the module implementation, including the guard near _TILE_V and the later guard around the decode logic. Preserve the module-level ImportError and unconditional DSL imports so unavailable CUTLASS DSL still fails explicitly at import time.
🤖 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 `@cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu`:
- Around line 1437-1443: Update the unsupported-case guard in the attention
forward function around attn_res_fwd_grid_size so failures and unmatched
configuration branches do not return silently after outputs are allocated. Raise
an assertion or otherwise propagate an explicit error to the caller for invalid
num_sm, N > N_MAX, or unsupported H, while preserving normal kernel launches for
supported configurations.
- Around line 982-989: Replace the unsynchronized function-local attrs_set
guards in launch_fwd, launch_s1_splitk, and launch_n1_ttile with per-device
std::call_once guards, ensuring each CUDA device independently initializes the
dynamic shared-memory attribute; launch_fwd should retain its conditional
initialization, while the other two always set it. Check and assert the
cudaFuncSetAttribute return status at all three sites. Apply these changes in
cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu at lines 982-989,
1332-1337, and 1378-1385.
- Around line 1490-1493: Update the `launch_fwd<7168, 4, false, true>` call in
the `N == 12 && T == 1024` branch to clamp the `num_sm - 1` CTA count to a
minimum of 1, preserving the existing value for devices with more than one SM.
- Around line 700-755: Insert a __syncwarp() after the chunk loop that writes
plan.logits_all and before the final block where logits_out and probs_out read
plan.logits_all. Keep the existing writer condition and output logic unchanged,
ensuring all consumer-warp lanes observe the current chunk’s logits before
reading them.
- Around line 42-56: Run clang-format on the CUDA source containing
warp_reduce_sum and the surrounding anonymous namespace, ensuring the formatting
complies with the repository’s .clang-format configuration, including
Allman-style braces.
In `@cpp/tensorrt_llm/thop/attnResOp.cpp`:
- Around line 71-93: In the attn_res_fwd validation flow, move the all-inputs
CUDA TORCH_CHECK before constructing CUDAGuard, then create the guard using
layer_residual.device(), and only afterward call check_attn_res_contract.
Preserve the existing dimension and type/contiguity/shape validations while
ensuring is_sm100_family() observes the tensor’s device.
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py`:
- Around line 1208-1212: Add an explicit non-None validation for A_log in
KdaPrefillRunner.forward alongside the existing chunk_size check, raising a
clear ValueError before either launcher runs. Preserve the current chunk-size
and CUTLASS DSL validations.
- Around line 285-330: Apply the existing _BUF_CACHE_MAX_ENTRIES LRU policy to
_g_sentinel_cache and _padded_input_cache: refresh accessed entries, evict the
least-recently-used entry when each cache exceeds the bound, and preserve
current buffer creation and return behavior in _get_g_sentinel_buffer and
_get_padded_input_buffers.
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py`:
- Around line 515-517: The stage_timing argument must not alias the bf16 output
when profiling can be enabled. In
tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py:515-517, replace
stage_timing_arg with a dedicated int64 buffer containing at least HV * N * 4
elements, or explicitly assert that PROFILE_STAGES remains disabled if retaining
the alias. In
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py:755-764,
update the kernel docstring to require an int64 stage_timing tensor of that size
and document indexing as (i_hv * grid_n + i_n) * 4.
- Line 741: Update the fake kernel’s empty-output allocation near the return
statement to derive its dtype from x_q rather than x_v, matching the eager
implementation’s output dtype while preserving x_v.shape.
- Around line 564-637: Update the `_compiled_cache[key]` launch call to use
`_from_dlpack_arg(...)` for every operand whose compile-time descriptor uses
that wrapper, matching the descriptors passed to `cute.compile`; preserve
`_dlpack_arg(...)` for operands compiled with it. Ensure the launch descriptors
for `h0_arg`, weights, state/cache tensors, and other corresponding arguments
are identical to their compile-time forms.
In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py`:
- Around line 21-48: Update the module docstring’s execution and barrier counts
to match the implementation constants: describe 992 threads and 31 warps, assign
warps 16–26 to MMA and 27–30 to store/inversion, and change both
stage_reuse_mbars and mma_done_mbars counts to 320. Leave the remaining pipeline
description unchanged.
- Around line 1059-1071: Update the SMEM zero-initialization loop around
_warp_id_in_cta and _row_base so every thread cooperatively covers all BT rows,
regardless of the 31-warp block size. Iterate row bases with a warp-count stride
(or equivalent coverage logic) so rows 62 and 63 in both sAqk and sAkk are
explicitly zeroed before the barrier, while preserving the existing per-lane
column initialization.
In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py`:
- Around line 63-70: Remove the module-level CuteExperimentalDSL and
CuTeDSL._get_pipeline monkey-patches, and pass the --uumn PTX option only
through this kernel’s compilation path. Replace the broad exception handling
around these compatibility imports with AttributeError and ImportError only, and
report failures via tensorrt_llm.logger rather than print. Keep unrelated CuTe
DSL compilations unaffected by this module.
- Around line 680-709: Before entering the K4 scheduler loop around
work.is_valid_tile, validate that every sequence represented by cu_seqlens has a
strictly positive length, rejecting mixed batches containing zero-length entries
rather than relying on the NT < 4 check. Alternatively, explicitly skip empty
scheduler tiles while preserving barrier participation for all processed tiles,
so empty sequences cannot bypass the mma6_done_mbar commits and deadlock state
barriers.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py`:
- Around line 569-585: Pre-declare t_stage1 alongside t_stage0 and the other
tracing variables before the dynamic run_precompute branch, then retain the
existing assignments in each PROFILE_STAGES path. Ensure the dynamic branch no
longer contains t_stage1’s first assignment, preserving its current profiling
behavior.
- Around line 477-524: Restrict the zero-accepted fast path to the shape it
implements by updating the host wrapper’s use_zero_accepted assignment near the
zero_accepted_hint handling to require num_spec == 2. This must prevent
USE_ZERO_ACCEPTED from reaching the unrolled branch in the KDA decode kernel for
any other num_spec, allowing those cases to use the generic loop.
- Around line 132-141: Clamp the device-loaded value in num_accepted_tokens
before assigning commit_len so it cannot exceed NUM_SPEC. Preserve the
USE_ZERO_ACCEPTED behavior, and ensure the clamped commit_len continues to
determine T_loop while t_max remains safely sized for all subsequent
shared-memory and cache accesses.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu`:
- Around line 48-137: Remove the unused constants ATTN_RES_BLOCK and
ATTN_RES_WARPS and the unreferenced helpers warp_reduce_sum, block_reduce_sum,
v_addr, tcgen05_after_thread_sync, umma_arrive_noelect, and float2_sub from this
file. Preserve float2_mul, float2_fma, and float2_add if they are used by the
remaining implementation.
In `@cpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txt`:
- Around line 30-31: Validate whether --use_fast_math in the
kimi_k3_attn_res_src CUDA compile options is necessary by comparing parity
accuracy and measurable performance with and without it. If it provides no
speedup, remove the option; otherwise retain it and document the measured parity
tolerance in a nearby comment, preserving the explicit rsqrtf and exp2f fast
paths.
In `@cpp/tensorrt_llm/thop/attnResOp.cpp`:
- Around line 129-140: Register a Fake/Meta implementation for the
trtllm::attn_res_fwd operator alongside its schema and CUDA implementation,
defining output metadata for all four returned tensors from the input tensors
and parameters. Ensure compile and export fake-tensor dispatch can resolve
attn_res_fwd without invoking the CUDA kernel, and add the corresponding Python
caller or registration path if required by the existing operator integration.
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py`:
- Around line 445-449: Remove the unused S_in and use_fast_sync parameters from
_launch_k4_persistent, and remove cute_wrappers from _launch_fused_k123_inv and
its call sites. In _get_buffers, delete the unused _k123_fns and _akk_inv_fn
cache slots while retaining _k4_fn if it is still read; rely on the module-level
_fused_k123_cache and _akk_inv_cache for those functions.
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py`:
- Around line 320-328: Update the layout-key and decode launch flow around
_layout_key so each tensor’s DLPack argument is constructed once and retained
for reuse. Derive dynamic_shapes_mask, dynamic_strides_mask, and the cache-key
components from that local argument, then pass the same argument to the launch
block instead of rebuilding it at Line 610; preserve the existing dynamic_layout
selection and key behavior.
- Around line 178-208: Annotate _require_stride_layout so all data/state/cache
arguments are torch.Tensor, all shape/count parameters are int, and the function
returns None. Replace Optional-based annotations with the Python 3.10 form using
|, and replace Tuple annotations with built-in tuple[...] types; remove the
now-unused typing imports.
- Around line 64-81: Remove the redundant IS_CUTLASS_DSL_AVAILABLE conditional
guards surrounding the module implementation, including the guard near _TILE_V
and the later guard around the decode logic. Preserve the module-level
ImportError and unconditional DSL imports so unavailable CUTLASS DSL still fails
explicitly at import time.
In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.py`:
- Around line 253-451: Remove the unused `_invert_diag_neumann` function and its
entire implementation. Keep `akk_inv_kernel` using `_invert_diag_forward_fp32`
for diagonal inversion, and preserve the existing rationale by retaining or
adding a concise comment near `_invert_diag_forward_fp32` that diagonal blocks
use FP32 forward substitution instead of BF16 Neumann inversion.
In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py`:
- Around line 81-92: Remove the module-level DEBUG_K4_LEVEL and DEBUG_SKIP_MMA5
switches and all branches they control in the K4 and MMA5 execution paths,
preserving the normal full computation without debug bypasses. Correct the
NUM_WARPS comment to reflect THREADS // WARP_SIZE evaluating to 20.
In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py`:
- Around line 77-82: Remove the unused Akk inversion helpers, including
mma_tf32_m16n8k8, inv_internal_barrier, and the listed inversion functions, plus
the unused register and temporary constants. Remove NUM_REGS_WG0, NUM_REGS_WG1,
NUM_REGS_WG2, MAX_REGS, and the nv_a_sl parameter from k4_persistent_kernel and
its call sites; also delete related commented-out usage. Preserve only symbols
required by the kernel, since Akk inversion is handled by akk_inverse.py.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py`:
- Around line 281-293: Add a compile-time static assertion near the top of the
kernel that requires KERNEL_WIDTH to equal 4, using the existing KERNEL_WIDTH
constexpr parameter and assertion mechanism. Place the guard at the kernel
boundary before the unrolled Q, K, state-shift, and V-weight logic so
unsupported widths fail during kernel compilation.
- Around line 471-473: Confirm that leaving warp 7 idle in the V convolution
stage is intentional, then add a concise explanatory comment beside the _v_idx =
tidx - 96 mapping describing the 96-thread offset and resulting inactive threads
for NUM_THREADS = 256 and V = 128.
🪄 Autofix (Beta)
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: ec8f9aae-3bdb-4434-9ca1-3bfd220db777
📒 Files selected for processing (16)
cpp/tensorrt_llm/CMakeLists.txtcpp/tensorrt_llm/kernels/CMakeLists.txtcpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txtcpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cucpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/attnResOp.cpptensorrt_llm/_torch/custom_ops/__init__.pytensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.pytensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/__init__.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py
|
PR_Github #63560 [ run ] triggered by Bot. Commit: |
|
PR_Github #63559 [ run ] completed with state |
…tn-res op - attnResFwd.cu: publish plan.logits_all with __syncwarp before the tail cross-lane reads; per-device cudaFuncSetAttribute via std::once_flag in all three launchers; replace silent returns with TLLM_CHECK_WITH_INFO (SM query, N > N_MAX, unsupported H); clamp num_sm - 1 grid to >= 1 - attnResOp.cpp: check CUDA placement, then set the device guard, then run the contract check so is_sm100_family() sees the tensors' device - cute_dsl_kimi_k3_custom_ops.py: LRU-bound the padded-input and g-sentinel scratch caches; validate A_log; document the K4 positive sequence-length precondition - cute_dsl_kimi_k3_kda_mtp_ops.py: gate the zero-accepted fast path on num_spec == 2; assert profiling stays off while stage_timing aliases out; fake kernel output dtype from x_q - kda_mtp_decode.py: clamp commit_len to NUM_SPEC on device; pre-declare t_stage1 for the profiling path; document stage_timing requirements - fused_k123.py: cover all 64 SMEM rows in the A_qk/A_kk zero-init (31 warps left rows 62-63 stale for the transposed A_kk store path); fix stale warp/thread counts in comments - k4_persistent.py: narrow monkey-patch exception handling, route messages through tensorrt_llm.logger, document the global ptx-options effect Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #63573 [ run ] triggered by Bot. Commit: |
|
PR_Github #63560 [ run ] completed with state |
|
Test results for this PR beyond the L0 pipelines, since the kernels added here have no in-tree callers yet (the KDA module wrappers, unit suites, and model integration arrive in follow-up PRs): Multi-architecture compile check —
Import / collection check on main + this PR only (no follow-up branches): End-to-end pipeline smoke — on an integration branch carrying the follow-up KDA modules and model on top of this PR's kernels (byte-identical kernel content as of c73ed72): a 4-GPU TP4 generation run on a layer-truncated checkpoint (bf16, reference MoE backend) passes, exercising the CuTe DSL prefill kernels through the full LLM API path on a Blackwell (sm_100 family) node. Kernel-vs-reference parity unit suites (they import the follow-up module wrappers, so they ship with the module PR): green on a Blackwell node against this PR's kernel content as of c73ed72. The device-code adjustments from the review feedback (2b2978c) are queued for a parity rerun; results will be reflected in the follow-up module PR's test coverage. |
|
/bot run |
|
PR_Github #63596 [ run ] triggered by Bot. Commit: |
|
PR_Github #63573 [ run ] completed with state |
|
Regression-wise this is clean: everything is What I'd want addressed before it merges: The kernels land with nothing exercising them in CI. The parity/cache-soundness suites ship with the follow-up module PR, so between this merging and that one, 11.8k lines of Blackwell kernels sit in the tree unbuilt-against and untested. That window is where the two fixes #17230 calls out — the Two unrelated drops in one PR. The CuTe DSL KDA prefill/MTP kernels (9 files) and the Holding my approval — 11.8k lines of new kernels with no other approval yet isn't one I want to be first on. |
|
PR_Github #63596 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63630 [ run ] triggered by Bot. Commit: |
|
PR_Github #63630 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63698 [ run ] triggered by Bot. Commit: |
|
PR_Github #63698 [ run ] completed with state
|
The review-round change narrowing the k4_persistent monkey-patch handler
to `except (AttributeError, ImportError)` breaks every cute.compile of
the module: the nvidia-cutlass-dsl 4.5.0 AST preprocessor cannot parse
tuple except handlers ("'Tuple' object has no attribute 'id'"). Revert
to a single bare Exception with a comment explaining the constraint.
Also gate the new A_log validation on a non-empty token batch: zero-token
calls take the early return in _chunk_kda_fwd and never touch A_log, and
the runtime emits such batches under the overlap scheduler + logprobs
flows.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py`:
- Around line 477-482: The bare except Exception handler around line 477-482
suppresses all exceptions including unexpected compatibility failures, masking
real issues until a later cute.compile call. Update this handler to only catch
and suppress the specific expected exceptions (AttributeError and ImportError
from the DSL parser bug), and re-raise all other exceptions. If the
nvidia-cutlass-dsl parser truly requires a bare except syntax, isolate this
compatibility workaround in a separate helper function that can use narrow
exception handling, then call that helper from the main code path.
🪄 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: 9c7be6cc-6496-4212-bc72-48ed2d6b4a94
📒 Files selected for processing (2)
tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py
|
PR_Github #63750 [ run ] triggered by Bot. Commit: |
|
PR_Github #63750 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63784 [ run ] triggered by Bot. Commit: |
|
PR_Github #63784 [ run ] completed with state
|
Review follow-up: the patch's bare `except Exception` swallowed every failure kind, so an unexpected error (e.g. a changed CUTLASS API) would leave the module silently unpatched and surface later as a ptxas register-budget failure at the first K4 compile, with no trail back to the patch. Re-raise anything that is not the expected AttributeError / ImportError. The handler clause itself stays a single `except Exception` with an isinstance check in the body: the nvidia-cutlass-dsl 4.5.0 AST preprocessor cannot parse tuple except handlers anywhere in a kernel module (see the previous commit's revert). Add CPU-only contract tests pinning both constraints (no tuple except handlers in the module; benign kinds soft-continue while unexpected errors fail the import loudly) and wire them into the CPU pre-merge lists. Nothing else in CI compiles these kernels yet, so the structural check is the only automated guard against reintroducing the parser breakage. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py`:
- Around line 118-134: Update test_import_applies_patch to replace
CuTeDSL._get_pipeline with a fake returning "cubin-format=bin", invoke the
patched method, and assert the result contains ptx-options='--uumn' rather than
only checking the wrapper name. Preserve the subprocess-based import test and
include the required coverage verdict and test-list check for this test change.
- Around line 43-49: Update the dependency probe around the cutlass and
flashinfer imports to catch only ImportError, allowing other runtime failures to
propagate instead of being converted into skips. In the fallback logic around
find_spec(), catch only ModuleNotFoundError so unexpected discovery failures
remain visible; preserve the existing source-tree fallback behavior for the
expected exception.
🪄 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: 0641806c-a88b-4966-aef1-58bfcc69365d
📒 Files selected for processing (4)
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.pytests/integration/test_lists/test-db/l0_cpu_arm.ymltests/integration/test_lists/test-db/l0_cpu_x86.ymltests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py
| def test_import_applies_patch(): | ||
| proc = subprocess.run( | ||
| [ | ||
| sys.executable, | ||
| "-c", | ||
| _LOAD_K4.format(k4=str(_k4_path())) | ||
| + """ | ||
| from cutlass.cutlass_dsl.cutlass import CuTeDSL | ||
| assert CuTeDSL._get_pipeline.__name__ == "_patched_get_pipeline", \\ | ||
| f"ptx-options patch not applied: {CuTeDSL._get_pipeline}" | ||
| """, | ||
| ], | ||
| capture_output=True, | ||
| text=True, | ||
| timeout=600, | ||
| ) | ||
| assert proc.returncode == 0, proc.stderr[-2000:] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Test coverage summary — insufficient: verify the pipeline transformation.
Added tests cover the AST constraint, expected soft failure, and unexpected-failure propagation. The file is registered in both CPU CI lists: tests/integration/test_lists/test-db/l0_cpu_arm.yml and tests/integration/test_lists/test-db/l0_cpu_x86.yml.
test_import_applies_patch only checks that _patched_get_pipeline is assigned. It passes if the wrapper returns the unmodified pipeline string. Use a fake CuTeDSL._get_pipeline() that returns "cubin-format=bin", then assert that the wrapped method returns a string containing ptx-options='--uumn'.
As per path instructions, test-code changes require a coverage verdict and test-list check.
🧰 Tools
🪛 ast-grep (0.45.0)
[error] 118-132: Command coming from incoming request
Context: subprocess.run(
[
sys.executable,
"-c",
_LOAD_K4.format(k4=str(_k4_path()))
+ """
from cutlass.cutlass_dsl.cutlass import CuTeDSL
assert CuTeDSL._get_pipeline.name == "_patched_get_pipeline", \
f"ptx-options patch not applied: {CuTeDSL._get_pipeline}"
""",
],
capture_output=True,
text=True,
timeout=600,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
🪛 Ruff (0.16.1)
[error] 119-119: subprocess call: check for execution of untrusted input
(S603)
🤖 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 `@tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py` around lines
118 - 134, Update test_import_applies_patch to replace CuTeDSL._get_pipeline
with a fake returning "cubin-format=bin", invoke the patched method, and assert
the result contains ptx-options='--uumn' rather than only checking the wrapper
name. Preserve the subprocess-based import test and include the required
coverage verdict and test-list check for this test change.
Source: Path instructions
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Description
Kernel-layer groundwork for the upcoming KimiLinear (Kimi K3) model support. Two independent kernel drops, combined because they share the same wave and reviewers:
trtllm::attn_res_fwdCUDA kernel for the K3 fused attention-residual path (sm_100 family).No model-code dependency: all files are new except one-line CMake registrations. Part of the K3 kernel enablement series alongside #17190, #17054, and #15138; the K3 modules and model land in follow-up PRs stacked on these.
Test Coverage
Unit tests exercising these kernels (KDA prefill op parity, prefill state parity, cache soundness, MTP decode parity, attn-res op parity) import the K3 module wrappers, so they ship with the follow-up module PR. The kernels in this PR were validated with those suites on a Blackwell (sm_103) node.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.
Dev Engineer Review
A_{kk}inversion, and KDA MTP speculative decode.trtllm::attn_res_fwdCUDA operator for Kimi K3 attention-residual computation on sm_100.QA Engineer Review
tests/integration/test_lists/test-db/l0_cpu_arm.ymlandtests/integration/test_lists/test-db/l0_cpu_x86.yml.tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.pyto both CPU pre-merge test lists.test_no_tuple_except_handlers(),test_import_applies_patch(),test_expected_patch_failure_is_soft(), andtest_unexpected_patch_failure_raises().