tpu-inference 0.12.0.dev20251213__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

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

@@ -1,24 +1,39 @@
-from typing import Union
+# 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.
 
 import jax
-import jax.numpy as jnp
 import torch
-import torch.nn.functional as F
 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 jax.sharding import Mesh
 from torch.nn.parameter import Parameter
-from torchax.interop import call_jax, torch_view
+from torchax.interop import 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.quantization.common import JaxCommonConfig
+from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.vllm.fused_moe import (FusedMoEBackend,
+                                                 fused_moe_apply,
+                                                 select_moe_backend)
+from tpu_inference.layers.vllm.process_weights.fused_moe_weights import (
+    FusedMoEWeights, process_moe_weights, shard_moe_weights)
+from tpu_inference.layers.vllm.quantization.configs import VllmQuantConfig
 from tpu_inference.layers.vllm.quantization.unquantized import \
     VllmUnquantizedFusedMoEMethod
+from tpu_inference.logger import init_logger
+from tpu_inference.utils import get_mesh_shape_product
 
 logger = init_logger(__name__)
 
@@ -31,7 +46,6 @@ class VllmCompressedTensorsMoEMethod(CompressedTensorsMoEMethod):
         layer: torch.nn.Module,
         layer_name: str,
     ) -> CompressedTensorsMoEMethod:
-
         assert isinstance(layer, FusedMoE)
 
         # FusedMoE was made by combining multiple Linears so need to
@@ -68,15 +82,44 @@ class VllmCompressedTensorsMoEMethod(CompressedTensorsMoEMethod):
 
 
 class VllmCompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsW8A8Fp8MoEMethod,
-                                            JaxCommonConfig):
+                                            VllmQuantConfig):
 
-    def __init__(self, weight_quant: QuantizationArgs,
-                 input_quant: QuantizationArgs, moe: FusedMoEConfig,
-                 mesh: Mesh):
+    def __init__(
+        self,
+        weight_quant: QuantizationArgs,
+        input_quant: QuantizationArgs,
+        moe: FusedMoEConfig,
+        mesh: Mesh,
+    ):
         super().__init__(weight_quant, input_quant, moe)
+
         self.mesh = mesh
+        self.moe_backend = select_moe_backend(self.moe)
+
+        self.extra_backend_kwargs = {}
+        if self.moe_backend == FusedMoEBackend.FUSED_MOE:
+            raise NotImplementedError(
+                "Per-channel quantization is not supported in FusedMoE kernel."
+            )
 
     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)
 
         w13_weight = t2j(layer.w13_weight, use_dlpack=False)
@@ -84,129 +127,73 @@ class VllmCompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsW8A8Fp8MoEMethod,
         w2_weight = t2j(layer.w2_weight, use_dlpack=False)
         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)
-
-        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)
-
-        w1_weight, w3_weight = jnp.split(w13_weight, 2, 1)
-        w1_weight_scale, w3_weight_scale = jnp.split(w13_weight_scale, 2, 1)
-
-        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)
+        if self.moe.has_bias:
+            w13_bias = t2j(layer.w13_bias, use_dlpack=False)
+            w2_bias = t2j(layer.w2_bias, use_dlpack=False)
         else:
-            n_shards = self.mesh.shape["model"]
-            assert intermediate_size % n_shards == 0
-
-            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"))),
+            w13_bias = w2_bias = None
+
+        @jax.jit
+        def process_fp8_moe_weights(
+            w13_weight: jax.Array,
+            w13_weight_scale: jax.Array,
+            w13_bias: jax.Array | None,
+            w2_weight: jax.Array,
+            w2_weight_scale: jax.Array,
+            w2_bias: jax.Array | None,
+        ) -> FusedMoEWeights:
+            w13_interleave = layer.activation == "swigluoai"
+            w13_reorder_size = get_mesh_shape_product(
+                self.mesh, ShardingAxisName.MLP_TENSOR)
+
+            return process_moe_weights(
+                weights=FusedMoEWeights(
+                    w13_weight=w13_weight,
+                    w13_weight_scale=w13_weight_scale,
+                    w13_bias=w13_bias,
+                    w2_weight=w2_weight,
+                    w2_weight_scale=w2_weight_scale,
+                    w2_bias=w2_bias,
+                ),
+                moe_backend=self.moe_backend,
+                w13_reorder_size=w13_reorder_size,
+                w13_interleave=w13_interleave,
             )
