tpu-inference 0.11.1.dev202511150811__py3-none-any.whl

tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

@@ -0,0 +1,208 @@
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+import torch
+from compressed_tensors.quantization import (QuantizationArgs,
+                                             QuantizationStrategy)
+from jax.sharding import NamedSharding, PartitionSpec
+from torchax.interop import jax_view, torch_view
+from torchax.ops.mappings import t2j
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import \
+    CompressedTensorsW8A8Fp8
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import \
+    per_tensor_dequantize
+
+from tpu_inference.layers.vllm.linear_common import (
+    sharded_quantized_matmul, slice_sharded_tensor_for_concatenation,
+    torch_to_jax_param)
+from tpu_inference.layers.vllm.quantization.common import JaxCommonLinearConfig
+
+P = PartitionSpec
+
+
+def requantize_with_max_scale(
+        weight: torch.Tensor, weight_scale: torch.Tensor,
+        logical_widths: list[int]) -> tuple[torch.Tensor, torch.Tensor]:
+    dtype = weight.dtype
+    dtype_info = torch.finfo(dtype)
+    maxval = float(dtype_info.max)
+    minval = float(dtype_info.min)
+
+    max_w_scale = weight_scale.max()
+
+    unfused_module_in_checkpoint = (weight_scale[-1]
+                                    > torch.finfo(torch.float8_e4m3fn).min)
+
+    # If unfused checkpoint, need requanize with the single scale.
+    if unfused_module_in_checkpoint:
+        start = 0
+        for idx, logical_width in enumerate(logical_widths):
+            # Skip any component with zero width.
+            if logical_width == 0:
+                continue
+            end = start + logical_width
+            weight_dq = per_tensor_dequantize(weight[start:end, :],
+                                              weight_scale[idx])
+            weight_q = weight_dq / max_w_scale
+            weight[start:end, :] = weight_q.clamp(minval, maxval).to(dtype)
+            start = end
+
+    return max_w_scale, weight
+
+
+class VllmCompressedTensorsW8A8Fp8(CompressedTensorsW8A8Fp8):
+
+    def __init__(
+        self,
+        weight_quant: QuantizationArgs,
+        is_static_input_scheme: bool,
+        jax_config: JaxCommonLinearConfig,
+    ):
+        super().__init__(weight_quant, is_static_input_scheme)
+
+        self.jax_config = jax_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight
+        weight_scale = layer.weight_scale
+
+        if self.is_static_input_scheme:
+            # In static quant, all input_scales share the same value.
+            assert layer.input_scale.min() == layer.input_scale.max()
+            input_scale_first = layer.input_scale[0]
+
+            input_scale = jax.device_put(
+                t2j(input_scale_first, use_dlpack=False),
+                NamedSharding(self.jax_config.mesh, P()))
+            input_scale = torch.nn.Parameter(torch_view(input_scale),
+                                             requires_grad=False)
+            delattr(layer, "input_scale")
+            layer.input_scale = input_scale
+
+            # TODO(kyuyeunk): Investigate performance gain from merging scales.
+            # By merging input and weight scales, we reduce the number of muls
+            # required for dequantization from 2 (for each scales) to 1.
+            # weight_scale *= input_scale_first
+
+        if self.strategy == QuantizationStrategy.TENSOR:
+            weight_scale, weight = requantize_with_max_scale(
+                weight, weight_scale, self.jax_config.output_sizes)
+            weight_scale = jax.device_put(
+                t2j(weight_scale, use_dlpack=False),
+                NamedSharding(self.jax_config.mesh, P()))
+            weight_scale = torch.nn.Parameter(torch_view(weight_scale),
+                                              requires_grad=False)
+        else:
+            weight_scale = weight_scale.squeeze(-1)
+            weight_scale = torch_to_jax_param(
+                weight_scale,
+                NamedSharding(self.jax_config.mesh,
+                              self.jax_config.bias_sharding),
+                self.jax_config.output_sizes, self.jax_config.n_shards,
+                self.jax_config.fuse_matmuls)
+        delattr(layer, "weight_scale")
+        layer.weight_scale = weight_scale
+
+        weight = torch_to_jax_param(
+            layer.weight,
+            NamedSharding(self.jax_config.mesh,
+                          self.jax_config.weight_sharding),
+            self.jax_config.output_sizes, self.jax_config.n_shards,
+            self.jax_config.fuse_matmuls)
+        delattr(layer, "weight")
+        layer.weight = weight
+
+        if layer.bias is not None:
+            bias = torch_to_jax_param(
+                layer.bias,
+                NamedSharding(self.jax_config.mesh,
+                              self.jax_config.bias_sharding),
+                self.jax_config.output_sizes, self.jax_config.n_shards,
+                self.jax_config.fuse_matmuls)
+            delattr(layer, "bias")
+            layer.bias = bias
+
+    def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor,
+                      bias: Optional[torch.Tensor]) -> torch.Tensor:
+        with jax.named_scope(layer._get_name()):
+            if self.jax_config.fuse_matmuls:
+                return self._apply_fused(layer, x, bias)
+            else:
+                return self._apply_split(layer, x, bias)
+
+    def _apply_fused(self, layer: torch.nn.Module, x: torch.Tensor,
+                     bias: Optional[torch.Tensor]) -> torch.Tensor:
+        x_jax = jax_view(x)
+        weight_jax = jax_view(layer.weight)
+        weight_scale_jax = jax_view(layer.weight_scale)
+
+        if self.is_static_input_scheme:
+            # TODO(kyuyeunk): Add kernel support for static quant
+            input_scale = jax_view(layer.input_scale)
+            dtype_info = jnp.finfo(weight_jax.dtype)
+            maxval = float(dtype_info.max)
+            minval = float(dtype_info.min)
+            x_q = jnp.clip(x_jax / input_scale.astype(x_jax.dtype), minval,
+                           maxval).astype(weight_jax.dtype)
+
+            outs = jax.lax.dot_general(
+                x_q,
+                weight_jax,
+                (((1, ), (1, )), ((), ())),
+                preferred_element_type=jnp.float32,
+            )
+            outs *= weight_scale_jax
+            outs = outs.astype(x_jax.dtype)
+        else:
+            outs = sharded_quantized_matmul(x_jax, weight_jax,
+                                            weight_scale_jax,
+                                            self.jax_config.mesh,
+                                            self.jax_config.weight_sharding)
+
+        if bias is not None and not layer.skip_bias_add:
+            outs += jax_view(bias)
+        outs = slice_sharded_tensor_for_concatenation(
+            outs, self.jax_config.output_sizes, self.jax_config.n_shards)
+        return torch_view(jnp.concatenate(outs, axis=-1))
+
+    def _apply_split(self, layer: torch.nn.Module, x: torch.Tensor,
+                     bias: Optional[torch.Tensor]) -> torch.Tensor:
+        assert isinstance(layer.weight, torch.nn.ParameterList)
+
+        x_jax = jax_view(x)
+        outs = []
+        for i, (weight, weight_scale) in enumerate(
+                zip(layer.weight, layer.weight_scale)):
+            weight_jax = jax_view(weight)
+            weight_scale_jax = jax_view(weight_scale)
+
+            if self.is_static_input_scheme:
+                # TODO(kyuyeunk): Add kernel support for static quant
+                input_scale = jax_view(layer.input_scale)
+                dtype_info = jnp.finfo(weight_jax.dtype)
+                maxval = float(dtype_info.max)
+                minval = float(dtype_info.min)
+                x_q = jnp.clip(x_jax / input_scale.astype(x_jax.dtype), minval,
+                               maxval).astype(weight_jax.dtype)
+
+                out = jax.lax.dot_general(
+                    x_q,
+                    weight_jax,
+                    (((1, ), (1, )), ((), ())),
+                    preferred_element_type=jnp.float32,
+                )
+                # TODO(kyuyeunk): Investigate performance gain from merging scales.
+                # out *= weight_scale_jax
+                out *= weight_scale_jax * input_scale
+                out = out.astype(x_jax.dtype)
+            else:
+                out = sharded_quantized_matmul(x_jax, weight_jax,
+                                               weight_scale_jax,
+                                               self.jax_config.mesh,
+                                               self.jax_config.weight_sharding)
+
+            if bias is not None and not layer.skip_bias_add:
+                out += jax_view(bias[i])
+            outs.append(out)
+        return torch_view(jnp.concatenate(outs, axis=-1))

tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py

@@ -0,0 +1,136 @@
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+import torch
+from compressed_tensors.quantization import QuantizationStrategy
+from jax.sharding import NamedSharding, PartitionSpec
+from torchax.interop import jax_view, torch_view
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_int8 import \
+    CompressedTensorsW8A8Int8
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import \
+    convert_to_channelwise
+
+from tpu_inference.layers.vllm.linear_common import (
+    sharded_quantized_matmul, slice_sharded_tensor_for_concatenation,
+    torch_to_jax_param)
+from tpu_inference.layers.vllm.quantization.common import JaxCommonLinearConfig
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+class VllmCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
+
+    def __init__(self, strategy: str, is_static_input_scheme: bool,
+                 input_symmetric: bool, jax_config: JaxCommonLinearConfig):
+        super().__init__(strategy, is_static_input_scheme, input_symmetric)
+
+        self.jax_config = jax_config
+        self.is_channelwise = (self.strategy == QuantizationStrategy.CHANNEL),
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = torch_to_jax_param(
+            layer.weight,
+            NamedSharding(self.jax_config.mesh,
+                          self.jax_config.weight_sharding),
+            self.jax_config.output_sizes,
+            self.jax_config.n_shards,
+            self.jax_config.fuse_matmuls,
+        )
+        delattr(layer, "weight")
+        layer.weight = weight
+
+        weight_scale = layer.weight_scale
+        is_fused_module = len(layer.logical_widths) > 1
+        if is_fused_module and not self.is_channelwise:
+            weight_scale = convert_to_channelwise(weight_scale,
+                                                  layer.logical_widths)
+        weight_scale = weight_scale.squeeze(-1)
+
+        weight_scale = torch_to_jax_param(
+            weight_scale,
+            NamedSharding(self.jax_config.mesh, self.jax_config.bias_sharding),
+            self.jax_config.output_sizes,
+            self.jax_config.n_shards,
+            self.jax_config.fuse_matmuls,
+        )
+        delattr(layer, "weight_scale")
+        layer.weight_scale = weight_scale
+
+        if layer.bias is not None and not layer.skip_bias_add:
+            if layer.return_bias:
+                logger.warning_once("Bias might return incorrect value.")
+
+            bias = torch_to_jax_param(
+                layer.bias,
+                NamedSharding(self.jax_config.mesh,
+                              self.jax_config.bias_sharding),
+                self.jax_config.output_sizes,
+                self.jax_config.n_shards,
+                self.jax_config.fuse_matmuls,
+            )
+            delattr(layer, "bias")
+            layer.bias = bias
+
+        # TODO(kyuyeunk): Support static range input quantization.
+        assert getattr(layer, "input_scale", None) is None
+        assert getattr(layer, "input_zero_point", None) is None
+        assert getattr(layer, "azp_adj", None) is None
+
+    def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor,
+                      bias: Optional[torch.Tensor]) -> torch.Tensor:
+        with jax.named_scope(layer._get_name()):
+            if self.jax_config.fuse_matmuls:
+                out = self._apply_fused(layer, x, bias)
+            else:
+                out = self._apply_split(layer, x, bias)
+
+        return out
+
+    def _apply_fused(self, layer: torch.nn.Module, x: torch.Tensor,
+                     bias: Optional[torch.Tensor]) -> torch.Tensor:
+        x_jax = jax_view(x)
+        weight_jax = jax_view(layer.weight)
+        weight_scale_jax = jax_view(layer.weight_scale)
+
+        outs = sharded_quantized_matmul(
+            x_jax,
+            weight_jax,
+            weight_scale_jax,
+            self.jax_config.mesh,
+            self.jax_config.weight_sharding,
+        )
+        if bias is not None and not layer.skip_bias_add:
+            outs += jax_view(bias)
+
+        outs = slice_sharded_tensor_for_concatenation(
+            outs, self.jax_config.output_sizes, self.jax_config.n_shards)
+        out = jnp.concatenate(outs, axis=-1)
+        return torch_view(out)
+
+    def _apply_split(self, layer: torch.nn.Module, x: torch.Tensor,
+                     bias: Optional[torch.Tensor]) -> torch.Tensor:
+        assert isinstance(layer.weight, torch.nn.ParameterList)
+
+        x_jax = jax_view(x)
+        outs = []
+        for i, (weight, weight_scale) in enumerate(
+                zip(layer.weight, layer.weight_scale)):
+            weight_jax = jax_view(weight)
+            weight_scale_jax = jax_view(weight_scale)
+
+            out = sharded_quantized_matmul(
+                x_jax,
+                weight_jax,
+                weight_scale_jax,
+                self.jax_config.mesh,
+                self.jax_config.weight_sharding,
+            )
+            if bias is not None and not layer.skip_bias_add:
+                out += jax_view(bias[i])
+
+            outs.append(out)
+        out = jnp.concatenate(outs, axis=-1)
+        return torch_view(out)

