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

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

@@ -0,0 +1,135 @@
+# 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 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 \
+    VllmCompressedTensorsMoEMethod
+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):
+            layer.moe_config = self.get_moe_config(layer)
+            return VllmCompressedTensorsMoEMethod.get_moe_method(
+                self, layer, layer_name=prefix)
+        if isinstance(layer, Attention):
+            return CompressedTensorsKVCacheMethod(self)
+        return None

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

@@ -0,0 +1,266 @@
+# 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 Union
+
+import jax
+import jax.numpy as jnp
+import torch
+from compressed_tensors.quantization import QuantizationArgs
+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 jax_view, 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_moe import (  # noqa: E501
+    CompressedTensorsMoEMethod, CompressedTensorsW8A8Fp8MoEMethod)
+
+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 \
+    VllmUnquantizedFusedMoEMethod
+
+logger = init_logger(__name__)
+
+
+class VllmCompressedTensorsMoEMethod(CompressedTensorsMoEMethod):
+
+    @staticmethod
+    def get_moe_method(
+        quant_config: "VllmCompressedTensorsConfig",  # type: ignore # noqa E501
+        layer: torch.nn.Module,
+        layer_name: str,
+    ) -> CompressedTensorsMoEMethod:
+        assert isinstance(layer, FusedMoE)
+
+        # FusedMoE was made by combining multiple Linears so need to
+        # make sure quantization config for Linear can target it
+        quant_config._add_fused_moe_to_target_scheme_map()
+        unfused_names = [
+            layer_name + proj_name
+            for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
+        ]
+        # TODO: refactor this to use expert_mapping and check all layer numbers
+        all_scheme_dicts = [
+            quant_config.get_scheme_dict(layer, name) for name in unfused_names
+        ]
+        scheme_dict = all_scheme_dicts.pop()
+
+        # multiple schemes found
+        if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]):
+            raise ValueError("All MoE projections need to have same "
+                             "quantization scheme but found multiple")
+
+        if scheme_dict is None:
+            return VllmUnquantizedFusedMoEMethod(layer.moe_config,
+                                                 quant_config.mesh)
+
+        weight_quant = scheme_dict.get("weights")
+        input_quant = scheme_dict.get("input_activations")
+
+        if quant_config._is_fp8_w8a8(weight_quant, input_quant):
+            return VllmCompressedTensorsW8A8Fp8MoEMethod(
+                weight_quant, input_quant, layer.moe_config, quant_config.mesh)
+        else:
+            raise RuntimeError(
+                f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}")
+
+
+class VllmCompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsW8A8Fp8MoEMethod,
+                                            JaxCommonConfig):
+
+    def __init__(
+        self,
+        weight_quant: QuantizationArgs,
+        input_quant: QuantizationArgs,
+        moe: FusedMoEConfig,
+        mesh: Mesh,
+    ):
+        super().__init__(weight_quant, input_quant, moe)
+        self.mesh = mesh
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """
+        Docstring for process_weights_after_loading
+
+        :param self: Description
+        :param layer: Description
+        :type layer: torch.nn.Module
+
+        Steps:
+        1. Read weights from layer object and convert to jax arrays
+        2. Interleave concat w13 weights
+        3. Shard weights for tp (rowwise w13, colwise w2)
+        4. Initialize Params as torch.nn.Parameter
+            a. w13_weight - float8_e4m3fn shape: (num_experts, 2 x intermediate_size, input_size)
+            b. w2_weight - float8_e4m3fn shape: (num_experts, output_size, intermediate_size)
+            c. w13_weight_scale - FP32 shape: (num_experts, 2 x intermediate_size, 1)
+            d. w2_weight_scale - FP32shape: (num_experts, output_size, 1)
+        """
+        assert isinstance(layer, FusedMoE)
+
+        # Read weights from layer object
+        w13_weight = t2j(
+            layer.w13_weight, use_dlpack=False
+        )  # float8_e4m3fn shape: (num_experts, 2 x intermediate_size, input_size)
+        w13_weight_scale = t2j(
+            layer.w13_weight_scale, use_dlpack=False
+        )  # FP32 shape: (num_experts, 2 x intermediate_size, 1)
+        w2_weight = t2j(
+            layer.w2_weight, use_dlpack=False
+        )  # float8_e4m3fn shape: (num_experts, output_size, intermediate_size)
+        w2_weight_scale = t2j(layer.w2_weight_scale, use_dlpack=False)
+        w13_weight_scale = w13_weight_scale.astype(jnp.bfloat16)
+        w2_weight_scale = w2_weight_scale.astype(jnp.bfloat16)
+        intermediate_size = layer.w13_weight.shape[1] // 2
+        assert intermediate_size == w2_weight.shape[-1]
+        n_shards = self.mesh.shape["model"]
+        assert intermediate_size % n_shards == 0
+        num_experts, hidden_size, intermediate_size = w2_weight.shape
+        assert w2_weight_scale.shape == (num_experts, hidden_size, 1)
+        assert w13_weight.shape == (num_experts, 2 * intermediate_size,
+                                    hidden_size)
+        assert w13_weight_scale.shape == (num_experts, 2 * intermediate_size,
+                                          1)
+
+        if not layer.use_ep:
+            # Interleave concat w13 weights
+            w13_weight = reorder_concatenated_tensor_for_sharding(
+                w13_weight,
+                split_sizes=(intermediate_size, intermediate_size),
+                dim=1,
+                n_shards=n_shards,
+            )
+            # Interleave concat w13 weight scales
+            w13_weight_scale = reorder_concatenated_tensor_for_sharding(
+                w13_weight_scale,
+                split_sizes=(intermediate_size, intermediate_size),
+                dim=1,
+                n_shards=n_shards,
+            )
+
+        # 160,5120,1 -> 160,1,5120
+        w13_weight_scale = jnp.swapaxes(w13_weight_scale, 1, 2)
+        # 160,1,5120 -> 160, 1, 1, 5120   (num_experts, num_blocks, 1, outer_dim)
+        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)
+
+        if layer.use_ep:
+            # Apply EP sharding
+            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)
+
+        else:
+            # Shard weights for tp (rowwise w13, colwise w2)
+            w13_format = Format(
+                Layout((0, 1, 2)),  # expert, 2xintermed, input
+                NamedSharding(self.mesh, P(None, "model", None)),
+            )  # rowwise sharding on intermed dim
+
+            w13_scale_format = Format(
+                Layout(
+                    (0, 1, 2, 3)),  #  (num_experts, num_blocks, 1, outer_dim)
+                NamedSharding(self.mesh, P(None, None, None, "model")),
+            )  # col wise GMM sharding on intermed dim
+
+            # Local shard shape: (num_experts, 2 x (intermediate_size // n_shards), input_size)
+            w13_weight = jax.lax.with_sharding_constraint(
+                w13_weight, w13_format)
+            # Local shard shape: (num_experts, (intermediate_size // n_shards), 1)
+            w13_weight_scale = jax.lax.with_sharding_constraint(
+                w13_weight_scale, w13_scale_format)
+
+            # Shard weights for tp (colwise w2)
+            w2_format = Format(
+                Layout((0, 1, 2)),  # expert, intermed, hidden
+                NamedSharding(self.mesh, P(None, None, "model")),
+            )
+            # Local shard shape: (num_experts, hidden, (intermediate_size // n_shards))
+            # #  (num_experts, num_blocks, 1, outer_dim)
+            w2_weight = jax.lax.with_sharding_constraint(w2_weight, w2_format)
+
+            w2_scale_format = Format(
+                Layout((0, 1, 2, 3)),  # expert, intermed, 1
+                NamedSharding(self.mesh, P(None, None, None, None)),
+            )
+            # Local shard shape: (num_experts, intermediate_size // n_shards, 1)
+            w2_weight_scale = jax.lax.with_sharding_constraint(
+                w2_weight_scale, w2_scale_format)
+
+        w13_weight = Parameter(torch_view(w13_weight), requires_grad=False)
+        w13_weight_scale = Parameter(torch_view(w13_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)
+
+        layer.w13_weight = w13_weight
+        layer.w13_weight_scale = w13_weight_scale
+        layer.w2_weight = w2_weight
+        layer.w2_weight_scale = w2_weight_scale
+
+    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.activation != "silu":
+            raise NotImplementedError(
+                "Only silu is supported for activation function.")
+        if layer.scoring_func != "softmax":
+            raise NotImplementedError(
+                "Only softmax is supported for scoring_func")
+
+        # TODO: Use MoE kernel when it supports fp8
+        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)
+        gating_output = jax_view(router_logits)
+        out = torch_view(
+            fused_moe_func(
+                hidden_states=x,
+                w1=w13_weight,
+                w2=w2_weight,
+                w1_scale=w13_weight_scale,
+                w2_scale=w2_weight_scale,
+                w1_bias=None,
+                w2_bias=None,
+                gating_output=gating_output,
+                topk=layer.top_k,
+                renormalize=layer.renormalize,
+                mesh=self.mesh,
+                use_ep=layer.use_ep,
+                activation=layer.activation,
+            ))
+
+        return out

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

@@ -0,0 +1,13 @@
+# 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.

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

@@ -0,0 +1,222 @@
+# 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 (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))