-            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
+
+        weights = process_fp8_moe_weights(
+            w13_weight,
+            w13_weight_scale,
+            w13_bias,
+            w2_weight,
+            w2_weight_scale,
+            w2_bias,
+        )
+        weights = torch_view(
+            shard_moe_weights(weights, self.moe_backend, self.mesh))
+
+        layer.w13_weight = Parameter(weights.w13_weight, requires_grad=False)
+        layer.w2_weight = Parameter(weights.w2_weight, requires_grad=False)
+
+        layer.w13_weight_scale = Parameter(weights.w13_weight_scale,
+                                           requires_grad=False)
+        layer.w2_weight_scale = Parameter(weights.w2_weight_scale,
+                                          requires_grad=False)
+
+        if self.moe.has_bias:
+            layer.w13_bias = Parameter(weights.w13_bias, requires_grad=False)
+            layer.w2_bias = Parameter(weights.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.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
-        seqlen = x.shape[0]
-
-        expert_weights = F.softmax(router_logits, dim=-1)
-        expert_weights, expert_indices = torch.topk(expert_weights,
-                                                    layer.top_k,
-                                                    dim=-1)
-        if layer.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
+    ) -> torch.Tensor:
+        return fused_moe_apply(
+            layer,
+            x,
+            router_logits,
+            self.moe_backend,
+            self.mesh,
+            self.extra_backend_kwargs,
+        )

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

@@ -1,3 +1,17 @@
+# 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
@@ -6,49 +20,27 @@ import torch
 from compressed_tensors.quantization import (QuantizationArgs,
                                              QuantizationStrategy)
 from jax.sharding import NamedSharding, PartitionSpec
+from torch.nn.parameter import Parameter
 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
+from tpu_inference.layers.common.quantization import (dequantize_tensor,
+                                                      quantize_tensor)
+from tpu_inference.layers.common.utils import \
+    slice_sharded_tensor_for_concatenation
+from tpu_inference.layers.vllm.linear import sharded_quantized_matmul
+from tpu_inference.layers.vllm.process_weights.linear_weights import (
+    LinearWeights, process_lienar_weights, shard_linear_weights,
+    to_parameter_list)
+from tpu_inference.layers.vllm.quantization.configs import \
+    VllmQuantLinearConfig
+from tpu_inference.logger import init_logger
 
 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
+logger = init_logger(__name__)
 
 
 class VllmCompressedTensorsW8A8Fp8(CompressedTensorsW8A8Fp8):
@@ -57,15 +49,86 @@ class VllmCompressedTensorsW8A8Fp8(CompressedTensorsW8A8Fp8):
         self,
         weight_quant: QuantizationArgs,
         is_static_input_scheme: bool,
-        jax_config: JaxCommonLinearConfig,
+        linear_config: VllmQuantLinearConfig,
     ):
         super().__init__(weight_quant, is_static_input_scheme)
 
