tpu-inference 0.12.0.dev20251222__py3-none-any.whl

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

@@ -0,0 +1,150 @@
+# 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
+
+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/fp8.py

@@ -0,0 +1,118 @@
+# 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 torch
+from jax.sharding import PartitionSpec
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
+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.fp8 import (Fp8Config,
+                                                         Fp8LinearMethod)
+from vllm.model_executor.layers.quantization.utils.quant_utils import \
+    is_layer_skipped
+
+from tpu_inference.layers.common.quant_methods import FP8, get_tpu_quant_method
+from tpu_inference.layers.vllm.quantization.common import (
+    JaxCommonConfig, JaxCommonLinearConfig)
+from tpu_inference.layers.vllm.quantization.unquantized import \
+    VllmUnquantizedLinearMethod
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+@register_quantization_config(get_tpu_quant_method(FP8))
+class VllmFp8Config(Fp8Config, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls):
+        return FP8
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        return [torch.bfloat16]
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[Union["LinearMethodBase", "QuantizeMethodBase"]]:
+        if isinstance(layer, LinearBase):
+            linear_config = self.get_linear_config(layer)
+            if is_layer_skipped(prefix, self.ignored_layers):
+                return VllmUnquantizedLinearMethod(linear_config)
+            return VllmFp8LinearMethod(self, linear_config)
+        elif isinstance(layer, FusedMoE):
+            raise NotImplementedError(
+                "FP8 FusedMoE is currently not supported in torchax-jax")
+        return None
+
+
+class VllmFp8LinearMethod(Fp8LinearMethod):
+
+    def __init__(self, quant_config: VllmFp8Config,
+                 jax_config: JaxCommonLinearConfig):
+        super().__init__(quant_config)
+        self.jax_config = jax_config
+        self._configure_sharding()
+
+    def _configure_sharding(self) -> None:
+
+        raise NotImplementedError(
+            "Configure PartitionSpec for weight_sharding and scale_sharding "
+            "based on layer type (RowParallel/ColumnParallel)")
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        raise NotImplementedError(
+            "Convert layer.weight, layer.weight_scale, and optionally "
+            "layer.input_scale and layer.bias from torch tensors to JAX arrays "
+            "using torch_to_jax_param() with appropriate sharding")
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> 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] = None) -> torch.Tensor:
+
+        raise NotImplementedError(
+            "Implement single matmul for fused outputs: "
+            "quantize input to fp8, perform fp8 matmul with weight and scales, "
+            "dequantize output, and add bias if present")
+
+    def _apply_split(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        raise NotImplementedError(
+            "Implement separate matmuls per output partition: "
+            "split weight/scale by output_sizes, perform fp8 matmul for each, "
+            "concatenate results, and add bias if present")

tpu_inference/layers/vllm/quantization/mxfp4.py

@@ -0,0 +1,396 @@
+# 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
+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, mxfp4_w4a16_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 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.common.quantization import (
+    dequantize_tensor_from_mxfp4_packed, quantize_tensor)
+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
+
+REQUANTIZED_BLOCK_SIZE = 512
+
+P = PartitionSpec
+
+logger = init_logger(__name__)
+
+
+@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)
+            logger.warning_once(
+                "MXFP4 linear layer is not implemented - falling back to "
+                "UnquantizedLinearMethod.")
+            return VllmUnquantizedLinearMethod(linear_config)
+        elif isinstance(layer, FusedMoE):
+            moe_config = self.get_moe_config(layer)
+            return VllmMxfp4MoEMethod(moe_config, self.mesh)
+        elif isinstance(layer, 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,
+                 ep_axis_name: str = 'model'):
+        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
+        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": 256,
+            "bf": 1024,
+            "bd1": 1024,
+            "bd2": 1024,
+            "btc": 256,
+            "bfc": 1024,
+            "bd1c": 1024,
+            "bd2c": 1024,
+        }
+
+    def get_fused_moe_quant_config(
+            self, layer: torch.nn.Module) -> FusedMoEQuantConfig | None:
+        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 = 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 = 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)
+
+        # 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
+                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:
+                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("model"))
+
+                    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 = self.mesh.shape["model"]
+                    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, "model", None)))
+                    w2_weight = jax.lax.with_sharding_constraint(
+                        w2_weight,
+                        NamedSharding(self.mesh, P(None, None, "model")))
+                    w13_weight_scale = jax.lax.with_sharding_constraint(
+                        w13_weight_scale,
+                        NamedSharding(self.mesh, P(None, None, None, "model")))
+                    w2_weight_scale = jax.lax.with_sharding_constraint(
+                        w2_weight_scale,
+                        NamedSharding(self.mesh, P(None, "model", None, None)))
+                    w13_bias = jax.lax.with_sharding_constraint(
+                        w13_bias,
+                        NamedSharding(self.mesh, P(None, None, "model")))
+                    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.w2_weight = Parameter(torch_view(w2_weight), requires_grad=False)
+
+        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,
+    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        assert isinstance(layer, FusedMoE)
+        if layer.scoring_func != "softmax":
+            raise NotImplementedError(
+                "Only softmax is supported for scoring_func")
+
+        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=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=layer.renormalize,
+                act_fn=layer.activation,
+                **self.block_size,
+            )[:, :actual_hidden_size]
+        else:
+            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=layer.activation,
+            )
+
+        return torch_view(output)