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

tpu_inference/layers/jax/attention_interface.py

@@ -0,0 +1,356 @@
+import functools
+import math
+from typing import Any, Callable, Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+from jax.experimental import shard_map
+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
+from tpu_inference.kernels.flash_attention.kernel import flash_attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+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
+
+
+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(
+        shard_map.shard_map(_flash_attention,
+                            mesh=mesh,
+                            in_specs=in_specs,
+                            out_specs=out_specs,
+                            check_rep=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(
+        shard_map.shard_map(
+            _paged_attention_fn,
+            mesh=mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_rep=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(
+        shard_map.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_rep=False,
+        ))
+
+
+def sharded_ragged_paged_attention(
+    sm_scale: float,
+    mesh: Mesh,
+    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(None, "model", None)
+    kv_cache_spec = P(None, None, "model")
+    in_specs = (
+        qkv_spec,  # q
+        qkv_spec,  # k
+        qkv_spec,  # v
+        kv_cache_spec,  # kv cache
+        P(),  # kv_lens
+        P(),  # page_indices
+        P(),  # cu_q_lens
+        P(),  # distribution
+    )
+    out_specs = (qkv_spec, kv_cache_spec)
+
+    def _ragged_paged_attention(*args):
+        return ragged_paged_attention(
+            *args,
+            sm_scale=sm_scale,
+            sliding_window=attention_chunk_size,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
+        )
+
+    return jax.jit(
+        shard_map.shard_map(
+            _ragged_paged_attention,
+            mesh=mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_rep=False,
+        ))
+
+
+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,
+) -> 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(
+        head_dim_original**-0.5, mesh, attention_chunk_size, q_scale, k_scale,
+        v_scale)(
+            q,
+            k,
+            v,
+            kv_cache,
+            md.seq_lens,
+            md.block_tables,
+            md.query_start_loc,
+            md.request_distribution,
+        )
+
+    return kv_cache, output

tpu_inference/layers/jax/base.py

@@ -0,0 +1,151 @@
+import dataclasses
+from dataclasses import dataclass, fields
+from typing import Any, Callable, Mapping
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from flax.typing import Sharding
+from jax.sharding import PartitionSpec as P
+
+from tpu_inference.logger import init_logger
+
+# Type alias for Initializer for cleaner type hints
+Initializer = Callable[..., jax.Array]
+logger = init_logger(__name__)
+
+# Define singleton initializers to avoid re-compilation.
+_scale_initializer = nnx.initializers.ones
+_sharded_initializer = nnx.initializers.xavier_normal()
+_init_fn = nnx.initializers.uniform()
+
+
+@dataclass
+class Config:
+    """Base configuration class with a robust factory method.
+
+    This class provides a `from_cfg` classmethod that allows creating a config
+    instance from a dictionary, ensuring that all required fields are present
+    and ignoring any extraneous keys.
+    """
+
+    @classmethod
+    def from_cfg(cls, cfg: dict[str, Any] | None = None, **kwargs):
+        """Creates a config instance from a dictionary and/or keyword arguments.
+
+        This factory method validates that all fields without default values
+        are provided in the input dictionary or keyword arguments.
+
+        Args:
+            cfg: A dictionary of configuration parameters.
+            **kwargs: Additional configuration parameters passed as keyword arguments.
+
+        Returns:
+            An instance of the configuration class.
+
+        Raises:
+            ValueError: If any required parameters are missing.
+        """
+        if cfg is None:
+            cfg = {}
+        cfg.update(kwargs)
+
+        required_params = {
+            f.name
+            for f in fields(cls) if f.default is dataclasses.MISSING
+            and f.default_factory is dataclasses.MISSING
+        }
+
+        # Check if any of the truly required parameters are missing from the provided config.
+        missing_params = required_params - set(cfg.keys())
+        if missing_params:
+            raise ValueError(
+                f"Missing required parameters for {cls.__name__}: {', '.join(sorted(list(missing_params)))}"
+            )
+
+        known_params = {f.name for f in fields(cls)}
+        filtered_cfg = {k: v for k, v in cfg.items() if k in known_params}
+
+        return cls(**filtered_cfg)
+
+    # TODO: check logic with some unit tests.
+    def maybe_apply_overrides(self):
+        """Update the args with additional_configs, hf_overrides, and override_generation_config settings.
+        If there is overlap in overrides between the configs, then print a warning declaring which
+        overrides will take precedent."""
+
+        if not getattr(self, "vllm_config"):
+            return
+
+        def _overrides_str(original: str, original_val: Any,
+                           new_val: Any) -> str:
+            return f"{original}: {original_val} ---> {new_val}"
+
+        def _get_overrides_dict(self) -> Mapping[str, Any]:
+            """Return the overrides from all of the possible vllm sections."""
+            overrides_dict = {}
+            vllm_model_config = self.vllm_config.model_config
+
+            for override_type in ordered_override_types:
+                if override_type == "additional_config":
+                    overrides_dict[
+                        override_type] = self.vllm_config.additional_config
+                else:
+                    overrides_dict[override_type] = getattr(
+                        vllm_model_config, override_type)
+            return overrides_dict
+
+        ordered_override_types = [
+            "additional_config", "hf_overrides", "override_generation_config"
+        ]
+
+        overrides_dict = _get_overrides_dict(self)
+
+        # Override the config values using the vLLM sections with highest
+        # precedence first.
+        for field in fields(self):
+            selected_type = None
+            for override_type in reversed(ordered_override_types):
+                if field.name in overrides_dict[override_type]:
+                    setattr(self, field.name,
+                            overrides_dict[override_type][field.name])
+                    selected_type = override_type
+                    break
+            if selected_type is None:
+                continue
+
+            # If multiple vLLM sections contain overrides, print a warning.
+            for override_type in ordered_override_types:
+                if override_type == selected_type:
+                    break
+                else:
+                    if field.name in overrides_dict[override_type]:
+                        overriden_keys_str = _overrides_str(
+                            field.name,
+                            overrides_dict[override_type][field.name],
+                            overrides_dict[selected_type][field.name])
+                        logger.warning(
+                            f"Overriding {override_type} arguments with the following {selected_type} args: {overriden_keys_str}"
+                        )
+
+    def __post_init__(self):
+        self.maybe_apply_overrides()
+
+
+def create_param(rngs: nnx.Rngs,
+                 shape: tuple[int, ...],
+                 sharding: Sharding = (),
+                 dtype: Any = jnp.float32,
+                 random_init=False) -> nnx.Param:
+    key = rngs.params()
+    if random_init:
+        initializer = _scale_initializer if len(
+            shape) == 1 else _sharded_initializer
+
+        jitted_initializer = jax.jit(initializer,
+                                     static_argnames=('shape', 'dtype'),
+                                     out_shardings=P(*sharding))
+        param_data = jitted_initializer(key, shape, dtype)
+        return nnx.Param(param_data, sharding=sharding)
+    else:
+        return nnx.Param(_init_fn(key, shape, dtype), sharding=sharding)