tpu-inference 0.11.1rc1__py3-none-any.whl → 0.11.1rc3__py3-none-any.whl

tpu_inference/kernels/collectives/all_gather_matmul_tuned_block_sizes.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""All-gather matmul kernel's tuned block sizes."""
+
+import jax
+
+# key:
+#    - tpu_version
+#    - m
+#    - n
+#    - k
+#    - dtype
+#    - tp_size
+# value:
+#    - bn
+#    - bk
+TUNED_BLOCK_SIZES = {
+    # go/keep-sorted start
+    (6, 1024, 51200, 5120, 'bfloat16', 8): (6400, 2560),
+    (6, 1024, 57344, 8192, 'bfloat16', 8): (7168, 8192),
+    (6, 2048, 51200, 5120, 'bfloat16', 8): (1280, 5120),
+    (6, 2048, 57344, 8192, 'bfloat16', 8): (1024, 8192),
+    (6, 4096, 51200, 5120, 'bfloat16', 8): (3200, 5120),
+    (6, 8192, 51200, 5120, 'bfloat16', 8): (1280, 5120),
+    # go/keep-sorted end
+}
+
+
+def get_tpu_version() -> int:
+    """Returns the numeric version of the TPU, or -1 if not on TPU."""
+    kind = jax.devices()[0].device_kind
+    if 'TPU' not in kind:
+        return -1
+    if kind.endswith(' lite'):
+        kind = kind[:-len(' lite')]
+    assert kind[:-1] == 'TPU v', kind
+    return int(kind[-1])
+
+
+def get_key(
+    m,
+    n,
+    k,
+    dtype,
+    tp_size,
+):
+    """Returns the key for the given parameters."""
+    return (
+        get_tpu_version(),
+        m,
+        n,
+        k,
+        dtype,
+        tp_size,
+    )
+
+
+def get_tuned_block_sizes(m, n, k, dtype_name, tp_size):
+    """Returns the tuned block sizes for the given parameters."""
+    key = get_key(m, n, k, dtype_name, tp_size)
+    return TUNED_BLOCK_SIZES.get(key, (None, None))

tpu_inference/kernels/collectives/util.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+"""utilities for collective kernels."""
+
+import functools
+
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+
+
+def local_barrier(left_neighbor, right_neighbor, double_barrier=True):
+    """Performs a barrier with neighbors on the global barrier semaphore.
+
+  Optionally performs a second barrier, which prevents a potential race
+  when reusing the same collective_id across kernel invocations.
+
+  Args:
+    left_neighbor: Left neighbor device id.
+    right_neighbor: Right neighbor device id.
+    double_barrier: Whether to perform a second barrier.
+  """
+    barrier_sem = pltpu.get_barrier_semaphore()
+    for neighbor in [left_neighbor, right_neighbor]:
+        pltpu.semaphore_signal(
+            barrier_sem,
+            inc=1,
+            device_id=(neighbor, ),
+            device_id_type=pltpu.DeviceIdType.MESH,
+        )
+    pltpu.semaphore_wait(barrier_sem, 2)
+    if double_barrier:
+        # The double-barrier prevents a race condition where one neighbor can
+        # re-enter the kernel again on a subsequent call and increment the
+        # barrier semaphore a second time. This would unblock the current device
+        # even if the other neighbor is not ready yet.
+        # To implement a double-barrier, we stack-allocate a second REGULAR
+        # semaphore using run_scoped.
+        @functools.partial(pl.run_scoped,
+                           second_barrier=pltpu.SemaphoreType.REGULAR)
+        def _(second_barrier):
+            for neighbor in [left_neighbor, right_neighbor]:
+                pltpu.semaphore_signal(
+                    second_barrier,
+                    inc=1,
+                    device_id=(neighbor, ),
+                    device_id_type=pltpu.DeviceIdType.MESH,
+                )
+            pltpu.semaphore_wait(second_barrier, 2)

tpu_inference/layers/common/attention_metadata.py

@@ -0,0 +1,34 @@
+import functools
+from dataclasses import dataclass, field
+from typing import Any
+
+import jax
+
+
+@functools.partial(
+    jax.tree_util.register_dataclass,
+    data_fields=[
+        "input_positions",
+        "block_tables",
+        "seq_lens",
+        "query_start_loc",
+        "request_distribution",
+    ],
+    meta_fields=[],
+    drop_fields=["query_start_loc_cpu", "seq_lens_cpu"],
+)
+@dataclass
+class AttentionMetadata(object):
+    # (padded_total_num_scheduled_tokens,)
+    input_positions: jax.Array
+    # (max_num_seqs * max_num_blocks_per_req,)
+    block_tables: jax.Array = None
+    # (max_num_seqs,)
+    seq_lens: jax.Array = None
+    # (max_num_seqs + 1,)
+    query_start_loc: jax.Array = None
+    # (3,)
+    request_distribution: jax.Array = None
+
+    query_start_loc_cpu: Any = field(init=False)
+    seq_lens_cpu: Any = field(init=False)

