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

195-198: Inconsistent top_k guard removal.

The comment on line 194 indicates "alltoall without allgather only supports top_k % 4 == 0", but the condition was not fully removed. This creates inconsistency between the comment and actual implementation. The enable_alltoall property now checks if experts_per_token % 4 == 0 rather than supporting generic top_k values as the broader changes suggest.

Either update the comment to reflect the actual constraint or fully remove the % 4 == 0 check:

-                and self.routing_method.experts_per_token % 4 ==
-                0  # alltoall without allgather only supports top_k % 4 == 0
+                # Generic top_k values now supported

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

141-209: Missing validation for expert count divisibility.

Line 150-151 checks that expertCount + 1 <= 512, but this check seems incomplete. The error message is also not descriptive enough about why this constraint exists.

     CHECK_INPUT(expertsIds, torch::kInt32);
     TORCH_CHECK(expertCount % 4 == 0, "expertCount must be divisible by 4");
     TORCH_CHECK(slotCount % 4 == 0, "slotCount must be divisible by 4");
-    TORCH_CHECK(expertCount + 1 <= 512, 
-        "expertCount + 1 is larger than 512");
+    TORCH_CHECK(expertCount + 1 <= 512, 
+        "expertCount + 1 (" + std::to_string(expertCount + 1) + 
+        ") exceeds maximum supported value of 512");

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

18-18: Prefer C++ headers: use instead of <stdint.h>.

Use C++ headers to align with coding standards and avoid pulling symbols into the global namespace unexpectedly.

-#include <stdint.h>
+#include <cstdint>

---

31-33: Potential duplication/confusion of WARP constants across the codebase.

WARP_SIZE and WARP_MASK are also defined/used elsewhere (e.g., FusedMoeCommunicator, kernels .cu). Consider centralizing these definitions in this common header and consuming them via this header only to avoid divergence.

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

1322-1326: Magic number for 48 KiB dynamic shared memory threshold.

Prefer a named constexpr (e.g., kDEFAULT_SMEM_LIMIT_PER_BLOCK) instead of 48 * 1024 to document intent. It also eases future adjustments.

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

109-121: Variable shadowing inside getFieldUncompactSize; use a distinct name or the constant directly.

alignedUnitBytes is re-declared as a constexpr within the if block, shadowing the outer variable. This is confusing and error-prone. Use a different name or the constant directly.

-    if (alignedUnitBytes != 16)
+    if (alignedUnitBytes != 16)
     {
-        constexpr int alignedUnitBytes = BYTES_PER_16B_BLOCK;
-        currentFieldSize = currentFieldSize / alignedUnitBytes * alignedUnitBytes;
-        currentFieldSize += alignedUnitBytes * 2;
+        currentFieldSize = (currentFieldSize / BYTES_PER_16B_BLOCK) * BYTES_PER_16B_BLOCK;
+        currentFieldSize += BYTES_PER_16B_BLOCK * 2;
     }

---

260-318: Duplicate WARP_SIZE definition vs common header.

FusedMoeCommunicator defines its own WARP_SIZE = 32, while moeCommKernelsCommon.h also defines WARP_SIZE and WARP_MASK. Prefer a single source of truth (e.g., the common header) to reduce risk of divergence.

cpp/tensorrt_llm/thop/moeCommOp.cpp (2)

31-36: Consider adding alignment checks for field info.

The function accepts a 2D tensor and extracts element size and stride information. Consider validating that the tensor's memory layout meets any alignment requirements that the fused MOE kernels might expect.

 void setMoeCommFieldInfo(tensorrt_llm::kernels::MoeCommFieldInfo& fieldInfo, torch::Tensor const& tensor)
 {
     TORCH_CHECK(tensor.dim() == 2, "tensor must be a 2D tensor");
     int eltSize = tensor.dtype().itemsize();
+    // Validate that the tensor is contiguous or meets alignment requirements
+    TORCH_CHECK(tensor.is_contiguous(), "tensor must be contiguous for optimal MOE communication");
     fieldInfo.fillFieldInfo(static_cast<uint8_t*>(tensor.data_ptr()), eltSize, tensor.size(1), tensor.stride(0));
 }

---

231-231: Add namespace closing comment.

According to the coding guidelines, namespaces should be closed with a comment naming the namespace.

-} // namespace torch_ext
+} // namespace torch_ext

cpp/tests/unit_tests/kernels/fusedMoeCommKernelTest.cpp (2)

314-314: Use EXPECT_FLOAT_EQ for float comparisons.

For floating-point comparisons, prefer EXPECT_FLOAT_EQ over EXPECT_NEAR with a fixed epsilon when comparing exact values.

-                        EXPECT_NEAR(expected, actual, 1e-6f) << "Scale mismatch at warp=" << tokenId << ", k=" << k;
+                        EXPECT_FLOAT_EQ(expected, actual) << "Scale mismatch at warp=" << tokenId << ", k=" << k;

---

1280-1289: Consistent cleanup pattern needed.

The cleanup section has inconsistent patterns for checking null pointers before deletion. Array deletion is safe with null pointers in C++.

         // Cleanup temporary arrays
-        if (resultRecvTokenSlots)
-            delete[] resultRecvTokenSlots;
-        if (resultRecvScales)
-            delete[] resultRecvScales;
+        delete[] resultRecvTokenSlots;  // Safe even if null
+        delete[] resultRecvScales;      // Safe even if null
         for (int i = 0; i < fieldCount; i++)
         {
-            if (resultRecvFields[i])
-                delete[] resultRecvFields[i];
+            delete[] resultRecvFields[i];  // Safe even if null
         }

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_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

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.

tests/unittest/_torch/thop/test_moe_alltoall.py (1)

319-435: Duplicate FP8 test class — remove the second definition

TestMoeAlltoAllFP8SingleGPU is defined twice. This will shadow the first definition and can lead to confusion and flakiness in test discovery. Keep only one copy.

Apply this diff to remove the duplicate block starting at Line 919:

- class TestMoeAlltoAllFP8SingleGPU(unittest.TestCase):
-
-     def setUp(self):
-         torch.manual_seed(0x1234)
-         tllm.logger.set_level('error')
-
-     def test_moe_alltoall_fp8_with_indices(self):
-         """Test fp8 alltoall with properly constructed indices"""
-         torch.cuda.set_device(0)
- 
-         # Match dimensions from the error
-         input_entry_count = 16384
-         output_entry_count = 16384
-         vector_dim = 2944
-         sf_vector_dim = 92  # Scaling factor dimension from error
-         send_recv_count = 1000  # Number of entries to send/receive
- 
-         # Create input tensors - first as float16, then convert
-         input_tensor_fp16 = torch.randn(input_entry_count,
-                                         vector_dim,
-                                         dtype=torch.float16,
-                                         device='cuda')
-         input_tensor_fp8 = input_tensor_fp16.to(torch.float8_e4m3fn)
- 
-         # Scaling factor tensor
-         input_sf_tensor = torch.randint(1,
-                                         255, (input_entry_count, sf_vector_dim),
-                                         dtype=torch.uint8,
-                                         device='cuda')
- 
-         # Expert selection tensors
-         input_experts = torch.randint(0,
-                                       64, (input_entry_count, 4),
-                                       dtype=torch.int32,
-                                       device='cuda')
-         input_scales = torch.rand(input_entry_count,
-                                   4,
-                                   dtype=torch.float32,
-                                   device='cuda')
- 
-         # Output tensors
-         output_tensor_fp8 = torch.zeros(output_entry_count,
-                                         vector_dim,
-                                         dtype=torch.float8_e4m3fn,
-                                         device='cuda')
-         output_sf_tensor = torch.zeros(output_entry_count,
-                                        sf_vector_dim,
-                                        dtype=torch.uint8,
-                                        device='cuda')
-         output_experts = torch.zeros(output_entry_count,
-                                      4,
-                                      dtype=torch.int32,
-                                      device='cuda')
-         output_scales = torch.zeros(output_entry_count,
-                                     4,
-                                     dtype=torch.float32,
-                                     device='cuda')
- 
-         # Construct send/recv indices
-         send_cumsum = torch.tensor([send_recv_count],
-                                    dtype=torch.int32,
-                                    device='cuda')
-         recv_cumsum = torch.tensor([send_recv_count],
-                                    dtype=torch.int32,
-                                    device='cuda')
- 
-         # Random indices for sending
-         send_indices = torch.randperm(input_entry_count,
-                                       dtype=torch.int32,
-                                       device='cuda')[:send_recv_count]
-         recv_indices = torch.randperm(output_entry_count,
-                                       dtype=torch.int32,
-                                       device='cuda')[:send_recv_count]
- 
-         # Create workspace
-         workspace_size = torch.ops.trtllm.get_moe_commworkspace_size_per_rank(1)
-         all_workspaces = torch.zeros(1,
-                                      workspace_size,
-                                      dtype=torch.uint64,
-                                      device='cuda')
- 
-         print(f"Test configuration:")
-         print(f"  Input entries: {input_entry_count}")
-         print(f"  Vector dim: {vector_dim}")
-         print(f"  SF vector dim: {sf_vector_dim}")
-         print(f"  Send/recv count: {send_recv_count}")
-         print(f"  FP8 tensor shape: {input_tensor_fp8.shape}")
-         print(f"  SF tensor shape: {input_sf_tensor.shape}")
- 
-         try:
-             # Test with all 4 tensors
-             torch.ops.trtllm.moe_comm([
-                 input_tensor_fp8, input_sf_tensor, input_experts, input_scales
-             ], send_cumsum, send_indices, [
-                 output_tensor_fp8, output_sf_tensor, output_experts,
-                 output_scales
-             ], recv_cumsum, recv_indices, all_workspaces, 0, 1)
- 
-             torch.cuda.synchronize()
-             print("FP8 alltoall test PASSED!")
- 
-             # Verify outputs
-             print(f"\nOutput verification:")
-             print(f"  Output FP8 shape: {output_tensor_fp8.shape}")
-             print(f"  Output SF shape: {output_sf_tensor.shape}")
-             print(
-                 f"  Non-zero FP8 elements: {(output_tensor_fp8 != 0).sum().item()}"
-             )
-             print(
-                 f"  Non-zero SF elements: {(output_sf_tensor != 0).sum().item()}"
-             )
- 
-         except Exception as e:
-             print(f"FP8 alltoall test FAILED: {e}")
-             print(f"Error type: {type(e)}")
-             raise

