cpp/tensorrt_llm/kernels/topkLastDim.cu (1)

1535-1552: Duplicate implementation of nextPowerOfTwo.

This function is duplicated from customMoeRoutingKernels.cu. As mentioned in the review of that file, consider moving this to a common utility header.

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py (1)

535-541: Additive fake op looks correct; consider deduplicating with renorm’s fake.

The new fake returns the right shapes/dtypes and mirrors the renorm variant. Minor: you can DRY the two fakes to a shared helper to avoid divergence.

Apply this refactor to reduce duplication:

@@
-    @torch.library.register_fake("trtllm::renorm_moe_routing_op")
-    def _(router_logits, topk):
-        num_tokens = router_logits.shape[0]
-        sz = (num_tokens, topk)
-        return router_logits.new_empty(
-            sz, dtype=torch.int32), router_logits.new_empty(sz,
-                                                            dtype=torch.float32)
+    def _fake_moe_routing(router_logits, topk):
+        num_tokens = router_logits.shape[0]
+        sz = (num_tokens, topk)
+        return router_logits.new_empty(sz, dtype=torch.int32), router_logits.new_empty(sz, dtype=torch.float32)
+
+    @torch.library.register_fake("trtllm::renorm_moe_routing_op")
+    def _(router_logits, topk):
+        return _fake_moe_routing(router_logits, topk)
@@
-    @torch.library.register_fake("trtllm::default_moe_routing_op")
-    def _(router_logits, topk):
-        num_tokens = router_logits.shape[0]
-        sz = (num_tokens, topk)
-        return router_logits.new_empty(
-            sz, dtype=torch.int32), router_logits.new_empty(sz,
-                                                            dtype=torch.float32)
+    @torch.library.register_fake("trtllm::default_moe_routing_op")
+    def _(router_logits, topk):
+        return _fake_moe_routing(router_logits, topk)

cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h (2)

2-2: Copyright header year.

Using 2025 is acceptable for new headers; ensure consistency across the newly added/updated kernel files in this PR.

---

26-29: Prefer const-correct input pointer in the launcher API.

The API takes InputT* for routerLogits, but the op treats router logits as read-only. Making it InputT const* clarifies intent and avoids forcing callers to request mutable pointers.

Proposed signature change:

-template <typename InputT, typename OutputT, typename IdxT, bool DoSoftmaxBeforeTopK>
-void invokeRenormMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens,
+template <typename InputT, typename OutputT, typename IdxT, bool DoSoftmaxBeforeTopK>
+void invokeRenormMoeRouting(InputT const* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens,
     int64_t const numExperts, int64_t const topK, cudaStream_t const stream);

If you take this, adjust the .cu implementation and the call sites to use data_ptr() instead of mutable_data_ptr().

tensorrt_llm/_torch/modules/fused_moe/routing.py (1)

67-75: Centralize dispatch thresholds for CUDA kernels
To keep Python logic in sync with the C++ TORCH_CHECK limits (num_experts ≤ 128, top_k ≤ 8), introduce module‐level constants—e.g.:

# tensorrt_llm/_torch/modules/fused_moe/routing.py
MAX_EXPERTS_FOR_CUDA = 128
MAX_TOPK_FOR_CUDA     = 8

Then update both apply implementations to use these constants:

• DefaultMoeRoutingMethod.apply (currently lines 67–75)
• RenormalizeMoeRoutingMethod.apply (currently lines 117–125)

Also add a comment linking back to the checks in customMoeRoutingOp.cpp, for example:

// Enforced in customMoeRoutingOp.cpp: TORCH_CHECK(topk <= 8 && num_experts <= 128)

Optional: extend existing unit tests to cover cases where num_experts > MAX_EXPERTS_FOR_CUDA and top_k > MAX_TOPK_FOR_CUDA, ensuring both the PyTorch fallback and the CUDA path are exercised.

cpp/tensorrt_llm/thop/customMoeRoutingOp.cpp (1)

48-66: Const-correctness and pointer access.

Inputs are read-only; prefer passing const pointers to the kernel and using data_ptr() over mutable_data_ptr(). This also avoids implying in-place modification of router_logits.

If you adopt the header change to take InputT const*, update calls like:

-        tk::invokeRenormMoeRouting<float, float, int32_t, DoSoftmaxBeforeTopK>(
-            reinterpret_cast<float*>(router_logits.mutable_data_ptr()),
+        tk::invokeRenormMoeRouting<float, float, int32_t, DoSoftmaxBeforeTopK>(
+            router_logits.data_ptr<float>(),
             reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
             reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()), num_tokens, num_experts, topk, stream);

Repeat similarly for bfloat16 and half.

cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh (2)

115-121: Consider providing a more descriptive macro name.

The macro TOPK_SWAP is very generic and could potentially conflict with other macros in the codebase. Consider using a more specific name that indicates its purpose within the topk reduction context.

Apply this diff to use a more descriptive name:

-#define TOPK_SWAP(I, J)                                                                                                \
+#define TOPK_REDUCTION_SWAP(I, J)                                                                                      \
     {                                                                                                                  \
         auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx);                                                    \
         auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx);                                                    \
         topK[I].compValIdx = pairMax;                                                                                  \
         topK[J].compValIdx = pairMin;                                                                                  \
     }

And update all usages accordingly:

-        TOPK_SWAP(0, 1);
+        TOPK_REDUCTION_SWAP(0, 1);

---

193-194: Complex ternary operator could benefit from clarification.

The ternary expression on Line 193 is correct but complex. Consider adding a comment to explain the logic: when updating (update == true) and we're at the last element (nn == N - 1), we replace it with minValue to prepare for the next iteration.

Add a clarifying comment:

 #pragma unroll
         for (int nn = 0; nn < N; ++nn)
         {
+            // When updating: shift all elements left by one, and fill the last position with minValue
             topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} : update ? topK[nn + 1] : topK[nn];
         }

cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu (3)

44-47: Redundant conditional check in calcSoftmax.

The condition score >= maxScore on Line 47 is redundant since maxScore is initialized to -INFINITY, making the condition always true for any valid score.

Simplify the code:

     T maxScore = T{-INFINITY};
     if (laneIdx < NumTopExperts)
     {
-        maxScore = score >= maxScore ? score : maxScore;
+        maxScore = score;
     }

---

231-232: Improve error message for kernel selection failure.

The error message "Can not find corresponding kernel instance" doesn't provide enough context about what configuration failed.

Provide more informative error message:

-        TLLM_CHECK_WITH_INFO(kernelInstance != nullptr, "Can not find corresponding kernel instance.");
+        TLLM_CHECK_WITH_INFO(kernelInstance != nullptr, 
+            "Cannot find kernel instance for maxNumExperts=" + std::to_string(maxNumExperts) + 
+            ", maxNumTopExperts=" + std::to_string(maxNumTopExperts));

---

168-185: Consolidate nextPowerOfTwo in common math utilities

The nextPowerOfTwo helper is currently defined in multiple places (and risks bit-level inconsistency):

• cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu (lines 168–185)
• cpp/tensorrt_llm/kernels/topkLastDim.cu (around line 1535)
• cpp/tensorrt_llm/thop/thUtils.h / thUtils.cpp
• cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.h

To DRY-up the code, please:

• Add an inline __host__ __device__
int nextPowerOfTwo(int) implementation to
  cpp/tensorrt_llm/common/mathUtils.h
  (in namespace tensorrt_llm::common).
• Replace the local definitions in each .cu/.h file with
  #include "cpp/tensorrt_llm/common/mathUtils.h"
  and call the shared utility.
• Remove the duplicated implementations from those modules.

This will centralize the logic and prevent divergence.

cpp/tensorrt_llm/kernels/topkLastDim.cu (2)

1564-1565: Consider extracting the max length calculation.

The pattern nextPowerOfTwo(len) < 32 ? 32 : nextPowerOfTwo(len) ensures a minimum of 32, but it calls nextPowerOfTwo twice. Consider computing it once.

-    uint32_t max_len = nextPowerOfTwo(len) < 32 ? 32 : nextPowerOfTwo(len);
+    uint32_t power_of_two_len = nextPowerOfTwo(len);
+    uint32_t max_len = power_of_two_len < 32 ? 32 : power_of_two_len;

---

1569-1612: Consider using a lookup table for kernel selection.

The nested switch statements for kernel selection could be simplified using a lookup table or a more compact approach.

Consider refactoring to use a template-based approach or lookup table to reduce code duplication:

+template<int MaxLen, int MaxTopK>
+struct KernelSelector {
+    static constexpr auto value = &moe_topk::moe_topk_kernel<InputT, OutputT, IdxT, MaxLen, MaxTopK>;
+};
+
+template<typename InputT, typename OutputT, typename IdxT>
+auto selectKernel(uint32_t max_len, uint32_t moe_topk) {
+    // Use a 2D lookup table or similar structure
+    // This is a sketch - actual implementation would need proper template metaprogramming
+}

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: sklevtsov-nvidia
PR: NVIDIA/TensorRT-LLM#3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1198-1209
Timestamp: 2025-08-08T22:03:40.707Z
Learning: In the CUTLASS MoE kernels (cpp/tensorrt_llm/cutlass_extensions), when `layout_info.fusion` is set to `TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE`, the `router_scales` parameter must be non-null by design. The fused finalize kernel epilogue does not perform nullptr checks and requires valid router scales to function correctly. This is an implicit contract that callers must satisfy when enabling the FINALIZE fusion mode.

Learnt from: sklevtsov-nvidia
PR: NVIDIA/TensorRT-LLM#3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1198-1209
Timestamp: 2025-08-08T22:03:40.707Z
Learning: In the CUTLASS MoE kernels (cpp/tensorrt_llm/cutlass_extensions), when `layout_info.fusion` is set to `TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE`, the `router_scales` parameter must be non-null by design. The fused finalize kernel epilogue does not perform nullptr checks and requires valid router scales to function correctly. This is an implicit contract that callers must satisfy when enabling the FINALIZE fusion mode.

Learnt from: MatthiasKohl
PR: NVIDIA/TensorRT-LLM#6904
File: cpp/tensorrt_llm/kernels/mlaKernels.cu:436-439
Timestamp: 2025-08-14T15:36:37.576Z
Learning: CUDA kernels prioritize performance and should avoid runtime bounds checking or conditional operations that cause branching/warp divergence. Input validation should be done at the host level before kernel launch, not per-thread in the kernel.

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.cuh (1)

381-387: Arch-based PDL gating switch looks correct; optionally restrict trigger to a single thread and add an #endif comment.

Compilation gating on SM90+ via CUDA_ARCH is the right direction. To avoid redundant programmatic-launch triggers from every thread, gate the trigger to a single elected thread. Also align the #endif style with the rest of the file.

Apply this diff:

-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    if constexpr (KernelParams::UsePdl)
-    {
-        cudaTriggerProgrammaticLaunchCompletion();
-    }
-#endif
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    if constexpr (KernelParams::UsePdl)
+    {
+        // Trigger once per cluster to minimize redundant triggers.
+        if (clusterBlockRank == 0 && warpIdx == 0 && threadIdx.x == 0)
+        {
+            cudaTriggerProgrammaticLaunchCompletion();
+        }
+    }
+#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py (1)

535-541: New fake op wiring looks correct and consistent with renorm_moe_routing_op.

• Output shapes/dtypes align with the routing path expectations: (num_tokens, topk) int32 indices and float32 values.
• Name matches the Torch op usage in routing.py.

Minor:

• Consider factoring the duplicated body with renorm_moe_routing_op into a tiny helper to avoid drift.
• Ensure the NVIDIA copyright header is present at the top of this Python file per repo guidelines.

tensorrt_llm/_torch/modules/fused_moe/routing.py (2)

54-58: Expose and document the new force_enable_pytorch_op in init.

The new switch is helpful for forcing the PyTorch path. Please add a brief docstring for the class (or the init signature) describing top_k and force_enable_pytorch_op to keep external usage self-explanatory.

---

67-75: Heuristic matches the CUDA fast-path constraints; keep thresholds in sync.

The gating (num_experts > 128 or top_k > 8) correctly routes large cases to the PyTorch path and small cases to the custom CUDA op. Please ensure this stays consistent with kernels/topkLastDim.cu fast-path constraints (len <= 128 and k <= 8). If those limits change, update this condition accordingly.

cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h (1)

26-29: Template API change is reasonable; consider naming or docs to reflect the new toggle.

Adding DoSoftmaxBeforeTopK is the right abstraction. Since invokeRenormMoeRouting now also covers the “Default” flow (softmax-before-topk), consider documenting this in a brief comment or renaming in a future refactor to avoid confusion at call sites.