-        self.jax_config = jax_config
+        self.linear_config = linear_config
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        weight = layer.weight
-        weight_scale = layer.weight_scale
+        weight = t2j(layer.weight, use_dlpack=False)
+        delattr(layer, "weight")
+        weight_scale = t2j(layer.weight_scale, use_dlpack=False)
+        delattr(layer, "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 = t2j(layer.bias, use_dlpack=False)
+            delattr(layer, "bias")
+        else:
+            bias = None
+
+        per_tensor = self.strategy == QuantizationStrategy.TENSOR
+
+        @jax.jit
+        def process_fp8_linear_weights(
+            weight: jax.Array,
+            weight_scale: jax.Array,
+            bias: jax.Array | None,
+        ) -> LinearWeights:
+            if per_tensor:
+                weights = []
+                start = 0
+                # Multiple weights may have been concatenated. Loop through
+                # each weight and perform dequantization.
+                for i, output_size in enumerate(
+                        self.linear_config.output_sizes):
+                    end = start + output_size
+                    weights.append(
+                        dequantize_tensor(weight[start:end], weight_scale[i]))
+                    start = end
+                weight = jnp.concat(weights, axis=0)
+                # Requantize into per-tensor.
+                weight, weight_scale = quantize_tensor(jnp.float8_e4m3fn,
+                                                       weight, None)
+            else:
+                weight_scale = jnp.squeeze(weight_scale, -1)
+
+            return process_lienar_weights(
+                LinearWeights(
+                    weight=weight,
+                    weight_scale=weight_scale,
+                    zero_point=None,
+                    bias=bias,
+                ),
+                fused=self.linear_config.fuse_matmuls,
+                output_sizes=self.linear_config.output_sizes,
+                reorder_size=self.linear_config.n_shards,
+                per_tensor=per_tensor,
+            )
+
+        weights = process_fp8_linear_weights(weight, weight_scale, bias)
+        weights = torch_view(
+            shard_linear_weights(
+                weights,
+                mesh=self.linear_config.mesh,
+                weight_p_spec=self.linear_config.weight_sharding,
+                bias_p_spec=self.linear_config.bias_sharding,
+                per_tensor=per_tensor,
+            ))
+
+        if self.linear_config.fuse_matmuls:
+            layer.weight = Parameter(weights.weight, requires_grad=False)
+            layer.weight_scale = Parameter(weights.weight_scale,
+                                           requires_grad=False)
+            if bias is not None:
+                layer.bias = Parameter(weights.bias, requires_grad=False)
+        else:
+            layer.weight = to_parameter_list(weights.weight)
+            layer.weight_scale = to_parameter_list(weights.weight_scale)
+            if bias is not None:
+                layer.bias = to_parameter_list(weights.bias)
 
         if self.is_static_input_scheme:
             # In static quant, all input_scales share the same value.
@@ -74,59 +137,16 @@ class VllmCompressedTensorsW8A8Fp8(CompressedTensorsW8A8Fp8):
 
             input_scale = jax.device_put(
                 t2j(input_scale_first, use_dlpack=False),
-                NamedSharding(self.jax_config.mesh, P()))
+                NamedSharding(self.linear_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:
+            if self.linear_config.fuse_matmuls:
                 return self._apply_fused(layer, x, bias)
             else:
                 return self._apply_split(layer, x, bias)
@@ -157,13 +177,13 @@ class VllmCompressedTensorsW8A8Fp8(CompressedTensorsW8A8Fp8):
         else:
             outs = sharded_quantized_matmul(x_jax, weight_jax,
                                             weight_scale_jax,
-                                            self.jax_config.mesh,
-                                            self.jax_config.weight_sharding)
+                                            self.linear_config.mesh,
+                                            self.linear_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)
+            outs, self.linear_config.output_sizes, self.linear_config.n_shards)
         return torch_view(jnp.concatenate(outs, axis=-1))
 
     def _apply_split(self, layer: torch.nn.Module, x: torch.Tensor,
@@ -197,10 +217,10 @@ class VllmCompressedTensorsW8A8Fp8(CompressedTensorsW8A8Fp8):
                 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)
+                out = sharded_quantized_matmul(
+                    x_jax, weight_jax, weight_scale_jax,
+                    self.linear_config.mesh,
+                    self.linear_config.weight_sharding)
 
             if bias is not None and not layer.skip_bias_add:
                 out += jax_view(bias[i])