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

tpu_inference/layers/common/binary_search.py

@@ -0,0 +1,295 @@
+# Copyright 2024 The T5X Authors.
+#
+# 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.
+"""Binary search over float32 bits.
+
+Includes fast algorithms top-k masking and top-p masking on probability
+distributions.
+"""
+
+from typing import Callable, Sequence
+
+import jax
+from jax import lax
+from jax import numpy as jnp
+
+
+def int32_bsearch(batch_shape: Sequence[int],
+                  predicate: Callable[[jnp.ndarray], jnp.ndarray]):
+    """Batched binary search over int32 values.
+
+    For each element of the batch, search for the largest int32 (closest to
+    positive infinity) for which the predicate is False. If the predicate is
+    always True, returns the minimum int32 value.
+
+    Args:
+      batch_shape: Shape of the search that we're batching over.
+      predicate: the query we're searching for. For every batch element, this is
+        required to be a monotonic function from int32 to bool. In other words,
+        the predicate must return False for all numbers <= some threshold and then
+        return True for all numbers > that threshold. The threshold may be
+        different for different elements of the batch.
+
+    Returns:
+      For each element of the batch, the largest int32 for which the predicate
+      returns False. Shape: batch_shape.
+    """
+    current_bits = jnp.zeros(batch_shape, dtype=jnp.int32)
+
+    # bit 31 is special, because it compares in the opposite order of all other
+    # bits. we use uint32 due to numpy promotion/casting rules.
+    midpoint = current_bits
+    predicate_satisfied = predicate(midpoint)
+    current_bits = current_bits | jnp.where(predicate_satisfied,
+                                            jnp.uint32(1 << 31), jnp.uint32(0))
+    del midpoint, predicate_satisfied
+
+    def loop_body(i, current_bits):
+        bit_index = 30 - i
+        bit = jnp.int32(1 << bit_index)
+        midpoint = current_bits | bit
+        predicate_satisfied = predicate(midpoint)
+        current_bits = current_bits | jnp.where(predicate_satisfied,
+                                                jnp.int32(0), bit)
+        return current_bits
+
+    current_bits = lax.fori_loop(0, 31, loop_body, current_bits)
+    return current_bits
+
+
+def _monotonic_int32_to_float32_bit_pattern(x: int) -> int:
+    """Converts an int32 to a float32 bit pattern with consistent ordering.
+
+    This function is the unique function that is monotonic with respect to the
+    floating point total order, see
+    https://en.wikipedia.org/wiki/IEEE_754#Total-ordering_predicate. Note that
+    this function returns an int32, not a float32. For the function that returns
+    float32, see `monotonic_int32_to_float32`.
+
+    Args:
+      x: int bit pattern.
+
+    Returns:
+      Bit pattern of a float32 number.
+    """
+    non_sign_bits = jnp.int32((1 << 31) - 1)
+    # See
+    # https://stackoverflow.com/questions/20097380/iee-754-total-order-in-standard-c11
+    # for the relationship between int32 order and f32 total order, including
+    # the "xor trick".
+
+    # Flip the sort order for numbers where the sign bit is set. On int32,
+    # the bit pattern with sign bit set and all other bits clear is the most
+    # negative bit pattern (it's int32::MIN), whereas on float32 it's the least
+    # negative bit pattern (it's -0.0). Flipping all the non-sign bits makes the
+    # int32 sort order consistent with the float32 sort order.
+    x = x ^ jnp.where(x < 0, non_sign_bits, jnp.int32(0))
+    return x
+
+
+def _monotonic_int32_to_float32(x: int) -> jax.Array:
+    """Converts an int32 to a float32 with consistent ordering.
+
+    This function is the unique function that is monotonic with respect to the
+    floating point total order, see
+    https://en.wikipedia.org/wiki/IEEE_754#Total-ordering_predicate.
+
+    Args:
+      x: int bit pattern.
+
+    Returns:
+      float32 number with consistent ordering.
+    """
+    x = _monotonic_int32_to_float32_bit_pattern(x)
+    return lax.bitcast_convert_type(x, jnp.float32)
+
+
+def float32_bsearch(batch_shape, predicate):
+    """Binary search on finite float32 numbers.
+
+    For each element of the batch, this function searches for the largest finite
+    non-NaN float32 for which the predicate is False.
+
+    Args:
+      batch_shape: Shape of the search that we're batching over.
+      predicate: the query we're searching for. This is required to be monotonic
+        with respect to the floating point order, i.e. it must be False for all
+        numbers <= a threshold, and then True for all numbers > the threshold. The
+        threshold may be different for different elements of the batch.
+
+    Returns:
+      For each element of the batch, the largest float32 for which the predicate
+      returns False. Shape: f32[batch_shape].
+    """
+    exponent_bits = jnp.int32((1 << 31) - (1 << (31 - 8)))
+
+    def int32_predicate(x):
+        x = _monotonic_int32_to_float32_bit_pattern(x)
+        is_finite = (x & exponent_bits) != exponent_bits
+
+        # Non-finite numbers (infinity and NaN) are at the very extremes of the
+        # int32 range, i.e. they include int32::MAX and int32::MIN, plus the numbers
+        # adjacent to them. For the nonfinite numbers touching int32::MIN, we
+        # arrange for them to return False from the predicate, and for the nonfinite
+        # numbers touching int32::MAX, we arrange for them to return True from the
+        # predicate. x>=0 is an easy way to achieve that.
+        predicate_on_nonfinite = x >= 0
+        x_float32 = lax.bitcast_convert_type(x, jnp.float32)
+        return jnp.where(is_finite, predicate(x_float32),
+                         predicate_on_nonfinite)
+
+    # We search over bit patterns, which requires bit shifting and ordering of bit
+    # patterns. This is natively supported on int32 but not on float32.
+    # Additionally, it's more common to reason about int32 bit arithmetic and
+    # ordering than float32 bit arithmetic and ordering, so we do the core of our
+    # search in int32. Additionally, this allows us to test the underlying binary
+    # search on int32 values.
+    #
+    # The function _monotonic_int32_to_float32 encapsulates all of the knowledge
+    # we need about float32 bit patterns.
+    result = int32_bsearch(batch_shape, int32_predicate)
+    return _monotonic_int32_to_float32(result)
+
+
+def topk_mask(x: jnp.ndarray, k: int, replace_val: jnp.ndarray) -> jnp.ndarray:
+    """Sets everything to replace_val, except the top k values per batch element.
+
+    Sharding considerations: this function does 32 reductions over the vocab_size
+    axis of the input array. To avoid excessive latency from these reductions, you
+    should ensure that the vocab_size axis is unsharded on input to this function.
+    Prefer to shard the batch axes instead.
+
+    Scratchpad memory considerations: this function is most efficient if the
+    entire input array can fit in a fast memory tier. To help ensure this, you may
+    wish to split the batch axes into microbatches and the microbatches in a
+    sequential loop.
+
+    Args:
+      x: Values before masking. [batch..., vocab_size]
+      k: Number of masked values to return. In presence of ties, more than k
+        values might be returned.
+      replace_val: For the masked values of x, what to overwrite them with.
+
+    Returns:
+      masked version of x. [batch..., vocab_size]
+    """
+    batch_shape = tuple(list(x.shape)[:-1])  # [batch...]
+
+    x_for_loop = x
+    reduce_axis = x.ndim - 1
+    if x.ndim > 1:
+        # We're going to be doing 32 reductions over 'reduce_axis'. Generally,
+        # reductions over the last dimension are the most expensive, because they
+        # involve reducing across vector lanes, which is often not efficient. So
+        # we transpose the reduce_axis to be the second-last dimension, to avoid
+        # this inefficiency.
+        #
+        # Normaly the XLA compiler would automatically perform this optimization,
+        # but it doesn't yet see through loops to do so. So we do it ourselves.
+        x_for_loop = jnp.swapaxes(x_for_loop, -1, -2)
+        reduce_axis = x.ndim - 2
+
+    # x: [batch..., vocab_size, batch]
+    def predicate(threshold):
+        # threshold: [batch...]
+
+        # Since we've negated, we now want a predicate that is True for small
+        # numbers and False for large numbers. The result of the bsearch is the
+        # smallest float32 for which the predicate is False.
+        threshold = -threshold
+
+        threshold = lax.expand_dims(threshold, (reduce_axis, ))
+        # threshold: [batch..., 1, last_batch]
+
+        # count_ge: [batch...]
+        count_gt = jnp.sum(x_for_loop > threshold, axis=reduce_axis)
+
+        return count_gt >= k
+
+    # cutoff: [batch...]
+    cutoff = float32_bsearch(batch_shape, predicate)
+    cutoff = -cutoff
+    # cutoff: [batch..., 1]
+    cutoff = lax.expand_dims(cutoff, (cutoff.ndim, ))
+    return jnp.where(x >= cutoff, x, jnp.full_like(x, replace_val))
+
+
+def topp_mask(logits: jnp.ndarray, p: float,
+              replace_val: jnp.ndarray) -> jnp.ndarray:
+    """Applies top-p masking to logits.
+
+    Masks logits down to the smallest set of choices, such that the total
+    probability mass is >= p. Values in this set are left as they are. All other
+    values are set with `replace_val`.
+
+    Sharding considerations: this function does 33 reductions over the vocab_size
+    axis of the input array. To avoid excessive latency from these reductions, you
+    should ensure that the vocab_size axis is unsharded on input to this function.
+    Prefer to shard the batch axes instead.
+
+    Scratchpad memory considerations: this function is most efficient if the
+    entire input array can fit in a fast memory tier. To help ensure this, you may
+    wish to split the batch axes into microbatches and the microbatches in a
+    sequential loop.
+
+    Args:
+      logits: Logits before masking. [batch..., vocab_size]
+      p: Minimum probability mass requested.
+      replace_val: For the masked values of logits, what to overwrite them with.
+
+    Returns:
+      masked version of x. [batch..., vocab_size]
+    """
+    batch_shape = tuple(list(logits.shape)[:-1])  # [batch...]
+
+    probs = jax.nn.softmax(logits, axis=-1)
+
+    probs_for_reduction = probs
+    reduce_axis = probs_for_reduction.ndim - 1
+    if probs_for_reduction.ndim > 1:
+        # We're going to be doing 33 reductions over 'reduce_axis'. Generally,
+        # reductions over the last dimension are the most expensive, because they
+        # involve reducing across vector lanes, which is often not efficient. So
+        # we transpose the reduce_axis to be the second-last dimension, to avoid
+        # this inefficiency.
+        probs_for_reduction = jnp.swapaxes(probs_for_reduction, -1, -2)
+        reduce_axis = probs_for_reduction.ndim - 2
+
+    # As we increase the threshold, the probability mass decreases, and the number
+    # selected decreases.
+    #
+    # We want the largest threshold with the probability mass >= p. Binary search
+    # searches for when the predicate is False, so we negate the output of the
+    # predicate, i.e. probability mass < p.
+
+    # probs_for_reduction: [batch..., vocab_size, batch]
+    def predicate(threshold):
+        # threshold: [batch...]
+        threshold = lax.expand_dims(threshold, (reduce_axis, ))
+        # threshold: [batch..., 1, last_batch]
+
+        # count_ge: [batch...]
+        probability_mass = jnp.sum(
+            jnp.where(probs_for_reduction >= threshold, probs_for_reduction,
+                      0.0),
+            axis=reduce_axis,
+        )
+
+        return probability_mass < p
+
+    # threshold: [batch...]
+    threshold = float32_bsearch(batch_shape, predicate)
+    # threshold: [batch..., 1]
+    threshold = lax.expand_dims(threshold, (threshold.ndim, ))
+    return jnp.where(probs >= threshold, logits,
+                     jnp.full_like(logits, replace_val))

tpu_inference/layers/common/quant_methods.py

@@ -0,0 +1,23 @@
+# 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.
+
+UNQUANTIZED = "unquantized"
+MXFP4 = "mxfp4"
+AWQ = "awq"
+COMPRESSED_TENSORS = "compressed-tensors"
+FP8 = "fp8"
+
+
+def get_tpu_quant_method(quant_method: str) -> str:
+    return "tpu-" + quant_method

tpu_inference/layers/common/quantization.py

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