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

tpu_inference/layers/jax/attention/gpt_oss_attention.py

@@ -0,0 +1,275 @@
+# 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.
+
+from dataclasses import InitVar, dataclass
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from flax.typing import Sharding
+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_hd64 import \
+    ragged_paged_attention_hd64
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.jax.base import create_param
+from tpu_inference.layers.jax.rope import GptOssRotaryEmbedding
+
+KVCache = Tuple[jax.Array, jax.Array]
+
+
+@dataclass(kw_only=True)
+class GptOssAttention(nnx.Module):
+    """
+    JAX implementation of the GPT-OSS Attention block
+    """
+    hidden_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    dtype: jnp.dtype
+    rngs: InitVar[nnx.Rngs]
+
+    rope_theta: float
+    initial_context_length: int = 4096
+    rope_scaling_factor: float = 32.0
+    rope_ntk_alpha: float = 1.0
+    rope_ntk_beta: float = 32.0
+    kv_cache_dtype: str
+
+    query_tnh: P = P()
+    keyvalue_skh: P = P()
+    attn_o_tnh: P = P()
+    dnh_sharding: Sharding = ()
+    dkh_sharding: Sharding = ()
+    nhd_sharding: Sharding = ()
+    n_sharding: Sharding = ()
+    nh_sharding: Sharding = ()
+    kh_sharding: Sharding = ()
+    d_sharding: Sharding = ()
+
+    random_init: bool = False
+    mesh: Mesh
+
+    _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 weights, biases, and RoPE module."""
+
+        self.sm_scale = 1.0 / (self.head_dim**0.5)
+
+        self.sinks_N = create_param(
+            rngs,
+            shape=(self.num_attention_heads, ),
+            dtype=jnp.float32,
+            sharding=self.n_sharding,
+            random_init=self.random_init,
+        )
+
+        # Q, K, V projection kernels
+        self.kernel_q_DNH = create_param(
+            rngs,
+            shape=(self.hidden_size, self.num_attention_heads, self.head_dim),
+            dtype=self.dtype,
+            sharding=self.dnh_sharding,
+            random_init=self.random_init,
+        )
+        self.bias_q_NH = create_param(
+            rngs,
+            shape=(self.num_attention_heads, self.head_dim),
+            dtype=self.dtype,
+            sharding=self.nh_sharding,
+            random_init=self.random_init,
+        )
+        self.kernel_k_DKH = create_param(
+            rngs,
+            shape=(self.hidden_size, self.num_key_value_heads, self.head_dim),
+            dtype=self.dtype,
+            sharding=self.dkh_sharding,
+            random_init=self.random_init,
+        )
+        self.bias_k_KH = create_param(
+            rngs,
+            shape=(self.num_key_value_heads, self.head_dim),
+            dtype=self.dtype,
+            sharding=self.kh_sharding,
+            random_init=self.random_init,
+        )
+        self.kernel_v_DKH = create_param(
+            rngs,
+            shape=(self.hidden_size, self.num_key_value_heads, self.head_dim),
+            dtype=self.dtype,
+            sharding=self.dkh_sharding,
+            random_init=self.random_init,
+        )
+        self.bias_v_KH = create_param(
+            rngs,
+            shape=(self.num_key_value_heads, self.head_dim),
+            dtype=self.dtype,
+            sharding=self.kh_sharding,
+            random_init=self.random_init,
+        )
+        # Output projection kernel
+        self.kernel_o_proj_NHD = create_param(
+            rngs,
+            shape=(self.num_attention_heads, self.head_dim, self.hidden_size),
+            dtype=self.dtype,
+            sharding=self.nhd_sharding,
+            random_init=self.random_init,
+        )
+        self.bias_o_D = create_param(
+            rngs,
+            shape=(self.hidden_size, ),
+            dtype=self.dtype,
+            sharding=self.d_sharding,
+            random_init=self.random_init,
+        )
+
+        # RoPE Module
+        self.rope = GptOssRotaryEmbedding(
+            head_dim=self.head_dim,
+            rope_theta=self.rope_theta,
+            dtype=self.dtype,
+            initial_context_length=self.initial_context_length,
+            rope_scaling_factor=self.rope_scaling_factor,
+            rope_ntk_alpha=self.rope_ntk_alpha,
+            rope_ntk_beta=self.rope_ntk_beta)
+
+        if self.kv_cache_dtype != "auto":
+            self.kv_cache_quantized_dtype = utils.get_jax_dtype_from_str_dtype(
+                self.kv_cache_dtype)
+
+    def attention(
+        self,
+        kv_cache: KVCache,
+        q_TNH: jax.Array,
+        k_SKH: jax.Array,
+        v_SKH: jax.Array,
+        sinks: 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 by calling the ragged_paged_attention kernel."""
+        md = attention_metadata
+        kv_cache_spec = P("data", None, "model")
+
+        in_specs = (
+            self.query_tnh,  # q
+            self.keyvalue_skh,  # k
+            self.keyvalue_skh,  # v
+            kv_cache_spec,  # kv_cache
+            P("data"),  # md.seq_lens
+            P("data"),  # page_indices_flat
+            P("data"),  # query_start_loc
+            P("data"),  # distribution
+            P(('model')),  # sinks
+        )
+        out_specs = (self.attn_o_tnh, kv_cache_spec)
+
+        def _ragged_paged_attention_wrapper(*args):
+            # Pass the GPT-OSS specific parameters to the kernel
+            return ragged_paged_attention_hd64(
+                *args,
+                sm_scale=self.sm_scale,
+                sliding_window=md.sliding_window,
+                q_scale=q_scale,
+                k_scale=k_scale,
+                v_scale=v_scale,
+            )
+
+        output_TNH, kv_cache = jax.jit(
+            jax.shard_map(
+                _ragged_paged_attention_wrapper,
+                mesh=mesh,
+                in_specs=in_specs,
+                out_specs=out_specs,
+                check_vma=False,
+            ))(
+                q_TNH,
+                k_SKH,
+                v_SKH,
+                kv_cache,
+                md.seq_lens,
+                md.block_tables,
+                md.query_start_loc,
+                md.request_distribution,
+                sinks,
+            )
+        return kv_cache, output_TNH
+
+    def __call__(self,
+                 x_TD,
+                 is_prefill,
+                 kv_cache: KVCache,
+                 attention_metadata: AttentionMetadata,
+                 use_attention_rope: bool = True):
+        """Forward pass for the Attention module using 3D kernels."""
+        md = attention_metadata
+        x_TD = jnp.asarray(x_TD, self.dtype)
+
+        with jax.named_scope("q_proj"):
+            q_TNH = jnp.einsum("TD,DNH->TNH", x_TD, self.kernel_q_DNH.value)
+            q_TNH += self.bias_q_NH.value
+
+        with jax.named_scope("k_proj"):
+            k_TKH = jnp.einsum("TD,DKH->TKH", x_TD, self.kernel_k_DKH.value)
+            k_TKH += self.bias_k_KH.value
+
+        with jax.named_scope("v_proj"):
+            v_TKH = jnp.einsum("TD,DKH->TKH", x_TD, self.kernel_v_DKH.value)
+            v_TKH += self.bias_v_KH.value
+
+        if use_attention_rope:
+            q_TNH, k_TKH = self.rope(q_TNH, k_TKH, md.input_positions)
+
+        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_TKH, v_TKH = utils.quantize_kv(k_TKH, v_TKH,
+                                             self.kv_cache_quantized_dtype,
+                                             k_scale, v_scale)
+
+        with jax.named_scope("attn_op"):
+            new_kv_cache, attn_out_TNH = self.attention(
+                kv_cache=kv_cache,
+                q_TNH=q_TNH,
+                k_SKH=k_TKH,
+                v_SKH=v_TKH,
+                sinks=self.sinks_N.value,
+                attention_metadata=md,
+                mesh=self.mesh,
+                q_scale=q_scale,
+                k_scale=k_scale,
+                v_scale=v_scale,
+            )
+            attn_out_TNH = attn_out_TNH[..., :self.head_dim]
+
+        with jax.named_scope("o_proj"):
+            output_TD = jnp.einsum("TNH,NHD->TD", attn_out_TNH,
+                                   self.kernel_o_proj_NHD.value)
+            output_TD += self.bias_o_D.value
+
+        return new_kv_cache, output_TD

