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

tpu_inference/layers/common/quantization.py

@@ -0,0 +1,282 @@
+# 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 itertools
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+
+MXFP4_BLOCK_SIZE = 32
+
+
+def quantize_tensor_to_mxfp4_packed(
+    tensor: jax.Array,
+    axis: int | tuple = -1,
+) -> Tuple[jax.Array, jax.Array]:
+    """Quantize a tensor to mxfp4 and pack it into uint8."""
+
+    # Perform regular block quantization.
+    tensor_q, scale = quantize_tensor(
+        jnp.float4_e2m1fn,
+        tensor,
+        axis,
+        MXFP4_BLOCK_SIZE,
+    )
+
+    # last two e2m1 elements will be packed into a single uint8 element.
+    bitcast_shape = tensor_q.shape[:-1] + (-1, 2)
+    tensor_q = tensor_q.reshape(bitcast_shape)
+    tensor_q_packed = jax.lax.bitcast_convert_type(tensor_q, jnp.uint8)
+
+    # Since TPU does not have native support for e8m0, we convert scale into
+    # e8m0 manually and store it as uint8.
+    e8m0_finfo = jnp.finfo(jnp.float8_e8m0fnu)
+    _, scale_exp = jnp.frexp(scale)
+    # Subtract exponents by one since e8m0 has no decimal.
+    scale_exp -= 1
+    scale_exp = (scale_exp - e8m0_finfo.minexp).astype(jnp.uint8)
+
+    return tensor_q_packed, scale_exp
+
+
+def u8_unpack_e2m1(u8_packed_e2m1: jax.Array) -> jax.Array:
+    """Unpack e2m1 tensor that was packed into u8."""
+    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:
+    """Convert e8m0 (that was bitcasted to u8) into fp32."""
+    assert u8.dtype == jnp.uint8
+
+    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 awq_u32_unpack_u4(awq_u32_packed: jax.Array) -> jax.Array:
+    """Unpack u4 tensor that was packed into u32 in awq ordering."""
+
+    awq_u4 = jax.lax.bitcast_convert_type(awq_u32_packed, jnp.uint4)
+
+    # AWQ packs 8 uint4 into 32-bits in this order: (0, 2, 4, 6, 1, 3, 5, 7).
+    # Following list maps the order used by AWQ into an ascending order.
+    reverse_awq_order = (0, 4, 1, 5, 2, 6, 3, 7)
+    u4 = awq_u4[..., reverse_awq_order]
+    return jnp.reshape(u4, u4.shape[:-2] + (-1, ))
+
+
+def dequantize_tensor(
+    tensor_q: jax.Array,
+    scale: jax.Array,
+    axis: int | None | tuple = -1,
+    out_dtype: jnp.dtype = jnp.bfloat16,
+) -> jax.Array:
+    """Dequantize a quantized tensor
+
+    Args:
+        tensor_q: Quantized tensor.
+        scale: Quantization scale.
+        axis: The axis tensor was quantized. None denotes per-tensor.
+        out_dtype: Dtype of the output.
+
+    Returns:
+        Dequantized tensor_q.
+    """
+    if axis is None:
+        # Perform per-tensor quantization.
+        axis = [i for i in range(tensor_q.ndim)]
+    if isinstance(axis, int):
+        axis = [axis]
+
+    orig_shape = tensor_q.shape
+    if tensor_q.ndim == scale.ndim:
+        # Indicates the tensor was block quantized.
+        blocked_shape = [[i] for i in orig_shape]
+        for i in axis:
+            num_blocks = scale.shape[i]
+            if tensor_q.shape[i] % num_blocks:
+                raise ValueError(
+                    f"Unable to perform block dequantization. axis={i} of "
+                    f"{tensor_q.shape=} is not divisible by {num_blocks=}", )
+            block_size = tensor_q.shape[i] // num_blocks
+
+            blocked_shape[i] = (num_blocks, block_size)
+
+        # Convert all axis into positive values.
+        axis = sorted([(i + tensor_q.ndim) % tensor_q.ndim for i in axis])
+        # Shift axis by 1 since its original position is now occupied by
+        # num_blocks dim. Also, if n axes before an axis was also quantized,
+        # shift its position by n.
+        axis = [1 + n + i for n, i in enumerate(axis)]
+
+        # Flatten list of lists that contains (num_blocks, block).
+        blocked_shape = list(itertools.chain(*blocked_shape))
+        tensor_q = tensor_q.reshape(blocked_shape)
+
+    scale = jnp.expand_dims(scale, axis)
+
+    tensor = (tensor_q.astype(jnp.float32) * scale).astype(out_dtype)
+
+    return tensor.reshape(orig_shape)
+
+
+def dequantize_tensor_from_mxfp4_packed(
+    tensor_q: jax.Array,
+    scale: jax.Array,
+    axis: int | tuple = -1,
+    out_dtype: jnp.dtype = jnp.bfloat16,
+) -> jax.Array:
+    """Dequantize packed mxfp4 tensor.
+
+    Args:
+        tensor_q: fp4 tensor packed into uint8.
+        scale: e8m0 scale packed into uint8.
+        axis: The axis tensor was quantized.
+        out_dtype: Dtype of the output.
+
+    Returns:
+        Dequantized tensor_q.
+    """
+    tensor_e2m1 = u8_unpack_e2m1(tensor_q)
+    scale_fp32 = e8m0_to_fp32(scale)
+
+    return dequantize_tensor(
+        tensor_e2m1,
+        scale_fp32,
+        axis,
+        out_dtype,
+    )
+
+
+def quantize_tensor(
+    dtype: jnp.dtype,
+    tensor: jax.Array,
+    axis: int | tuple | None = -1,
+    block_size: int | None = None,
+    pad_tensor: bool = False,
+) -> tuple[jax.Array, jax.Array]:
+    """Quantize tensor.
+
+    Args:
+        dtype: dtype to perform quantization.
+        tensor: Unquantized tensor
+        axis: Axis to perform quantization. None denotes per-tensor.
+        block_size: Specify block quantization size.
+        pad_tensor: Whether to pad the axis along block size.
+
+    Returns:
+        Tensor quantized to dtype.
+    """
+    if axis is None:
+        # Perform per-tensor quantization.
+        axis = [i for i in range(tensor.ndim)]
+    if isinstance(axis, int):
+        axis = [axis]
+
+    orig_shape = tensor.shape
+    mask = jnp.ones_like(tensor, jnp.int32)
+
+    if block_size is not None:
+        if isinstance(block_size, int):
+            block_size = [block_size] * len(axis)
+
+        blocked_shape = [[i] for i in orig_shape]
+        pad_width = [[0, 0] for _ in range(tensor.ndim)]
+        for i, block in zip(axis, block_size):
+            num_blocks = (tensor.shape[i] + block - 1) // block
+            padding_size = num_blocks * block - tensor.shape[i]
+            if padding_size and not pad_tensor:
+                raise ValueError(
+                    f"Unable to perform block quantization. axis={i} of "
+                    f"{tensor.shape=} is not divisible by {block=}")
+
+            # Pad the tensor to align with block size.
+            pad_width[i][1] = padding_size
+
+            blocked_shape[i] = (num_blocks, block)
+
+        # In order to avoid padded values affecting scale value, we pad it
+        # using edge value of the tensor.
+        tensor = jnp.pad(tensor, pad_width, "edge")
+        mask = jnp.pad(mask, pad_width)
+
+        orig_shape = tensor.shape
+        # Convert all axis into positive values.
+        axis = sorted([i % tensor.ndim for i in axis])
+        # Shift axis by 1 since its original position is now occupied by
+        # num_blocks dim. Also, if n axes before an axis was also quantized,
+        # shift its position by n.
+        axis = [1 + n + i for n, i in enumerate(axis)]
+
+        # Flatten list of lists that contains (num_blocks, block).
+        blocked_shape = list(itertools.chain(*blocked_shape))
+        tensor = tensor.reshape(blocked_shape)
+
+    if jnp.issubdtype(dtype, jnp.integer):
+        dtype_info = jnp.iinfo(dtype)
+    else:
+        dtype_info = jnp.finfo(dtype)
+
+    dtype_max = float(dtype_info.max)
+    dtype_min = float(dtype_info.min)
+
+    abs_max = jnp.max(jnp.abs(tensor), axis=axis, keepdims=True)
+    scale = abs_max / dtype_max
+
+    tensor_q = jnp.clip(tensor / scale, dtype_min, dtype_max)
+    tensor_q = tensor_q.reshape(orig_shape)
+    tensor_q = tensor_q.astype(dtype)
+
+    # To avoid padded values affecting output of quantized matmul, we mask them
+    # out with 0s.
+    tensor_q = jnp.where(mask, tensor_q, 0)
+
+    scale = jnp.squeeze(scale, axis).astype(jnp.float32)
+
+    return tensor_q, scale
+
+
+def static_per_tensor_quantize_tensor(
+    dtype: jnp.dtype,
+    tensor: jax.Array,
+    scale: float,
+) -> jax.Array:
+    if jnp.issubdtype(dtype, jnp.integer):
+        dtype_info = jnp.iinfo(dtype)
+    else:
+        dtype_info = jnp.finfo(dtype)
+
+    dtype_max = float(dtype_info.max)
+    dtype_min = float(dtype_info.min)
+
+    return jnp.clip(tensor / scale, dtype_min, dtype_max).astype(dtype)
+
+
+def quantize_kv(
+    dtype: jnp.dtype,
+    key: jax.Array,
+    value: jax.Array,
+    k_scale: float,
+    v_scale: float,
+) -> Tuple[jax.Array, jax.Array]:
+    """Static quantize key and value tensors."""
+    key = static_per_tensor_quantize_tensor(dtype, key, k_scale)
+    value = static_per_tensor_quantize_tensor(dtype, value, v_scale)
+    return key, value