tpu_inference/layers/jax/attention/attention.py

@@ -0,0 +1,254 @@
+from dataclasses import InitVar, dataclass
+from typing import Any, Tuple
+
+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
+
+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.jax.base import create_param
+from tpu_inference.layers.jax.rope_interface import apply_rope
+
+KVCache = Tuple[jax.Array, jax.Array]
+
+
+@dataclass(kw_only=True)
+class Attention(nnx.Module):
+    """An implementation of attention.
+
+    This module performs the attention mechanism for a transformer model,
+    including query, key, and value projections, application of Rotary
+    Position Embeddings (RoPE), and management of a KV cache for efficient
+    autoregressive generation. It supports both prefill and generation
+    (decode) modes and handles tensor sharding for distributed computation.
+
+    Attributes:
+        mesh: The JAX device mesh for distributed computation.
+    """
+    hidden_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    rope_theta: float
+    rope_scaling: dict[str, Any]
+    dtype: jnp.dtype
+    mesh: Mesh
+    kv_cache_dtype: str
+
+    dnh_sharding: Sharding = ()
+    dkh_sharding: Sharding = ()
+    nhd_sharding: Sharding = ()
+
+    activation_q_td: Sharding = ()
+    query_tnh: P = P()
+    keyvalue_skh: P = P()
+
+    attn_o_tnh: P = P()
+    rngs: InitVar[nnx.Rngs]
+
+    random_init: bool = False
+    attention_chunk_size: int | None = None
+    rope_input_ordering: str = "split"
+
+    _q_scale: float = 1.0
+    _k_scale: float = 1.0
+    _v_scale: float = 1.0
+
+    kv_cache_quantized_dtype = None
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        """Initializes the weight kernels for Q, K, V, and O projections."""
+        N = self.num_attention_heads
+        K = self.num_key_value_heads
+        D = self.hidden_size
+        H = self.head_dim
+
+        self.kernel_q_proj_DNH = create_param(rngs, (D, N, H),
+                                              self.dnh_sharding,
+                                              self.dtype,
+                                              random_init=self.random_init)
+        self.kernel_k_proj_DKH = create_param(rngs, (D, K, H),
+                                              self.dkh_sharding,
+                                              self.dtype,
+                                              random_init=self.random_init)
+        self.kernel_v_proj_DKH = create_param(rngs, (D, K, H),
+                                              self.dkh_sharding,
+                                              self.dtype,
+                                              random_init=self.random_init)
+        self.kernel_o_proj_NHD = create_param(rngs, (N, H, D),
+                                              self.nhd_sharding,
+                                              self.dtype,
+                                              random_init=self.random_init)
+
+        if self.kv_cache_dtype != "auto":
+            self.kv_cache_quantized_dtype = utils.get_jax_dtype_from_str_dtype(
+                self.kv_cache_dtype)
+
+    def __call__(self,
+                 x,
+                 is_prefill,
+                 kv_cache: KVCache,
+                 attention_metadata: AttentionMetadata,
+                 use_attention_rope: bool = True):
+        """Performs the forward pass of the attention module.
+
+        This method computes the attention output by projecting the input `x`
+        to queries, keys, and values, applying RoPE, performing scaled
+        dot-product attention, and projecting the result back to the model
+        dimension. It updates and utilizes a KV cache.
+
+        Args:
+            x: The input tensor of shape `(seq_len, d_model)`.
+            is_prefill: Whether the operation mode is prefill (otherwise it is generate).
+            kv_cache: The key-value cache for storing past attention states.
+            attention_metadata: Metadata for attention, such as input positions.
+            use_attention_rope: Whether to use RoPE.
+
+        Returns:
+            A tuple containing:
+                - The updated KV cache.
+                - The attention output tensor of shape
+                  `(batch_size, seq_len, d_model)`.
+        """
+        md = attention_metadata
+        x_SD = jnp.asarray(x, self.dtype)
+        x_q_TD = nnx.with_sharding_constraint(x, self.activation_q_td)
+        H = self.head_dim
+        with jax.named_scope("q_proj"):
+            q_TNH = jnp.einsum('TD,DNH -> TNH', x_q_TD,
+                               self.kernel_q_proj_DNH.value)
+            if use_attention_rope:
+                q_TNH = apply_rope(q_TNH, md.input_positions, H,
+                                   self.rope_theta, self.rope_scaling,
+                                   self.rope_input_ordering)
+            q_TNH = nnx.with_sharding_constraint(q_TNH, self.query_tnh)
+        with jax.named_scope("k_proj"):
+            k_SKH = jnp.einsum('SD,DKH -> SKH', x_SD,
+                               self.kernel_k_proj_DKH.value)
+            if use_attention_rope:
+                k_SKH = apply_rope(k_SKH, md.input_positions, H,
+                                   self.rope_theta, self.rope_scaling,
+                                   self.rope_input_ordering)
+            k_SKH = nnx.with_sharding_constraint(k_SKH, self.keyvalue_skh)
+
+        with jax.named_scope("v_proj"):
+            v_SKH = jnp.einsum('SD,DKH -> SKH', x_SD,
+                               self.kernel_v_proj_DKH.value)
+
+        q_scale = k_scale = v_scale = None
+        if self.kv_cache_quantized_dtype:
+            # TODO(kyuyeunk/jacobplatin): Enable w8a8 when VREG spill issue is resolved.
+            # 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)
+
+        with jax.named_scope("attn_op"):
+            new_kv_cache, outputs_TNH = self.attention(
+                is_prefill,
+                kv_cache,
+                q_TNH,
+                k_SKH,
+                v_SKH,
+                attention_metadata,
+                self.mesh,
+                q_scale=q_scale,
+                k_scale=k_scale,
+                v_scale=v_scale,
+            )
+
+        with jax.named_scope("o_proj"):
+            o_TD = jnp.einsum('TNH,NHD -> TD', outputs_TNH,
+                              self.kernel_o_proj_NHD.value)
+        return new_kv_cache, o_TD
+
+    def attention(
+        self,
+        is_prefill: bool,
+        kv_cache: KVCache,
+        q_TNH: jax.Array,
+        k_SKH: jax.Array,
+        v_SKH: jax.Array,
+        attention_metadata: AttentionMetadata,
+        mesh: Mesh,
+        q_scale: float | None = None,
+        k_scale: float | None = None,
+        v_scale: float | None = None,
+    ) -> Tuple[KVCache, jax.Array]:
+        """Performs scaled dot-product attention and updates the KV cache.
+
+        This function handles the core attention logic, which varies between
+        prefill and generation modes. In prefill, it computes self-attention
+        over the input sequence with a causal mask. In generation, it attends
+        to the full history of keys and values stored in the cache.
+
+        Args:
+            is_prefill: A boolean indicating if the mode is 'prefill'.
+            kv_cache: The key-value cache to be updated and used.
+            q_TNH: Query tensor of shape `(query_seq, num_attention_heads, head_dim)`.
+            k_SKH: Key tensor of shape `(kv_seq, num_key_value_heads, head_dim)`.
+            v_SKH: Value tensor of shape `(kv_seq, num_key_value_heads, head_dim)`.
+            attention_metadata: Metadata containing sequence lengths.
+            mesh: The JAX device mesh (unused in this specific function but
+                kept for potential future use or API consistency).
+            q_scale: Quantization scale for q.
+            k_scale: Quantization scale for k.
+            v_scale: Quantization scale for v.
+
+        Returns:
+            A tuple containing:
+                - The updated KV cache.
+                - The attention output tensor of shape
+                  `(seq, num_q_heads, head_dim)`.
+        """
+        md = attention_metadata
+        kv_cache_spec = P(None, None, "model")
+        in_specs = (
+            self.query_tnh,  # q
+            self.keyvalue_skh,  # k
+            self.keyvalue_skh,  # v
+            kv_cache_spec,  # kv_cache
+            P(),  # md.seq_lens: Replicated
+            P(),  # page_indices_flat: Replicated
+            P(),  # query_start_loc: Replicated
+            P(),  # distribution: Replicated
+        )
+
+        out_specs = (self.attn_o_tnh, kv_cache_spec)
+
+        def _ragged_paged_attention(*args):
+            return ragged_paged_attention(
+                *args,
+                sm_scale=q_TNH.shape[-1]**-0.5,
+                q_scale=q_scale,
+                k_scale=k_scale,
+                v_scale=v_scale,
+            )
+
+        output_TNH, kv_cache = jax.jit(
+            shard_map.shard_map(
+                _ragged_paged_attention,
+                mesh=mesh,
+                in_specs=in_specs,
+                out_specs=out_specs,
+                check_rep=False,
+            ))(
+                q_TNH,
+                k_SKH,
+                v_SKH,
+                kv_cache,
+                md.seq_lens,
+                md.block_tables,
+                md.query_start_loc,
+                md.request_distribution,
+            )
+        return kv_cache, output_TNH