MoE Refactor: Refactor fp8.py -> flashinfer_trllm.py (#15151)

Co-authored-by: Brayden Zhong <b8zhong@users.noreply.github.com>

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -1146,14 +1146,14 @@ class FlashInferFusedMoE(FusedMoE):
 
         else:
 
-            final_hidden_states = self.quant_method.apply_with_router_logits(
+            final_hidden_states = self.quant_method.apply(
                 layer=self,
                 dispatch_output=StandardDispatchOutput(
                     hidden_states=hidden_states,
                     hidden_states_scale=None,
                     topk_output=topk_output,
                 ),
-            )
+            ).hidden_states
 
         # NOTE for symmetric memory tagging:
         # We do not create the context in this function.

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, cast
+
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+from sglang.srt.distributed import get_tp_group
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.layers.dp_attention import is_allocation_symmetric
+from sglang.srt.layers.moe.moe_runner.base import (
+    MoeQuantInfo,
+    MoeRunnerConfig,
+    register_fused_func,
+)
+from sglang.srt.layers.quantization.fp8_kernel import (
+    per_token_group_quant_fp8,
+    scaled_fp8_quant,
+)
+from sglang.srt.utils.common import next_power_of_2
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import (
+        StandardCombineInput,
+        StandardDispatchOutput,
+    )
+
+
+def align_fp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
+    """Prepare FP8 MoE weights/scales for FlashInfer TRT-LLM kernels."""
+    from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a
+
+    # Note: No need to swap W13 halves, they are already in the correct order:
+    # [Gate, Up]
+    w13_weight = cast(torch.Tensor, layer.w13_weight)
+    w2_weight = cast(torch.Tensor, layer.w2_weight)
+    num_experts, two_n, hidden = w13_weight.shape
+
+    w13_interleaved_list = [
+        reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts)
+    ]
+    w13_interleaved: torch.Tensor = torch.stack(w13_interleaved_list).reshape(
+        num_experts, two_n, hidden
+    )
+
+    # Shuffle weights for transposed MMA output (both W13, W2)
+    epilogue_tile_m = 128
+    w13_shuffled = [
+        shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m)
+        for i in range(num_experts)
+    ]
+    w2_shuffled = [
+        shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m)
+        for i in range(num_experts)
+    ]
+
+    layer.w13_weight = Parameter(
+        torch.stack(w13_shuffled).view(torch.float8_e4m3fn),
+        requires_grad=False,
+    )
+    layer.w2_weight = Parameter(
+        torch.stack(w2_shuffled).view(torch.float8_e4m3fn),
+        requires_grad=False,
+    )
+
+    # Precompute and register per-expert output scaling factors for FI MoE.
+    # Note: w13_input_scale and w2_input_scale are scalar Parameters post-reduction.
+    assert hasattr(layer, "w13_input_scale") and layer.w13_input_scale is not None
+    assert hasattr(layer, "w2_input_scale") and layer.w2_input_scale is not None
+    assert hasattr(layer, "w13_weight_scale") and layer.w13_weight_scale is not None
+    assert hasattr(layer, "w2_weight_scale") and layer.w2_weight_scale is not None
+
+    input_scale = cast(torch.Tensor, layer.w13_input_scale).to(torch.float32)
+    activation_scale = cast(torch.Tensor, layer.w2_input_scale).to(torch.float32)
+    w13_weight_scale = cast(torch.Tensor, layer.w13_weight_scale).to(torch.float32)
+    w2_weight_scale = cast(torch.Tensor, layer.w2_weight_scale).to(torch.float32)
+
+    output1_scales_scalar = w13_weight_scale * input_scale * (1.0 / activation_scale)
+    output1_scales_gate_scalar = w13_weight_scale * input_scale
+    output2_scales_scalar = activation_scale * w2_weight_scale
+
+    layer.output1_scales_scalar = Parameter(output1_scales_scalar, requires_grad=False)
+    layer.output1_scales_gate_scalar = Parameter(
+        output1_scales_gate_scalar, requires_grad=False
+    )
+    layer.output2_scales_scalar = Parameter(output2_scales_scalar, requires_grad=False)
+
+
+@dataclass
+class FlashInferTrtllmFp8MoeQuantInfo(MoeQuantInfo):
+    """Quantization payload consumed by FlashInfer TRT-LLM FP8 MoE kernels."""
+
+    # Weights
+    w13_weight: torch.Tensor
+    w2_weight: torch.Tensor
+
+    # Expert-parallel metadata
+    global_num_experts: int
+    local_expert_offset: int
+    local_num_experts: int
+
+    routing_method_type: int
+
+    # Block-quant path
+    block_quant: bool
+    weight_block_k: int | None = None
+    w13_weight_scale_inv: torch.Tensor | None = None
+    w2_weight_scale_inv: torch.Tensor | None = None
+
+    # Per-tensor path
+    w13_input_scale: torch.Tensor | None = None
+    output1_scales_scalar: torch.Tensor | None = None
+    output1_scales_gate_scalar: torch.Tensor | None = None
+    output2_scales_scalar: torch.Tensor | None = None
+
+
+@register_fused_func("none", "flashinfer_trtllm")
+def fused_experts_none_to_flashinfer_trtllm_fp8(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: FlashInferTrtllmFp8MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+) -> StandardCombineInput:
+    from flashinfer.fused_moe import (
+        trtllm_fp8_block_scale_moe,
+        trtllm_fp8_per_tensor_scale_moe,
+    )
+
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.layers.moe.utils import RoutingMethodType
+
+    assert runner_config.activation == "silu", "Only silu is supported."
+    assert not runner_config.no_combine, "no_combine is not supported for flashinfer."
+
+    hidden_states = dispatch_output.hidden_states
+    topk_output = dispatch_output.topk_output
+    assert TopKOutputChecker.format_is_bypassed(topk_output)
+
+    router_logits = topk_output.router_logits
+    topk_config = topk_output.topk_config
+    correction_bias = (
+        None
+        if topk_config.correction_bias is None
+        else topk_config.correction_bias.to(hidden_states.dtype)
+    )
+
+    routing_method_type = quant_info.routing_method_type
+
+    if quant_info.block_quant:
+        assert quant_info.weight_block_k is not None
+        assert quant_info.w13_weight_scale_inv is not None
+        assert quant_info.w2_weight_scale_inv is not None
+
+        a_q, a_sf = per_token_group_quant_fp8(hidden_states, quant_info.weight_block_k)
+        a_sf_t = a_sf.t().contiguous()
+
+        with use_symmetric_memory(
+            get_tp_group(), disabled=not is_allocation_symmetric()
+        ):
+            # FIXME: there is a bug in the trtllm_fp8_block_scale_moe.
+            # It ignored the `output` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325
+            # so we put the whole function under the ``use_symmetric_memory`` context manager.
+            # If the bug is fixed, we can only put the output tensor allocation under the context manager.
+            output = trtllm_fp8_block_scale_moe(
+                routing_logits=(
+                    router_logits.to(torch.float32)
+                    if routing_method_type == RoutingMethodType.DeepSeekV3
+                    else router_logits
+                ),
+                routing_bias=correction_bias,
+                hidden_states=a_q,
+                hidden_states_scale=a_sf_t,
+                gemm1_weights=quant_info.w13_weight,
+                gemm1_weights_scale=quant_info.w13_weight_scale_inv,
+                gemm2_weights=quant_info.w2_weight,
+                gemm2_weights_scale=quant_info.w2_weight_scale_inv,
+                num_experts=quant_info.global_num_experts,
+                top_k=topk_config.top_k,
+                n_group=topk_config.num_expert_group,
+                topk_group=topk_config.topk_group,
+                intermediate_size=int(quant_info.w2_weight.shape[2]),
+                local_expert_offset=quant_info.local_expert_offset,
+                local_num_experts=quant_info.local_num_experts,
+                routed_scaling_factor=(
+                    runner_config.routed_scaling_factor
+                    if runner_config.routed_scaling_factor is not None
+                    else 1.0
+                ),
+                tile_tokens_dim=None,
+                routing_method_type=routing_method_type,
+                use_shuffled_weight=False,
+                tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+            )
+    else:
+        assert quant_info.w13_input_scale is not None
+        assert quant_info.output1_scales_scalar is not None
+        assert quant_info.output1_scales_gate_scalar is not None
+        assert quant_info.output2_scales_scalar is not None
+
+        a_q, _ = scaled_fp8_quant(hidden_states, quant_info.w13_input_scale)
+        routing_bias_cast = (
+            None if correction_bias is None else correction_bias.to(torch.bfloat16)
+        )
+
+        with use_symmetric_memory(
+            get_tp_group(), disabled=not is_allocation_symmetric()
+        ):
+            output = trtllm_fp8_per_tensor_scale_moe(
+                routing_logits=router_logits.to(torch.bfloat16),
+                routing_bias=routing_bias_cast,
+                hidden_states=a_q,
+                gemm1_weights=quant_info.w13_weight,
+                output1_scales_scalar=quant_info.output1_scales_scalar,
+                output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar,
+                gemm2_weights=quant_info.w2_weight,
+                output2_scales_scalar=quant_info.output2_scales_scalar,
+                num_experts=quant_info.global_num_experts,
+                top_k=topk_config.top_k,
+                n_group=topk_config.num_expert_group,
+                topk_group=topk_config.topk_group,
+                intermediate_size=int(quant_info.w2_weight.shape[2]),
+                local_expert_offset=quant_info.local_expert_offset,
+                local_num_experts=quant_info.local_num_experts,
+                routed_scaling_factor=(
+                    runner_config.routed_scaling_factor
+                    if runner_config.routed_scaling_factor is not None
+                    else 1.0
+                ),
+                use_routing_scales_on_input=False,
+                routing_method_type=routing_method_type,
+                tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+            )
+
+    return StandardCombineInput(hidden_states=output)