Also applies to: 919-1035

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

70-85: Bound field count to compile-time maximum to avoid OOB

Before filling fieldsInfo, enforce fieldCount <= MOE_COMM_FIELD_MAX_COUNT to avoid writing past the fixed-size array.

 int fieldCount = inputs.size();
+tensorrt_llm::kernels::FusedMoeFieldInfo sendFieldInfo, recvFieldInfo;
+TORCH_CHECK(
+    fieldCount <= tensorrt_llm::kernels::MOE_COMM_FIELD_MAX_COUNT,
+    "Number of fields (", fieldCount, ") exceeds maximum allowed (",
+    tensorrt_llm::kernels::MOE_COMM_FIELD_MAX_COUNT, ")");
-sendFieldInfo.isBasicInterleaved = false;
+sendFieldInfo.isBasicInterleaved = false;

Note: include or ensure fusedMoeCommKernels.h exposes MOE_COMM_FIELD_MAX_COUNT (already included at Line 19).

tests/unittest/_torch/thop/test_moe_alltoall.py (2)

117-126: Warmup path uses an uninitialized workspace

do_warmup allocates a fresh all_workspaces but does not call moe_initialize_workspace before moe_comm. Initialize it to mirror the main test path.

 all_workspaces = torch.zeros(1,
                              workspace_size // 8,
                              dtype=torch.uint64,
                              device=torch.device('cuda'))
+torch.ops.trtllm.moe_initialize_workspace(all_workspaces, 0, 1)
 torch.ops.trtllm.moe_comm(input_tensors, send_cumsum, send_indices,
                           output_tensors, recv_cumsum, recv_indices,
                           all_workspaces, 0, 1)

---

251-255: Reference assembly may be fragile across ranks; verify offset usage

output_indice_offset is accumulated across nested loops over other_rank. This is correct if each rank’s recv counts are appended contiguously per source rank, but it couples the reference construction to the same concatenation order you use later. If you change concatenation order, this will silently miscompare.

If you expect to change concatenation order in future, consider computing per-rank-per-source slices and then concatenating both actual and reference with the same order variable to avoid hidden coupling.

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

190-198: Top-k divisibility guard contradicts fused path capability

enable_alltoall still enforces experts_per_token % 4 == 0, but the fused multi-field alltoall path is intended to support arbitrary top_k. This condition will prevent the new path from engaging for non-multiples-of-4.

If the fused kernels support arbitrary top_k, drop this guard:

 @cached_property
 def enable_alltoall(self):
-    return (self.mapping.moe_ep_size > self.routing_method.experts_per_token
-            and self.routing_method.experts_per_token % 4 == 0  # alltoall without allgather only supports top_k % 4 == 0
+    return (self.mapping.moe_ep_size > self.routing_method.experts_per_token
             and self.mapping.enable_attention_dp
             and self.mapping.tp_size > 1
             and os.environ.get("TRTLLM_MOE_DISABLE_ALLTOALLV", "0") != "1"
             and MnnvlMemory.supports_mnnvl())

If there are known exceptions (e.g., legacy/non-fused kernels), gate them by kernel selection rather than a global condition here.

---

245-249: No-op assignment for output_dtype

output_dtype = output_dtype in the isinstance(x, Fp4QuantizedTensor) branch is redundant.

 if isinstance(x, Fp4QuantizedTensor):
-    assert output_dtype is not None
-    output_dtype = output_dtype
+    assert output_dtype is not None
 else:
     output_dtype = x.dtype

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

141-207: moePrepareOp shape/flow is sound; minor cleanups

The new API returns the cumsums and indices plus optional gathered statics, which aligns with the kernel changes. Two small cleanups:

• Remove unused preparedLocalExpertIds allocation.
• Consider asserting slotCount % epSize == 0 to match expertCountPerRank = slotCount / epSize assumption in the kernel.

-torch::Tensor preparedLocalExpertIds
-    = torch::empty({maxTokenCountPerRank * epSize, topK}, expertsIds.options().dtype(torch::kInt32));
+// preparedLocalExpertIds no longer used in the new API; removed.

+TORCH_CHECK(slotCount % epSize == 0, "slotCount must be divisible by epSize");

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: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.420Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

Learnt from: sklevtsov-nvidia
PR: NVIDIA/TensorRT-LLM#3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

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

60-77: Fix cross-CTA visibility and type mismatch in CounterCommunicator (use s32 acquire/release).

Plain volatile stores/loads across CTAs are not sufficient; readers can spin on stale zeros. Also, value is uint64_t while values[] is typically int (s32). Use release/acquire semantics on s32 and keep types consistent.

Apply this diff:

-    __forceinline__ __device__ void releaseValue(uint64_t value, int index)
+    __forceinline__ __device__ void releaseValue(int value, int index)
     {
         // Avoid block on 0
-        fifoConnInfo->values[index] = value + 1;
+        int v = value + 1;
+        // Release store to make writes visible across CTAs/SMs
+        asm volatile("st.release.sys.global.s32 [%0], %1;" :: "l"(&fifoConnInfo->values[index]), "r"(v) : "memory");
     }
 
-    __forceinline__ __device__ uint64_t acquireValue(int index)
+    __forceinline__ __device__ int acquireValue(int index)
     {
-        uint64_t localValue = 0;
+        int localValue = 0;
         do
         {
-            localValue = fifoConnInfo->values[index];
+            // Acquire load to observe prior release
+            asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(localValue) : "l"(&fifoConnInfo->values[index]) : "memory");
         } while (localValue == 0);
 
         fifoConnInfo->values[index] = 0; // reset the value
 
         return localValue - 1;
     }

Note: Call sites already operate on ints; this change aligns types and ordering.

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

46-51: Tighten architecture guards: require CUDA_ARCH presence and use parentheses.

You changed from the prior incorrect || to && (good), but the recommended guard is #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= X) rather than defined(__CUDACC__) && __CUDA_ARCH__ >= X. The latter is redundant and risks subtle host-compilation parsing.

Apply representative diffs (repeat for all guarded blocks):

