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

344-346: Non-OSS path passes a torch::Tensor where a raw pointer is expected (compile-time error).

In the non-OSS branch, workspace_info.workspace is a torch::Tensor, but it’s passed where a char* is expected. This will not compile since you’re casting a tensor object to a pointer.

Apply this diff to pass the underlying data pointer:

-            static_cast<char*>(workspace_info.workspace), output.data_ptr(),
+            static_cast<char*>(workspace_info.workspace.data_ptr()), output.data_ptr(),

---

474-476: Same raw pointer issue in min-latency non-OSS path.

Mirror the fix here as well: workspace_info.workspace is a torch::Tensor and needs .data_ptr().

-            static_cast<char*>(workspace_info.workspace), output.data_ptr(),
+            static_cast<char*>(workspace_info.workspace.data_ptr()), output.data_ptr(),

cpp/tensorrt_llm/thop/moeOp.cpp (4)

580-581: Member naming: use m-prefix and avoid shadowing.

The new member workspace_info violates the class’ member naming convention (members use mPrefix). It also shadows the local workspace_info variable used at call sites, reducing clarity.

Apply this diff to rename the member:

-    WorkspaceInfo workspace_info;
+    WorkspaceInfo mWorkspaceInfo;

Notes:

• The function below (getWorkspaceInfo) should be updated to refer to mWorkspaceInfo (see proposed diff in that function’s comment).
• Call-site locals can keep the name workspace_info if you prefer; the member’s m-prefix prevents ambiguity.

---

311-313: Thread device index through to getWorkspaceInfo to ensure correct device placement for the workspace.

getWorkspaceInfo allocates a CUDA tensor internally. Without using input.options() (not available in that scope) or a device guard, the allocation may land on the wrong device if current device doesn’t match input.get_device(). Passing the device index fixes this deterministically.

Apply this diff to pass the input device index:

-        WorkspaceInfo const& workspace_info = getWorkspaceInfo(num_rows, hidden_size, inter_size, num_experts_total,
-            static_cast<int>(experts_per_token), activation_type, parallelism_config, min_latency_mode, stream);
+        WorkspaceInfo const& workspace_info = getWorkspaceInfo(num_rows, hidden_size, inter_size, num_experts_total,
+            static_cast<int>(experts_per_token), activation_type, parallelism_config, min_latency_mode,
+            input.get_device(), stream);

---

442-444: Ditto: pass device index to ensure workspace allocation lands on the correct GPU for min-latency path.

-        WorkspaceInfo const& workspace_info = getWorkspaceInfo(num_rows, hidden_size, inter_size, num_experts_total,
-            static_cast<int>(experts_per_token), activation_type, parallelism_config, min_latency_mode, stream);
+        WorkspaceInfo const& workspace_info = getWorkspaceInfo(num_rows, hidden_size, inter_size, num_experts_total,
+            static_cast<int>(experts_per_token), activation_type, parallelism_config, min_latency_mode,
+            input.get_device(), stream);

---

626-660: Make workspace allocation device-correct, avoid member/local shadowing, and fix 64-bit logging specifiers.

Good call to persist and grow the workspace and to reallocate during CUDA Graph capture. A few targeted improvements:

• Ensure allocations are placed on the correct CUDA device by accepting a device_index parameter and using it in TensorOptions.
• Use the mWorkspaceInfo member to avoid shadowing and match the class’ naming convention.
• Fix %ld format specifiers for 64-bit values to avoid UB/incorrect output on platforms where long is 32-bit.

Apply this diff:

-    WorkspaceInfo const& getWorkspaceInfo(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size,
-        int num_experts, int experts_per_token, ActivationType activation_type,
-        kernels::MOEParallelismConfig const& parallelismConfig, bool min_latency_mode, cudaStream_t stream)
+    WorkspaceInfo const& getWorkspaceInfo(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size,
+        int num_experts, int experts_per_token, ActivationType activation_type,
+        kernels::MOEParallelismConfig const& parallelismConfig, bool min_latency_mode,
+        int device_index, cudaStream_t stream)
     {
         size_t moe_workspace_size = mKernelRunner->getWorkspaceSize(num_rows, hidden_size, inter_size, num_experts,
             experts_per_token, activation_type, parallelismConfig, /* use_lora */ false, mUseDeepSeekFP8BlockScaling,
             min_latency_mode, mUseW4A8GroupScaling);
         size_t src_to_dest_map_size = experts_per_token * num_rows * sizeof(int);

         std::vector<size_t> workspaces{moe_workspace_size, src_to_dest_map_size};

-        int64_t const total_workspace_size = common::calculateTotalWorkspaceSize(workspaces.data(), workspaces.size());
+        int64_t const total_workspace_size = common::calculateTotalWorkspaceSize(workspaces.data(), workspaces.size());

-        bool is_capturing = tensorrt_llm::common::isCapturing(stream);
+        bool is_capturing = tensorrt_llm::common::isCapturing(stream);
         // Always allocate workspace when capturing cuda graph to avoid illegal memory access during replay
-        if (is_capturing || workspace_info.workspace.numel() < total_workspace_size)
+        if (is_capturing || mWorkspaceInfo.workspace.numel() < total_workspace_size)
         {
             if (is_capturing)
             {
-                TLLM_LOG_DEBUG(
-                    "Allocating MoE workspace with %ld bytes size during cuda graph capture", total_workspace_size);
+                TLLM_LOG_DEBUG(
+                    "Allocating MoE workspace with %lld bytes size during cuda graph capture",
+                    static_cast<long long>(total_workspace_size));
             }
             else
             {
-                TLLM_LOG_DEBUG("MoE workspace size is not enough, increase the size from %ld bytes to %ld bytes",
-                    workspace_info.workspace.numel(), total_workspace_size);
+                TLLM_LOG_DEBUG("MoE workspace size is not enough, increase the size from %lld bytes to %lld bytes",
+                    static_cast<long long>(mWorkspaceInfo.workspace.numel()),
+                    static_cast<long long>(total_workspace_size));
             }
-            workspace_info.workspace = torch::empty({static_cast<long>(total_workspace_size)},
-                torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
+            mWorkspaceInfo.workspace = torch::empty({static_cast<int64_t>(total_workspace_size)},
+                torch::dtype(torch::kInt8).device(c10::Device(torch::kCUDA, device_index)).requires_grad(false));
         }
-        workspace_info.src_to_dest_map
-            = common::nextWorkspacePtr(static_cast<int8_t*>(workspace_info.workspace.data_ptr()), moe_workspace_size);
+        mWorkspaceInfo.src_to_dest_map
+            = common::nextWorkspacePtr(static_cast<int8_t*>(mWorkspaceInfo.workspace.data_ptr()), moe_workspace_size);

-        return workspace_info;
+        return mWorkspaceInfo;
     }

Optional consideration:

• If you foresee capturing multiple CUDA graphs with the same runner instance (on the same or different streams) and replaying all of them concurrently, please double-check lifetime semantics with the caching allocator’s graph pools to ensure reallocating mWorkspaceInfo.workspace for a later capture doesn’t free memory held by a previously captured graph. If needed, we can store per-graph workspaces keyed by stream/graph.

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.