tpu_inference/layers/vllm/quantization/mxfp4.py

@@ -0,0 +1,266 @@
+from typing import Callable, Optional, Union
+
+import jax
+import jax.numpy as jnp
+import torch
+from jax.experimental.layout import Format, Layout
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torch.nn.parameter import Parameter
+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)
+from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
+                                                        FusedMoEMethodBase)
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.quantization import \
+    register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import \
+    QuantizeMethodBase
+from vllm.model_executor.layers.quantization.mxfp4 import (Mxfp4Backend,
+                                                           Mxfp4Config,
+                                                           Mxfp4MoEMethod)
+from vllm.model_executor.layers.quantization.utils.quant_utils import \
+    is_layer_skipped
+
+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.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
+
+MXFP4_BLOCK_SIZE = 32
+
+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)
+
+
+@register_quantization_config(get_tpu_quant_method(MXFP4))
+class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls):
+        return MXFP4
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional["QuantizeMethodBase"]:
+        from vllm.attention.layer import Attention  # Avoid circular import
+
+        if isinstance(layer, LinearBase):
+            linear_config = self.get_linear_config(layer)
+            if self.ignored_layers and is_layer_skipped(
+                    prefix=prefix,
+                    ignored_layers=self.ignored_layers,
+                    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.")
+            return VllmUnquantizedLinearMethod(linear_config)
+        elif isinstance(layer, FusedMoE):
+            return VllmMxfp4MoEMethod(layer.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
+
+
+class VllmMxfp4MoEMethod(Mxfp4MoEMethod):
+
+    def __init__(self, moe: FusedMoEConfig, mesh: Mesh):
+        FusedMoEMethodBase.__init__(self, moe)
+
+        # We piggyback on triton implementation as it applies minimal hardware
+        # specific post processing to the weights.
+        self.mxfp4_backend = Mxfp4Backend.TRITON
+        self.mesh = mesh
+
+    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,
+        )
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        assert isinstance(layer, FusedMoE)
+
+        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_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_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)
+
+        # TODO(kyuyeunk): Add weight processing logic for the new kernel.
+        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))))
+
+        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))))
+
+        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
+
+    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":
+            raise NotImplementedError(
+                "Only softmax is supported for scoring_func")
+
+        # 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) if self.moe.has_bias else None,
+            jax_view(layer.w2_bias) if self.moe.has_bias else None,
+            jax_view(router_logits),
+            topk=top_k,
+            global_num_experts=global_num_experts,
+            renormalize=renormalize,
+            reduce_results=layer.reduce_results,
+            mesh=self.mesh,
+            use_ep=layer.use_ep,
+            activation=activation,
+        )
+
+        return torch_view(output)