tpu-inference 0.11.1rc1__py3-none-any.whl → 0.11.1rc3__py3-none-any.whl

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,
+                                                         is_layer_skipped_awq)
+from vllm.model_executor.layers.quantization.base_config import \
+    QuantizeMethodBase
+from vllm.model_executor.layers.quantization.utils.quant_utils import \
+    unpack_quantized_values_into_int32
+from vllm.scalar_type import scalar_types
+
+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("jax-awq")
+class VllmAWQConfig(AWQConfig, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "jax-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_awq(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,121 @@
+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, is_activation_quantization_format,
+    should_ignore_layer)
+
+from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
+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("jax-compressed-tensors")
+class VllmCompressedTensorsConfig(CompressedTensorsConfig, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "jax-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")
+            format = scheme_dict.get("format")
+
+        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
+        act_quant_format = is_activation_quantization_format(  # noqa: F841
+            format
+        ) if format is not None else is_activation_quantization_format(
+            self.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):
+            raise NotImplementedError(
+                "FusedMoE quantization is currently not supported.")
+        if isinstance(layer, Attention):
+            return CompressedTensorsKVCacheMethod(self)
+        return None

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))