tpu-inference 0.11.1.dev202512030818__py3-none-any.whl → 0.13.0rc2.post7__py3-none-any.whl

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

@@ -1,203 +1,266 @@
-from typing import Callable, Optional, 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.
+
+from typing import Union
 
 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 torch.nn.parameter import Parameter
-from torchax.interop import call_jax, torch_view
+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 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 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, quant_config: "CompressedTensorsConfig",
-                 moe: FusedMoEConfig, mesh: Mesh):
-        super().__init__(quant_config, moe)
+    def __init__(
+        self,
+        weight_quant: QuantizationArgs,
+        input_quant: QuantizationArgs,
+        moe: FusedMoEConfig,
+        mesh: Mesh,
+    ):
+        super().__init__(weight_quant, input_quant, 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:
+        """
+        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
-        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)
+        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:
-            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
+            # 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)
 
-            # 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_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)),
-                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"))),
+                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")),
             )
-            w2_weight_scale = jax.device_put(
-                w2_weight_scale,
-                Format(Layout((0, 1, 2)), NamedSharding(self.mesh, P())),
-            )  # replicate
+            # 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)
 
-        w1_weight = Parameter(torch_view(w1_weight), requires_grad=False)
-        w1_weight_scale = Parameter(torch_view(w1_weight_scale),
-                                    requires_grad=False)
+            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)
-        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.w13_weight = w13_weight
+        layer.w13_weight_scale = w13_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":
+        if layer.activation != "silu":
             raise NotImplementedError(
                 "Only silu is supported for activation function.")
-        if scoring_func != "softmax":
+        if layer.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)
+        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

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

tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.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

tpu_inference/layers/vllm/quantization/fp8.py

@@ -0,0 +1,118 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Union
+
+import jax
+import torch
+from jax.sharding import PartitionSpec
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
+from vllm.model_executor.layers.quantization import \
+    register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import \
+    QuantizeMethodBase
+from vllm.model_executor.layers.quantization.fp8 import (Fp8Config,
+                                                         Fp8LinearMethod)
+from vllm.model_executor.layers.quantization.utils.quant_utils import \
+    is_layer_skipped
+
+from tpu_inference.layers.common.quant_methods import FP8, get_tpu_quant_method
+from tpu_inference.layers.vllm.quantization.common import (
+    JaxCommonConfig, JaxCommonLinearConfig)
+from tpu_inference.layers.vllm.quantization.unquantized import \
+    VllmUnquantizedLinearMethod
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+@register_quantization_config(get_tpu_quant_method(FP8))
+class VllmFp8Config(Fp8Config, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls):
+        return FP8
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        return [torch.bfloat16]
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[Union["LinearMethodBase", "QuantizeMethodBase"]]:
+        if isinstance(layer, LinearBase):
+            linear_config = self.get_linear_config(layer)
+            if is_layer_skipped(prefix, self.ignored_layers):
+                return VllmUnquantizedLinearMethod(linear_config)
+            return VllmFp8LinearMethod(self, linear_config)
+        elif isinstance(layer, FusedMoE):
+            raise NotImplementedError(
+                "FP8 FusedMoE is currently not supported in torchax-jax")
+        return None
+
+
+class VllmFp8LinearMethod(Fp8LinearMethod):
+
+    def __init__(self, quant_config: VllmFp8Config,
+                 jax_config: JaxCommonLinearConfig):
+        super().__init__(quant_config)
+        self.jax_config = jax_config
+        self._configure_sharding()
+
+    def _configure_sharding(self) -> None:
+
+        raise NotImplementedError(
+            "Configure PartitionSpec for weight_sharding and scale_sharding "
+            "based on layer type (RowParallel/ColumnParallel)")
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        raise NotImplementedError(
+            "Convert layer.weight, layer.weight_scale, and optionally "
+            "layer.input_scale and layer.bias from torch tensors to JAX arrays "
+            "using torch_to_jax_param() with appropriate sharding")
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        with jax.named_scope(layer._get_name()):
+            if self.jax_config.fuse_matmuls:
+                out = self._apply_fused(layer, x, bias)
+            else:
+                out = self._apply_split(layer, x, bias)
+
+        return out
+
+    def _apply_fused(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        raise NotImplementedError(
+            "Implement single matmul for fused outputs: "
+            "quantize input to fp8, perform fp8 matmul with weight and scales, "
+            "dequantize output, and add bias if present")
+
+    def _apply_split(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        raise NotImplementedError(
+            "Implement separate matmuls per output partition: "
+            "split weight/scale by output_sizes, perform fp8 matmul for each, "
+            "concatenate results, and add bias if present")