python/sglang/srt/layers/moe/moe_runner/runner.py

@@ -39,6 +39,8 @@ class MoeRunner:
             self.runner_core = DeepGemmRunnerCore(config)
         elif runner_backend.is_marlin():
             self.runner_core = None  # Marlin only supports fused path
+        elif runner_backend.is_flashinfer_trtllm():
+            self.runner_core = None  # FlashInfer TRT-LLM only supports fused path
         else:
             raise NotImplementedError(f"Unsupported runner backend: {runner_backend}")
 
@@ -50,6 +52,12 @@ class MoeRunner:
             a2a_backend_name, runner_backend_name
         )
 
+        if self.runner_core is None and self.fused_func is None:
+            raise NotImplementedError(
+                f"Runner backend {runner_backend} requires a fused func for a2a backend "
+                f"{a2a_backend_name}, but none is registered."
+            )
+
         self.down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None
         self.meta_overlap_args: Optional[dict] = None
 
@@ -69,6 +77,7 @@ class MoeRunner:
         if self.fused_func is not None:
             return self.fused_func(dispatch_output, quant_info, self.config)
 
+        assert self.runner_core is not None
         dispatch_format = dispatch_output.format.value
         runner_format = self.runner_core.runner_backend.value
         self.pre_permute_func = PermuteMethodPool.get_pre_permute(

python/sglang/srt/layers/quantization/fp8.py

@@ -18,8 +18,11 @@ from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
 from sglang.srt.layers.dp_attention import is_allocation_symmetric
 from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo
+from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+    FlashInferTrtllmFp8MoeQuantInfo,
+)
 from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
