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

tpu_inference/layers/vllm/quantization/__init__.py

@@ -0,0 +1,39 @@
+import copy
+
+from jax.sharding import Mesh
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.quantization.base_config import \
+    QuantizationConfig
+
+from tpu_inference.layers.common import quant_methods
+from tpu_inference.layers.vllm.quantization.awq import VllmAWQConfig
+from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
+from tpu_inference.layers.vllm.quantization.compressed_tensors.compressed_tensors import \
+    VllmCompressedTensorsConfig  # noqa: E501
+from tpu_inference.layers.vllm.quantization.mxfp4 import VllmMxfp4Config
+from tpu_inference.layers.vllm.quantization.unquantized import \
+    VllmUnquantizedConfig
+
+
+def get_tpu_quantization_config(vllm_config: VllmConfig,
+                                mesh: Mesh) -> QuantizationConfig:
+    model_config = copy.deepcopy(vllm_config.model_config)
+    # TODO(kyuyeunk): Add support for "tpu_int8".
+    method_to_config: dict[str, str] = {
+        None: VllmUnquantizedConfig,
+        quant_methods.COMPRESSED_TENSORS: VllmCompressedTensorsConfig,
+        quant_methods.AWQ: VllmAWQConfig,
+        quant_methods.MXFP4: VllmMxfp4Config,
+    }
+    if model_config.quantization not in method_to_config:
+        raise NotImplementedError(
+            f"{model_config.quantization} quantization method not supported."
+            f" Supported methods are {method_to_config.keys()}")
+    quant_config = method_to_config[model_config.quantization]
+    assert issubclass(quant_config, JaxCommonConfig)
+    quant_config.set_configs(vllm_config, mesh)
+
+    model_config.quantization = quant_methods.get_tpu_quant_method(
+        quant_config.get_name())
+    return VllmConfig.get_quantization_config(model_config,
+                                              vllm_config.load_config)

tpu_inference/layers/vllm/quantization/awq.py