tpu_inference/layers/common/sharding.py

@@ -1,6 +1,19 @@
+# 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 json
 import math
-import os
 from dataclasses import asdict, dataclass
 from typing import TYPE_CHECKING, List, Optional
 
@@ -8,7 +21,7 @@ import jax.numpy as jnp
 import numpy as np
 from jax.sharding import Mesh
 
-from tpu_inference import utils
+from tpu_inference import envs, utils
 
 if TYPE_CHECKING:
     from vllm.v1.configs.vllm_config import VllmConfig
@@ -27,7 +40,7 @@ class ShardingAxisNameBase:
     MLP_TENSOR = ('attn_dp', 'model', 'expert')
     MOE_TENSOR = ('attn_dp', 'model')
     EXPERT = ('attn_dp', 'expert', 'model')
-    VOCAB = ('expert', 'model')
+    VOCAB = ('expert', 'attn_dp', 'model')
 
 
 class ShardingAxisName2D:
@@ -48,7 +61,7 @@ class ShardingAxisName2D:
 
 
 try:
-    _use_base_sharding = os.getenv("NEW_MODEL_DESIGN", False)
+    _use_base_sharding = envs.NEW_MODEL_DESIGN
     if _use_base_sharding:
         ShardingAxisName = ShardingAxisNameBase
     else:
@@ -120,10 +133,19 @@ class ShardingConfigManager:
                                                     False)
         if enable_dp_attention:
             # Replicate attention layer when num_kv_heads < TP
-            num_kv_heads = vllm_config.model_config.get_total_num_kv_heads()
+            num_kv_heads = 1 if vllm_config.model_config.use_mla else vllm_config.model_config.get_total_num_kv_heads(
+            )
+            cache_dtype = vllm_config.cache_config.cache_dtype
+            if cache_dtype == 'auto':
+                cache_dtype = vllm_config.model_config.dtype
             kv_dtype = utils.get_jax_dtype_from_str_dtype(
-                vllm_config.cache_config.cache_dtype) or jnp.bfloat16
+                cache_dtype) or jnp.bfloat16
             packing = 4 // jnp.dtype(kv_dtype).itemsize
+
+            # The default head dim is 128 but 64 is also supported as a special case.
+            if vllm_config.model_config.get_head_size() == 64:
+                packing *= 2
+
             # When num_kv_heads * 2 / packing < TP, tensor parallelism would
             # duplicate KV heads across devices, wasting kv cache memory.
             # Use attention DP instead to reduce per-device num_kv_heads and
@@ -166,10 +188,11 @@ class ShardingConfigManager:
                     f"LoRA is not supported with data parallelism "
                     f"(DP size: {total_dp_size}). Please disable LoRA or "
                     f"set data parallelism to 1.")
-            if not os.environ.get("NEW_MODEL_DESIGN", False):
+        if sharding_strategy.attention_data_parallelism > 1:
+            if not envs.NEW_MODEL_DESIGN:
                 raise ValueError(
-                    "Must run DP with NEW_MODEL_DESIGN enabled. Please set the "
-                    "NEW_MODEL_DESIGN=True.")
+                    "Must run Attention DP with NEW_MODEL_DESIGN enabled. Please set "
+                    "NEW_MODEL_DESIGN=True")
 
     @property
     def total_dp_size(self) -> int:

tpu_inference/layers/common/utils.py