cpp/tensorrt_llm/kernels/topkLastDim.cu (4)

25-36: Add to use std::max in host code.

This TU uses max() in host code. Prefer std::max and include to avoid macro collisions.

Apply:

 #include "topkLastDim.h"
+#include <algorithm>
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

---

1323-1324: Use std::max for clarity and to avoid macro pitfalls.

Switch to std::max since this is host code and is now included.

Apply:

-    temp_storage_bytes = max(temp_storage_bytes, temp_storage_bytes_sort);
+    temp_storage_bytes = std::max(temp_storage_bytes, temp_storage_bytes_sort);

---

1452-1452: Use std::max here as well.

Apply:

-    temp_storage_bytes = max(temp_storage_bytes, temp_storage_bytes_sort);
+    temp_storage_bytes = std::max(temp_storage_bytes, temp_storage_bytes_sort);

---

1555-1618: Kernel dispatch table: minor cleanups and edge-case note.

• Case 96 for max_len is unreachable with current nextPowerOfTwo() logic (it yields 32/64/128). Safe to remove or keep for readability, but it's dead code.
• Param greater is unused in moe_reduce_topk; either wire it (if you add “small-K min” in future) or remove it to reduce confusion.

cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh (4)

33-35: Verify architecture gating for redux.sync path

kTLLM_GEN_HAS_FAST_REDUX is tied to arch::is_major_v<10>. Please confirm this matches the first SM that supports redux.sync.max.u32 (SM90 on Hopper in practice). If not, gate on the correct SM to avoid illegal instruction on older GPUs.

Would you like me to update this to a helper constant like arch::hasFastRedux() and wire it to SM90+?

---

83-95: Use volatile inline asm for redux to prevent unwanted reordering

Minor but conventional; also documents that there are no memory side-effects.

-            asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(compValIdx));
+            asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(compValIdx));

---

115-121: Avoid relying on unqualified min/max in device code

Unqualified min/max can be problematic depending on includes. Being explicit removes ambiguity.

-        auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx);
-        auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx);
+        auto pairMin = topK[I].compValIdx < topK[J].compValIdx ? topK[I].compValIdx : topK[J].compValIdx;
+        auto pairMax = topK[I].compValIdx < topK[J].compValIdx ? topK[J].compValIdx : topK[I].compValIdx;

---

169-173: Fix static_assert condition/message mismatch

N is the per-lane candidate count (= MaxNumExperts / 32); supported N is up to 4. The message currently says “<= 128” which is misleading here.

-    static_assert(N < 5, "Only support candidates number less than or equal to 128");
+    static_assert(N <= 4, "Only supports up to 4 candidates per lane (MaxNumExperts <= 128)");

cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu (2)

220-227: Unreachable CASE(96) branch

maxNumExperts is nextPowerOfTwo(numExperts) floored to at least 32, so values are 32/64/128. CASE(96) will never be selected.

Remove the CASE(96) branch to reduce confusion.

---

168-186: Deduplicate nextPowerOfTwo implementation

nextPowerOfTwo is already implemented in topkLastDim.cu. Prefer a single utility (header) to avoid divergence.

I can extract nextPowerOfTwo into a small header (e.g., kernels/topkUtils.cuh) and refactor both call sites to include it.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: sklevtsov-nvidia
PR: NVIDIA/TensorRT-LLM#3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1198-1209
Timestamp: 2025-08-08T22:03:40.707Z
Learning: In the CUTLASS MoE kernels (cpp/tensorrt_llm/cutlass_extensions), when `layout_info.fusion` is set to `TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE`, the `router_scales` parameter must be non-null by design. The fused finalize kernel epilogue does not perform nullptr checks and requires valid router scales to function correctly. This is an implicit contract that callers must satisfy when enabling the FINALIZE fusion mode.

Learnt from: MatthiasKohl
PR: NVIDIA/TensorRT-LLM#6904
File: cpp/tensorrt_llm/kernels/mlaKernels.cu:436-439
Timestamp: 2025-08-14T15:36:37.576Z
Learning: CUDA kernels prioritize performance and should avoid runtime bounds checking or conditional operations that cause branching/warp divergence. Input validation should be done at the host level before kernel launch, not per-thread in the kernel.