-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
   asm("mbarrier.init.shared.b64 [%0], %1;" : : "r"(__as_ptr_smem(addr)), "r"(count) : "memory");
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   asm("mbarrier.expect_tx.relaxed.cta.shared::cta.b64 [%0], %1;" : : "r"(__as_ptr_smem(addr)), "r"(txCount) : "memory");
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
   uint64_t state;
   asm("mbarrier.arrive.shared.b64 %0, [%1];" : "=l"(state) : "r"(__as_ptr_smem(addr)) : "memory");
   return state;
 #else
   return 0;
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   uint64_t state;
   asm("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 %0, [%1], %2;" : "=l"(state) : "r"(__as_ptr_smem(addr)), "r"(txCount) : "memory");
   return state;
 #else
   return 0;
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   uint32_t waitComplete;
   asm("{\n\t .reg .pred P_OUT; \n\t"
       "mbarrier.try_wait.parity.shared::cta.b64  P_OUT, [%1], %2;\n\t"
       "selp.b32 %0, 1, 0, P_OUT; \n"
       "}" : "=r"(waitComplete) : "r"(__as_ptr_smem(addr)), "r"(phaseParity) : "memory");
   return static_cast<bool>(waitComplete);
 #else
   return false;
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
   asm volatile(
       "{\n"
       "  .reg .pred p;\n"
       "  setp.ne.b32 p, %0, 0;\n"
       "  @p cp.async.ca.shared.global [%1], [%2], %3;\n"
       "}\n" ::"r"((int) predGuard), "r"(__as_ptr_smem(dstShm)), "l"(__as_ptr_gmem(srcMem)), "n"(COPY_SIZE));
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
   asm volatile("cp.async.commit_group;" : : :);
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
   asm volatile("cp.async.wait_group %0;" : : "n"(N) : "memory");
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   asm("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];"
       : : "r"(__as_ptr_smem(dstMem)), "l"(__as_ptr_gmem(srcMem)), "r"(copySize), "r"( __as_ptr_smem(smemBar)) : "memory");
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   asm("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" : : "l"(__as_ptr_gmem(dstMem)), "r"(__as_ptr_smem(srcMem)), "r"(copySize) : "memory");
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   asm volatile("cp.async.bulk.commit_group;" : : :);
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   asm volatile("cp.async.bulk.wait_group %0;" : : "n"(N) : "memory");
 #endif
@@
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   asm volatile("cp.async.bulk.wait_group.read %0;" : : "n"(N) : "memory");
 #endif

Also applies to: 53-61, 63-73, 74-86, 88-103, 105-117, 119-124, 126-132, 134-142, 144-152, 154-159, 161-167, 169-175

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

93-94: Confirm slotCount divisibility by epSize.

expertCountPerRank = slotCount / epSize assumes even partitioning; any remainder will corrupt rank mapping when doing experts[...] / expertCountPerRank. Validate this at the host before launching the kernel.

Add a host-side precondition in computeCountAndIndice:

 void computeCountAndIndice( /* args... */ )
 {
+    TLLM_CHECK_WITH_INFO(slotCount % rankCount == 0, "slotCount (%d) must be divisible by epSize (%d).",
+                         slotCount, rankCount);
     // existing code ...
 }

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

1298-1303: Static caching of max dynamic SMEM may be device-dependent.

static int maxDynamicShmSize = computeMoeAlltoallMaxDynamicSharedMemorySize(); will cache the first device’s value. If the active device can change, recompute per call or per device.

Apply this diff:

-    static int maxDynamicShmSize = fused_moe_impl::computeMoeAlltoallMaxDynamicSharedMemorySize();
-    int groupCountPerCta = std::min(maxGroupCountPerCta, maxDynamicShmSize / warpShmSize);
+    int maxDynamicShmSize = fused_moe_impl::computeMoeAlltoallMaxDynamicSharedMemorySize();
+    int groupCountPerCta = std::min(maxGroupCountPerCta, maxDynamicShmSize / warpShmSize);

---

1322-1326: Optional: avoid magic number 48*1024.

Use a named constexpr (e.g., kDEFAULT_SMEM_PER_BLOCK = 48 * 1024) or query device attribute to decide when to opt-in.

-    if (groupCountPerCta * warpShmSize > 48 * 1024)
+    constexpr int kDEFAULT_SMEM_PER_BLOCK = 48 * 1024;
+    if (groupCountPerCta * warpShmSize > kDEFAULT_SMEM_PER_BLOCK)
     {
         TLLM_CUDA_CHECK(cudaFuncSetAttribute(
             kernelFn, cudaFuncAttributeMaxDynamicSharedMemorySize, groupCountPerCta * warpShmSize));
     }

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: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.420Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

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/thop/moeCommOp.cpp (1)

174-183: Validate expertsStatics when provided

If expertsStatics is present, ensure dtype and shape are as expected before passing raw pointers into kernels.

-    if (expertsStatics.has_value())
+    if (expertsStatics.has_value())
     {
+        CHECK_INPUT(expertsStatics.value(), torch::kInt32);
+        TORCH_CHECK(
+            expertsStatics.value().dim() == 1 && expertsStatics.value().size(0) == expertCount,
+            "experts_statics must be a 1D int32 tensor of length expert_count (", expertCount, ")");
         localExpertStaticsPtr = expertsStatics.value().data_ptr<int>();
         gatheredExpertStatics = torch::empty({epSize, expertCount}, expertsIds.options().dtype(torch::kInt32));
         gatheredExpertStaticsPtr = gatheredExpertStatics.value().data_ptr<int>();
     }

tests/unittest/_torch/thop/test_moe_alltoall.py (1)

319-435: Remove duplicate test class definition.

The TestMoeAlltoAllFP8SingleGPU class is defined twice (lines 319 and 919), which will cause a redefinition error. The second definition should be removed.

Remove the duplicate class definition starting at line 919:

-class TestMoeAlltoAllFP8SingleGPU(unittest.TestCase):
-
-    def setUp(self):
-        torch.manual_seed(0x1234)
-        tllm.logger.set_level('error')
-
-    def test_moe_alltoall_fp8_with_indices(self):
-        """Test fp8 alltoall with properly constructed indices"""
-        torch.cuda.set_device(0)
-        
-        # Match dimensions from the error
-        input_entry_count = 16384
-        output_entry_count = 16384
-        vector_dim = 2944
-        sf_vector_dim = 92  # Scaling factor dimension from error
-        send_recv_count = 1000  # Number of entries to send/receive
-        
-        # ... rest of the duplicate implementation ...

Also applies to: 919-1035

tensorrt_llm/_mnnvl_utils.py (1)

529-599: Fix return type inconsistency.

When x is a single tensor input, the function converts it to a list for unified handling. However, at line 597, when returning for a single tensor case, it returns [result[0]] which is still a list containing one element. According to the function signature -> Union[torch.Tensor, List[torch.Tensor]], it should return a single tensor when the input was a single tensor.

Apply this diff to fix the return type:

-        # If input was a single tensor, return a single tensor
+        # If input was a single tensor, return a single tensor
         if is_single_tensor:
-            result = [result[0]]
+            result = result[0]
 
         return result

Additionally, update the type hints to reflect support for None entries:

-def mnnvl_moe_alltoallv(
-        x: Union[torch.Tensor, List[torch.Tensor]],
+def mnnvl_moe_alltoallv(
+        x: Union[torch.Tensor, List[Optional[torch.Tensor]]],
         alltoall_info: MoEAlltoallInfo,
         workspace: torch.Tensor,
         ep_rank: int,
         ep_size: int,
-    ) -> Union[torch.Tensor, List[torch.Tensor]]:
+    ) -> Union[torch.Tensor, List[Optional[torch.Tensor]]]:

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

60-77: Memory ordering concerns with volatile operations.

The CounterCommunicator class uses plain volatile stores/loads for inter-CTA communication. While the diff shows type consistency with uint64_t values and int storage, the lack of proper memory ordering could lead to visibility issues where CTAs spin on stale values.

Consider using proper release/acquire semantics:

 __forceinline__ __device__ void releaseValue(uint64_t value, int index)
 {
     // Avoid block on 0
-    fifoConnInfo->values[index] = value + 1;
+    int v = static_cast<int>(value + 1);
+    // Release store to make writes visible across CTAs
+    asm volatile("st.release.sys.global.s32 [%0], %1;" :: "l"(&fifoConnInfo->values[index]), "r"(v) : "memory");
 }

 __forceinline__ __device__ uint64_t acquireValue(int index)
 {
-    uint64_t localValue = 0;
+    int localValue = 0;
     do
     {
-        localValue = fifoConnInfo->values[index];
+        // Acquire load to observe prior release
+        asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(localValue) : "l"(&fifoConnInfo->values[index]));
     } while (localValue == 0);

     fifoConnInfo->values[index] = 0; // reset the value
-    return localValue - 1;
+    return static_cast<uint64_t>(localValue - 1);
 }

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

46-175: Fix architecture guards to prevent unsupported PTX on lower compute capabilities.

