fix autotune and apply fp8 auto tune to test

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

Author: jiahanc

flashinfer/fused_moe/core.py

@@ -926,6 +926,15 @@ def __init__(
             self.gated_act_type = GatedActType(gated_act_type)
             self.use_shuffled_weight = use_shuffled_weight
             self.weight_layout = WeightLayout(weight_layout)
+            if (
+                not self.use_shuffled_weight
+                or self.weight_layout != WeightLayout.MajorK
+            ):
+                assert (
+                    self.use_deepseek_fp8 and self.dtype_weights == DtypeTrtllmGen.E4m3
+                ), (
+                    "use_shuffled_weight is False or weight_layout is not MajorK is only supported for FP8 block scale"
+                )
 
         def get_valid_tactics(
             self,
@@ -1022,7 +1031,7 @@ def forward(
                         dtype=torch.float,
                         device=hidden_states.device,
                     )
-                    return moe_op.trtllm_fp8_block_scale_moe(
+                    moe_op.trtllm_fp8_block_scale_moe(
                         routing_logits,
                         kwargs["routing_bias"],
                         hidden_states,
@@ -1031,6 +1040,7 @@ def forward(
                         kwargs["gemm1_weights_scale"],
                         kwargs["gemm2_weights"],
                         kwargs["gemm2_weights_scale"],
+                        output,
                         kwargs["num_experts"],
                         self.top_k,
                         kwargs["n_group"],
@@ -1047,7 +1057,7 @@ def forward(
                     )
                 else:
                     # FP8 per tensor scale
-                    return moe_op.trtllm_fp8_per_tensor_scale_moe(
+                    moe_op.trtllm_fp8_per_tensor_scale_moe(
                         routing_logits,
                         kwargs["routing_bias"],
                         hidden_states,
@@ -1056,6 +1066,7 @@ def forward(
                         kwargs["output1_scales_gate_scalar"],
                         kwargs["gemm2_weights"],
                         kwargs["output2_scales_scalar"],
+                        output,
                         kwargs["num_experts"],
                         self.top_k,
                         kwargs["n_group"],
@@ -1188,6 +1199,8 @@ def trtllm_fp8_per_tensor_scale_moe_op(
             use_deepseek_fp8=False,  # per_tensor mode
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
+            weight_layout=WeightLayout.MajorK,
+            use_shuffled_weight=True,
         )
 
         inputs = [output, routing_logits, topk_ids, expert_weights, hidden_states]
@@ -1203,6 +1216,7 @@ def trtllm_fp8_per_tensor_scale_moe_op(
             output1_scales_gate_scalar=output1_scales_gate_scalar,
             gemm2_weights=gemm2_weights,
             output2_scales_scalar=output2_scales_scalar,
+            num_experts=num_experts,
             n_group=n_group,
             topk_group=topk_group,
             local_expert_offset=local_expert_offset,
@@ -1325,6 +1339,8 @@ def trtllm_fp8_block_scale_moe_op(
             use_deepseek_fp8=True,  # block_scale mode
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
+            weight_layout=weight_layout,
+            use_shuffled_weight=use_shuffled_weight,
         )
 
         inputs = [
@@ -1346,6 +1362,7 @@ def trtllm_fp8_block_scale_moe_op(
             gemm1_weights_scale=gemm1_weights_scale,
             gemm2_weights=gemm2_weights,
             gemm2_weights_scale=gemm2_weights_scale,
+            num_experts=num_experts,
             n_group=n_group,
             topk_group=topk_group,
             local_expert_offset=local_expert_offset,
@@ -1498,9 +1515,8 @@ def trtllm_fp4_block_scale_moe_op(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
             gated_act_type=gated_act_type,
-            # NOTE(siyuan): do not fix the tile_tokens_dim to let tunnable runner decide the tile_tokens_dim itself.
-            # however, when the user chooses a different heuristic for tile_tokens_dim, the autotuner will fail to find the correct cached tactics.
-            # tile_tokens_dim=tile_tokens_dim,
+            weight_layout=WeightLayout.MajorK,
+            use_shuffled_weight=True,
         )
         tunning_config = (
             MoERunner.tuning_config_no_hidden_states_scales

tests/moe/test_trtllm_gen_fused_moe.py

@@ -769,29 +769,31 @@ def call_moe(
             "NaN detected in hidden_states_fp8"
         )
 
-        output = trtllm_fp8_block_scale_moe(
-            expert_logits,
-            routing_bias,
-            hidden_states_fp8,
-            hidden_states_scale,
-            static_data["gemm1_weights"],
-            static_data["gemm1_scales"],
-            static_data["gemm2_weights"],
-            static_data["gemm2_scales"],
-            num_experts,
-            top_k,
-            n_groups,
-            top_k_groups,
-            intermediate_size,
-            0,
-            num_experts,
-            routed_scaling,
-            None,
-            routing_method_type,
-            use_shuffled_weight=static_data["use_shuffled_weight"],
-            weight_layout=static_data["weight_layout"],
-            enable_pdl=enable_pdl,
-        )
+        # Use autotuner for optimal kernel selection
+        with autotune(True):
+            output = trtllm_fp8_block_scale_moe(
+                expert_logits,
+                routing_bias,
+                hidden_states_fp8,
+                hidden_states_scale,
+                static_data["gemm1_weights"],
+                static_data["gemm1_scales"],
+                static_data["gemm2_weights"],
+                static_data["gemm2_scales"],
+                num_experts,
+                top_k,
+                n_groups,
+                top_k_groups,
+                intermediate_size,
+                0,
+                num_experts,
+                routed_scaling,
+                None,
+                routing_method_type,
+                use_shuffled_weight=static_data["use_shuffled_weight"],
+                weight_layout=static_data["weight_layout"],
+                enable_pdl=enable_pdl,
+            )
 
         return output.to(torch.float)
 
@@ -940,32 +942,34 @@ def call_moe(
             hidden_states_orig, hidden_states_scale_global
         )
 
-        output = trtllm_fp8_per_tensor_scale_moe(
-            (
-                expert_logits.to(torch.bfloat16)
-                if routing_method_type == RoutingMethodType.Llama4
-                else expert_logits
-            ),
-            routing_bias,
-            hidden_states_fp8,
-            static_data["gemm1_weights"],
-            static_data["scale_c_fc1"],
-            static_data["scale_gate_fc1"],
-            static_data["gemm2_weights"],
-            static_data["scale_c_fc2"],
-            num_experts,
-            top_k,
-            n_groups,
-            top_k_groups,
-            intermediate_size,
-            0,
-            num_experts,
-            routed_scaling,
-            routing_method_type
-            == RoutingMethodType.Llama4,  # Use_routing_scales_on_input
-            None,
-            routing_method_type,
-        )
+        # Use autotuner for optimal kernel selection
+        with autotune(True):
+            output = trtllm_fp8_per_tensor_scale_moe(
+                (
+                    expert_logits.to(torch.bfloat16)
+                    if routing_method_type == RoutingMethodType.Llama4
+                    else expert_logits
+                ),
+                routing_bias,
+                hidden_states_fp8,
+                static_data["gemm1_weights"],
+                static_data["scale_c_fc1"],
+                static_data["scale_gate_fc1"],
+                static_data["gemm2_weights"],
+                static_data["scale_c_fc2"],
+                num_experts,
+                top_k,
+                n_groups,
+                top_k_groups,
+                intermediate_size,
+                0,
+                num_experts,
+                routed_scaling,
+                routing_method_type
+                == RoutingMethodType.Llama4,  # Use_routing_scales_on_input
+                None,
+                routing_method_type,
+            )
 
         return output.to(torch.float)