-from sglang.srt.layers.moe.utils import get_moe_runner_backend
+from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend
 from sglang.srt.layers.parameter import (
     BlockQuantScaleParameter,
     ModelWeightParameter,
@@ -69,18 +72,13 @@ from sglang.srt.utils import (
     is_sm90_supported,
     is_sm100_supported,
     log_info_on_rank0,
-    next_power_of_2,
     print_warning_once,
     set_weight_attrs,
     use_intel_amx_backend,
 )
 
 if TYPE_CHECKING:
-    from sglang.srt.layers.moe.token_dispatcher import (
+    from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
-        CombineInput,
-        DispatchOutput,
-        StandardDispatchOutput,
-    )
     from sglang.srt.layers.moe.topk import TopKOutput
     from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
 
@@ -947,83 +945,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
 
             # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled
             if get_moe_runner_backend().is_flashinfer_trtllm():
-                from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a
+                from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
-
+                    align_fp8_moe_weights_for_flashinfer_trtllm,
-                # Note: No need to swap W13 halves, they are already in the correct order: [Gate, Up]
-                num_experts, two_n, hidden = layer.w13_weight.shape
-
-                # 2) Reorder rows for fused gated activation (W13)
-                w13_interleaved = [
-                    reorder_rows_for_gated_act_gemm(layer.w13_weight[i])
-                    for i in range(num_experts)
-                ]
-                w13_interleaved = torch.stack(w13_interleaved).reshape(
-                    num_experts, two_n, hidden
                 )
 
-                # 3) Shuffle weights for transposed MMA output (both W13, W2)
+                align_fp8_moe_weights_for_flashinfer_trtllm(layer)
-                epilogue_tile_m = 128
-                w13_shuffled = [
-                    shuffle_matrix_a(
-                        w13_interleaved[i].view(torch.uint8), epilogue_tile_m
-                    )
-                    for i in range(num_experts)
-                ]
-                w2_shuffled = [
-                    shuffle_matrix_a(
-                        layer.w2_weight[i].view(torch.uint8), epilogue_tile_m
-                    )
-                    for i in range(num_experts)
-                ]
-
-                layer.w13_weight = Parameter(
-                    torch.stack(w13_shuffled).view(torch.float8_e4m3fn),
-                    requires_grad=False,
-                )
-                layer.w2_weight = Parameter(
-                    torch.stack(w2_shuffled).view(torch.float8_e4m3fn),
-                    requires_grad=False,
-                )
-
-                # Precompute and register per-expert output scaling factors for FI MoE
-                # Note: w13_input_scale and w2_input_scale are scalar Parameters post-reduction
-                assert (
-                    hasattr(layer, "w13_input_scale")
-                    and layer.w13_input_scale is not None
-                )
-                assert (
-                    hasattr(layer, "w2_input_scale")
-                    and layer.w2_input_scale is not None
-                )
-                assert (
-                    hasattr(layer, "w13_weight_scale")
-                    and layer.w13_weight_scale is not None
-                )
-                assert (
-                    hasattr(layer, "w2_weight_scale")
-                    and layer.w2_weight_scale is not None
-                )
-
-                input_scale = layer.w13_input_scale.to(torch.float32)
-                activation_scale = layer.w2_input_scale.to(torch.float32)
-                w13_weight_scale = layer.w13_weight_scale.to(torch.float32)
-                w2_weight_scale = layer.w2_weight_scale.to(torch.float32)
-
-                output1_scales_scalar = (
-                    w13_weight_scale * input_scale * (1.0 / activation_scale)
-                )
-                output1_scales_gate_scalar = w13_weight_scale * input_scale
-                output2_scales_scalar = activation_scale * w2_weight_scale
-
-                layer.output1_scales_scalar = Parameter(
-                    output1_scales_scalar, requires_grad=False
-                )
-                layer.output1_scales_gate_scalar = Parameter(
-                    output1_scales_gate_scalar, requires_grad=False
-                )
-                layer.output2_scales_scalar = Parameter(
-                    output2_scales_scalar, requires_grad=False
-                )
             return
 
     def process_weights_hip_int4(self, layer: Module):
@@ -1119,7 +1045,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 moe_runner_backend = MoeRunnerBackend.DEEP_GEMM
             else:
                 moe_runner_backend = MoeRunnerBackend.TRITON
-        if moe_runner_backend.is_deep_gemm() or moe_runner_backend.is_triton():
+        if (
+            moe_runner_backend.is_deep_gemm()
+            or moe_runner_backend.is_triton()
+            or moe_runner_backend.is_flashinfer_trtllm()
+        ):
             self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
         else:
             # TODO(cwan): refactor other backends
@@ -1245,6 +1175,51 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 w2_scale=w2_scale,
                 block_shape=block_shape,
             )