@@ -0,0 +1,94 @@
+# 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
+
+
+def reorder_concatenated_tensor_for_sharding(concatenated_tensor: jax.Array,
+                                             split_sizes: list[int],
+                                             n_shards: int, dim: int):
+    """
+    Reorder a replicated concatenated tensor such that when sharded on multiple chips, each shard is a concatenation of the shards of the individual tensors.
+    For example, let the concatenated_tensor be:
+        AAAAAAAAAAAABBBBBBBBCCCC
+            12 As     8 Bs  4 Cs
+    and let the split_sizes = [12, 8, 4] and n_shards = 4.
+    The output is:
+        AAABBCAAABBCAAABBCAAABBC
+    In other words, it reorders the input tensor into 4 segements, with each segment corresponding to a shard and being AAABBC.
+    Args:
+        concatenated_tensor: the tensor, concatenated on the dimension specified by `dim`.
+        split_sizes: each individual tensor's size on the dimension specified by `dim`.
+        n_shards: num of shards.
+        dim: the dimension on which the concatenated_tensor is concatenated.
+    """
+    # Split the concatenated tensor into individual tensors.
+    if dim < 0:
+        dim += concatenated_tensor.ndim
+    split_tensors = []
+    start_offset = 0
+    old_shape = concatenated_tensor.shape
+    # New shape ensures each split_tensor[i] maps to a tensor in ith shards
+    new_shape = old_shape[:dim] + (n_shards, -1) + old_shape[dim + 1:]
+    for split_size in split_sizes:
+        split_tensor = jax.lax.slice_in_dim(concatenated_tensor,
+                                            start_offset,
+                                            start_offset + split_size,
+                                            axis=dim)
+        split_tensors.append(split_tensor.reshape(new_shape))
+        start_offset += split_size
+    # While maintaining 0th dim as a shard dim, we concatenate along 1th dim to
+    # to create concatenated tnensor where 0th dim maps to shard dim.
+    reordered_tensor = jnp.concatenate(split_tensors, axis=dim + 1)
+    return reordered_tensor.reshape(old_shape)
+
+
+def slice_sharded_tensor_for_concatenation(sharded_tensor: jax.Array,
+                                           split_sizes: list[int],
+                                           n_shards: int):
+    """
+    Slice the input tensor which is sharded on multiple chips (on the last dim) into individual tensors with the same sharding.
+    For example, let the sharded_tensor be:
+        AAABBC | AAABBC | AAABBC | AAABBC
+        Shard0   Shard1   Shard2   Shard3
+    and let the split_sizes = [12, 8, 4] and n_shards = 4.
+    The output is a list of 3 tensors:
+         AAA   |  AAA   |  AAA   |  AAA
+          BB   |   BB   |   BB   |   BB
+           C   |    C   |    C   |    C
+        Shard0   Shard1   Shard2   Shard3
+    In other words, each individual tensor is a slice of the input tensor with the same sharding.
+    Args:
+        sharded_tensor: the input tensor, sharded on the last dim.
+        split_sizes: each individual tensor's size on the last dim.
+        n_shards: num of shards.
+    """
+    new_shape = sharded_tensor.shape[:-1] + (n_shards, -1)
+    # New shape ensures each sharded_tensor[:, i] maps to a tensor in ith shards
+    sharded_tensor = sharded_tensor.reshape(new_shape)
+
+    split_tensors = []
+    start_offset = 0
+    for split_size in split_sizes:
+        assert split_size % n_shards == 0
+        sz = split_size // n_shards  # size of this split tensor per shard
+        end_offset = start_offset + sz
+        # Because we are slicing over last dim, sharding dim remains intact.
+        # Therefore, splitting happens locally.
+        split_tensor = sharded_tensor[..., start_offset:end_offset]
+        split_tensors.append(split_tensor.reshape(new_shape[:-2] + (-1, )))
+        start_offset = end_offset
+
+    return split_tensors

tpu_inference/layers/jax/__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/jax/attention/__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/jax/attention/attention.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 dataclasses import InitVar, dataclass
 from typing import Any, Tuple
 
@@ -5,7 +19,6 @@ import jax
 import jax.numpy as jnp
 from flax import nnx
 from flax.typing import Sharding
-from jax.experimental import shard_map
 from jax.sharding import Mesh
 from jax.sharding import PartitionSpec as P
 
@@ -13,6 +26,7 @@ from tpu_inference import utils
 from tpu_inference.kernels.ragged_paged_attention.v3.kernel import \
     ragged_paged_attention
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
 from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference.layers.jax.base import create_param
 from tpu_inference.layers.jax.rope_interface import apply_rope
@@ -149,9 +163,8 @@ class Attention(nnx.Module):
             # q_scale = self._q_scale
             k_scale = self._k_scale
             v_scale = self._v_scale
-            k_SKH, v_SKH = utils.quantize_kv(k_SKH, v_SKH,
-                                             self.kv_cache_quantized_dtype,
-                                             k_scale, v_scale)
+            k_SKH, v_SKH = quantize_kv(self.kv_cache_quantized_dtype, k_SKH,
+                                       v_SKH, k_scale, v_scale)
 
         with jax.named_scope("attn_op"):
             new_kv_cache, outputs_TNH = self.attention(
@@ -236,12 +249,12 @@ class Attention(nnx.Module):
             )
 
         output_TNH, kv_cache = jax.jit(
-            shard_map.shard_map(
+            jax.shard_map(
                 _ragged_paged_attention,
                 mesh=mesh,
                 in_specs=in_specs,
                 out_specs=out_specs,
-                check_rep=False,
+                check_vma=False,
             ))(
                 q_TNH,
                 k_SKH,