@@ -0,0 +1,207 @@
+from typing import Optional, Union
+
+import jax
+import jax.numpy as jnp
+import torch
+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.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.awq import (AWQConfig,
+                                                         AWQLinearMethod)
+from vllm.model_executor.layers.quantization.base_config import \
+    QuantizeMethodBase
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    is_layer_skipped, unpack_quantized_values_into_int32)
+from vllm.scalar_type import scalar_types
+
+from tpu_inference.layers.common.quant_methods import AWQ, get_tpu_quant_method
+from tpu_inference.layers.vllm.linear_common import (
+    slice_sharded_tensor_for_concatenation, torch_to_jax_param)
+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(AWQ))
+class VllmAWQConfig(AWQConfig, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls):
+        return AWQ
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        # NOTE: AWQ checkpoint was quantized with float16. But on TPUs, using
+        # bfloat16 is signifcantly preferred over foat16. This might lead to
+        # some numeric output change.
+        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.modules_to_not_convert):
+                return VllmUnquantizedLinearMethod(linear_config)
+            return VllmAWQLinearMethod(self, linear_config)
+        elif isinstance(layer, FusedMoE):
+            raise NotImplementedError(
+                "AWQ FusedMoE is currently not supported in torchax-jax")
+        return None
+
+
+class VllmAWQLinearMethod(AWQLinearMethod):
+
+    def __init__(self, quant_config: VllmAWQConfig,
+                 jax_config: JaxCommonLinearConfig):
+        super().__init__(quant_config)
+        self.jax_config = jax_config
+
+        out_sharding, in_sharding = self.jax_config.weight_sharding[:]
+        self.jax_config.weight_sharding = P(in_sharding, None, out_sharding)
+        self.jax_config.scale_sharding = P(in_sharding, out_sharding)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        qweight = layer.qweight
+        qweight = unpack_awq_weight(qweight, qweight.packed_dim)
+
+        group_size = self.quant_config.group_size
+        # Reshape so that each qweight[i] were quantized with same scales[i].
+        qweight = qweight.reshape((-1, group_size, layer.output_size))
+        qweight = torch_to_jax_param(qweight,
+                                     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,
+                                     dim=2,
+                                     jax_dtype=jnp.uint4)
+        delattr(layer, "qweight")
+        layer.qweight = qweight
+
+        qzeros = layer.qzeros
+        qzeros = unpack_awq_weight(qzeros, qzeros.packed_dim)
+        qzeros = torch_to_jax_param(qzeros,
+                                    NamedSharding(
+                                        self.jax_config.mesh,
+                                        self.jax_config.scale_sharding),
+                                    self.jax_config.output_sizes,
+                                    self.jax_config.n_shards,
+                                    self.jax_config.fuse_matmuls,
+                                    dim=1,
+                                    jax_dtype=jnp.uint4)
+        delattr(layer, "qzeros")
+        layer.qzeros = qzeros
+
+        scales = torch_to_jax_param(layer.scales,
+                                    NamedSharding(
+                                        self.jax_config.mesh,
+                                        self.jax_config.scale_sharding),
+                                    self.jax_config.output_sizes,
+                                    self.jax_config.n_shards,
+                                    self.jax_config.fuse_matmuls,
+                                    dim=1)
+        delattr(layer, "scales")
+        layer.scales = scales
+
+        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
+
+    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:
+        x_jax = jax_view(x)
+
+        qweight = jax_view(layer.qweight)
+        qzeros = jnp.expand_dims(jax_view(layer.qzeros), 1)
+        scales = jnp.expand_dims(jax_view(layer.scales), 1)
+
+        qweight = qweight.astype(jnp.int8)
+        qzeros = qzeros.astype(jnp.int8)
+
+        weight = (qweight - qzeros) * scales
+        weight = weight.reshape((-1, weight.shape[-1]))
+        outs = jnp.einsum("bd,df->bf", x_jax, weight)
+
+        if bias is not None and not layer.skip_bias_add:
+            outs += bias.jax()
+
+        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] = None) -> torch.Tensor:
+        assert isinstance(layer.qweight, torch.nn.ParameterList)
+
+        x_jax = jax_view(x)
+        params = zip(layer.qweight, layer.qzeros, layer.scales)
+        outs = []
+        for i, (qweight, qzeros, scales) in enumerate(params):
+            qweight = jax_view(qweight)
+            scales = jnp.expand_dims(jax_view(scales), 1)
+            qzeros = jnp.expand_dims(jax_view(qzeros), 1)
+
+            qweight = qweight.astype(jnp.int8)
+            qzeros = qzeros.astype(jnp.int8)
+
+            weight = (qweight - qzeros) * scales
+            weight = weight.reshape((-1, weight.shape[-1]))
+            out = jnp.einsum("bd,df->bf", x_jax, weight)
+
+            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)
+
+
+def unpack_awq_weight(weight: torch.Tensor, packed_dim: int):
+    weight = unpack_quantized_values_into_int32(weight, scalar_types.uint4,
+                                                packed_dim)
+
+    # AWQ packs 8 uint4 into 32-bits in this order: (0, 2, 4, 6, 1, 3, 5, 7).
+    # Following list maps the order used by AWQ into an ascending order.
+    reverse_awq_order = (0, 4, 1, 5, 2, 6, 3, 7)
+
+    orig_shape = weight.shape
+    weight = weight.reshape(orig_shape[:-1] + (-1, 8))
+    return weight[..., reverse_awq_order].reshape(orig_shape)

tpu_inference/layers/vllm/quantization/common.py

