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

tpu_inference/layers/common/attention_interface.py

@@ -0,0 +1,403 @@
+# 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 functools
+import math
+from typing import Any, Callable, Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+from jax.experimental.pallas.ops.tpu.paged_attention import paged_attention
+from jax.experimental.pallas.ops.tpu.splash_attention import \
+    splash_attention_kernel as splash
+from jax.experimental.pallas.ops.tpu.splash_attention import \
+    splash_attention_mask as mask_lib
+from jax.sharding import Mesh
+from jax.sharding import PartitionSpec as P
+
+import tpu_inference.kernels.ragged_paged_attention.v3.kernel as rpa
+import tpu_inference.kernels.ragged_paged_attention.v3.kernel_hd64 as rpa_hd64
+from tpu_inference.kernels.flash_attention.kernel import flash_attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.utils import get_megacore
+
+MAX_ALLOWED_PAGE_INDICES_N = (
+    128 * 1024
+)  # Based on experiments on v5e, 256x1024 results in smem oom but 128x1024 not. TODO: Adjust this based on TPU version.
+
+ragged_paged_attention = rpa.ragged_paged_attention
+get_kv_cache_shape = rpa.get_kv_cache_shape
+
+ragged_paged_attention_hd64 = rpa_hd64.ragged_paged_attention_hd64
+get_kv_cache_shape_hd64 = rpa_hd64.get_kv_cache_shape
+
+
+def sharded_flash_attention(
+    mesh: Mesh,
+    causal: bool = True,
+    sm_scale: Optional[float] = None,
+    vmem_limit_bytes: int | None = None,
+) -> Callable[..., Any]:
+    in_specs = (
+        P("data", "model", None, None),  # q
+        P("data", "model", None, None),  # k
+        P("data", "model", None, None),  # v
+        P(),  # segment_ids
+    )
+    out_specs = P("data", "model", None, None)
+
+    def _flash_attention(q, k, v, segment_ids):
+        return flash_attention(q,
+                               k,
+                               v,
+                               segment_ids=segment_ids,
+                               sm_scale=sm_scale,
+                               causal=causal,
+                               vmem_limit_bytes=vmem_limit_bytes)
+
+    return jax.jit(
+        jax.shard_map(_flash_attention,
+                      mesh=mesh,
+                      in_specs=in_specs,
+                      out_specs=out_specs,
+                      check_vma=False))
+
+
+def sharded_paged_attention(
+    mesh: Mesh,
+    attn_logits_soft_cap: Optional[float] = None,
+) -> Callable[..., Any]:
+    """Shards GQA PagedAttention along KV heads."""
+    in_specs = (
+        P(None, "model", None),  # q
+        P("model", None, None, None),  # k
+        P("model", None, None, None),  # v
+        P(),  # lengths
+        P(),  # page_indices
+    )
+    out_specs = P(None, "model", None)
+
+    def _paged_attention_fn(q, k, v, lengths, page_indices):
+        if page_indices.size > MAX_ALLOWED_PAGE_INDICES_N:
+            raise ValueError(
+                "This will result in smem OOM. Use `paged_attention_with_guarded_smem` to run with minibatches."
+            )
+        return paged_attention(
+            q,
+            k,
+            v,
+            lengths,
+            page_indices,
+            attn_logits_soft_cap=attn_logits_soft_cap,
+            pages_per_compute_block=min(
+                16, page_indices.shape[1]),  # 512 / page_size:32,
+            megacore_mode="kv_head" if get_megacore() else None,
+        )
+
+    return jax.jit(
+        jax.shard_map(
+            _paged_attention_fn,
+            mesh=mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_vma=False,
+        ))
+
+
+# TODO(xiangxu): merge this with sharded_paged_attention
+@functools.partial(jax.jit, static_argnums=[0])
+def paged_attention_with_guarded_smem(
+    paged_attention_kernel: Callable,
+    q: jax.Array,
+    k_pages: jax.Array,
+    v_pages: jax.Array,
+    lengths: jax.Array,
+    page_indices: jax.Array,
+):
+    # Addresses b/336316706. Summary:
+    # Paged attention kernel stores `lengths` (batch_size * 4 bytes) and `page_indices` (batch_size * num_blocks_per_seq * 4 bytes) in SMEM.
+    # Capacity of SMEM is quite limited which is also TPU version dependent. Models with higher context length or higher batch size, can cause OOM in SMEM.
+    # There are two solutions:
+    # 1. Reduce blocks per seq by increasing page size.
+    # 2. Splitting the batch into several minibatches (Higher perf based on my benchmark).
+
+    batch_size, blocks_per_seq = page_indices.shape
+
+    if page_indices.size <= MAX_ALLOWED_PAGE_INDICES_N:
+        return paged_attention_kernel(q, k_pages, v_pages, lengths,
+                                      page_indices)
+
+    mini_batch_size = MAX_ALLOWED_PAGE_INDICES_N // blocks_per_seq
+
+    # If batch_size is not disible by mini_batch_size,
+    # we set mini_batch_size to a smaller value, i.e GCD,
+    # which will trigger more kernel launches but it's fine.
+    # TODO: Fix --decode_seqs_padding with this limitation.
+    mini_batch_size = math.gcd(batch_size, mini_batch_size)
+
+    num_kernel_launches = batch_size // mini_batch_size
+
+    outputs = jnp.zeros_like(q).reshape(
+        (num_kernel_launches, mini_batch_size, *q.shape[1:]))
+    q = q.reshape((num_kernel_launches, mini_batch_size, *q.shape[1:]))
+    seq_lens = lengths.reshape((num_kernel_launches, mini_batch_size))
+    block_indices = page_indices.reshape(
+        (num_kernel_launches, mini_batch_size, page_indices.shape[1]))
+
+    for i in range(num_kernel_launches):
+        outputs = outputs.at[i].set(
+            paged_attention_kernel(q[i], k_pages, v_pages, seq_lens[i],
+                                   block_indices[i]))
+
+    outputs = outputs.reshape((batch_size, *outputs.shape[2:]))
+
+    return outputs
+
+
+# ruff: noqa: E741
+def update_cache(
+    is_prefill,
+    cache,
+    indices,
+    operand,
+    prefill_seq_len=None,
+    sliding_window=None,
+) -> jax.Array:
+
+    # (8, 55640, 32, 128) (1, 8, 256, 128) -> K (8, 8, 32, 128)
+    # I = B * T // S
+    # k cache, operand
+
+    B, K, T, H = operand.shape
+    K_c, L, S, H = cache.shape
+    assert K == K_c
+    # NOTE: The cache updating is pretty tricky:
+    # 1. The random access updating cache is not as performant as the slice updating.
+    #    If the random access is necessary, make sure the indexing count is as small as possible.
+    # 2. The random access updating may trigger extra tranpose (memory copy) of cache,
+    #    which is a disaster because the cache is huge. This is a data formatting op inserted by
+    #    the XLA compiler and not well documented.
+    # To mitigate the issues above:
+    # For prefill:
+    # We reshape the operand so that we can update the cache in block wise, which only requires the block indices.
+    # For decode:
+    # We reshape the cache so that we can update the cache in token wise, which only requires the token indices (block_id + offset).
+    if is_prefill:
+        # In the case of sliding window, we should select sliding_window tokens from actual prompt, not from the padded tokens.
+        if sliding_window and T > sliding_window:
+            assert B == 1
+            start_index = jax.lax.max(0, prefill_seq_len - sliding_window)
+            operand = jax.lax.dynamic_slice_in_dim(
+                operand, start_index, sliding_window,
+                axis=2)  # TODO: @pooyam Perf check this.
+            T = sliding_window
+
+        I = B * T // S
+        # cache: (K, L, S, H)
+        # operand: (B, K, T, H) -> (K, I, S, H)
+        # indices: (B, T // S) -> (I,)
+        operand = jnp.swapaxes(operand, 0, 1).reshape(K, I, S, H)
+        indices = indices.reshape(I)
+        cache = cache.at[:, indices, :, :].set(operand)
+    else:
+        # cache: (K, L, S, H) -> (K, L * S, H)
+        # operand: (B, K, 1, H) -> (K, B, H)
+        # indices: (B,)
+        cache = cache.reshape(K, L * S, H)
+        operand = jnp.swapaxes(operand, 0, 1).reshape(K, B, H)
+        # NOTE: `cache.[:, indices, :].set()` will trigger the extra tranpose of the cache.
+        # The `jnp.arange(K)[..., None]` trick is to avoid it. WTF?
+        cache = cache.at[jnp.arange(K)[..., None], indices, :].set(operand)
+        cache = cache.reshape(K, L, S, H)
+    return cache
+
+
+@functools.partial(
+    jax.jit, static_argnames=["window_size", "attn_logits_soft_cap", "is_mqa"])
+def apply_splash(q, k, v, window_size, attn_logits_soft_cap,
+                 is_mqa) -> jax.Array:
+    # q: (batch_size, num_heads, seq_len, head_dim)
+    num_heads = q.shape[1]
+    q_seq_len = q.shape[2]
+    kv_seq_len = k.shape[2]
+    assert kv_seq_len >= q_seq_len
+
+    masks = [
+        mask_lib.LocalMask((q_seq_len, kv_seq_len), (window_size, 0),
+                           kv_seq_len - q_seq_len) for _ in range(num_heads)
+    ]
+    mask = mask_lib.MultiHeadMask(tuple((m for m in masks)))
+    block_sizes = splash.BlockSizes.get_default()
+
+    if is_mqa:
+        attn = splash.make_splash_mqa_single_device(
+            mask,
+            block_sizes=block_sizes,
+            attn_logits_soft_cap=attn_logits_soft_cap)
+    else:
+        attn = splash.make_splash_mha_single_device(
+            mask,
+            block_sizes=block_sizes,
+            attn_logits_soft_cap=attn_logits_soft_cap)
+    attn = jax.vmap(attn)
+    outputs = attn(q, k, v, None)
+
+    return outputs
+
+
+def sharded_splash_attention(
+    mesh: Mesh,
+    window_size: Optional[int] = None,
+    attn_logits_soft_cap: Optional[float] = None,
+    is_mqa: bool = False,
+) -> Callable[..., Any]:
+    in_specs = (
+        P("data", "model", None, None),  # q
+        P("data", "model", None, None),  # k
+        P("data", "model", None, None),  # vx
+    )
+    out_specs = P("data", "model", None, None)
+    return jax.jit(
+        jax.shard_map(
+            functools.partial(
+                apply_splash,
+                window_size=window_size,
+                attn_logits_soft_cap=attn_logits_soft_cap,
+                is_mqa=is_mqa,
+            ),
+            mesh=mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_vma=False,
+        ))
+
+
+def sharded_ragged_paged_attention(
+    mesh: Mesh,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    kv_cache: jax.Array,
+    kv_lens: jax.Array,
+    page_indices: jax.Array,
+    cu_q_lens: jax.Array,
+    distribution: jax.Array,
+    attention_sink: jax.Array | None,
+    sm_scale: float,
+    attention_chunk_size: int | None = None,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+):
+    """Shards along KV heads."""
+
+    qkv_spec = P(ShardingAxisName.ATTN_DATA, ShardingAxisName.ATTN_HEAD, None)
+    kv_cache_spec = P(ShardingAxisName.ATTN_DATA, None,
+                      ShardingAxisName.ATTN_HEAD, None, None)
+    in_specs = (
+        qkv_spec,  # q
+        qkv_spec,  # k
+        qkv_spec,  # v
+        kv_cache_spec,  # kv cache
+        P(ShardingAxisName.ATTN_DATA),  # kv_lens
+        P(ShardingAxisName.ATTN_DATA),  # page_indices
+        P(ShardingAxisName.ATTN_DATA),  # cu_q_lens
+        P(ShardingAxisName.ATTN_DATA),  # distribution
+    )
+    out_specs = (qkv_spec, kv_cache_spec)
+
+    args = (q, k, v, kv_cache, kv_lens, page_indices, cu_q_lens, distribution)
+
+    use_hd64 = q.shape[-1] == 64
+    func = ragged_paged_attention_hd64 if use_hd64 else ragged_paged_attention
+
+    if attention_sink is not None:
+        if not use_hd64:
+            raise NotImplementedError(
+                "Attention sink support is only available when head_dim==64")
+
+        in_specs += (P(ShardingAxisName.ATTN_HEAD), )
+        args += (attention_sink, )
+
+    def _ragged_paged_attention(*args):
+        return func(
+            *args,
+            sm_scale=sm_scale,
+            sliding_window=attention_chunk_size,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
+        )
+
+    return jax.shard_map(
+        _ragged_paged_attention,
+        mesh=mesh,
+        in_specs=in_specs,
+        out_specs=out_specs,
+        check_vma=False,
+    )(*args)
+
+
+def attention(
+    kv_cache: jax.Array,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    attention_metadata: AttentionMetadata,
+    mesh: Mesh,
+    head_dim_original: int | None = None,  # before padding,
+    attention_chunk_size: int | None = None,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+    sinks: jax.Array | None = None,
+) -> Tuple[jax.Array, jax.Array]:
+    # T: seq_len
+    # N: num_heads
+    # K: num_kv_heads
+    # D: hidden_size
+    # H: head_dim
+    # L: num_blocks
+    # S: block_size
+
+    # TODO(jevinjiang, cuiq): transpose q weight offline.
+    # q: (T, N, H)
+    # k,v: (T, K, H)
+
+    if head_dim_original is None:
+        head_dim_original = q.shape[-1]
+
+    md = attention_metadata
+
+    # (T, N, H)
+    output, kv_cache = sharded_ragged_paged_attention(
+        mesh,
+        q,
+        k,
+        v,
+        kv_cache,
+        md.seq_lens,
+        md.block_tables,
+        md.query_start_loc,
+        md.request_distribution,
+        sinks,
+        sm_scale=head_dim_original**-0.5,
+        attention_chunk_size=attention_chunk_size,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+
+    return kv_cache, output

tpu_inference/layers/common/attention_metadata.py

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