tpu-inference 0.11.1.dev202512030818__py3-none-any.whl → 0.13.0rc2.post7__py3-none-any.whl

tpu_inference/layers/vllm/quantization/mxfp4.py

@@ -1,4 +1,18 @@
-from typing import Callable, Optional, Union
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Union
 
 import jax
 import jax.numpy as jnp
@@ -10,7 +24,7 @@ from torchax.interop import jax_view, torch_view
 from torchax.ops.mappings import t2j
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEConfig, FusedMoEQuantConfig, biased_moe_quant_config)
+    FusedMoEConfig, FusedMoEQuantConfig, mxfp4_w4a16_moe_quant_config)
 from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
                                                         FusedMoEMethodBase)
 from vllm.model_executor.layers.linear import LinearBase
@@ -28,44 +42,22 @@ from tpu_inference import envs
 from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe
 from tpu_inference.layers.common.quant_methods import (MXFP4,
                                                        get_tpu_quant_method)
-from tpu_inference.layers.vllm.fused_moe import fused_moe_func_padded
+from tpu_inference.layers.common.quantization import (
+    dequantize_tensor_from_mxfp4_packed, quantize_tensor)
+from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.vllm.fused_moe import fused_moe_func
 from tpu_inference.layers.vllm.linear_common import \
     reorder_concatenated_tensor_for_sharding
 from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
 from tpu_inference.layers.vllm.quantization.unquantized import \
     VllmUnquantizedLinearMethod
+from tpu_inference.utils import get_mesh_shape_product
 
-MXFP4_BLOCK_SIZE = 32
+REQUANTIZED_BLOCK_SIZE = 512
 
 P = PartitionSpec
-logger = init_logger(__name__)
-
-
-# TODO(kyuyeunk): Move these functions into a common utility file.
-def u8_unpack_e2m1(u8_packed_e2m1: jax.Array) -> jax.Array:
-    assert u8_packed_e2m1.dtype == jnp.uint8
-    e2m1 = jax.lax.bitcast_convert_type(u8_packed_e2m1, jnp.float4_e2m1fn)
-    # bitcast creates one more dimension that splits 8 bits into two e2m1.
-    # we flatten them with the last dim.
-    return jnp.reshape(e2m1, e2m1.shape[:-2] + (-1, ))
-
-
-def e8m0_to_fp32(u8: jax.Array) -> jax.Array:
-    e8_finfo = jnp.finfo(jnp.float8_e8m0fnu)
-    exponents = u8.astype(jnp.int32) + e8_finfo.minexp
-    ones = jnp.ones_like(u8, dtype=jnp.float32)
-    return jnp.ldexp(ones, exponents)
 
-
-def dequantize_block_weight(weight: jax.Array,
-                            scale: jax.Array,
-                            block_size: int,
-                            out_dtype: jnp.dtype = jnp.bfloat16) -> jax.Array:
-    orig_shape = weight.shape
-    weight_block = weight.reshape(orig_shape[:-1] + (-1, block_size))
-    weight_dequantized = weight_block.astype(jnp.float32) * jnp.expand_dims(
-        scale, -1)
-    return weight_dequantized.reshape(orig_shape).astype(out_dtype)
+logger = init_logger(__name__)
 
 
 @register_quantization_config(get_tpu_quant_method(MXFP4))
@@ -87,9 +79,6 @@ class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
                     fused_mapping=self.packed_modules_mapping,
             ):
                 return VllmUnquantizedLinearMethod(linear_config)
-            # TODO: Add support for MXFP4 Linear Method.
-            # MXFP4 LinearMethod is available in AMD-Quark, refer to that
-            # implementation if you are interested in enabling MXFP4 here.
             logger.warning_once(
                 "MXFP4 linear layer is not implemented - falling back to "
                 "UnquantizedLinearMethod.")
@@ -98,7 +87,6 @@ class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
             moe_config = self.get_moe_config(layer)
             return VllmMxfp4MoEMethod(moe_config, self.mesh)
         elif isinstance(layer, Attention):
-            # TODO: Add support for MXFP4 Attention.
             logger.warning_once("MXFP4 attention layer is not implemented. "
                                 "Skipping quantization for this layer.")
         return None
@@ -117,225 +105,306 @@ class VllmMxfp4MoEMethod(Mxfp4MoEMethod):
         self.mxfp4_backend = Mxfp4Backend.TRITON
 
         self.mesh = mesh
