[TRTLLM-14702][feat] Integrate Kimi K3 KDA decode kernel - #17054
Conversation
Add a fused decode kernel for Kimi Delta Attention (KDA) recurrent-state update at generation time, with a thop custom op (kda_decode) exposing it to the PyTorch flow. Supports the K3 head geometry and per-request state indexing into the mamba cache pool. Co-authored-by: Pengbo Wang <sephw@nvidia.com> Co-authored-by: Brian Nguyen (TensorRT) <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>
pengbowang-nv
changed the title
WalkthroughAdds a fused CUDA KDA single-token decode kernel with compact and many-head backends. Integrates it into the build and PyTorch custom-operator layers. Adds meta shape inference and parity tests against ChangesKDA decode
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Test
participant KdaDecodeOperator
participant invokeKdaDecode
participant CUDADecodeKernel
participant StateStorage
Test->>KdaDecodeOperator: Submit tensors and decode options
KdaDecodeOperator->>KdaDecodeOperator: Validate shapes, dtypes, strides, and alignment
KdaDecodeOperator->>invokeKdaDecode: Pack KdaDecodeParams and launch
invokeKdaDecode->>CUDADecodeKernel: Select compact or many-head backend
CUDADecodeKernel->>StateStorage: Update recurrent and convolution state
CUDADecodeKernel-->>KdaDecodeOperator: Produce decoded output
KdaDecodeOperator-->>Test: Return output and mutated state
🚥 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: 5
🧹 Nitpick comments (6)
cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu (4)
1395-1399: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
cudaFuncSetAttributeruns on every launch.This is a driver call on the per-token decode path; the attribute only needs to be set once per kernel instantiation. Hoist it behind a function-local static so repeated decodes skip it.
⚡ Set the attribute once per instantiation
+ auto* kernel = kda_decode_fusion_compact_heads_kernel<kApplyOnorm, true, kUseStaticDecodeLayout, kHeads, kHeads, + false, false, false, true, false, true, kUpdateConvState, kUseLowerBound, kApplyBetaSigmoid>; + static bool const smemConfigured = [&] + { + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kStageDynamicSmemBytes)); + return true; + }(); + (void) smemConfigured; - TLLM_CUDA_CHECK(cudaFuncSetAttribute( - kda_decode_fusion_compact_heads_kernel<kApplyOnorm, true, kUseStaticDecodeLayout, kHeads, kHeads, false, false, - false, true, false, true, kUpdateConvState, kUseLowerBound, kApplyBetaSigmoid>, - cudaFuncAttributeMaxDynamicSharedMemorySize, kStageDynamicSmemBytes));🤖 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/kdaDecode/kdaDecode.cu` around lines 1395 - 1399, Update the setup around kStageDynamicSmemBytes and cudaFuncSetAttribute for kda_decode_fusion_compact_heads_kernel to guard the attribute call with a function-local static, ensuring it executes only once for each kernel instantiation while preserving the existing CUDA error check.
122-136: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDead code:
cp_async_wait_oldestis never called;Bis an unused kernel parameter.
cp_async_wait_oldesthas no call site (both kernels inline their ownwait_group_*selection), andint Bat lines 374 and 911 is never read inside either kernel body — the batch size is only used for the grid dimensions in the launch wrappers. Dropping both shrinks the surface.As per coding guidelines, "do not leave defined functions unused".
♻️ Remove the unused helper
-__device__ __forceinline__ void cp_async_wait_oldest(int outstanding_groups) -{ - if (outstanding_groups >= 3) - { - cp_async_wait_group_2(); - } - else if (outstanding_groups == 2) - { - cp_async_wait_group_1(); - } - else - { - cp_async_wait_group_0(); - } -} -🤖 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/kdaDecode/kdaDecode.cu` around lines 122 - 136, Remove the unused device helper cp_async_wait_oldest and delete the unused int B parameters from both kernel definitions. Update all corresponding kernel launch sites and argument lists to match, while preserving batch-size usage in the launch wrapper grid dimensions.Source: Coding guidelines
43-64: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffNaming does not follow the repository C++ conventions.
The whole file uses snake_case for functions and locals (
bf16_load,store_state_float4,cp_async_*,kda_decode_fusion_compact_heads_kernel,s_state,dot_hk_a, …). The C++ guidelines require lower camelCase for functions and local variables (k-prefixed camelCase constants are already correct here). Worth normalizing while the file is new, before it accumulates call sites.As per coding guidelines, "Use lower-camel-case filenames, types in PascalCase, local variables/functions/namespaces in lower camelCase".
🤖 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/kdaDecode/kdaDecode.cu` around lines 43 - 64, Normalize function and local-variable naming throughout the new kdaDecode implementation to lower camelCase, including bf16_load, bf16_store, store_state_float4, and existing snake_case helpers, kernels, and locals such as cp_async_*, kda_decode_fusion_compact_heads_kernel, s_state, and dot_hk_a. Preserve behavior and keep k-prefixed constants in their existing camelCase form.Source: Coding guidelines
1625-1640: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: template fan-out cost, and
case 2is unreachable.The 12 head values multiply with the layout/onorm/lower-bound/beta/conv-update booleans into ~384 instantiations of a large kernel, which shows up in build time and binary size; several of the boolean flags are currently pinned to one value at every call site and could be dropped. Also,
invokeKdaDecoderoutesH == HV == 2to the compact backend, socase 2here can never be taken.🤖 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/kdaDecode/kdaDecode.cu` around lines 1625 - 1640, The KDA decode dispatch in invokeKdaDecode includes unnecessary template fan-out and an unreachable H == 2 specialization. Remove case 2 from dispatch_kda_decode_layout, and reduce template boolean parameters that are fixed at all call sites, including their corresponding dispatch arguments and kernel signatures, while preserving the supported runtime combinations and compact-backend routing.tests/unittest/_torch/thop/parallel/test_kda_decode.py (1)
16-16: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winScale
NUM_CACHE_SLOTSwith head count to avoid multi-GB allocations.Slot count is independent of
num_heads, so thenum_heads=96indexed cases allocate514 * 96 * 128 * 128fp32 ≈ 3.2 GB of state, plus another ~3.2 GB for thestate_beforeclone (line 367) — repeated across 3 batch sizes and 2 indexed configs. Indexing coverage does not need 514 slots for wide-head cases.♻️ Cap the pool by total elements
- num_cache_slots = NUM_CACHE_SLOTS if use_state_indices else batch_size + if use_state_indices: + # Keep the indexed pool a few slots wider than the batch without letting + # per-slot state (num_heads * 128 * 128 floats) blow up device memory. + num_cache_slots = min(NUM_CACHE_SLOTS, max(batch_size + 8, 8 * 1024 // num_heads)) + else: + num_cache_slots = batch_sizeAlso applies to: 76-77, 99-103
🤖 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/thop/parallel/test_kda_decode.py` at line 16, Scale NUM_CACHE_SLOTS with num_heads in the test configurations, capping the cache pool by a fixed total-element budget so wide-head cases use fewer slots while preserving adequate coverage for smaller head counts. Apply the same calculation consistently to the related NUM_CACHE_SLOTS definitions and keep the indexed cases within a memory-safe allocation limit.cpp/tensorrt_llm/thop/kdaDecodeOp.cpp (1)
39-43: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winTake the tensor parameters by
const&.Both helpers copy 22
at::Tensorhandles per call (44 atomic refcount ops per decode step) for read-only use.♻️ Signature change
-void validate_kda_decode_fusion_inputs(at::Tensor x_q, at::Tensor x_k, at::Tensor x_v, at::Tensor w_q_t, - at::Tensor w_k_t, at::Tensor w_v_t, at::Tensor bias_q, at::Tensor bias_k, at::Tensor bias_v, at::Tensor cs_q, - at::Tensor cs_k, at::Tensor cs_v, at::Tensor a_log, at::Tensor g, at::Tensor dt_bias, at::Tensor beta, - at::Tensor onorm_g, at::Tensor onorm_weight, std::optional<at::Tensor> const& ssm_state_indices, - at::Tensor cu_seqlens, at::Tensor state, at::Tensor out, bool apply_onorm, bool update_conv_cache) +void validate_kda_decode_fusion_inputs(at::Tensor const& x_q, at::Tensor const& x_k, at::Tensor const& x_v, + at::Tensor const& w_q_t, at::Tensor const& w_k_t, at::Tensor const& w_v_t, at::Tensor const& bias_q, + at::Tensor const& bias_k, at::Tensor const& bias_v, at::Tensor const& cs_q, at::Tensor const& cs_k, + at::Tensor const& cs_v, at::Tensor const& a_log, at::Tensor const& g, at::Tensor const& dt_bias, + at::Tensor const& beta, at::Tensor const& onorm_g, at::Tensor const& onorm_weight, + std::optional<at::Tensor> const& ssm_state_indices, at::Tensor const& cu_seqlens, at::Tensor const& state, + at::Tensor const& out, bool apply_onorm, bool update_conv_cache)Apply the same to
launch_selected_kernel.Also applies to: 182-187
🤖 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/kdaDecodeOp.cpp` around lines 39 - 43, Update the read-only tensor parameters in both validate_kda_decode_fusion_inputs and launch_selected_kernel to accept const at::Tensor& instead of copying at::Tensor handles; preserve the existing parameter order and behavior, including the optional ssm_state_indices reference.
🤖 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/kdaDecode/kdaDecode.cu`:
- Around line 654-701: Protect shared-memory state-stage reuse with block-wide
barriers: in cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu lines 654-701, add
__syncthreads() after consuming the current state_stage and before issuing the
chunk+2 prefetch; in lines 1288-1294, add __syncthreads() before the prefetch.
Keep the existing cp.async waits and __syncwarp() logic, ensuring both kernels
cannot overwrite a stage while another warp is still reading it.
- Around line 445-449: Widen convolution-state base offset calculations from int
to int64_t throughout the affected decode paths, including the kdaDecode offsets
near hk, hvv, and the many-heads equivalents. Ensure the full slot *
head-dimension * kConvStateWidth expression is evaluated in 64-bit arithmetic
before assignment, matching the existing recurrent-state offsets.
- Around line 91-95: Guard the cp_async_cg_16b device helper so the
cp.async.cg.shared.global assembly is only emitted when __CUDA_ARCH__ is at
least 800. Provide a compatible fallback for pre-SM80 architectures, or
otherwise ensure this target excludes those architectures while preserving the
existing 16-byte copy behavior.
In `@cpp/tensorrt_llm/thop/kdaDecodeOp.cpp`:
- Around line 202-204: In the KDA decode launch path, add a c10::cuda::CUDAGuard
using x_q.device() before obtaining the stream and calling invokeKdaDecode.
Ensure at::cuda::getCurrentCUDAStream() and the kernel launch execute on x_q’s
device.
- Around line 214-219: Move the rank validation for x_q and x_v ahead of the
x_q.size(1)/x_v.size(2) accesses and out allocation in the kda decode flow,
reusing validate_kda_decode_fusion_inputs or an equivalent rank-only pre-check.
Ensure rank-2/3 inputs produce the op’s “must be rank-4 tensors” validation
message before any dimension dereference, while retaining the existing full
validation before execution.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu`:
- Around line 1395-1399: Update the setup around kStageDynamicSmemBytes and
cudaFuncSetAttribute for kda_decode_fusion_compact_heads_kernel to guard the
attribute call with a function-local static, ensuring it executes only once for
each kernel instantiation while preserving the existing CUDA error check.
- Around line 122-136: Remove the unused device helper cp_async_wait_oldest and
delete the unused int B parameters from both kernel definitions. Update all
corresponding kernel launch sites and argument lists to match, while preserving
batch-size usage in the launch wrapper grid dimensions.
- Around line 43-64: Normalize function and local-variable naming throughout the
new kdaDecode implementation to lower camelCase, including bf16_load,
bf16_store, store_state_float4, and existing snake_case helpers, kernels, and
locals such as cp_async_*, kda_decode_fusion_compact_heads_kernel, s_state, and
dot_hk_a. Preserve behavior and keep k-prefixed constants in their existing
camelCase form.
- Around line 1625-1640: The KDA decode dispatch in invokeKdaDecode includes
unnecessary template fan-out and an unreachable H == 2 specialization. Remove
case 2 from dispatch_kda_decode_layout, and reduce template boolean parameters
that are fixed at all call sites, including their corresponding dispatch
arguments and kernel signatures, while preserving the supported runtime
combinations and compact-backend routing.
In `@cpp/tensorrt_llm/thop/kdaDecodeOp.cpp`:
- Around line 39-43: Update the read-only tensor parameters in both
validate_kda_decode_fusion_inputs and launch_selected_kernel to accept const
at::Tensor& instead of copying at::Tensor handles; preserve the existing
parameter order and behavior, including the optional ssm_state_indices
reference.
In `@tests/unittest/_torch/thop/parallel/test_kda_decode.py`:
- Line 16: Scale NUM_CACHE_SLOTS with num_heads in the test configurations,
capping the cache pool by a fixed total-element budget so wide-head cases use
fewer slots while preserving adequate coverage for smaller head counts. Apply
the same calculation consistently to the related NUM_CACHE_SLOTS definitions and
keep the indexed cases within a memory-safe allocation limit.
🪄 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: 86aaac3a-bf6f-4dad-94fc-45fcbeb7df1a
📒 Files selected for processing (10)
cpp/tensorrt_llm/CMakeLists.txtcpp/tensorrt_llm/kernels/CMakeLists.txtcpp/tensorrt_llm/kernels/kdaDecode/CMakeLists.txtcpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cucpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/kdaDecodeOp.cpprequirements-dev.txttensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytests/unittest/_torch/thop/parallel/test_kda_decode.py
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>
Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu (1)
1416-1422: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the positional boolean template packs. Eight/nine consecutive bare
true/falsetemplate arguments select kernel behavior; inserting or reordering a flag in the kernel signature shifts them all with no compile error. Naming them makes the intent auditable.As per coding guidelines: "document non-obvious call arguments with
/*paramName=*/".♻️ Example annotation
- kda_decode_fusion_compact_heads_kernel<kApplyOnorm, true, kUseStaticDecodeLayout, kHeads, kHeads, false, false, - false, true, false, true, kUpdateConvState, kUseLowerBound, kApplyBetaSigmoid> + kda_decode_fusion_compact_heads_kernel<kApplyOnorm, /*kAccumulateOnormSumsq=*/true, kUseStaticDecodeLayout, + /*kFixedHeads=*/kHeads, /*kFixedValueHeads=*/kHeads, /*kUseHeadGrid=*/false, /*kUseCacheGlobalStore=*/false, + /*kComputeOutputBeforeStore=*/false, /*kPreloadOnormParams=*/true, /*kIssueThirdStatePrefetchEarly=*/false, + /*kUseActiveQkReduction=*/true, kUpdateConvState, kUseLowerBound, kApplyBetaSigmoid>Also applies to: 1441-1443
🤖 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/kdaDecode/kdaDecode.cu` around lines 1416 - 1422, Annotate each positional boolean template argument in both the cudaFuncSetAttribute and kernel launch instantiations of kda_decode_fusion_compact_heads_kernel with /*paramName=*/ comments, using the kernel template parameter names. Keep argument order and values unchanged while making every boolean flag auditable.Source: Coding guidelines
cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.h (1)
31-35: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSupported head counts have no single source of truth. The set
{1,2,3,4,6,8,12,16,24,32,48,96}is now spelled out in the header predicate, in the CUDA dispatch switch, and again assupportedHeadsincpp/tensorrt_llm/thop/kdaDecodeOp.cpp(line 81). Adding a head count to one site without the others either aborts at thedefaultbranch or silently rejects a supported shape.
cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.h#L31-L35: make this the single source, e.g. aconstexpr std::arrayof supported counts thatisSupportedHeadCountscans.cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu#L1544-L1557: generate the dispatch cases from that list (fold expression over the array) instead of restating the values, and reuseisSupportedHeadCountfor the validation incpp/tensorrt_llm/thop/kdaDecodeOp.cpp.🤖 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/kdaDecode/kdaDecode.h` around lines 31 - 35, Make the supported head-count list in isSupportedHeadCount the single source of truth by defining a shared constexpr array and scanning it for validation. In cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu lines 1544-1557, generate dispatch cases from that array via a fold expression instead of duplicating values; update supportedHeads validation in cpp/tensorrt_llm/thop/kdaDecodeOp.cpp line 81 to reuse isSupportedHeadCount.
🤖 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.
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu`:
- Around line 1416-1422: Annotate each positional boolean template argument in
both the cudaFuncSetAttribute and kernel launch instantiations of
kda_decode_fusion_compact_heads_kernel with /*paramName=*/ comments, using the
kernel template parameter names. Keep argument order and values unchanged while
making every boolean flag auditable.
In `@cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.h`:
- Around line 31-35: Make the supported head-count list in isSupportedHeadCount
the single source of truth by defining a shared constexpr array and scanning it
for validation. In cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu lines
1544-1557, generate dispatch cases from that array via a fold expression instead
of duplicating values; update supportedHeads validation in
cpp/tensorrt_llm/thop/kdaDecodeOp.cpp line 81 to reuse isSupportedHeadCount.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1db00dcb-1c9c-4f6c-9594-f075ab1073e4
📒 Files selected for processing (3)
cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cucpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.hcpp/tensorrt_llm/thop/kdaDecodeOp.cpp
tburt-nv
left a comment
tburt-nv
left a comment
There was a problem hiding this comment.
MIT licensed, no new transitive dependencies.
Confirmed tests/unittest/_torch/thop/parallel is scheduled in CI.
|
/bot run --disable-fail-fast |
|
PR_Github #62932 [ run ] triggered by Bot. Commit: |
|
PR_Github #62932 [ run ] completed with state |
…-SM80 builds Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu (2)
1557-1583: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPut each
breakinside a compound case.The head dispatch uses one-line cases. The coding rule requires compound case braces with
breakinside the block.Proposed structure
-case 1: dispatch_kda_decode_layout<kCompact, 1>(p); break; +case 1: +{ + dispatch_kda_decode_layout<kCompact, 1>(p); + break; +}As per coding guidelines: “Structure switch statements clearly; ... place
breakinside compound case braces.”🤖 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/kdaDecode/kdaDecode.cu` around lines 1557 - 1583, Update the switch cases in dispatch_kda_decode_heads so each supported head-count case uses a compound brace block, with its dispatch_kda_decode_layout call and break inside the block. Preserve the existing head mappings and default error handling.Source: Coding guidelines
43-100: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffAlign new declarations with the required C++ naming and const style.
bf16_load,bf16_store,cp_async_cg_16b, and related helpers use snake_case.smem_addris unmodified but lacksconst. Declarations such asconst __nv_bfloat16 q_stateandconst Sum4 dot_hkviolate east-const. Apply the required naming and declaration style consistently.As per coding guidelines: “local variables/functions/namespaces in lower camel case”; “declare unmodified variables
const”; and “use east-const style.”Also applies to: 466-467, 481-482, 748-749
🤖 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/kdaDecode/kdaDecode.cu` around lines 43 - 100, Rename the affected helper functions and local variables to lower camel case, including bf16_load, bf16_store, store_state_float4, sigmoid_fast, silu_fast, softplus_fast, warp_reduce_sum, and cp_async_cg_16b. Apply the same naming style to related declarations and uses, including q_state and smem_addr. Mark unmodified locals such as smem_addr and the referenced reduction variables const, using east-const syntax for declarations such as __nv_bfloat16 const qState and Sum4 const dotHk.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu`:
- Around line 1557-1583: Update the switch cases in dispatch_kda_decode_heads so
each supported head-count case uses a compound brace block, with its
dispatch_kda_decode_layout call and break inside the block. Preserve the
existing head mappings and default error handling.
- Around line 43-100: Rename the affected helper functions and local variables
to lower camel case, including bf16_load, bf16_store, store_state_float4,
sigmoid_fast, silu_fast, softplus_fast, warp_reduce_sum, and cp_async_cg_16b.
Apply the same naming style to related declarations and uses, including q_state
and smem_addr. Mark unmodified locals such as smem_addr and the referenced
reduction variables const, using east-const syntax for declarations such as
__nv_bfloat16 const qState and Sum4 const dotHk.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cb1d1932-769b-4929-83d5-29fa8805f77f
📒 Files selected for processing (1)
cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu
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 --disable-fail-fast |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
cpp/tensorrt_llm/thop/kdaDecodeOp.cpp (3)
136-138: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftEnforce the
cu_seqlensvalue contract.
KdaDecodeParamsrequirescuSeqlensto equalarange(batchSize + 1), but this wrapper checks only dtype, layout, rank, and length. Invalid values reachinvokeKdaDecodeunchanged and can produce incorrect per-batch offsets or state updates.Construct the required sequence internally, or validate its contents before launch. Add a negative test for non-arange values.
🤖 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/kdaDecodeOp.cpp` around lines 136 - 138, Update the cu_seqlens validation in the KdaDecode wrapper to enforce that its values equal arange(B + 1), not only its CUDA dtype, contiguity, rank, and length; validate the contents before invoking invokeKdaDecode or construct the required sequence internally, and add a negative test covering non-arange values.
39-64: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winEnforce one CUDA device before launching KDA decode. Compare every tensor, including
ssm_state_indiceswhen present, withx_q.device()before output allocation and dispatch. Use a stream for that device instead of the implicit current-device stream.🤖 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/kdaDecodeOp.cpp` around lines 39 - 64, Update validate_kda_decode_fusion_inputs and the KDA decode dispatch to require every tensor, including present ssm_state_indices, to be on x_q.device() before allocating outputs or launching work. Select the CUDA stream associated with x_q.device() rather than relying on the implicit current-device stream, and preserve the existing dtype validation.
128-135: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate state-slot indices before launching the kernel.
Negative or out-of-range
ssm_state_indicesvalues cause invalid global-memory accesses tostateand, whenupdate_conv_cacheis enabled, the convolution-cache pools. Add bounds validation againststate.size(0)and negative tests. Reject duplicate indices only if the state-update contract requires unique slots.🤖 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/kdaDecodeOp.cpp` around lines 128 - 135, Extend validation in the ssm_state_indices checks before the kernel launch to verify every index is nonnegative and less than state.size(0), rejecting invalid values before any state or convolution-cache access. Preserve the existing CUDA, int32, contiguous, one-dimensional, and [B] checks, and only add duplicate-index rejection if the state-update contract requires unique slots.
🧹 Nitpick comments (3)
cpp/tensorrt_llm/thop/kdaDecodeOp.cpp (3)
123-127: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winName the alignment constants.
Replace the literal
4and16values with named constants such askFloat4ElementsandkStateAlignmentBytes.As per coding guidelines: avoid magic literals except
0,nullptr,true, andfalse; initialize named constants instead, usingk-prefixed camelCase names.🤖 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/kdaDecodeOp.cpp` around lines 123 - 127, In the validation checks near the state alignment logic, replace the literal 4 and 16 values with locally defined k-prefixed camelCase constants such as kFloat4Elements and kStateAlignmentBytes. Use those constants consistently for the stride divisibility and data-pointer alignment checks while preserving the existing error messages and behavior.Source: Coding guidelines
189-196: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the
KdaDecodeParamsaggregate initialization.This initializer passes many pointers, strides, dimensions, and flags positionally. A future field insertion or reordering can compile while binding the wrong value.
Add
/*xQ=*/,/*logA=*/,/*stateSlotStride=*/, and equivalent comments for every field, as required by the C++ guidelines.As per coding guidelines: document non-obvious call arguments with
/*paramName=*/.🤖 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/kdaDecodeOp.cpp` around lines 189 - 196, Annotate every argument in the KdaDecodeParams aggregate initialization with inline /*paramName=*/ comments, including pointers, dimensions, strides, and flags. Update only the initializer for tensorrt_llm::kernels::kdaDecode::KdaDecodeParams, using the exact corresponding field names such as xQ, logA, and stateSlotStride, and ensure all positional values are documented.Source: Coding guidelines
39-43: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse lower-camel-case names for the new C++ identifiers.
Rename
validate_kda_decode_fusion_inputs,launch_selected_kernel,kda_decode_fusion_forward, and their snake_case parameters to lower camel case, such asvalidateKdaDecodeFusionInputs,launchSelectedKernel,kdaDecodeFusionForward,xQ, andssmStateIndices.Keep the external PyTorch schema name unchanged if it is part of the public API.
As per coding guidelines: use lower-camel-case filenames, types in PascalCase, and local variables/functions/namespaces in lower camel case.
Also applies to: 178-183, 203-208
🤖 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/kdaDecodeOp.cpp` around lines 39 - 43, Rename the new functions validate_kda_decode_fusion_inputs, launch_selected_kernel, and kda_decode_fusion_forward to validateKdaDecodeFusionInputs, launchSelectedKernel, and kdaDecodeFusionForward, and convert their snake_case parameters and local identifiers to lower camel case such as xQ and ssmStateIndices. Apply the same naming updates throughout their declarations, definitions, and call sites while preserving any external PyTorch schema name unchanged.Source: Coding guidelines
🤖 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.
Outside diff comments:
In `@cpp/tensorrt_llm/thop/kdaDecodeOp.cpp`:
- Around line 136-138: Update the cu_seqlens validation in the KdaDecode wrapper
to enforce that its values equal arange(B + 1), not only its CUDA dtype,
contiguity, rank, and length; validate the contents before invoking
invokeKdaDecode or construct the required sequence internally, and add a
negative test covering non-arange values.
- Around line 39-64: Update validate_kda_decode_fusion_inputs and the KDA decode
dispatch to require every tensor, including present ssm_state_indices, to be on
x_q.device() before allocating outputs or launching work. Select the CUDA stream
associated with x_q.device() rather than relying on the implicit current-device
stream, and preserve the existing dtype validation.
- Around line 128-135: Extend validation in the ssm_state_indices checks before
the kernel launch to verify every index is nonnegative and less than
state.size(0), rejecting invalid values before any state or convolution-cache
access. Preserve the existing CUDA, int32, contiguous, one-dimensional, and [B]
checks, and only add duplicate-index rejection if the state-update contract
requires unique slots.
---
Nitpick comments:
In `@cpp/tensorrt_llm/thop/kdaDecodeOp.cpp`:
- Around line 123-127: In the validation checks near the state alignment logic,
replace the literal 4 and 16 values with locally defined k-prefixed camelCase
constants such as kFloat4Elements and kStateAlignmentBytes. Use those constants
consistently for the stride divisibility and data-pointer alignment checks while
preserving the existing error messages and behavior.
- Around line 189-196: Annotate every argument in the KdaDecodeParams aggregate
initialization with inline /*paramName=*/ comments, including pointers,
dimensions, strides, and flags. Update only the initializer for
tensorrt_llm::kernels::kdaDecode::KdaDecodeParams, using the exact corresponding
field names such as xQ, logA, and stateSlotStride, and ensure all positional
values are documented.
- Around line 39-43: Rename the new functions validate_kda_decode_fusion_inputs,
launch_selected_kernel, and kda_decode_fusion_forward to
validateKdaDecodeFusionInputs, launchSelectedKernel, and kdaDecodeFusionForward,
and convert their snake_case parameters and local identifiers to lower camel
case such as xQ and ssmStateIndices. Apply the same naming updates throughout
their declarations, definitions, and call sites while preserving any external
PyTorch schema name unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1b8c607b-41be-4bf0-9de8-b2a9193682ee
📒 Files selected for processing (2)
cpp/tensorrt_llm/thop/kdaDecodeOp.cpptests/unittest/_torch/thop/parallel/test_kda_decode.py
💤 Files with no reviewable changes (1)
- tests/unittest/_torch/thop/parallel/test_kda_decode.py
|
PR_Github #63679 [ run ] triggered by Bot. Commit: |
|
PR_Github #63679 [ run ] completed with state
|
Mgluhovskoi
left a comment
Mgluhovskoi
left a comment
There was a problem hiding this comment.
LGTM. All concerns were addressed.
|
/bot run |
|
PR_Github #63850 [ run ] triggered by Bot. Commit: |
|
PR_Github #63850 [ run ] completed with state |
Squash of the KimiLinear model integration on top of the kernel PRs (NVIDIA#17190, NVIDIA#17054, NVIDIA#17266, NVIDIA#17225): - KimiLinear model (modeling_kimi_k3) and KimiLinearConfig registration - Kimi K3 support modules: KDA mixer, K3 MoE, K3 MLA, fused attention-residual wrapper - K3 MLA module refactored onto the general MLA path (TRTLLM-14811) - fused_moe: SiTu activation and communication_method support - SiTu in the ActType_TrtllmGen python enum (python mirror of the C++ enum from the MoE kernel drop) - kda_decode: accept an optional out tensor (to be folded into NVIDIA#17054) - KDA kernel/runtime unit tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Squash of the KimiLinear model integration on top of the kernel PRs (NVIDIA#17190, NVIDIA#17054, NVIDIA#17266, NVIDIA#17225): - KimiLinear model (modeling_kimi_k3) and KimiLinearConfig registration - Kimi K3 support modules: KDA mixer, K3 MoE, K3 MLA, fused attention-residual wrapper - K3 MLA module refactored onto the general MLA path (TRTLLM-14811) - fused_moe: SiTu activation and communication_method support - SiTu in the ActType_TrtllmGen python enum (python mirror of the C++ enum from the MoE kernel drop) - kda_decode: accept an optional out tensor (to be folded into NVIDIA#17054) - KDA kernel/runtime unit tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Squash of the KimiLinear model integration on top of the kernel PRs (NVIDIA#17190, NVIDIA#17054, NVIDIA#17266, NVIDIA#17225): - KimiLinear model (modeling_kimi_k3) and KimiLinearConfig registration - Kimi K3 support modules: KDA mixer, K3 MoE, K3 MLA, fused attention-residual wrapper - K3 MLA module refactored onto the general MLA path (TRTLLM-14811) - fused_moe: SiTu activation and communication_method support - SiTu in the ActType_TrtllmGen python enum (python mirror of the C++ enum from the MoE kernel drop) - kda_decode: accept an optional out tensor (to be folded into NVIDIA#17054) - KDA kernel/runtime unit tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
8 participants
Dev Engineer Review
int64_taddress calculations and pre-SM80 guards forcp.asyncpaths.trtllm::kda_decodetorch operation and its fake implementation.flash-linear-attention==0.5.2for reference validation.QA Engineer Review
test_kda_decode_matches_fla(...)intests/unittest/_torch/thop/parallel/test_kda_decode.py.torch.ops.trtllm.kda_decodewithfla.ops.kda.fused_recurrent_kda.Description
Integrate kimi k3 kda decode kernel as a thop and added related unittest. Note that it requires FlashLinearAttention package as a dependency for accuracy validation, it has been added into
requirements-dev.txt. The tests will be included in CI automatically.Test Coverage
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.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.