@@ -0,0 +1,105 @@
+import torchax
+from jax.sharding import Mesh, PartitionSpec
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEConfig
+# yapf: disable
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
+                                               LinearBase,
+                                               MergedColumnParallelLinear,
+                                               QKVParallelLinear,
+                                               ReplicatedLinear,
+                                               RowParallelLinear)
+
+from tpu_inference.layers.vllm.linear_common import \
+    get_model_matmul_fusion_assignment
+from tpu_inference.utils import TPU_SECOND_LAST_MINOR
+
+# yapf: enable
+
+P = PartitionSpec
+
+logger = init_logger(__name__)
+
+
+class JaxCommonLinearConfig:
+
+    def __init__(self, vllm_config: VllmConfig, mesh: Mesh, layer: LinearBase):
+        assert isinstance(layer, LinearBase)
+
+        self.mesh = mesh
+        self.output_sizes = [layer.output_size]
+        self.weight_sharding = P(None, None)
+        self.fuse_matmuls = True
+        self.enable_sequence_parallelism = vllm_config.compilation_config.pass_config.enable_sequence_parallelism
+        self.input_sharding = None
+        self.output_sharding = None
+
+        if isinstance(layer, RowParallelLinear):
+            self.weight_sharding = P(None, "model")
+            if self.enable_sequence_parallelism:
+                self.output_sharding = P("model", None)
+        elif isinstance(layer, ColumnParallelLinear):
+            self.weight_sharding = P("model", None)
+            if self.enable_sequence_parallelism:
+                self.input_sharding = P("model", None)
+
+            if isinstance(layer, MergedColumnParallelLinear) or isinstance(
+                    layer, QKVParallelLinear):
+                self.output_sizes = layer.output_sizes
+
+            self.fuse_matmuls = get_model_matmul_fusion_assignment(
+                vllm_config.model_config.model,
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                vllm_config.parallel_config.tensor_parallel_size,
+                layer._get_name())
+        elif isinstance(layer, ReplicatedLinear):
+            self.weight_sharding = P(None, None)
+        else:
+            logger.warning(
+                "Unsupported linear layer type of %s. Can potentially yield "
+                " bad performance.", type(layer))
+
+        self.bias_sharding = P(self.weight_sharding[0])
+        self.n_shards = self.mesh.shape.get(self.weight_sharding[0], 1)
+
+    def get_input_sharding(self, x: torchax.tensor.Tensor):
+        if self.enable_sequence_parallelism:
+            token_num = x.shape[0]
+            # NOTE(chengjiyao): make sure the sharded token_num is larger than TPU_SECOND_LAST_MINOR
+            if token_num // self.mesh.shape["model"] >= TPU_SECOND_LAST_MINOR:
+                return self.input_sharding
+            else:
+                return None
+        return self.input_sharding
+
+    def get_output_sharding(self, x: torchax.tensor.Tensor):
+        if self.enable_sequence_parallelism:
+            token_num = x.shape[0]
+            # NOTE(chengjiyao): make sure the sharded token_num is larger than TPU_SECOND_LAST_MINOR
+            if token_num // self.mesh.shape["model"] >= TPU_SECOND_LAST_MINOR:
+                return self.output_sharding
+            else:
+                return None
+        return self.output_sharding
+
+
+class JaxCommonConfig:
+    vllm_config: VllmConfig
+    mesh: Mesh
+
+    @classmethod
+    def set_configs(cls, vllm_config: VllmConfig, mesh: Mesh):
+        cls.vllm_config = vllm_config
+        cls.mesh = mesh
+
+    def get_linear_config(self, layer: LinearBase) -> JaxCommonLinearConfig:
+        assert isinstance(layer, LinearBase)
+        return JaxCommonLinearConfig(self.vllm_config, self.mesh, layer)
+
+    def get_moe_config(self, layer: FusedMoE) -> FusedMoEConfig:
+        assert isinstance(layer, FusedMoE)
+        moe_config = layer.moe_config
+        use_ep = self.vllm_config.parallel_config.enable_expert_parallel
+        moe_config.moe_parallel_config.use_ep = use_ep
+        return moe_config

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