-        self.use_kernel = envs.USE_MOE_EP_KERNEL
+        self.use_kernel = envs.USE_MOE_EP_KERNEL and moe.use_ep
         self.ep_axis_name = ep_axis_name
         # TODO: Use autotune table once we have it.
         self.block_size = {
-            "bt": 64,
+            "bt": 256,
             "bf": 1024,
-            "bd1": 1536,
-            "bd2": 1536,
-            "btc": 64,
+            "bd1": 1024,
+            "bd2": 1024,
+            "btc": 256,
             "bfc": 1024,
-            "bd1c": 1536,
-            "bd2c": 1536,
+            "bd1c": 1024,
+            "bd2c": 1024,
         }
 
     def get_fused_moe_quant_config(
             self, layer: torch.nn.Module) -> FusedMoEQuantConfig | None:
-        # Because we have dequantized weights, we only need biased moe config.
-        # TODO(kyuyeunk): Add native support for MXFP4.
-        return biased_moe_quant_config(
-            layer.w13_bias,
-            layer.w2_bias,
+        return mxfp4_w4a16_moe_quant_config(
+            w1_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            w1_bias=layer.w13_bias,
+            w2_bias=layer.w2_bias,
         )
 
     def process_weights_after_loading(self, layer: torch.nn.Module):
         assert isinstance(layer, FusedMoE)
         assert layer.moe_config.has_bias, "mxfp4 quantization alwyas use bias."
 
-        w13_weight = u8_unpack_e2m1(t2j(layer.w13_weight, use_dlpack=False))
-        w13_weight_scale = e8m0_to_fp32(
-            t2j(layer.w13_weight_scale, use_dlpack=False))
+        w13_weight = t2j(layer.w13_weight, use_dlpack=False)
+        w13_weight_scale = t2j(layer.w13_weight_scale, use_dlpack=False)
         w13_bias = t2j(layer.w13_bias, use_dlpack=False)
 
-        w2_weight = u8_unpack_e2m1(t2j(layer.w2_weight, use_dlpack=False))
-        w2_weight_scale = e8m0_to_fp32(
-            t2j(layer.w2_weight_scale, use_dlpack=False))
+        w2_weight = t2j(layer.w2_weight, use_dlpack=False)
+        w2_weight_scale = t2j(layer.w2_weight_scale, use_dlpack=False)
         w2_bias = t2j(layer.w2_bias, use_dlpack=False)
 
-        # We dequantize fp4 weights into bf16.
-        # TODO(kyuyeunk): Add native support for MXFP4.
-        w13_weight = dequantize_block_weight(w13_weight, w13_weight_scale,
-                                             MXFP4_BLOCK_SIZE, jnp.bfloat16)
-        w2_weight = dequantize_block_weight(w2_weight, w2_weight_scale,
-                                            MXFP4_BLOCK_SIZE, jnp.bfloat16)
-
-        # Because we have dequantized weights, scales are not used anymore.
-        delattr(layer, "w13_weight_scale")
-        delattr(layer, "w2_weight_scale")
-
-        if layer.activation == "swigluoai":
-            # When using swigluoai, vLLM splits gmm output in a interleaved way.
-            # However, interleaved split is not performant on TPU. Therefore,
-            # we preprocess the weight so that splitting gmm output by middle
-            # can still get the same result.
-            w1_weight = w13_weight[:, ::2, :]
-            w3_weight = w13_weight[:, 1::2, :]
-            w13_weight = jnp.concat([w1_weight, w3_weight], axis=1)
-
-            w1_bias = w13_bias[:, ::2]
-            w3_bias = w13_bias[:, 1::2]
-            w13_bias = jnp.concat([w1_bias, w3_bias], axis=1)
-
-        if self.use_kernel and layer.use_ep:
-            # Kernel expects:
-            # w13: (num_experts, 2, hidden_size, intermediate_size)
-            # w2: (num_experts, intermediate_size, hidden_size)
-            # Current format:
-            # w13_weight: (num_experts, 2*intermediate_size, hidden_size)
-            # w2_weight: (num_experts, hidden_size, intermediate_size)
-            num_experts = w13_weight.shape[0]
-            intermediate_size = w13_weight.shape[1] // 2
-            hidden_size = w13_weight.shape[2]
-
-            # Reshape and transpose w13_weight to (num_experts, 2, hidden_size, intermediate_size)
-            w13_reshaped = w13_weight.reshape(num_experts, 2,
-                                              intermediate_size, hidden_size)
-            w13_weight_transposed = jnp.transpose(w13_reshaped, (0, 1, 3, 2))
-
-            # Transpose w2_weight to (num_experts, intermediate_size, hidden_size)
-            w2_weight_transposed = jnp.transpose(w2_weight, (0, 2, 1))
-
-            # Apply EP sharding
-            w13_weight = jax.device_put(
-                w13_weight_transposed,
-                Format(Layout((0, 1, 2, 3)),
-                       NamedSharding(self.mesh, P("model", None, None, None))))
-            w2_weight = jax.device_put(
-                w2_weight_transposed,
-                Format(Layout((0, 1, 2)),
-                       NamedSharding(self.mesh, P("model", None, None))))
-
-            if self.moe.has_bias:
-                w13_bias = w13_bias.reshape(num_experts, 2, intermediate_size)
+        # Wrap functions in jit to speedup requantization.
+        @jax.jit
+        def wrapper(w13_weight, w13_weight_scale, w13_bias, w2_weight,
+                    w2_weight_scale, w2_bias):
+            # Dequantize fp4 weights into fp32.
+            w13_weight = dequantize_tensor_from_mxfp4_packed(
+                w13_weight, w13_weight_scale, 2)
+            w2_weight = dequantize_tensor_from_mxfp4_packed(
+                w2_weight, w2_weight_scale, 2)
+
+            num_experts, orig_hidden_size, orig_intermediate_size = w2_weight.shape
+
+            # Requantize the weights into TPU friendly block size.
+            w13_weight, w13_weight_scale = quantize_tensor(
+                jnp.float4_e2m1fn, w13_weight, 2, REQUANTIZED_BLOCK_SIZE, True)
+            w2_weight, w2_weight_scale = quantize_tensor(
+                jnp.float4_e2m1fn, w2_weight, 2, REQUANTIZED_BLOCK_SIZE, True)
+
+            intermediate_size = w2_weight.shape[-1]
+            hidden_size = w13_weight.shape[-1]
+
+            # Dims may have been padded to align with subchannel size during
+            # quantization. We pad the corresponding dim on other weight.
+            # NOTE: We perform padding after quantization as padding value can
+            # affect quantization numerics.
+            intermediate_padding_size = 2 * (intermediate_size -
+                                             orig_intermediate_size)
+            w13_weight = jnp.pad(w13_weight,
+                                 ((0, 0), (0, intermediate_padding_size),
+                                  (0, 0)))
+            w13_weight_scale = jnp.pad(w13_weight_scale,
+                                       ((0, 0), (0, intermediate_padding_size),
+                                        (0, 0)))
+            w13_bias = jnp.pad(w13_bias,
+                               ((0, 0), (0, intermediate_padding_size)))
+
+            hidden_padding_size = hidden_size - orig_hidden_size
+            w2_weight = jnp.pad(w2_weight,
+                                ((0, 0), (0, hidden_padding_size), (0, 0)))
+            w2_weight_scale = jnp.pad(w2_weight_scale,
+                                      ((0, 0), (0, hidden_padding_size),
+                                       (0, 0)))
+            w2_bias = jnp.pad(w2_bias, ((0, 0), (0, hidden_padding_size)))
+
+            if layer.activation == "swigluoai":
+                # When using swigluoai, vLLM splits gmm output in a interleaved way.
+                # However, interleaved split is not performant on TPU. Therefore,
+                # we preprocess the weight so that splitting gmm output by middle
+                # can still get the same result.
+                w1_weight = w13_weight[:, ::2, :]
+                w3_weight = w13_weight[:, 1::2, :]
+                w13_weight = jnp.concat([w1_weight, w3_weight], axis=1)
+
+                w1_weight_scale = w13_weight_scale[:, ::2, :]
+                w3_weight_scale = w13_weight_scale[:, 1::2, :]
+                w13_weight_scale = jnp.concat(
+                    [w1_weight_scale, w3_weight_scale], axis=1)
+
+                w1_bias = w13_bias[:, ::2]
+                w3_bias = w13_bias[:, 1::2]
+                w13_bias = jnp.concat([w1_bias, w3_bias], axis=1)
+
+            if self.use_kernel:
+                # Kernel expects:
+                # w13: (num_experts, 2, hidden_size, intermediate_size)
+                # w2: (num_experts, intermediate_size, hidden_size)
+                # Current format:
+                # w13_weight: (num_experts, 2*intermediate_size, hidden_size)
+                # w2_weight: (num_experts, hidden_size, intermediate_size)
+
+                w13_weight = w13_weight.reshape(num_experts, 2,
+                                                intermediate_size, hidden_size)
+
+                w13_weight_scale = w13_weight_scale.reshape(
+                    num_experts, 2, intermediate_size, 1, -1)
+                w2_weight_scale = w2_weight_scale.reshape(
+                    num_experts, hidden_size, 1, -1)
+
+                w13_bias = w13_bias.astype(jnp.float32).reshape(
+                    num_experts, 2, 1, intermediate_size)
+                w2_bias = w2_bias.astype(jnp.float32).reshape(
+                    num_experts, 1, hidden_size)
+
+                # Transpose non-constracting dim to right most dim
+                w13_weight = jnp.swapaxes(w13_weight, 2, 3)
+                w2_weight = jnp.swapaxes(w2_weight, 1, 2)
+
+                w13_weight_scale = jnp.swapaxes(w13_weight_scale, 2, 4)
+                w2_weight_scale = jnp.swapaxes(w2_weight_scale, 1, 3)
 
                 # Apply EP sharding
-                w13_bias = jax.device_put(
-                    w13_bias,
-                    Format(Layout((0, 1, 2)),
-                           NamedSharding(self.mesh, P("model", None, None))))
-                w2_bias = jax.device_put(
-                    w2_bias,
-                    Format(Layout((0, 1)),
-                           NamedSharding(self.mesh, P("model", None))))
-
-        else:
-            if layer.use_ep:
-                w13_weight = jax.device_put(
-                    w13_weight,
-                    Format(Layout((0, 1, 2)),
-                           NamedSharding(self.mesh, P("model", None, None))))
-                w2_weight = jax.device_put(
-                    w2_weight,
-                    Format(Layout((0, 1, 2)),
-                           NamedSharding(self.mesh, P("model", None, None))))
-
-                w13_bias = jax.device_put(
-                    w13_bias,
-                    Format(Layout((0, 1)),
-                           NamedSharding(self.mesh, P("model", None))))
-                w2_bias = jax.device_put(
-                    w2_bias,
-                    Format(Layout((0, 1)),
-                           NamedSharding(self.mesh, P("model", None))))
-
+                ep_sharding = NamedSharding(self.mesh, P("model"))
+
+                w13_weight = jax.lax.with_sharding_constraint(
+                    w13_weight, Format(Layout((0, 1, 2, 3)), ep_sharding))
+                w2_weight = jax.lax.with_sharding_constraint(
+                    w2_weight, Format(Layout((0, 1, 2)), ep_sharding))
+
+                w13_weight_scale = jax.lax.with_sharding_constraint(
+                    w13_weight_scale,
+                    Format(Layout((0, 1, 2, 3, 4)), ep_sharding))
+                w2_weight_scale = jax.lax.with_sharding_constraint(
+                    w2_weight_scale, Format(Layout((0, 1, 2, 3)), ep_sharding))
+
+                w13_bias = jax.lax.with_sharding_constraint(
+                    w13_bias, Format(Layout((0, 1, 2, 3)), ep_sharding))
+                w2_bias = jax.lax.with_sharding_constraint(
+                    w2_bias, Format(Layout((0, 1, 2)), ep_sharding))
             else:
-                intermediate_size = w13_weight.shape[1] // 2
-                assert intermediate_size == w2_weight.shape[-1]
-                output_sizes = [intermediate_size, intermediate_size]
-                n_shards = self.mesh.shape["model"]
-                assert intermediate_size % n_shards == 0
-                w13_weight = reorder_concatenated_tensor_for_sharding(
-                    w13_weight, output_sizes, n_shards, dim=1)
-                w13_weight = jax.device_put(
-                    w13_weight,
-                    Format(Layout((0, 1, 2)),
-                           NamedSharding(self.mesh, P(None, "model", None))))
-                w2_weight = jax.device_put(
-                    w2_weight,
-                    Format(Layout((0, 1, 2)),
-                           NamedSharding(self.mesh, P(None, None, "model"))))
-
-                w13_bias = reorder_concatenated_tensor_for_sharding(
-                    w13_bias, output_sizes, n_shards, dim=1)
-                w13_bias = jax.device_put(
-                    w13_bias,
-                    Format(Layout((0, 1)),
-                           NamedSharding(self.mesh, P(None, "model"))))
-                w2_bias = jax.device_put(
-                    w2_bias,
-                    Format(Layout((0, 1)),
-                           NamedSharding(self.mesh, P(None, None))))
+                w13_weight_scale = jnp.swapaxes(w13_weight_scale, 1, 2)
+                w13_weight_scale = jnp.expand_dims(w13_weight_scale, 2)
+                w2_weight_scale = jnp.swapaxes(w2_weight_scale, 1, 2)
+                w2_weight_scale = jnp.expand_dims(w2_weight_scale, 2)
+
+                w13_bias = jnp.expand_dims(w13_bias, 1)
+                w2_bias = jnp.expand_dims(w2_bias, 1)
+
+                if layer.use_ep:
+                    ep_sharding = NamedSharding(self.mesh,
+                                                P(ShardingAxisName.EXPERT))
+
+                    w13_weight = jax.lax.with_sharding_constraint(
+                        w13_weight, ep_sharding)
+                    w2_weight = jax.lax.with_sharding_constraint(
+                        w2_weight, ep_sharding)
+
+                    w13_weight_scale = jax.lax.with_sharding_constraint(
+                        w13_weight_scale, ep_sharding)
+                    w2_weight_scale = jax.lax.with_sharding_constraint(
+                        w2_weight_scale, ep_sharding)
+
+                    w13_bias = jax.lax.with_sharding_constraint(
+                        w13_bias, ep_sharding)
+                    w2_bias = jax.lax.with_sharding_constraint(
+                        w2_bias, ep_sharding)
+
+                else:
+                    output_sizes = [intermediate_size, intermediate_size]
+                    n_shards = get_mesh_shape_product(
+                        self.mesh, ShardingAxisName.MLP_TENSOR)
+                    assert intermediate_size % n_shards == 0
+
+                    # Reorder w13 weights so that splitting w1 and w3 output
+                    # can happen locally without any collective operations.
+                    w13_weight = reorder_concatenated_tensor_for_sharding(
+                        w13_weight,
+                        output_sizes,
+                        n_shards,
+                        dim=1,
+                    )
+                    w13_weight_scale = reorder_concatenated_tensor_for_sharding(
+                        w13_weight_scale,
+                        output_sizes,
+                        n_shards,
+                        dim=3,
+                    )
+                    w13_bias = reorder_concatenated_tensor_for_sharding(
+                        w13_bias,
+                        output_sizes,
+                        n_shards,
+                        dim=2,
+                    )
+
+                    w13_weight = jax.lax.with_sharding_constraint(
+                        w13_weight,
+                        NamedSharding(
+                            self.mesh,
+                            P(None, ShardingAxisName.MLP_TENSOR, None)))
+                    w2_weight = jax.lax.with_sharding_constraint(
+                        w2_weight,
+                        NamedSharding(
+                            self.mesh,
+                            P(None, None, ShardingAxisName.MLP_TENSOR)))
+                    w13_weight_scale = jax.lax.with_sharding_constraint(
+                        w13_weight_scale,
+                        NamedSharding(
+                            self.mesh,
+                            P(None, None, None, ShardingAxisName.MLP_TENSOR)))
+                    w2_weight_scale = jax.lax.with_sharding_constraint(
+                        w2_weight_scale,
+                        NamedSharding(
+                            self.mesh,
+                            P(None, ShardingAxisName.MLP_TENSOR, None, None)))
+                    w13_bias = jax.lax.with_sharding_constraint(
+                        w13_bias,
+                        NamedSharding(
+                            self.mesh,
+                            P(None, None, ShardingAxisName.MLP_TENSOR)))
+                    w2_bias = jax.lax.with_sharding_constraint(
+                        w2_bias, NamedSharding(self.mesh, P(None, None, None)))
+
+            return w13_weight, w13_weight_scale, w13_bias, w2_weight, w2_weight_scale, w2_bias
+
+        w13_weight, w13_weight_scale, w13_bias, w2_weight, w2_weight_scale, w2_bias = wrapper(
+            w13_weight, w13_weight_scale, w13_bias, w2_weight, w2_weight_scale,
+            w2_bias)
 
         layer.w13_weight = Parameter(torch_view(w13_weight),
                                      requires_grad=False)
-        layer.w13_bias = Parameter(torch_view(w13_bias), requires_grad=False)
-
         layer.w2_weight = Parameter(torch_view(w2_weight), requires_grad=False)
-        layer.w2_bias = Parameter(torch_view(w2_bias), requires_grad=False)
 
-        pass
+        layer.w13_weight_scale = Parameter(torch_view(w13_weight_scale),
+                                           requires_grad=False)
+        layer.w2_weight_scale = Parameter(torch_view(w2_weight_scale),
+                                          requires_grad=False)
+
+        layer.w13_bias = Parameter(torch_view(w13_bias), requires_grad=False)
+        layer.w2_bias = Parameter(torch_view(w2_bias), requires_grad=False)
 
     def apply(
         self,
         layer: torch.nn.Module,
         x: torch.Tensor,
         router_logits: torch.Tensor,
-        top_k: int,
-        renormalize: bool,
-        use_grouped_topk: bool = False,
-        topk_group: Optional[int] = None,
-        num_expert_group: Optional[int] = None,
-        global_num_experts: int = -1,
-        expert_map: Optional[torch.Tensor] = None,
-        custom_routing_function: Optional[Callable] = None,
-        scoring_func: str = "softmax",
-        routed_scaling_factor: float = 1.0,
-        e_score_correction_bias: Optional[torch.Tensor] = None,
-        apply_router_weight_on_input: bool = False,
-        activation: str = "silu",
-        enable_eplb: bool = False,
-        expert_load_view: Optional[torch.Tensor] = None,
-        logical_to_physical_map: Optional[torch.Tensor] = None,
-        logical_replica_count: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         assert isinstance(layer, FusedMoE)
-        if scoring_func != "softmax":
+        if layer.scoring_func != "softmax":
             raise NotImplementedError(
                 "Only softmax is supported for scoring_func")
 
-        if self.use_kernel and layer.use_ep:
+        x = jax_view(x)
+        w13_weight = jax_view(layer.w13_weight)
+        w2_weight = jax_view(layer.w2_weight)
+        w13_weight_scale = jax_view(layer.w13_weight_scale)
+        w2_weight_scale = jax_view(layer.w2_weight_scale)
+        w13_bias = jax_view(layer.w13_bias)
+        w2_bias = jax_view(layer.w2_bias)
+        gating_output = jax_view(router_logits)
+
+        if self.use_kernel:
+            actual_hidden_size = x.shape[-1]
+            padding_size = w13_weight.shape[-2] - actual_hidden_size
+            x = jnp.pad(x, ((0, 0), (0, padding_size)))
             output = fused_ep_moe(
                 mesh=self.mesh,
-                tokens=jax_view(x),
-                w1=jax_view(layer.w13_weight),
-                w2=jax_view(layer.w2_weight),
-                b1=jax_view(layer.w13_bias),
-                b2=jax_view(layer.w2_bias),
-                gating_output=jax_view(router_logits),
-                top_k=top_k,
+                tokens=x,
+                w1=w13_weight,
+                w2=w2_weight,
+                w1_scale=w13_weight_scale,
+                w2_scale=w2_weight_scale,
+                b1=w13_bias,
+                b2=w2_bias,
+                gating_output=gating_output,
+                subc_quant_wsz=REQUANTIZED_BLOCK_SIZE,
+                top_k=layer.top_k,
                 ep_axis_name=self.ep_axis_name,
-                renormalize_topk_logits=renormalize,
-                act_fn=activation,
+                renormalize_topk_logits=layer.renormalize,
+                act_fn=layer.activation,
                 **self.block_size,
-            )
+            )[:, :actual_hidden_size]
         else:
-            # Use the original implementation
-            output = fused_moe_func_padded(
-                jax_view(x),
-                jax_view(layer.w13_weight),
-                jax_view(layer.w2_weight),
-                jax_view(layer.w13_bias),
-                jax_view(layer.w2_bias),
-                jax_view(router_logits),
-                topk=top_k,
-                global_num_experts=global_num_experts,
-                renormalize=renormalize,
-                reduce_results=layer.reduce_results,
+            output = fused_moe_func(
+                hidden_states=x,
+                w1=w13_weight,
+                w2=w2_weight,
+                w1_scale=w13_weight_scale,
+                w2_scale=w2_weight_scale,
+                w1_bias=w13_bias,
+                w2_bias=w2_bias,
+                gating_output=gating_output,
+                topk=layer.top_k,
+                renormalize=layer.renormalize,
                 mesh=self.mesh,
                 use_ep=layer.use_ep,
-                activation=activation,
+                activation=layer.activation,
             )
 
         return torch_view(output)