tpu-inference 0.11.1.dev202511220812__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tpu_inference/layers/vllm/quantization/mxfp4.py

@@ -1,16 +1,29 @@
-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 Optional
 
 import jax
 import jax.numpy as jnp
 import torch
-from jax.experimental.layout import Format, Layout
-from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from jax.sharding import Mesh, PartitionSpec
 from torch.nn.parameter import Parameter
-from torchax.interop import jax_view, torch_view
+from torchax.interop import torch_view
 from torchax.ops.mappings import t2j
-from vllm.logger import init_logger
+from vllm.attention.layer import Attention
 from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEConfig, FusedMoEQuantConfig, biased_moe_quant_config)
+    FusedMoEConfig, FusedMoEQuantConfig, mxfp4_w4a16_moe_quant_config)
 from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
                                                         FusedMoEMethodBase)
 from vllm.model_executor.layers.linear import LinearBase
@@ -26,48 +39,30 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import \
 
 from tpu_inference.layers.common.quant_methods import (MXFP4,
                                                        get_tpu_quant_method)
-from tpu_inference.layers.vllm.fused_moe import fused_moe_func_padded
-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.common.quantization import \
+    dequantize_tensor_from_mxfp4_packed
+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, quantize_moe_weights,
+    shard_moe_weights)
+from tpu_inference.layers.vllm.quantization.configs import VllmQuantConfig
 from tpu_inference.layers.vllm.quantization.unquantized import \
     VllmUnquantizedLinearMethod
+from tpu_inference.logger import init_logger
+from tpu_inference.utils import get_mesh_shape_product
 
-MXFP4_BLOCK_SIZE = 32
+REQUANTIZED_BLOCK_SIZE = 512
 
 P = PartitionSpec
-logger = init_logger(__name__)
-
-
-# TODO(kyuyeunk): Move these functions into a common utility file.
-def u8_unpack_e2m1(u8_packed_e2m1: jax.Array) -> jax.Array:
-    assert u8_packed_e2m1.dtype == jnp.uint8
-    e2m1 = jax.lax.bitcast_convert_type(u8_packed_e2m1, jnp.float4_e2m1fn)
-    # bitcast creates one more dimension that splits 8 bits into two e2m1.
-    # we flatten them with the last dim.
-    return jnp.reshape(e2m1, e2m1.shape[:-2] + (-1, ))
-
-
-def e8m0_to_fp32(u8: jax.Array) -> jax.Array:
-    e8_finfo = jnp.finfo(jnp.float8_e8m0fnu)
-    exponents = u8.astype(jnp.int32) + e8_finfo.minexp
-    ones = jnp.ones_like(u8, dtype=jnp.float32)
-    return jnp.ldexp(ones, exponents)
 
-
-def dequantize_block_weight(weight: jax.Array,
-                            scale: jax.Array,
-                            block_size: int,
-                            out_dtype: jnp.dtype = jnp.bfloat16) -> jax.Array:
-    orig_shape = weight.shape
-    weight_block = weight.reshape(orig_shape[:-1] + (-1, block_size))
-    weight_dequantized = weight_block.astype(jnp.float32) * jnp.expand_dims(
-        scale, -1)
-    return weight_dequantized.reshape(orig_shape).astype(out_dtype)
+logger = init_logger(__name__)
 
 
 @register_quantization_config(get_tpu_quant_method(MXFP4))
-class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
+class VllmMxfp4Config(Mxfp4Config, VllmQuantConfig):
 
     @classmethod
     def get_name(cls):
@@ -75,7 +70,6 @@ class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
 
     def get_quant_method(self, layer: torch.nn.Module,
                          prefix: str) -> Optional["QuantizeMethodBase"]:
-        from vllm.attention.layer import Attention  # Avoid circular import
 
         if isinstance(layer, LinearBase):
             linear_config = self.get_linear_config(layer)
@@ -85,17 +79,14 @@ class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
                     fused_mapping=self.packed_modules_mapping,
             ):
                 return VllmUnquantizedLinearMethod(linear_config)
-            # TODO: Add support for MXFP4 Linear Method.
-            # MXFP4 LinearMethod is available in AMD-Quark, refer to that
-            # implementation if you are interested in enabling MXFP4 here.
             logger.warning_once(
                 "MXFP4 linear layer is not implemented - falling back to "
                 "UnquantizedLinearMethod.")
             return VllmUnquantizedLinearMethod(linear_config)
         elif isinstance(layer, FusedMoE):
-            return VllmMxfp4MoEMethod(layer.moe_config, self.mesh)
+            moe_config = self.get_moe_config(layer)
+            return VllmMxfp4MoEMethod(moe_config, self.mesh)
         elif isinstance(layer, Attention):
-            # TODO: Add support for MXFP4 Attention.
             logger.warning_once("MXFP4 attention layer is not implemented. "
                                 "Skipping quantization for this layer.")
         return None
