fix bf16 api

Signed-off-by: jiahanc <[email protected]>

Author: jiahanc

csrc/trtllm_fused_moe_kernel_launcher.cu

@@ -1241,10 +1241,11 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
 Tensor trtllm_bf16_moe(TensorView const& routing_logits, Optional<TensorView> const& routing_bias,
                        TensorView const& hidden_states, TensorView const& gemm1_weights,
                        TensorView const& gemm2_weights, int64_t num_experts, int64_t top_k,
-                       int64_t n_group, int64_t topk_group, int64_t intermediate_size,
-                       int64_t local_expert_offset, int64_t local_num_experts,
-                       int64_t routing_method_type, bool use_shuffled_weight, int64_t weight_layout,
-                       bool enable_pdl, Array<int64_t> moe_tactic) {
+                       Optional<int64_t> n_group, Optional<int64_t> topk_group,
+                       int64_t intermediate_size, int64_t local_expert_offset,
+                       int64_t local_num_experts, int64_t routing_method_type,
+                       bool use_shuffled_weight, int64_t weight_layout, bool enable_pdl,
+                       Array<int64_t> moe_tactic) {
   // Just some basic type validation first and leave more checks to the launcher
   TVM_FFI_ICHECK(routing_logits.dtype() == dl_float32 || routing_logits.dtype() == dl_bfloat16)
       << "BF16 MoE: routing_logits must be bfloat16 or float.";
@@ -1275,8 +1276,9 @@ Tensor trtllm_bf16_moe(TensorView const& routing_logits, Optional<TensorView> co
     args->hidden_size = hidden_size;
     args->hidden_size_output = args->hidden_size;
     args->top_k = top_k;
-    args->n_group = n_group;
-    args->topk_group = topk_group;
+    args->n_group = n_group.value_or(0);
+    args->topk_group = topk_group.value_or(0);
+    ;
     args->local_expert_offset = local_expert_offset;
     args->local_num_experts = local_num_experts;
     args->intermediate_size = intermediate_size;

flashinfer/fused_moe/core.py

@@ -1168,8 +1168,8 @@ def trtllm_bf16_moe_op(
         gemm2_weights: torch.Tensor,
         num_experts: int,
         top_k: int,
-        n_group: int,
-        topk_group: int,
+        n_group: Optional[int],
+        topk_group: Optional[int],
         intermediate_size: int,
         local_expert_offset: int,
         local_num_experts: int,
@@ -1268,8 +1268,8 @@ def _fake_trtllm_bf16_moe(
         gemm2_weights: torch.Tensor,
         num_experts: int,
         top_k: int,
-        n_group: int,
-        topk_group: int,
+        n_group: Optional[int],
+        topk_group: Optional[int],
         intermediate_size: int,
         local_expert_offset: int,
         local_num_experts: int,
@@ -1808,8 +1808,8 @@ def trtllm_bf16_moe(
     gemm2_weights: torch.Tensor,
     num_experts: int,
     top_k: int,
-    n_group: int,
-    topk_group: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
     intermediate_size: int,
     local_expert_offset: int,
     local_num_experts: int,
@@ -1867,8 +1867,8 @@ def trtllm_bf16_moe(
         gemm2_weights,
         num_experts,
         top_k,
-        n_group or 0,  # may receive None from test configs, convert to 0
-        topk_group or 0,
+        n_group,
+        topk_group,
         intermediate_size,
         local_expert_offset,
         local_num_experts,