fix perf

Signed-off-by: Enwei Zhu <[email protected]>

Author: syuoni

cpp/tensorrt_llm/kernels/cuteDslKernels/moeUtils.cu

@@ -142,7 +142,7 @@ void moePermute(InputType const* input, InputType* permuted_output, SFType const
 #endif
 
     static int32_t const smCount = tensorrt_llm::common::getMultiProcessorCount();
-    int32_t const blocks = std::min(smCount, max_num_permuted_tokens);
+    int32_t const blocks = std::min(smCount * 8, max_num_permuted_tokens);
     int32_t const threads = kThreadsPerBlock;
 
     auto kernel = &moePermuteKernel<InputType, SFType, kSFVecSize, kThreadsPerBlock>;
@@ -383,7 +383,7 @@ void moeActivation(InputType const* input, OutputType* output, float const* glob
 #endif
 
     static int32_t const smCount = tensorrt_llm::common::getMultiProcessorCount();
-    int32_t const blocks = std::min(smCount, max_num_permuted_tokens);
+    int32_t const blocks = std::min(smCount * 8, max_num_permuted_tokens);
     int32_t const threads = kThreadsPerBlock;
 
     auto get_act_kernel = [](ActivationType activation_type) -> void (*)(InputType const* input, OutputType* output,

cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp

@@ -120,8 +120,8 @@ std::vector<torch::Tensor> moe_sort(torch::Tensor const& token_selected_experts,
     TORCH_CHECK(token_final_scales.size(0) == num_tokens, "token_final_scales.size(0) must be num_tokens.");
     TORCH_CHECK(token_final_scales.size(1) == top_k, "token_final_scales.size(1) must be top_k.");
     return moe_topk_sort_impl(std::nullopt, std::nullopt, token_selected_experts, token_final_scales, num_experts,
-        top_k, std::nullopt, std::nullopt, local_expert_offset, local_num_experts, std::nullopt, tile_tokens_dim,
-        RoutingMethodType::Renormalize);
+        top_k, 1, 1, local_expert_offset, local_num_experts, std::nullopt, tile_tokens_dim,
+        RoutingMethodType::DeepSeekV3);
 }
 
 // Permute

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -676,6 +676,10 @@ def forward(self, inputs: List[torch.Tensor],
                     mma_tiler_mn=mma_tiler_mn,
                     cluster_shape_mn=cluster_shape_mn,
                 )
+                # Compute max active clusters on current device
+                hardware_info = cutlass.utils.HardwareInfo()
+                max_active_clusters = hardware_info.get_max_active_clusters(
+                    cluster_shape_mn[0] * cluster_shape_mn[1])
 
                 compiled_gemm = cute.compile(
                     gemm.wrapper,
@@ -693,7 +697,7 @@ def forward(self, inputs: List[torch.Tensor],
                     l,
                     tile_size=self.tile_size,
                     scaling_vector_size=self.scaling_vector_size,
-                    max_active_clusters=16,
+                    max_active_clusters=max_active_clusters,
                     stream=stream,
                 )
                 self.__class__.kernel_cache[cache_key] = compiled_gemm
@@ -970,6 +974,10 @@ def forward(self, inputs: List[torch.Tensor],
                     mma_tiler_mn=mma_tiler_mn,
                     cluster_shape_mn=cluster_shape_mn,
                 )
+                # Compute max active clusters on current device
+                hardware_info = cutlass.utils.HardwareInfo()
+                max_active_clusters = hardware_info.get_max_active_clusters(
+                    cluster_shape_mn[0] * cluster_shape_mn[1])
 
                 compiled_gemm = cute.compile(
                     gemm.wrapper,
@@ -992,7 +1000,7 @@ def forward(self, inputs: List[torch.Tensor],
                     self.top_k,
                     tile_size=self.tile_size,
                     scaling_vector_size=self.scaling_vector_size,
-                    max_active_clusters=16,
+                    max_active_clusters=max_active_clusters,
                     stream=stream,
                 )
                 self.__class__.kernel_cache[cache_key] = compiled_gemm

tensorrt_llm/tools/layer_wise_benchmarks/deepseekv3_runner.py

@@ -209,7 +209,13 @@ def __init__(
         tensorrt_llm._torch.models.modeling_deepseekv3.DeepseekV3Gate = gate_cls_orig
 
     def replace_routing_method(self, balance_method: BalanceMethod, balance_ratio: float):
-        if self.model_config.moe_backend not in ["CUTLASS", "DEEPGEMM", "TRTLLM", "WIDEEP"]:
+        if self.model_config.moe_backend not in [
+            "CUTLASS",
+            "DEEPGEMM",
+            "TRTLLM",
+            "WIDEEP",
+            "CUTEDSL",
+        ]:
             raise NotImplementedError(
                 f'Not support replace routing method for moe_backend "{self.model_config.moe_backend}",'
                 f' please set balance_method to "NotModified"'