@@ -103,164 +94,132 @@ class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
 
 class VllmMxfp4MoEMethod(Mxfp4MoEMethod):
 
-    def __init__(self, moe: FusedMoEConfig, mesh: Mesh):
+    def __init__(
+        self,
+        moe: FusedMoEConfig,
+        mesh: Mesh,
+        ep_axis_name: str = "model",
+    ):
         FusedMoEMethodBase.__init__(self, moe)
 
         # We piggyback on triton implementation as it applies minimal hardware
         # specific post processing to the weights.
         self.mxfp4_backend = Mxfp4Backend.TRITON
+
         self.mesh = mesh
+        self.moe_backend = select_moe_backend(self.moe)
+
+        self.extra_backend_kwargs = {}
+        if self.moe_backend == FusedMoEBackend.FUSED_MOE:
+            # When fused moe kernle is used, we pass extra arguments like
+            # tuned block sizes to the kernel.
+            self.extra_backend_kwargs = dict(
+                subc_quant_wsz=REQUANTIZED_BLOCK_SIZE,
+                ep_axis_name=ep_axis_name,
+                # TODO: Use autotune table once we have it.
+                bt=256,
+                bf=1024,
+                bd1=1024,
+                bd2=1024,
+                btc=256,
+                bfc=1024,
+                bd1c=1024,
+                bd2c=1024,
+            )
 
     def get_fused_moe_quant_config(
             self, layer: torch.nn.Module) -> FusedMoEQuantConfig | None:
