tpu-inference 0.11.1.dev202511150811__py3-none-any.whl

tpu_inference/layers/jax/transformer_block.py

@@ -0,0 +1,107 @@
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+# Flax and JAX sharding imports
+import jax
+from flax import nnx
+
+from tpu_inference.layers.jax.attention.attention import (AttentionMetadata,
+                                                          KVCache)
+from tpu_inference.layers.jax.layers import DenseFFW
+from tpu_inference.layers.jax.moe.moe import MoE
+
+
+@dataclass(kw_only=True)
+class TransformerBlock(nnx.Module):
+    """
+    A heavy weight module which serves as the stateful live blocks in serving
+
+    custom_module can be either a dense module (i.e., DenseFFW) or MoE.
+    """
+    pre_attention_norm: nnx.Module
+    pre_mlp_norm: nnx.Module
+    custom_module: Optional[nnx.Module] = None
+    attn: nnx.Module
+    use_attention_rope: bool = True
+    quant: Any | None = None
+
+    def __call__(
+            self, x_TD: jax.Array, is_prefill: bool, kv_cache: KVCache,
+            attention_metadata: AttentionMetadata
+    ) -> Tuple[KVCache, jax.Array]:
+        # Attn Block
+        attn_residual_TD = x_TD
+        x_TD = self.pre_attention_norm(x_TD)
+        new_cache, attn_output_TD = self.attn(x_TD, is_prefill, kv_cache,
+                                              attention_metadata,
+                                              self.use_attention_rope)
+        attn_output_TD += attn_residual_TD
+
+        # FFW Block
+        ffw_residual_TD = attn_output_TD
+        normed_ffw_input_TD = self.pre_mlp_norm(attn_output_TD)
+        logits_TD = self.custom_module(normed_ffw_input_TD)
+        logits_TD += ffw_residual_TD
+        return new_cache, logits_TD
+
+
+@dataclass(kw_only=True)
+class SharedExpertsTransformerBlock(TransformerBlock):
+    """Create a modified TransformerBlock that sums MoE layer output with shared expert output.
+
+    Users can provide the FFW layer in two ways:
+    1.  Pass the module (either `MoE` or `DenseFFW`) to the `custom_module`
+        attribute.
+    2.  Specify the `moe_ffw` or `dense_ffw` attributes
+        (e.g., for passing quantized modules).
+
+    Attributes:
+        moe_ffw: Optional MoE layer.
+        dense_ffw: Optional DFF layer.
+        shared_experts: Optional shared experts module, used if MoE is enabled.
+
+    If an `MoE` layer is used (from either path), its output is summed
+    with the `shared_experts` module.
+    """
+
+    moe_ffw: Optional[MoE] = None
+    dense_ffw: Optional[DenseFFW] = None
+    shared_experts: Optional[DenseFFW] = None
+
+    def __call__(self, x_TD, is_prefill, kv_cache, attention_metadata):
+        # Attn Block
+        attn_residual_TD = x_TD
+        x_TD = self.pre_attention_norm(x_TD)
+        new_cache, attn_output_TD = self.attn(x_TD, is_prefill, kv_cache,
+                                              attention_metadata,
+                                              self.use_attention_rope)
+        attn_output_TD += attn_residual_TD
+
+        # FFW Block
+        ffw_residual_TD = attn_output_TD
+        normed_ffw_input_TD = self.pre_mlp_norm(attn_output_TD)
+
+        if isinstance(self.custom_module, MoE):
+            moe_layer = self.custom_module
+        else:
+            moe_layer = self.moe_ffw
+
+        if isinstance(self.custom_module, DenseFFW):
+            dense_layer = self.custom_module
+        else:
+            dense_layer = self.dense_ffw
+
+        if moe_layer is not None:
+            logits_TD = moe_layer(normed_ffw_input_TD)
+            # Add the shared expert outputs to the MoE outputs.
+            shared_expert_output_TD = self.shared_experts(normed_ffw_input_TD)
+            logits_TD += shared_expert_output_TD
+        elif dense_layer is not None:
+            logits_TD = dense_layer(normed_ffw_input_TD)
+        else:
+            raise ValueError(
+                "Neither custom_module, moe_ffw nor dense_ffw attribute is set for this SharedExpertsTransformerBlock!"
+            )
+
+        logits_TD += ffw_residual_TD
+        return new_cache, logits_TD