The current preprocessor guards use #if defined(__CUDACC__) && __CUDA_ARCH__ >= XXX, but this is incorrect. The defined(__CUDACC__) check is for host-side compilation detection, not device-side. This could allow unsupported PTX instructions to be compiled for lower compute capabilities.

Apply these fixes to ensure proper architecture guards:

 __device__ __forceinline__ void mbarrier_init(uint64_t* addr, uint32_t const& count)
 {
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
     asm("mbarrier.init.shared.b64 [%0], %1;" : : "r"(__as_ptr_smem(addr)), "r"(count) : "memory");
 #endif
 }

 __device__ __forceinline__ void mbarrier_expect_tx(uint64_t* addr, const uint32_t txCount)
 {
-#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
     asm("mbarrier.expect_tx.relaxed.cta.shared::cta.b64 [%0], %1;"
         :
         : "r"(__as_ptr_smem(addr)), "r"(txCount)
         : "memory");
 #endif
 }

Apply similar fixes to all other architecture-gated functions (mbarrier_arrive, mbarrier_arrive_expect_tx, mbarrier_try_wait_parity, ldgsts, cp_async_commit_group, cp_async_wait_group, cp_async_bulk_g2s, cp_async_bulk_s2g, cp_async_bulk_commit_group, cp_async_bulk_wait_group, cp_async_bulk_wait_group_read).

cpp/tensorrt_llm/thop/moeCommOp.cpp (3)

209-227: Fix misleading error text and strengthen guards in memsetExpertIds