-        # Because we have dequantized weights, we only need biased moe config.
-        # TODO(kyuyeunk): Add native support for MXFP4.
-        return biased_moe_quant_config(
-            layer.w13_bias,
-            layer.w2_bias,
+        return mxfp4_w4a16_moe_quant_config(
+            w1_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            w1_bias=layer.w13_bias,
+            w2_bias=layer.w2_bias,
         )
 
     def process_weights_after_loading(self, layer: torch.nn.Module):
         assert isinstance(layer, FusedMoE)
+        assert layer.moe_config.has_bias, "mxfp4 quantization alwyas use bias."
 
-        w13_weight = u8_unpack_e2m1(t2j(layer.w13_weight, use_dlpack=False))
-        w13_weight_scale = e8m0_to_fp32(
-            t2j(layer.w13_weight_scale, use_dlpack=False))
+        w13_weight = t2j(layer.w13_weight, use_dlpack=False)
+        w13_weight_scale = t2j(layer.w13_weight_scale, use_dlpack=False)
         w13_bias = t2j(layer.w13_bias, use_dlpack=False)
 
-        w2_weight = u8_unpack_e2m1(t2j(layer.w2_weight, use_dlpack=False))
-        w2_weight_scale = e8m0_to_fp32(
-            t2j(layer.w2_weight_scale, use_dlpack=False))
+        w2_weight = t2j(layer.w2_weight, use_dlpack=False)
+        w2_weight_scale = t2j(layer.w2_weight_scale, use_dlpack=False)
         w2_bias = t2j(layer.w2_bias, use_dlpack=False)
 
-        # We dequantize fp4 weights into bf16.
-        # TODO(kyuyeunk): Add native support for MXFP4.
-        w13_weight = dequantize_block_weight(w13_weight, w13_weight_scale,
-                                             MXFP4_BLOCK_SIZE, jnp.bfloat16)
-        w2_weight = dequantize_block_weight(w2_weight, w2_weight_scale,
-                                            MXFP4_BLOCK_SIZE, jnp.bfloat16)
-
-        # Because we have dequantized weights, scales are not used anymore.
-        delattr(layer, "w13_weight_scale")
-        delattr(layer, "w2_weight_scale")
-
-        if layer.activation == "swigluoai":
-            # When using swigluoai, vLLM splits gmm output in a interleaved way.
-            # However, interleaved split is not performant on TPU. Therefore,
-            # we preprocess the weight so that splitting gmm output by middle
-            # can still get the same result.
-            w1_weight = w13_weight[:, ::2, :]
-            w3_weight = w13_weight[:, 1::2, :]
-            w13_weight = jnp.concat([w1_weight, w3_weight], axis=1)
-
-            w1_bias = w13_bias[:, ::2]
-            w3_bias = w13_bias[:, 1::2]
-            w13_bias = jnp.concat([w1_bias, w3_bias], axis=1)
-
-        # TODO(kyuyeunk): Add weight processing logic for the new kernel.
-        if layer.use_ep:
-            w13_weight = jax.device_put(
-                w13_weight,
-                Format(Layout((0, 1, 2)),
-                       NamedSharding(self.mesh, P("model", None, None))))
-            w2_weight = jax.device_put(
-                w2_weight,
-                Format(Layout((0, 1, 2)),
-                       NamedSharding(self.mesh, P("model", None, None))))
-
-            w13_bias = jax.device_put(
-                w13_bias,
-                Format(Layout((0, 1)),
-                       NamedSharding(self.mesh, P("model", None))))
-            w2_bias = jax.device_put(
-                w2_bias,
-                Format(Layout((0, 1)),
-                       NamedSharding(self.mesh, P("model", None))))
-
-        else:
-            intermediate_size = w13_weight.shape[1] // 2
-            assert intermediate_size == w2_weight.shape[-1]
-            output_sizes = [intermediate_size, intermediate_size]
-            n_shards = self.mesh.shape["model"]
-            assert intermediate_size % n_shards == 0
-            w13_weight = reorder_concatenated_tensor_for_sharding(w13_weight,
-                                                                  output_sizes,
-                                                                  n_shards,
-                                                                  dim=1)
-            w13_weight = jax.device_put(
-                w13_weight,
-                Format(Layout((0, 1, 2)),
-                       NamedSharding(self.mesh, P(None, "model", None))))
-            w2_weight = jax.device_put(
-                w2_weight,
-                Format(Layout((0, 1, 2)),
-                       NamedSharding(self.mesh, P(None, None, "model"))))
-
-            w13_bias = reorder_concatenated_tensor_for_sharding(w13_bias,
-                                                                output_sizes,
-                                                                n_shards,
-                                                                dim=1)
-            w13_bias = jax.device_put(
-                w13_bias,
-                Format(Layout((0, 1)),
-                       NamedSharding(self.mesh, P(None, "model"))))
-            w2_bias = jax.device_put(
-                w2_bias,
-                Format(Layout((0, 1)), NamedSharding(self.mesh, P(None,
-                                                                  None))))
+        @jax.jit
+        def process_mxfp4_moe_weights(
+            w13_weight: jax.Array,
+            w13_weight_scale: jax.Array,
+            w13_bias: jax.Array,
+            w2_weight: jax.Array,
+            w2_weight_scale: jax.Array,
+            w2_bias: jax.Array,
+        ) -> FusedMoEWeights:
+            # Dequantize fp4 weights into fp32.
+            w13_weight = dequantize_tensor_from_mxfp4_packed(
+                w13_weight, w13_weight_scale, 2)
+            w2_weight = dequantize_tensor_from_mxfp4_packed(
+                w2_weight, w2_weight_scale, 2)
+
+            w13_interleave = layer.activation == "swigluoai"
+            w13_reorder_size = get_mesh_shape_product(
+                self.mesh, ShardingAxisName.MLP_TENSOR)
+
+            weights = quantize_moe_weights(
+                FusedMoEWeights(
+                    w13_weight=w13_weight,
+                    w13_weight_scale=None,
+                    w13_bias=w13_bias,
+                    w2_weight=w2_weight,
+                    w2_weight_scale=None,
+                    w2_bias=w2_bias,
+                ),
+                jnp.float4_e2m1fn,
+                REQUANTIZED_BLOCK_SIZE,
+            )
+            return process_moe_weights(
+                weights,
+                moe_backend=self.moe_backend,
+                w13_reorder_size=w13_reorder_size,
+                w13_interleave=w13_interleave,
+            )
+
+        weights = process_mxfp4_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(torch_view(w13_weight),
-                                     requires_grad=False)
-        layer.w13_bias = Parameter(torch_view(w13_bias), requires_grad=False)
+        layer.w13_weight = Parameter(weights.w13_weight, requires_grad=False)
+        layer.w2_weight = Parameter(weights.w2_weight, requires_grad=False)
 
-        layer.w2_weight = Parameter(torch_view(w2_weight), requires_grad=False)
-        layer.w2_bias = Parameter(torch_view(w2_bias), 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)
 
-        pass
+        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,
-        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 scoring_func != "softmax":
-            raise NotImplementedError(
-                "Only softmax is supported for scoring_func")
-
-        # Use the original implementation
-        output = fused_moe_func_padded(
-            jax_view(x),
-            jax_view(layer.w13_weight),
-            jax_view(layer.w2_weight),
-            jax_view(layer.w13_bias) if self.moe.has_bias else None,
-            jax_view(layer.w2_bias) if self.moe.has_bias else None,
-            jax_view(router_logits),
-            topk=top_k,
-            global_num_experts=global_num_experts,
-            renormalize=renormalize,
-            reduce_results=layer.reduce_results,
-            mesh=self.mesh,
-            use_ep=layer.use_ep,
-            activation=activation,
+    ) -> torch.Tensor:
+
+        return fused_moe_apply(
+            layer,
+            x,
+            router_logits,
+            self.moe_backend,
+            self.mesh,
+            self.extra_backend_kwargs,
         )
-
-        return torch_view(output)