tpu_inference/layers/vllm/attention.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools
+from typing import Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+import torch
+from jax.sharding import Mesh
+from torchax.interop import jax_view, torch_view
+from torchax.ops.mappings import t2j
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionLayer, AttentionType)
+
+from tpu_inference import utils
+from tpu_inference.layers.common.attention_interface import attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.logger import init_logger
+from tpu_inference.models.vllm.vllm_model_wrapper_context import \
+    get_vllm_model_wrapper_context
+
+logger = init_logger(__name__)
+
+
+class PallasAttentionBackend(AttentionBackend):
+
+    @staticmethod
+    def get_name() -> str:
+        return "PALLAS"
+
+    @staticmethod
+    def get_impl_cls() -> type["PallasAttentionBackendImpl"]:
+        return PallasAttentionBackendImpl
+
+
+class PallasAttentionBackendImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None = None,
+        attn_type: AttentionType = AttentionType.DECODER,
+        kv_sharing_target_layer_name: str | None = None,
+        sinks: torch.Tensor | None = None,
+    ) -> None:
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        self.sliding_window = sliding_window
+        self.logits_soft_cap = logits_soft_cap
+        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
+
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+        if alibi_slopes is not None:
+            raise NotImplementedError("Alibi slopes is not supported.")
+        self.kv_cache_quantized_dtype = None
+        if kv_cache_dtype != "auto":
+            self.kv_cache_quantized_dtype = utils.get_jax_dtype_from_str_dtype(
+                kv_cache_dtype)
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError("Encoder self-attention and "
+                                      "encoder/decoder cross-attention "
+                                      "are not implemented for "
+                                      "PallasAttentionBackendImpl")
+
+        self.sinks = sinks
+        if self.sinks is not None:
+            assert self.sinks.shape[0] == num_heads, (
+                "Sinks must have the same number of heads as the number of "
+                "heads in the layer")
+
+    def process_weights_after_loading(self, act_dtype: torch.dtype):
+        #TODO (kyuyeunk): Shard the sinks along num_heads dim
+        if self.sinks is not None:
+            sinks = t2j(self.sinks, use_dlpack=False)
+            sinks = torch_view(sinks.astype(jnp.float32))
+            self.sinks = torch.nn.Parameter(sinks, requires_grad=False)
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        output: Optional[torch.Tensor] = None,
+        output_scale: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if output_scale is not None:
+            raise NotImplementedError(
+                "fused output quantization is not yet supported for "
+                "PallasAttentionBackendImpl")
+
+        if kv_cache.numel():
+            raise RuntimeError(
+                "KV cache from vLLM Attention layer should be empty but has "
+                "the size of %s.", kv_cache.numel())
+
+        del kv_cache  # Use kv_cache from vllm wrapper context values instead.
+
+        vllm_model_wrapper_context = get_vllm_model_wrapper_context()
+        kv_cache_index = vllm_model_wrapper_context.layer_name_to_kvcache_index[
+            layer.layer_name]
+        kv_cache = vllm_model_wrapper_context.kv_caches[kv_cache_index]
+
+        mesh = vllm_model_wrapper_context.mesh
+
+        query, key, value = jax_view(query), jax_view(key), jax_view(value)
+        q_scale = k_scale = v_scale = None
+        if self.kv_cache_quantized_dtype:
+            key, value = utils.quantize_kv(key, value,
+                                           self.kv_cache_quantized_dtype,
+                                           layer._k_scale_float,
+                                           layer._v_scale_float)
+            # TODO(kyuyeunk): Enable w8a8 when VREG spill issue is resolved.
+            # q_scale = layer._q_scale_float
+            k_scale = layer._k_scale_float
+            v_scale = layer._v_scale_float
+
+        sinks = jax_view(self.sinks)
+
+        new_kv_cache, outputs = _jax_attn_func(
+            kv_cache,
+            query,
+            key,
+            value,
+            sinks,
+            attn_metadata,
+            mesh,
+            self.scale,
+            self.head_size,
+            self.num_heads,
+            self.num_kv_heads,
+            q_scale,
+            k_scale,
+            v_scale,
+            self.sliding_window,
+        )
+        vllm_model_wrapper_context.kv_caches[kv_cache_index] = new_kv_cache
+
+        return torch_view(outputs)
+
+
+@functools.partial(
+    jax.jit,
+    static_argnames=(
+        "mesh",
+        "scale",
+        "head_size",
+        "num_heads",
+        "num_kv_heads",
+        "q_scale",
+        "k_scale",
+        "v_scale",
+        "sliding_window",
+    ),
+    donate_argnames=("kv_cache"),
+)
+def _jax_attn_func(
+    kv_cache: jax.Array,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    sinks: jax.Array | None,
+    attention_metadata: AttentionMetadata,
+    mesh: Mesh,
+    scale: float,
+    head_size: int,
+    num_heads: int,
+    num_kv_heads: int,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+    sliding_window: int | None = None,
+) -> Tuple[jax.Array, jax.Array]:
+    del scale  # Unused for now, as the attention function applies a default scale.
+
+    # Get shapes from vllm
+    q_len, q_compute_dim = q.shape
+    k_len, k_compute_dim = k.shape
+    assert k.shape == v.shape
+    assert q_compute_dim == head_size * num_heads
+    assert k_compute_dim == head_size * num_kv_heads
+
+    # Convert the shapes from vLLM's convetion to what the attention function expects
+    # bs, num_heads, q_len, head_size
+    q = q.reshape(q_len, num_heads, head_size)
+    # bs, num_kv_heads, k_len, head_size
+    k = k.reshape(k_len, num_kv_heads, head_size)
+    v = v.reshape(k_len, num_kv_heads, head_size)
+
+    new_kv_cache, outputs = attention(
+        kv_cache,
+        q,
+        k,
+        v,
+        attention_metadata,
+        mesh,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+        sinks=sinks,
+        attention_chunk_size=sliding_window,
+    )
+
+    # Convert the shape back to vLLM's convention
+    assert outputs.shape[0] == q_len
+    assert outputs.shape[1] == num_heads
+    assert outputs.shape[2] == head_size
+    outputs = outputs.reshape(q_len, q_compute_dim)
+
+    return new_kv_cache, outputs