The dim check says “1D” while enforcing 2D. Also add basic param guards.

 void memsetExpertIds(torch::Tensor expertsIds, torch::Tensor recvRankCountCumSum, int64_t maxTokenCountPerRank,
     int64_t topK, int64_t slotCount, int64_t epSize)
 {
     CHECK_INPUT(expertsIds, torch::kInt32);
-    TORCH_CHECK(expertsIds.dim() == 2, "expertsIds must be a 1D tensor");
+    TORCH_CHECK(expertsIds.dim() == 2, "expertsIds must be a 2D tensor [max_token_count_per_rank*ep_size, top_k]");
     TORCH_CHECK(
         expertsIds.size(0) == maxTokenCountPerRank * epSize, "expertsIds must have maxTokenCountPerRank * epSize rows");
     TORCH_CHECK(expertsIds.size(1) == topK, "expertsIds must have topK columns");
 
     CHECK_INPUT(recvRankCountCumSum, torch::kInt32);
     TORCH_CHECK(recvRankCountCumSum.dim() == 1, "recvRankCountCumSum must be a 1D tensor");
     TORCH_CHECK(recvRankCountCumSum.size(0) == epSize, "recvRankCountCumSum must have epSize elements");
+    TORCH_CHECK(maxTokenCountPerRank > 0 && topK > 0 && slotCount > 0 && epSize > 0,
+        "max_token_count_per_rank, top_k, slot_count, ep_size must be positive");
 
     auto stream = at::cuda::getCurrentCUDAStream();

---

70-76: Enforce fieldCount ≤ MOE_COMM_FIELD_MAX_COUNT to avoid out-of-bounds on fieldsInfo[]

fieldsInfo has a compile-time max (8). Guard fieldCount to prevent OOB writes.

-    int fieldCount = inputs.size();
+    int fieldCount = inputs.size();
+    TORCH_CHECK(
+        fieldCount <= tensorrt_llm::kernels::MOE_COMM_FIELD_MAX_COUNT,
+        "Number of fields (", fieldCount, ") exceeds maximum allowed (",
+        tensorrt_llm::kernels::MOE_COMM_FIELD_MAX_COUNT, ")");

---

95-99: topK is not initialized; pass 0 explicitly (no basic fields) to fillMetaInfo

params.expertParallelInfo.topK is never set. Passing its default to fillMetaInfo is fragile; for fused comm here we don’t use basic fields. Pass 0 explicitly and set the struct for clarity.

     params.recvFieldInfo = recvFieldInfo;
     // Do not need expertParallelInfo for fused moe comm now
+    params.expertParallelInfo.topK = 0;
 
-    params.sendFieldInfo.fillMetaInfo(&(params.sendCommMeta), params.expertParallelInfo.topK, false, false);
-    params.recvFieldInfo.fillMetaInfo(&(params.recvCommMeta), params.expertParallelInfo.topK, false, false);
+    params.sendFieldInfo.fillMetaInfo(&(params.sendCommMeta), /*topK=*/0, false, false);
+    params.recvFieldInfo.fillMetaInfo(&(params.recvCommMeta), /*topK=*/0, false, false);

cpp/tests/unit_tests/kernels/fusedMoeCommKernelTest.cpp (2)

1-15: Switch to SPDX header format (consistent with project headers)

Replace the legacy Apache header with SPDX lines as used elsewhere.

-/*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */

---

148-151: random_shuffle removed in C++17; use std::shuffle with a real RNG

This will not compile under C++17. Replace with std::shuffle and drop the extra std::srand.

         // Shuffle unused indices for random assignment
-        std::srand(42); // Fixed seed for reproducible tests
-        std::random_shuffle(unusedIndices.begin(), unusedIndices.end());
+        std::mt19937 gen(42); // Fixed seed for reproducible tests
+        std::shuffle(unusedIndices.begin(), unusedIndices.end(), gen);

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

465-475: Clarify statistics gathering on the last call

The implementation accumulates local expert statistics on every chunk (and repeat) via update_local_statistic (when updates_enabled), storing counts in local_statistic_tensor, and only calls get_local_statistic_tensor when is_last_call to retrieve the fully aggregated tensor for the subsequent AlltoAll. This avoids redundant communications while ensuring a complete view across all repeats. Static routing (is_static_routing == True) correctly skips both updates and gathering.

• File tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py, around line 465:
Add a comment above the if is_last_call and … block explaining that

1. local_statistic_tensor has accumulated stats from all prior calls, and
2. gathering is deferred to the final call by design to minimize communication overhead.

No changes to the logic itself are required.

cpp/tensorrt_llm/thop/moeCommOp.cpp (3)

81-88: Optional: Validate each input/output tensor is CUDA and contiguous

setMoeCommFieldInfo asserts 2D shape, but neither CUDA placement nor contiguity are asserted. A bad placement/stride can miscompute element sizes and lead to memory faults at kernel time.

     for (int i = 0; i < fieldCount; i++)
     {
-        setMoeCommFieldInfo(sendFieldInfo.fieldsInfo[i], inputs[i]);
-        setMoeCommFieldInfo(recvFieldInfo.fieldsInfo[i], outputs[i]);
+        TORCH_CHECK(inputs[i].is_cuda() && outputs[i].is_cuda(), "inputs/outputs must be CUDA tensors");
+        TORCH_CHECK(inputs[i].is_contiguous() && outputs[i].is_contiguous(), "inputs/outputs must be contiguous");
+        setMoeCommFieldInfo(sendFieldInfo.fieldsInfo[i], inputs[i]);
+        setMoeCommFieldInfo(recvFieldInfo.fieldsInfo[i], outputs[i]);
     }

---

153-155: Remove unused preparedLocalExpertIds allocation

This tensor is allocated but never used. Drop it to save memory and avoid confusion.

-    torch::Tensor preparedLocalExpertIds
-        = torch::empty({maxTokenCountPerRank * epSize, topK}, expertsIds.options().dtype(torch::kInt32));
+    // preparedLocalExpertIds removed: not used in fused prepare workflow

---

157-158: Nit: Use lowerCamelCase for local variable names

RecvRankCountCumSum uses uppercased initial letter, while other locals are lowerCamelCase. Consider recvRankCountCumSum for consistency.

cpp/tests/unit_tests/kernels/fusedMoeCommKernelTest.cpp (6)

49-58: Avoid C RNG in tests; don’t seed global state

Use a local RNG to avoid global state and improve reproducibility. Remove std::srand here.

     void SetUp() override
     {
         if (shouldSkip())
         {
             skipped = true;
             GTEST_SKIP() << "Skipping due to no/unsupported GPU";
         }
         TLLM_CUDA_CHECK(cudaStreamCreate(&stream));
-        std::srand(42); // Initialize random seed
     }

---

68-70: Add a per-fixture RNG

Use a member engine to power all random data generation paths.

     bool skipped = false;
     cudaStream_t stream = nullptr;
+    std::mt19937 rng{42}; // Fixed seed for reproducible tests

---

88-104: Replace rand() with distributions

rand() is non-deterministic across platforms and discouraged. Use the fixture RNG with type-appropriate distributions.

         else
         {
             // Default initialization with random values
-            for (size_t i = 0; i < count; i++)
-            {
-                if constexpr (std::is_same_v<T, float>)
-                {
-                    (*hostPtr)[i] = static_cast<float>(rand()) / RAND_MAX * 10.0f;
-                }
-                else if constexpr (std::is_same_v<T, int>)
-                {
-                    (*hostPtr)[i] = rand() % 1000;
-                }
-                else
-                {
-                    (*hostPtr)[i] = static_cast<T>(rand() % 100);
-                }
-            }
+            std::uniform_real_distribution<float> floatDist(0.0f, 10.0f);
+            std::uniform_int_distribution<int> intDist(0, 999);
+            std::uniform_int_distribution<int> byteDist(0, 99);
+            for (size_t i = 0; i < count; i++)
+            {
+                if constexpr (std::is_same_v<T, float>)
+                {
+                    (*hostPtr)[i] = floatDist(rng);
+                }
+                else if constexpr (std::is_same_v<T, int>)
+                {
+                    (*hostPtr)[i] = intDist(rng);
+                }
+                else
+                {
+                    (*hostPtr)[i] = static_cast<T>(byteDist(rng));
+                }
+            }
         }

---

153-163: Handle exhausted unusedIndices and fix signed/unsigned mismatch

Prevent leaving -1 in fullMapping if unusedIndices runs out; also use size_t for the index to avoid signed/unsigned warnings.

-        int unusedIdx = 0;
+        size_t unusedIdx = 0;
         for (int i = 0; i < totalSize; i++)
         {
             if (i < providedSize && fullMapping[i] == -1)
             {
                 // Fix invalid mapping
-                if (unusedIdx < unusedIndices.size())
+                if (unusedIdx < unusedIndices.size())
                 {
                     fullMapping[i] = unusedIndices[unusedIdx++];
                 }
+                else
+                {
+                    // No more unused indices available, use identity as fallback
+                    fullMapping[i] = i;
+                }
             }
             else if (i >= providedSize)
             {
                 // Extend mapping
                 if (unusedIdx < unusedIndices.size())
                 {
                     fullMapping[i] = unusedIndices[unusedIdx++];
                 }
                 else
                 {
                     // Fallback: identity mapping for remaining
                     fullMapping[i] = i;
                 }
             }

---

473-476: Avoid strict-aliasing UB when packing float into int

Use memcpy instead of pointer punning.

             else if (hasScales && offsetInWarp < topK * 2)
             {
                 // Scales area
-                float scale
-                    = 1.0f + static_cast<float>(warpIdx) * 0.1f + static_cast<float>(offsetInWarp - topK) * 0.01f;
-                return *reinterpret_cast<int*>(&scale);
+                float scale = 1.0f + static_cast<float>(warpIdx) * 0.1f
+                    + static_cast<float>(offsetInWarp - topK) * 0.01f;
+                int packed;
+                std::memcpy(&packed, &scale, sizeof(packed));
+                return packed;
             }

---

757-767: Optional: Prefer RAII for host buffers to simplify cleanup

Using std::vector or smart pointers prevents leaks on early returns and is clearer.

-        if (resultRecvTokenSlots)
-            delete[] resultRecvTokenSlots;
-        if (resultRecvScales)
-            delete[] resultRecvScales;
-        for (int i = 0; i < fieldCount; i++)
-        {
-            if (resultRecvFields[i])
-                delete[] resultRecvFields[i];
-        }
+        std::unique_ptr<int[]> resultRecvTokenSlotsPtr(resultRecvTokenSlots);
+        std::unique_ptr<float[]> resultRecvScalesPtr(resultRecvScales);
+        for (int i = 0; i < fieldCount; i++)
+        {
+            std::unique_ptr<uint8_t[]> fieldGuard(resultRecvFields[i]);
+        }

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: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.420Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

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.

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-12T10:28:57.320Z
Learning: Applies to **/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py} : Prepend NVIDIA copyright header (current year) to all source files

Learnt from: sklevtsov-nvidia
PR: NVIDIA/TensorRT-LLM#3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

cpp/tensorrt_llm/thop/moeCommOp.cpp (6)

38-46: Validate dtype for allWorkspaces before data_ptr<uint64_t>().

We dereference allWorkspaces as 64-bit elements later. Add a dtype check to prevent UB if callers pass a different dtype.

Apply this diff:

     CHECK_INPUT(sendRankCumSum, torch::kInt32);
     CHECK_INPUT(sendIndiceTensor, torch::kInt32);
     CHECK_INPUT(recvRankCumSum, torch::kInt32);
     CHECK_INPUT(recvIndiceTensor, torch::kInt32);
+    CHECK_INPUT(allWorkspaces, torch::kInt64);

---

126-139: Initialize workspace: validate allWorkspaces dtype.

Same precondition as moeCommOp; ensure allWorkspaces holds 64-bit words before taking data_ptr<uint64_t>().

Apply this diff:

 void initializeMoeWorkspace(torch::Tensor allWorkspaces, int64_t epRank, int64_t epSize)
 {
     TORCH_CHECK(allWorkspaces.dim() == 2, "allWorkspaces must be a 2D tensor");
     TORCH_CHECK(epRank >= 0 && epRank < epSize, "epRank must be in the range [0, epSize)");
+    CHECK_INPUT(allWorkspaces, torch::kInt64);

---

141-151: Add shape/dtype guards for expertsIds and allWorkspaces.

computeCountAndIndice expects [tokenCount, topK] expertsIds and 64-bit workspaces.

Apply this diff:

     CHECK_INPUT(expertsIds, torch::kInt32);
+    TORCH_CHECK(expertsIds.dim() == 2, "experts_ids must be a 2D tensor [token_count, top_k]");
+    TORCH_CHECK(
+        expertsIds.size(1) == topK,
+        "experts_ids second dimension must equal top_k");
+    CHECK_INPUT(allWorkspaces, torch::kInt64);

---

209-221: Fix misleading tensor-dim error message.

The check enforces 2D but the message says 1D. Correct the message to avoid confusion.

Apply this diff:

-    TORCH_CHECK(expertsIds.dim() == 2, "expertsIds must be a 1D tensor");
+    TORCH_CHECK(expertsIds.dim() == 2, "expertsIds must be a 2D tensor [maxTokenCountPerRank*epSize, topK]");

---

70-76: Guard against overflowing the fieldsInfo array.

fieldCount must not exceed the compile-time capacity of FusedMoeFieldInfo::fieldsInfo. Add a check against MOE_COMM_FIELD_MAX_COUNT.

Apply this diff:

-    int fieldCount = inputs.size();
+    int fieldCount = static_cast<int>(inputs.size());
+    TORCH_CHECK(
+        fieldCount <= tensorrt_llm::kernels::MOE_COMM_FIELD_MAX_COUNT,
+        "Number of fields exceeds MOE_COMM_FIELD_MAX_COUNT");

---

97-99: params.expertParallelInfo.topK is uninitialized before use.

You’re passing an indeterminate value into fillMetaInfo. Either set it or pass a known constant (0 if unused in fused comm meta).

Apply this diff:

-    params.sendFieldInfo.fillMetaInfo(&(params.sendCommMeta), params.expertParallelInfo.topK, false, false);
-    params.recvFieldInfo.fillMetaInfo(&(params.recvCommMeta), params.expertParallelInfo.topK, false, false);
+    params.sendFieldInfo.fillMetaInfo(&(params.sendCommMeta), /*topK=*/0, false, false);
+    params.recvFieldInfo.fillMetaInfo(&(params.recvCommMeta), /*topK=*/0, false, false);

Optionally also initialize the struct field explicitly before the calls:

+    params.expertParallelInfo.topK = 0;

cpp/tensorrt_llm/thop/moeCommOp.cpp (5)

2-2: Update copyright year to include 2025.

Repo guideline says to prepend the current year; header still shows 1993-2024.

Apply this diff:

- * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

---

86-87: Clarify magic arguments with inline comments.

Use inline C comments to document the meaning of positional booleans/integers per style guide.

Apply this diff:

-    sendFieldInfo.fillFieldPlacementInfo(0, false);
-    recvFieldInfo.fillFieldPlacementInfo(0, false);
+    sendFieldInfo.fillFieldPlacementInfo(/*groupOffset=*/0, /*isBroadcast=*/false);
+    recvFieldInfo.fillFieldPlacementInfo(/*groupOffset=*/0, /*isBroadcast=*/false);

---

157-205: Naming consistency: rename RecvRankCountCumSum to lowerCamelCase.

Local variables should be lowerCamelCase. Rename and update the uses in kernel calls and return tuple.

Apply this diff:

-    torch::Tensor RecvRankCountCumSum = torch::empty({epSize}, expertsIds.options().dtype(torch::kInt32));
+    torch::Tensor recvRankCountCumSum = torch::empty({epSize}, expertsIds.options().dtype(torch::kInt32));
@@
-    tensorrt_llm::kernels::moe_prepare::computeCountAndIndice(expertsIds.data_ptr<int>(),
-        sendRankCountCumSum.data_ptr<int>(), RecvRankCountCumSum.data_ptr<int>(), sendRankIndices.data_ptr<int>(),
+    tensorrt_llm::kernels::moe_prepare::computeCountAndIndice(expertsIds.data_ptr<int>(),
+        sendRankCountCumSum.data_ptr<int>(), recvRankCountCumSum.data_ptr<int>(), sendRankIndices.data_ptr<int>(),
@@
-    tensorrt_llm::kernels::moe_prepare::computeCumsum(
-        sendRankCountCumSum.data_ptr<int>(), RecvRankCountCumSum.data_ptr<int>(), epRank, epSize, stream);
+    tensorrt_llm::kernels::moe_prepare::computeCumsum(
+        sendRankCountCumSum.data_ptr<int>(), recvRankCountCumSum.data_ptr<int>(), epRank, epSize, stream);
@@
-    tensorrt_llm::kernels::moe_prepare::moveIndice(sendRankCountCumSum.data_ptr<int>(),
-        RecvRankCountCumSum.data_ptr<int>(), sendRankIndices.data_ptr<int>(), gatherSendRankIndices.data_ptr<int>(),
+    tensorrt_llm::kernels::moe_prepare::moveIndice(sendRankCountCumSum.data_ptr<int>(),
+        recvRankCountCumSum.data_ptr<int>(), sendRankIndices.data_ptr<int>(), gatherSendRankIndices.data_ptr<int>(),
@@
-    return std::make_tuple(sendRankCountCumSum, gatherSendRankIndices, RecvRankCountCumSum, gatherRecvRankIndices,
+    return std::make_tuple(sendRankCountCumSum, gatherSendRankIndices, recvRankCountCumSum, gatherRecvRankIndices,
         gatherBackwardRecvRankIndices, gatheredExpertStatics);

---

276-279: Unify parameter naming style in schema (allWorkspace → all_workspaces).

Elsewhere you use snake_case; keep it consistent across op schemas.

Apply this diff:

-        "mnnvl_moe_alltoallv_prepare_without_allgather(Tensor experts_ids, Tensor? experts_statics, "
-        "Tensor allWorkspace, int max_token_count_per_rank, int ep_rank, int ep_size, int expert_count, int "
+        "mnnvl_moe_alltoallv_prepare_without_allgather(Tensor experts_ids, Tensor? experts_statics, "
+        "Tensor all_workspaces, int max_token_count_per_rank, int ep_rank, int ep_size, int expert_count, int "

---

146-149: Avoid magic numbers for expert/slot divisibility and max bound.

Replace 4 and 512 with named constexpr constants to document intent and ease maintenance.

You can add near the top of this file (namespace scope):

namespace {
constexpr int kEXPERT_SLOT_ALIGNMENT = 4;
constexpr int kMAX_EXPERT_PLUS_ONE = 512;
} // anonymous namespace

Then update:

-    TORCH_CHECK(expertCount % 4 == 0, "expertCount must be divisible by 4");
-    TORCH_CHECK(slotCount % 4 == 0, "slotCount must be divisible by 4");
-    TORCH_CHECK(expertCount + 1 <= 512, "expertCount + 1 is larger than 512");
+    TORCH_CHECK(expertCount % kEXPERT_SLOT_ALIGNMENT == 0, "expertCount must be divisible by 4");
+    TORCH_CHECK(slotCount % kEXPERT_SLOT_ALIGNMENT == 0, "slotCount must be divisible by 4");
+    TORCH_CHECK(expertCount + 1 <= kMAX_EXPERT_PLUS_ONE, "expertCount + 1 is larger than 512");

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: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.420Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

Learnt from: sklevtsov-nvidia
PR: NVIDIA/TensorRT-LLM#3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

cpp/tests/unit_tests/kernels/fusedMoeCommKernelTest.cpp (9)

49-58: Stop using global C RNG; use a local C++ RNG.

Remove std::srand(42) and add a deterministic engine as a member. This also supports switching rand() to distributions below.

     void SetUp() override
     {
         if (shouldSkip())
         {
             skipped = true;
             GTEST_SKIP() << "Skipping due to no/unsupported GPU";
         }
         TLLM_CUDA_CHECK(cudaStreamCreate(&stream));
-        std::srand(42); // Initialize random seed
     }
@@
-    bool skipped = false;
-    cudaStream_t stream = nullptr;
+    bool skipped = false;
+    cudaStream_t stream = nullptr;
+    std::mt19937 rng{42}; // deterministic engine for tests

Also applies to: 68-70

---

71-107: Replace rand() with distributions (deterministic, guideline-compliant).

Prefer C++ RNGs over C APIs and avoid implicit RAND_MAX scaling.

         else
         {
             // Default initialization with random values
-            for (size_t i = 0; i < count; i++)
-            {
-                if constexpr (std::is_same_v<T, float>)
-                {
-                    (*hostPtr)[i] = static_cast<float>(rand()) / RAND_MAX * 10.0f;
-                }
-                else if constexpr (std::is_same_v<T, int>)
-                {
-                    (*hostPtr)[i] = rand() % 1000;
-                }
-                else
-                {
-                    (*hostPtr)[i] = static_cast<T>(rand() % 100);
-                }
-            }
+            std::uniform_real_distribution<float> floatDist(0.0f, 10.0f);
+            std::uniform_int_distribution<int> intDist(0, 999);
+            std::uniform_int_distribution<int> byteDist(0, 99);
+            for (size_t i = 0; i < count; i++)
+            {
+                if constexpr (std::is_same_v<T, float>)
+                {
+                    (*hostPtr)[i] = floatDist(rng);
+                }
+                else if constexpr (std::is_same_v<T, int>)
+                {
+                    (*hostPtr)[i] = intDist(rng);
+                }
+                else
+                {
+                    (*hostPtr)[i] = static_cast<T>(byteDist(rng));
+                }
+            }
         }

---

758-767: Prefer RAII for temporary buffers to simplify cleanup.

Use std::vector or std::unique_ptr for resultRecv* arrays to avoid manual delete[] and reduce error surface.

Example with std::vector:

-        int* resultRecvTokenSlots = nullptr;
+        std::vector<int> resultRecvTokenSlots;
...
-            resultRecvTokenSlots = new int[tokenCount * topK];
-            TLLM_CUDA_CHECK(cudaMemcpy(
-                resultRecvTokenSlots, deviceRecvTokenSlots, tokenCount * topK * sizeof(int), cudaMemcpyDeviceToHost));
+            resultRecvTokenSlots.resize(tokenCount * topK);
+            TLLM_CUDA_CHECK(cudaMemcpy(
+                resultRecvTokenSlots.data(), deviceRecvTokenSlots, tokenCount * topK * sizeof(int), cudaMemcpyDeviceToHost));

Also applies to: 1285-1294

---

1-15: Switch to SPDX header (project standard).

Replace the legacy Apache header block with the SPDX-style header used elsewhere in this repo.

-/*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+// SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0

---

17-25: Missing headers for used APIs (compile issues + UB fix prerequisites).

• You call std::random_shuffle and std::shuffle replacements need .
• You use std::memcpy in recommended fixes; add .
• You use std::is_same_v but <type_traits> is not included (this is a hard compile error).

Add these includes.

 #include <atomic>
 #include <chrono>
 #include <functional>
 #include <gtest/gtest.h>
 #include <memory>
 #include <random>
+#include <algorithm>
+#include <cstring>
+#include <type_traits>
 #include <thread>
 #include <vector>

---

109-113: UB: deleting T[] via delete[] on a char; make cleanup type-aware and fix call sites.*

Current delete[] static_cast<char*>(hostPtr) is undefined behavior for non-char arrays. Provide a templated deleter and use the correct element type everywhere.

-    void cleanup(void* hostPtr, void* devicePtr)
-    {
-        delete[] static_cast<char*>(hostPtr);
-        TLLM_CUDA_CHECK(cudaFree(devicePtr));
-    }
+    template <typename T>
+    void cleanupArray(T* hostPtr, void* devicePtr)
+    {
+        delete[] hostPtr;
+        TLLM_CUDA_CHECK(cudaFree(devicePtr));
+    }

Apply across call sites (examples):

- cleanup(hostTokenSlots, deviceTokenSlots);
+ cleanupArray<int>(hostTokenSlots, deviceTokenSlots);

- cleanup(hostScales, deviceScales);
+ cleanupArray<float>(hostScales, deviceScales);

- cleanup(hostFieldPtrs[i], deviceFieldPtrs[i]);
+ cleanupArray<uint8_t>(static_cast<uint8_t*>(hostFieldPtrs[i]), deviceFieldPtrs[i]);

- cleanup(hostShmDump, deviceShmDump);
+ cleanupArray<int>(hostShmDump, deviceShmDump);

Run this to list remaining old calls that need updating:

#!/bin/bash
rg -nP 'cleanup\(' cpp/tests/unit_tests/kernels/fusedMoeCommKernelTest.cpp

---

148-151: std::random_shuffle was removed in C++17; use std::shuffle with a seeded engine.

This is a compile error with -std=c++17. Replace with std::shuffle.

-        // Shuffle unused indices for random assignment
-        std::srand(42); // Fixed seed for reproducible tests
-        std::random_shuffle(unusedIndices.begin(), unusedIndices.end());
+        // Shuffle unused indices for random assignment
+        std::mt19937 gen(42); // Fixed seed for reproducible tests
+        std::shuffle(unusedIndices.begin(), unusedIndices.end(), gen);

---

156-176: Fill invalid mappings when unusedIndices are exhausted to avoid -1 entries.

Prevent leaving -1 in fullMapping, which propagates invalid indices downstream.

             if (i < providedSize && fullMapping[i] == -1)
             {
                 // Fix invalid mapping
-                if (unusedIdx < unusedIndices.size())
+                if (unusedIdx < static_cast<int>(unusedIndices.size()))
                 {
                     fullMapping[i] = unusedIndices[unusedIdx++];
                 }
+                else
+                {
+                    // Fallback: identity mapping
+                    fullMapping[i] = i;
+                }
             }

---

472-476: Fix strict-aliasing UB when bit-packing float into int.

Replace reinterpret_cast-based type punning with std::memcpy.

-                float scale
-                    = 1.0f + static_cast<float>(warpIdx) * 0.1f + static_cast<float>(offsetInWarp - topK) * 0.01f;
-                return *reinterpret_cast<int*>(&scale);
+                float scale = 1.0f + static_cast<float>(warpIdx) * 0.1f
+                             + static_cast<float>(offsetInWarp - topK) * 0.01f;
+                int bits;
+                static_assert(sizeof(bits) == sizeof(scale), "size mismatch");
+                std::memcpy(&bits, &scale, sizeof(bits));
+                return bits;

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

16-16: Use project include guard instead of #pragma once.

Per coding guidelines, replace #pragma once with TRTLLM_FUSEDMOECOMMKERNELS_H include guard.

-#pragma once
+#ifndef TRTLLM_FUSEDMOECOMMKERNELS_H
+#define TRTLLM_FUSEDMOECOMMKERNELS_H
@@
-} // namespace kernels
-} // namespace tensorrt_llm
-
+#endif // TRTLLM_FUSEDMOECOMMKERNELS_H

Also applies to: 561-563

tests/unittest/_torch/thop/test_moe_alltoall.py (4)

75-81: Avoid silent truncation when sizing workspaces; compute u64 words with ceiling.

get_moe_commworkspace_size_per_rank() returns bytes. Using // 8 truncates if the size isn’t 8-byte aligned (future changes could break this). Compute words with a ceiling division in all three places and use a clearer variable name.

Apply this diff in the three blocks:

-        workspace_size = torch.ops.trtllm.get_moe_commworkspace_size_per_rank(1)
-        all_workspaces = torch.zeros(1,
-                                     workspace_size // 8,
+        workspace_size_bytes = torch.ops.trtllm.get_moe_commworkspace_size_per_rank(1)
+        workspace_words = (workspace_size_bytes + 7) // 8
+        all_workspaces = torch.zeros(1,
+                                     workspace_words,
                                      dtype=torch.uint64,
                                      device=torch.device('cuda'))

And similarly in do_warmup() and the multi-rank section:

-        workspace_size = torch.ops.trtllm.get_moe_commworkspace_size_per_rank(
-            world_size)
-        all_workspaces = torch.zeros(world_size,
-                                     workspace_size // 8,
+        workspace_size_bytes = torch.ops.trtllm.get_moe_commworkspace_size_per_rank(world_size)
+        workspace_words = (workspace_size_bytes + 7) // 8
+        all_workspaces = torch.zeros(world_size,
+                                     workspace_words,
                                      dtype=torch.uint64,
                                      device=torch.device('cuda'))

Also applies to: 113-119, 254-260

---

82-89: Clarify/standardize the optional boolean list arg to moe_comm.

Here you pass [True], while the multi-rank calls below omit the flag. If the last argument toggles “combine” or “gathered” mode, keep usage consistent in tests (or add a short comment on its semantics) to reduce reader confusion.

---

95-113: Warmup defeats stream isolation due to global synchronizes.

do_warmup() calls torch.cuda.synchronize() at entry/exit, so wrapping it in per-rank streams won’t avoid first-launch syncs. Either drop the leading sync or allow an optional stream arg and run the ops on that stream to preserve isolation.

Minimal tweak:

-    def do_warmup(self):
-        torch.cuda.synchronize()
+    def do_warmup(self):
         input_tensor = torch.randn(1,
                                    8,

Also applies to: 119-122

---

526-526: Wrap overlong lines flagged by Ruff (E501).

A few lines exceed 120 chars. Break them across lines (parentheses) to satisfy style without disabling linting.

Also applies to: 528-528, 606-607

cpp/tensorrt_llm/kernels/moePrepareKernels.h (4)

43-48: Avoid magic-numbered FIFO array size; introduce a named constant.

values[512] is opaque and hard to change. Prefer a constexpr and reuse it for the member.

 namespace moe_prepare
 {
 
-#define UNIT_PER_PIPELINE 128
+#define UNIT_PER_PIPELINE 128
 #define PIPELINE_PER_CTA 4
 #define CUMSUM_THREADS_PER_BLOCK 128
 
 static constexpr int THREADS_PER_PIPELINE = UNIT_PER_PIPELINE;
+static constexpr int kFIFO_VALUES_SIZE = 512;
@@
 struct ALIGN_256 MoeCommFifoConnInfo
 {
     volatile uint64_t head;   // write position
     volatile uint64_t tail;   // read position
-    int volatile values[512]; // for values
+    int volatile values[kFIFO_VALUES_SIZE]; // for values
 };

---

59-66: Document/validate alignment of workspacePtr before casting to MoeCommFifoConnInfo*.

Given ALIGN_256 on MoeCommFifoConnInfo, consider adding an assertion or a brief comment guaranteeing workspacePtr’s base and rankStrideInU64 keep each instance 256B-aligned to avoid misaligned accesses on device.

---

1-15: License years may need updating to include 2025.

Repo-wide policy often bumps to the current year on touched files. If that applies here, extend 2022-2024 to 2022-2025.

---

71-75: Align memsetExpertIds parameter names

The declaration in moePrepareKernels.h uses epSize, while the definition in moePrepareKernels.cu uses rankCount. The other kernels (computeCumsum, moveIndice) already match parameter names between header and implementation.

To resolve, pick one of the following:

• Rename the header parameter to rankCount
• Rename the .cu parameter to epSize

Example header change (rename to rankCount):

--- a/cpp/tensorrt_llm/kernels/moePrepareKernels.h
+++ b/cpp/tensorrt_llm/kernels/moePrepareKernels.h
@@ -80,7 +80,7 @@ void memsetExpertIds(int* expertIds, int* recvCountsCumsum, int maxTokenCountPe
-    int epSize, cudaStream_t stream);
+    int rankCount, cudaStream_t stream);

Files to update:

• cpp/tensorrt_llm/kernels/moePrepareKernels.h (declaration)
• cpp/tensorrt_llm/kernels/moePrepareKernels.cu (definition)

Verification script (already run) confirms only memsetExpertIds is mismatched:

rg -nP 'void\s+memsetExpertIds\s*\(' -C2 cpp/tensorrt_llm/kernels
rg -nP 'void\s+computeCumsum\s*\(' -C2 cpp/tensorrt_llm/kernels
rg -nP 'void\s+moveIndice\s*\(' -C2 cpp/tensorrt_llm/kernels

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

353-358: Use self.num_slots for slot_count for clarity (even if equal to num_experts).

In this module num_slots == num_experts, but conceptually they differ in other backends. Passing self.num_slots improves consistency with WideEP and reduces future diff churn.

-                max_num_token, self.ep_rank, self.ep_size, self.num_experts,
-                self.num_experts, top_k)
+                max_num_token, self.ep_rank, self.ep_size, self.num_experts,
+                self.num_slots, top_k)

---

373-376: Minor: maintain x_row only once.

You set x_row earlier and reassign from x_sf.shape[0] here. Both should match; consider asserting equality or dropping the reassignment to avoid hidden skew if upstream changes.

tensorrt_llm/_mnnvl_utils.py (1)

339-347: Type annotation mismatch: local_gather_indices can be None.

You assign None but the dataclass types it as torch.Tensor. Make it Optional to reflect reality and help static analyzers.

-from dataclasses import dataclass
-from typing import List, Optional, Union
+from dataclasses import dataclass
+from typing import List, Optional, Union
@@
-@dataclass
-class MoEAlltoallInfo:
-    local_gather_indices: torch.Tensor
+@dataclass
+class MoEAlltoallInfo:
+    local_gather_indices: Optional[torch.Tensor]

Also applies to: 433-444

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (3)

465-476: Two-stage prepare/dispatch API looks clean; minor naming nit.

alltoall_prepare returns token_selected_slots unchanged; consider returning only the new artifacts (gathered_*, alltoall_info) to reduce redundancy, or rename to prepare_slots_and_info if the intent is to emphasize slot-based routing.

Also applies to: 867-876

---

879-895: Dispatch + memset_expert_ids is correct; ensure scales can be None.

Downstream fused_moe accepts token_final_scales=None for the “apply-on-input” path. Your dispatch preserves that — good. Consider a brief comment here to clarify that token_final_scales may be None by design.

---

465-465: Ruff E501: one overlong line; wrap to meet 120 cols.

Wrap line 465 to satisfy linting.

cpp/tests/unit_tests/kernels/fusedMoeCommKernelTest.cpp (1)

296-321: Optional: avoid reinterpret_cast<float const>(int) when reading back scales.**

Reading floats through an int* via reinterpret_cast violates aliasing rules; consider memcpy per element for safety in both verifyG2SResults and verifyS2GResults.

Also applies to: 518-526

cpp/tensorrt_llm/kernels/fusedMoeCommKernels.h (4)

38-40: Clarify or remove “do we use this?” in public header.

Question-style comments in public headers age poorly. Either justify the member or remove the aside.

---

109-120: Avoid shadowing ‘alignedUnitBytes’ inside getFieldUncompactSize.

Shadowing the name with a constexpr of the same identifier harms readability. Use a distinct constant.

-        if (alignedUnitBytes != 16)
-        {
-            constexpr int alignedUnitBytes = BYTES_PER_16B_BLOCK;
-            currentFieldSize = currentFieldSize / alignedUnitBytes * alignedUnitBytes;
-            currentFieldSize += alignedUnitBytes * 2;
-        }
+        if (alignedUnitBytes != BYTES_PER_16B_BLOCK)
+        {
+            constexpr int kAlignedBytes = BYTES_PER_16B_BLOCK;
+            currentFieldSize = currentFieldSize / kAlignedBytes * kAlignedBytes;
+            currentFieldSize += kAlignedBytes * 2;
+        }

---

485-491: Typo: mappedMemoryrank → mappedMemoryRank.

Minor readability polish in pointer arithmetic helper.

-        int mappedMemoryrank = isSenderSideBuffer ? pairInfo.senderRank : pairInfo.receiverRank;
+        int mappedMemoryRank = isSenderSideBuffer ? pairInfo.senderRank : pairInfo.receiverRank;
-        auto* mappedMemory = reinterpret_cast<uint8_t*>(workspacePtr + mappedMemoryrank * rankStrideInU64);
+        auto* mappedMemory = reinterpret_cast<uint8_t*>(workspacePtr + mappedMemoryRank * rankStrideInU64);

---

292-300: Fix typos in diagnostics and variable name.

“at lease” → “at least”; “perferredChannel” → “preferredChannel”.

-        TLLM_CHECK_WITH_INFO(
-            blockCountPerChannel <= smCount, "GPU should support at lease one channel, usableSmCount=%d", smCount);
-        int perferredChannel = smCount / 2 / blockCountPerChannel; // use half SMs for communication
-        int channelCount = std::max(perferredChannel, 1);          // at lease one channel
+        TLLM_CHECK_WITH_INFO(
+            blockCountPerChannel <= smCount, "GPU should support at least one channel, usableSmCount=%d", smCount);
+        int preferredChannel = smCount / 2 / blockCountPerChannel; // use half SMs for communication
+        int channelCount = std::max(preferredChannel, 1);          // at least one channel

Configuration used: Path: .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_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

Learnt from: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

Learnt from: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4616-4626
Timestamp: 2025-08-19T03:35:20.866Z
Learning: In the MOE profiler TMA workspace preparation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu), the overlapping of TMA WS regions for NONE and FINALIZE variants is deliberate design to save memory space, as confirmed by djns99. The comment "reuse the same pointers to save space" reflects this intentional behavior.

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-12T10:28:57.320Z
Learning: Applies to **/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py} : Prepend NVIDIA copyright header (current year) to all source files

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-12T10:28:57.320Z
Learning: Applies to **/*.{h,hpp,hxx,hh,cuh} : Header files must use include guards named TRTLLM_<FILENAME>_H without underscores prefix/suffix (e.g., TRTLLM_FOO_BAR_HELLO_H)

Learnt from: sklevtsov-nvidia
PR: NVIDIA/TensorRT-LLM#3294
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:36-36
Timestamp: 2025-08-08T05:06:31.596Z
Learning: CUTLASS extension files (under cpp/tensorrt_llm/cutlass_extensions/) follow CUTLASS coding style conventions, including using #pragma once instead of TRTLLM_ prefixed header guards, even though they are .hpp files.

tests/unittest/_torch/thop/test_moe_alltoall.py (3)

125-139: Good parameter coverage for multi-rank smoke/large tests.

---

236-239: Reference reconstruction logic is correct but fragile to future shape changes.

The global output_indice_offset assumes each token contributes exactly once across ranks. If future tests exercise multi-lane routes (top-k > 1), this needs adjustment to accumulate per-rank-per-tensor counts.

Would you like me to extend the test to a top-k>1 variant and validate the reconstruction order?

---

271-286: Nice parallel a2a invocation and reconstruction with per-rank streams.

Also applies to: 290-306

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

35-36: Nice: replaced macro with a constexpr for threads-per-pipeline.

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

363-371: Good: unified multi-tensor alltoallv dispatch and post-population of expert IDs.

This matches the fused path and removes the padding hacks.

tensorrt_llm/_mnnvl_utils.py (3)

370-373: Nice: explicit workspace init + barrier avoids first-use races.

---

528-590: Multi-tensor a2a API: solid handling of None entries and single-tensor round-trip.

The validation and restoration logic look correct.

---

602-615: Combine path correctly switches to list-based moe_comm and unpacks.

The reshape-and-sum pattern matches the expected top-k combine semantics.

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

196-202: Good: allocate both a2a and prepare workspaces under MNNVL.

This mirrors the fused path expectations.

---

299-301: Always-on alltoall gating for MNNVL is reasonable.

This simplifies control flow for the fused path.

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

315-329: Confirm launch dimension policy.

The 2 in grid.z encodes both directions; ensure kernels assume grid.z==2 and no divergence when hasBasicFields=false. If not guaranteed, guard by template parameter or assert.

Would you like a quick audit script to grep kernel launch sites and ensure consistent interpretation of grid.z across kernels?

---

1-15: Adopt SPDX header.

Align with repository convention: replace the block header with SPDX lines.

-/*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+// SPDX-FileCopyrightText: Copyright (c) 2019-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0

⛔ Skipped due to learnings

Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-12T10:28:57.320Z
Learning: Applies to **/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py} : Prepend NVIDIA copyright header (current year) to all source files