tpu_inference/layers/jax/attention/llama4_attention.py

@@ -0,0 +1,167 @@
+# 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.
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jax.sharding import Sharding
+
+from tpu_inference import utils
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.jax.attention.attention import Attention, KVCache
+from tpu_inference.layers.jax.rope_interface import apply_rope
+from tpu_inference.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class L2Norm(nnx.Module):
+    """
+  Implementation of L2 Norm in JAX (taken from MaxText repo - maxtext/MaxText/layers/attentions.py).
+
+  Attributes:
+    eps: float, epsilon used for numerical stability (default value should be ok for most cases).
+  """
+
+    def __init__(self, eps: float = 1e-6):
+        self.eps = eps
+
+    def __call__(self, x):
+        return x * jax.lax.rsqrt(
+            jnp.mean(x**2, axis=-1, keepdims=True) + self.eps)
+
+
+@dataclass(kw_only=True)
+class Llama4Attention(Attention):
+    use_qk_norm: bool
+    temperature_tuning: bool
+    temperature_tuning_floor_scale: float
+    temperature_tuning_scale: float
+    activation_attention_td: Sharding
+    activation_attention_out_td: Sharding
+
+    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 and L2Norm if specified,
+        performing scaled dot-product attention, and projecting the results
+        back to the model dimension.
+        If no RoPE (NoPE) is specified, one can also perform temperature tuning
+        which is useful to combat dilution of attention scores in long-context attention.
+
+        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 = jnp.asarray(x, self.dtype)
+        x_SD = nnx.with_sharding_constraint(x, self.activation_attention_td)
+        x_q_TD = nnx.with_sharding_constraint(x, self.activation_q_td)
+        rope_scaling = self.rope_scaling
+        rope_theta = self.rope_theta
+        H = self.head_dim
+        l2_norm = L2Norm()
+
+        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, rope_theta,
+                                   rope_scaling, self.rope_input_ordering)
+
+                # Apply normaliation after RoPE
+                if self.use_qk_norm:
+                    q_TNH = l2_norm(q_TNH)
+            else:
+                if self.temperature_tuning:
+                    q_TNH = self.apply_temperature_tuning(md, q_TNH)
+
+            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, rope_theta,
+                                   rope_scaling, self.rope_input_ordering)
+
+                # Apply normaliation after RoPE
+                if self.use_qk_norm:
+                    k_SKH = l2_norm(k_SKH)
+            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)
+            v_SKH = nnx.with_sharding_constraint(v_SKH, self.keyvalue_skh)
+
+        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)
+            o_TD = nnx.with_sharding_constraint(
+                o_TD, self.activation_attention_out_td)
+        return new_kv_cache, o_TD
+
+    def apply_temperature_tuning(self, md: AttentionMetadata,
+                                 input_arr_TNH: jax.Array) -> jax.Array:
+        """Applies temperature tuning to the input array of shape (T, N, H).
+        Args:
+            md: AttentionMetadata object containing the input positions.
+            input_arr_TNH: Input array of shape (T, N, H) which will have scaled temperatures applied.
+        """
+        attn_scales = (jnp.log(
+            jnp.floor((md.input_positions.astype(self.dtype) + 1.0) /
+                      self.temperature_tuning_floor_scale) + 1.0) *
+                       self.temperature_tuning_scale + 1.0)
+        return input_arr_TNH * attn_scales[:, None, None]

tpu_inference/layers/jax/base.py

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