+        elif self.runner.runner_backend.is_flashinfer_trtllm():
+            # FlashInfer TRT-LLM backend only supports fused execution and consumes
+            # router logits directly (no separate apply_with_router_logits needed).
+            global_num_experts = int(getattr(layer, "num_experts"))
+            num_local_experts = int(getattr(layer, "num_local_experts"))
+            moe_ep_rank = int(getattr(layer, "moe_ep_rank"))
+
+            quant_info = FlashInferTrtllmFp8MoeQuantInfo(
+                w13_weight=layer.w13_weight,
+                w2_weight=layer.w2_weight,
+                global_num_experts=global_num_experts,
+                local_expert_offset=moe_ep_rank * num_local_experts,
+                local_num_experts=num_local_experts,
+                routing_method_type=int(
+                    getattr(layer, "routing_method_type", RoutingMethodType.DeepSeekV3)
+                ),
+                block_quant=self.block_quant,
+                weight_block_k=(
+                    None
+                    if self.quant_config.weight_block_size is None
+                    else self.quant_config.weight_block_size[1]
+                ),
+                w13_weight_scale_inv=(
+                    layer.w13_weight_scale_inv if self.block_quant else None
+                ),
+                w2_weight_scale_inv=(
+                    layer.w2_weight_scale_inv if self.block_quant else None
+                ),
+                w13_input_scale=layer.w13_input_scale if not self.block_quant else None,
+                output1_scales_scalar=(
+                    getattr(layer, "output1_scales_scalar", None)
+                    if not self.block_quant
+                    else None
+                ),
+                output1_scales_gate_scalar=(
+                    getattr(layer, "output1_scales_gate_scalar", None)
+                    if not self.block_quant
+                    else None
+                ),
+                output2_scales_scalar=(
+                    getattr(layer, "output2_scales_scalar", None)
+                    if not self.block_quant
+                    else None
+                ),
+            )
         elif self.runner.runner_backend.is_triton():
             quant_info = TritonMoeQuantInfo(
                 w13_weight=layer.w13_weight,
@@ -1316,123 +1291,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
 
         self._cutlass_buffers_ready = True
 
-    def apply_with_router_logits(
-        self,
-        layer: torch.nn.Module,
-        dispatch_output: StandardDispatchOutput,
-    ) -> torch.Tensor:
-        x = dispatch_output.hidden_states
-        topk_output = dispatch_output.topk_output
-
-        activation = self.moe_runner_config.activation
-        routed_scaling_factor = self.moe_runner_config.routed_scaling_factor
-
-        from flashinfer.fused_moe import (
-            trtllm_fp8_block_scale_moe,
-            trtllm_fp8_per_tensor_scale_moe,
-        )
-
-        from sglang.srt.layers.moe.topk import TopKOutputChecker
-        from sglang.srt.layers.moe.utils import RoutingMethodType
-
-        assert TopKOutputChecker.format_is_bypassed(topk_output)
-        router_logits = topk_output.router_logits
-        topk_config = topk_output.topk_config
-        assert (
-            activation == "silu"
-        ), "Only silu is supported for flashinfer blockscale fp8 moe"
-
-        if self.block_quant:
-            a_q, a_sf = per_token_group_quant_fp8(
-                x, self.quant_config.weight_block_size[1]
-            )
-            # NOTE: scales of hidden states have to be transposed!
-            a_sf_t = a_sf.t().contiguous()
-        else:
-            a_q, _ = scaled_fp8_quant(x, layer.w13_input_scale)
-
-        correction_bias = (
-            None
-            if topk_config.correction_bias is None
-            else topk_config.correction_bias.to(x.dtype)
-        )
-
-        routing_method_type = getattr(
-            layer, "routing_method_type", RoutingMethodType.DeepSeekV3
-        )
-
-        with use_symmetric_memory(
-            get_tp_group(), disabled=not is_allocation_symmetric()
-        ):
-
-            if self.block_quant:
-                # FIXME: there is a bug in the trtllm_fp8_block_scale_moe.
-                # It ignored the `output`` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325
-                # so we put the whole function under the ``use_symmetric_memory`` context manager.
-                # If the bug is fixed, we can only put the output tensor allocation under the context manager.
-                return trtllm_fp8_block_scale_moe(
-                    routing_logits=(
-                        router_logits.to(torch.float32)
-                        if routing_method_type == RoutingMethodType.DeepSeekV3
-                        else router_logits
-                    ),
-                    routing_bias=correction_bias,
-                    hidden_states=a_q,
-                    hidden_states_scale=a_sf_t,
-                    gemm1_weights=layer.w13_weight,
-                    gemm1_weights_scale=layer.w13_weight_scale_inv,
-                    gemm2_weights=layer.w2_weight,
-                    gemm2_weights_scale=layer.w2_weight_scale_inv,
-                    num_experts=layer.num_experts,
-                    top_k=topk_config.top_k,
-                    n_group=topk_config.num_expert_group,
-                    topk_group=topk_config.topk_group,
-                    intermediate_size=layer.w2_weight.shape[2],
-                    local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
-                    local_num_experts=layer.num_local_experts,
-                    routed_scaling_factor=(
-                        routed_scaling_factor
-                        if routed_scaling_factor is not None
-                        else 1.0
-                    ),
-                    tile_tokens_dim=None,
-                    routing_method_type=routing_method_type,
-                    use_shuffled_weight=False,
-                    tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
-                )
-            else:
-                routing_bias_cast = (
-                    None
-                    if correction_bias is None
-                    else correction_bias.to(torch.bfloat16)
-                )
-
-                return trtllm_fp8_per_tensor_scale_moe(
-                    routing_logits=router_logits.to(torch.bfloat16),
-                    routing_bias=routing_bias_cast,
-                    hidden_states=a_q,
-                    gemm1_weights=layer.w13_weight,
-                    output1_scales_scalar=layer.output1_scales_scalar,
-                    output1_scales_gate_scalar=layer.output1_scales_gate_scalar,
-                    gemm2_weights=layer.w2_weight,
-                    output2_scales_scalar=layer.output2_scales_scalar,
-                    num_experts=layer.num_experts,
-                    top_k=topk_config.top_k,
-                    n_group=topk_config.num_expert_group,
-                    topk_group=topk_config.topk_group,
-                    intermediate_size=layer.w2_weight.shape[2],
-                    local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
-                    local_num_experts=layer.num_local_experts,
-                    routed_scaling_factor=(
-                        routed_scaling_factor
-                        if routed_scaling_factor is not None
-                        else 1.0
-                    ),
-                    use_routing_scales_on_input=False,
-                    routing_method_type=routing_method_type,
-                    tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
-                )
-
     def maybe_apply_hip_fused_experts(
         self,
         layer: torch.nn.Module,