@@ -0,0 +1,120 @@
+from typing import Optional
+
+import torch
+from jax.sharding import PartitionSpec
+from vllm.attention.layer import Attention
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+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  # noqa: E501
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
+    CompressedTensorsConfig, CompressedTensorsKVCacheMethod,
+    CompressedTensorsLinearMethod, CompressedTensorsScheme)
+from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
+    find_matched_target, should_ignore_layer)
+
+from tpu_inference.layers.common.quant_methods import (COMPRESSED_TENSORS,
+                                                       get_tpu_quant_method)
+from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
+from tpu_inference.layers.vllm.quantization.compressed_tensors.compressed_tensors_moe import \
+    VllmCompressedTensorsW8A8Fp8MoEMethod
+from tpu_inference.layers.vllm.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import \
+    VllmCompressedTensorsW8A8Fp8
+from tpu_inference.layers.vllm.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_int8 import \
+    VllmCompressedTensorsW8A8Int8
+from tpu_inference.layers.vllm.quantization.unquantized import \
+    VllmUnquantizedConfig
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+@register_quantization_config(get_tpu_quant_method(COMPRESSED_TENSORS))
+class VllmCompressedTensorsConfig(CompressedTensorsConfig, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls) -> str:
+        return COMPRESSED_TENSORS
+
+    def get_scheme(self,
+                   layer: torch.nn.Module,
+                   layer_name: Optional[str] = None
+                   ) -> Optional["CompressedTensorsScheme"]:
+        """
+        compressed-tensors supports non uniform in the following way:
+
+        targets of config_groups: There can be N config_groups which each
+            have a quantization scheme. Each config_group has a list of targets
+            which can be a full layer_name, a regex for a layer_name, or
+            an nn.Module name.
+
+        Detect whether a layer_name is found in any target and
+        use the quantization scheme corresponding to the matched target
+        to select the CompressedTensorsScheme used for inference.
+        """
+
+        # Will be empty for models with only sparsity
+        weight_quant = input_quant = None
+        if self.target_scheme_map:
+            matched_target = find_matched_target(
+                layer_name=layer_name,
+                module=layer,
+                targets=self.target_scheme_map.keys(),
+                fused_mapping=self.packed_modules_mapping,
+            )
+
+            scheme_dict = self.target_scheme_map[matched_target]
+            weight_quant = scheme_dict.get("weights")
+            input_quant = scheme_dict.get("input_activations")
+
+        if weight_quant is None:
+            logger.warning_once("Acceleration for non-quantized schemes is "
+                                "not supported by Compressed Tensors. "
+                                "Falling back to UnquantizedLinearMethod")
+            return None
+
+        # TODO(kyuyeunk): Add support for different act_quant_format
+
+        linear_config = self.get_linear_config(layer)
+        if self._is_fp8_w8a8(weight_quant, input_quant):
+            is_static_input_scheme = input_quant and not input_quant.dynamic
+            return VllmCompressedTensorsW8A8Fp8(
+                weight_quant=weight_quant,
+                is_static_input_scheme=is_static_input_scheme,
+                jax_config=linear_config,
+            )
+        if self._is_dynamic_token_w8a8(weight_quant, input_quant):
+            return VllmCompressedTensorsW8A8Int8(
+                strategy=weight_quant.strategy,
+                is_static_input_scheme=False,
+                input_symmetric=input_quant.symmetric,
+                jax_config=linear_config,
+            )
+        raise NotImplementedError(
+            "No compressed-tensors compatible scheme was found.")
+
+    def get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+    ) -> Optional[QuantizeMethodBase]:
+        if should_ignore_layer(prefix,
+                               ignore=self.ignore,
+                               fused_mapping=self.packed_modules_mapping):
+            return VllmUnquantizedConfig.get_quant_method(self, layer, prefix)
+        if isinstance(layer, LinearBase):
+            scheme = self.get_scheme(layer=layer, layer_name=prefix)
+            if scheme is None:
+                return VllmUnquantizedConfig.get_quant_method(
+                    self, layer, prefix)
+            layer.scheme = scheme
+            return CompressedTensorsLinearMethod(self)
+        if isinstance(layer, FusedMoE):
+            return VllmCompressedTensorsW8A8Fp8MoEMethod(
+                self, layer.quant_config, self.mesh)
+        if isinstance(layer, Attention):
+            return CompressedTensorsKVCacheMethod(self)
+        return None

tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -0,0 +1,203 @@
+from typing import Callable, Optional, Union
+
+import jax
+import jax.numpy as jnp
+import torch
+import torch.nn.functional as F
+from jax.experimental.layout import Format, Layout
+from jax.sharding import Mesh, NamedSharding
+from jax.sharding import PartitionSpec as P
+from torch.nn.parameter import Parameter
+from torchax.interop import call_jax, torch_view
+from torchax.ops.mappings import t2j
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import FusedMoE, FusedMoEConfig
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import \
+    CompressedTensorsConfig
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import \
+    CompressedTensorsW8A8Fp8MoEMethod
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import (  # noqa
+    WNA16_SUPPORTED_BITS, WNA16_SUPPORTED_TYPES_MAP)
+
+from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
+
+logger = init_logger(__name__)
+
+
+class VllmCompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsW8A8Fp8MoEMethod,
+                                            JaxCommonConfig):
+
+    def __init__(self, quant_config: "CompressedTensorsConfig",
+                 moe: FusedMoEConfig, mesh: Mesh):
+        super().__init__(quant_config, moe)
+        self.mesh = mesh
+        self.quant_config = quant_config
+
+        # disable GPU paths
+        self.use_marlin = False
+        self.rocm_aiter_moe_enabled = False  # is_rocm_aiter_moe_enabled()
+        self.is_fp8_w8a8_sm100 = False
+        self.use_cutlass = False
+        self.disable_expert_map = False
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        assert isinstance(layer, FusedMoE)
+
+        intermediate_size = layer.w13_weight.shape[1] // 2
+        w1_weight = layer.w13_weight[:, :intermediate_size]
+        w3_weight = layer.w13_weight[:, intermediate_size:]
+        w1_weight_scale = layer.w13_weight_scale[:, :intermediate_size]
+        w3_weight_scale = layer.w13_weight_scale[:, intermediate_size:]
+
+        w2_weight = t2j(layer.w2_weight, use_dlpack=False)
+        w2_weight_scale = t2j(layer.w2_weight_scale.to(torch.bfloat16),
+                              use_dlpack=False)
+        w1_weight = t2j(w1_weight, use_dlpack=False)
+        w1_weight_scale = t2j(w1_weight_scale.to(torch.bfloat16),
+                              use_dlpack=False)
+        w3_weight = t2j(w3_weight, use_dlpack=False)
+        w3_weight_scale = t2j(w3_weight_scale.to(torch.bfloat16),
+                              use_dlpack=False)
+
+        if layer.use_ep:
+            format = Format(Layout((0, 1, 2)),
+                            NamedSharding(self.mesh, P("model", None, None)))
+            w1_weight = jax.device_put(w1_weight, format)
+            w1_weight_scale = jax.device_put(w1_weight_scale, format)
+            w3_weight = jax.device_put(w3_weight, format)
+            w3_weight_scale = jax.device_put(w3_weight_scale, format)
+            w2_weight = jax.device_put(w2_weight, format)
+            w2_weight_scale = jax.device_put(w2_weight_scale, format)
+        else:
+            assert intermediate_size == w2_weight.shape[-1]
+            n_shards = self.mesh.shape["model"]
+            assert intermediate_size % n_shards == 0
+
+            # TODO: enable this if using fused weights
+            # output_sizes = [intermediate_size, intermediate_size]
+            # w13_weight = reorder_concatenated_tensor_for_sharding(
+            #    w13_weight, output_sizes, n_shards, dim=1
+            # )
+
+            w13_format = Format(
+                Layout((0, 1, 2)),
+                NamedSharding(self.mesh, P(None, "model", None)))
+            w1_weight = jax.device_put(w1_weight, w13_format)
+            w1_weight_scale = jax.device_put(w1_weight_scale, w13_format)
+            w3_weight = jax.device_put(w3_weight, w13_format)
+            w3_weight_scale = jax.device_put(w3_weight_scale, w13_format)
+            w2_weight = jax.device_put(
+                w2_weight,
+                Format(Layout((0, 1, 2)),
+                       NamedSharding(self.mesh, P(None, None, "model"))),
+            )
+            w2_weight_scale = jax.device_put(
+                w2_weight_scale,
+                Format(Layout((0, 1, 2)), NamedSharding(self.mesh, P())),
+            )  # replicate
+
+        w1_weight = Parameter(torch_view(w1_weight), requires_grad=False)
+        w1_weight_scale = Parameter(torch_view(w1_weight_scale),
+                                    requires_grad=False)
+        w2_weight = Parameter(torch_view(w2_weight), requires_grad=False)
+        w2_weight_scale = Parameter(torch_view(w2_weight_scale),
+                                    requires_grad=False)
+        w3_weight = Parameter(torch_view(w3_weight), requires_grad=False)
+        w3_weight_scale = Parameter(torch_view(w3_weight_scale),
+                                    requires_grad=False)
+
+        # TODO dont reuse variable
+        layer.w13_weight = w1_weight
+        layer.w13_weight_scale = w1_weight_scale
+        layer.w2_weight = w2_weight
+        layer.w2_weight_scale = w2_weight_scale
+        layer.w3_weight = w3_weight
+        layer.w3_weight_scale = w3_weight_scale
+
+    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 activation != "silu":
+            raise NotImplementedError(
+                "Only silu is supported for activation function.")
+        if scoring_func != "softmax":
+            raise NotImplementedError(
+                "Only softmax is supported for scoring_func")
+
+        # import sys
+        # sys.stdin = open(0)
+        # breakpoint()
+
+        # TODO: Use MoE kernel when it supports fp8
+
+        seqlen = x.shape[0]
+
+        expert_weights = F.softmax(router_logits, dim=-1)
+        expert_weights, expert_indices = torch.topk(expert_weights,
+                                                    top_k,
+                                                    dim=-1)
+        if renormalize:
+            expert_weights /= expert_weights.sum(dim=-1, keepdim=True)
+
+        # cond ffn
+        # e = total num of exp = 160
+        # t = seqlen
+        # o = config.imtermediate size
+        # i = config.dim
+        #torch.einsum("ti, eoi -> teo", x, layer.w13_weight) * self.w13_weight_scale)
+        ux1 = call_jax(jax.lax.dot,
+                       x,
+                       layer.w13_weight,
+                       dimension_numbers=(((1, ), (2, )), ((), ())),
+                       preferred_element_type=jnp.bfloat16.dtype)
+        x1 = F.silu(ux1 * layer.w13_weight_scale.squeeze(2))
+
+        #x3 = torch.einsum("ti, eoi -> teo", x, layer.w3_weight) * self.w3_weight_scale
+        x3 = call_jax(jax.lax.dot,
+                      x,
+                      layer.w3_weight,
+                      dimension_numbers=(((1, ), (2, )), ((), ())),
+                      preferred_element_type=jnp.bfloat16.dtype
+                      ) * layer.w3_weight_scale.squeeze(2)
+
+        #expert_outs = torch.einsum("teo, eio -> tei", (x1 * x3), self.w2_weight) * self.w2_weight_scale
+        expert_outs = call_jax(
+            jax.lax.dot,
+            x1 * x3,
+            layer.w2_weight,
+            dimension_numbers=(((2, ), (2, )), ((1, ), (0, ))),
+            preferred_element_type=jnp.bfloat16.dtype).transpose(
+                0, 1) * layer.w2_weight_scale.squeeze(2)
+
+        seq_indexes = torch.arange(seqlen, device='jax').unsqueeze(1)
+        expert_outs = expert_outs[seq_indexes, expert_indices]
+
+        # out = torch.einsum("tai,ta -> ti", expert_outs, expert_weights)
+        out = call_jax(jax.lax.dot,
+                       expert_outs,
+                       expert_weights,
+                       dimension_numbers=(((1, ), (1, )), ((0, ), (0, ))),
+                       preferred_element_type=jnp.bfloat16.dtype)
+
+        return out