ai-edge-torch-nightly 0.3.0.dev20250123__py3-none-any.whl → 0.3.0.dev20250125__py3-none-any.whl

ai_edge_torch/generative/layers/experimental/attention.py

@@ -0,0 +1,269 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Common building blocks for a GPU-specific Attention layer.
+
+This is a temporary implemenation for the GPU. It is subject to change/removal
+at any time.
+"""
+
+from typing import Optional, Tuple, Union
+
+from ai_edge_torch.generative.layers import builder
+from ai_edge_torch.generative.layers import lora as lora_utils
+from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
+from ai_edge_torch.generative.layers.experimental import scaled_dot_product_attention as sdpa
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+import torch
+from torch import nn
+
+
+class TransformerBlock(nn.Module):
+
+  def __init__(
+      self,
+      config: cfg.TransformerBlockConfig,
+      model_config: cfg.ModelConfig,
+  ) -> None:
+    """Initialize an instance of the TransformerBlock.
+
+    Args:
+      config (cfg.TransformerBlockConfig): the configuration object for this
+        transformer block.
+      model_config (cfg.ModelConfig): the configuration object for the model
+        this transformer block belongs to.
+    """
+    super().__init__()
+    self.pre_atten_norm = builder.build_norm(
+        model_config.embedding_dim,
+        config.pre_attention_norm_config,
+    )
+    self.atten_func = CausalSelfAttention(
+        model_config.batch_size,
+        model_config.embedding_dim,
+        config.attn_config,
+        model_config.enable_hlfb,
+    )
+    self.post_atten_norm = builder.build_norm(
+        model_config.embedding_dim,
+        config.post_attention_norm_config,
+    )
+    self.ff = builder.build_ff(model_config.embedding_dim, config.ff_config)
+    self.config = config
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+      mask: Optional[torch.Tensor] = None,
+      input_pos: Optional[torch.Tensor] = None,
+      kv_cache: kv_utils.KVCacheEntryBase = None,
+      lora: Optional[lora_utils.LoRAEntry] = None,
+  ) -> Union[torch.Tensor, Tuple[torch.Tensor, kv_utils.KVCacheEntryBase]]:
+    """Forward function of the TransformerBlock.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+      kv_cache (KVCacheEntryBase): the optional kv cache entry.
+      lora (LoRAEntry): the optional lora entry.
+
+    Returns:
+      output activation from this transformer block, and updated kv cache (if
+      passed in).
+    """
+    kv = None
+    if self.config.parallel_residual:
+      x_norm = self.pre_atten_norm(x)
+      atten_func_out = self.atten_func(
+          x_norm, rope, mask, input_pos, kv_cache, lora
+      )
+      if kv_cache is None:
+        attn_out = atten_func_out
+      else:
+        attn_out, kv = atten_func_out
+      ff_out = self.ff(x_norm)
+      output = x + attn_out + ff_out
+    else:
+      x_norm = self.pre_atten_norm(x)
+      atten_func_out = self.atten_func(
+          x_norm, rope, mask, input_pos, kv_cache, lora
+      )
+      if kv_cache is None:
+        attn_out = atten_func_out
+      else:
+        attn_out, kv = atten_func_out
+      x = x + attn_out
+      x_norm = self.post_atten_norm(x)
+      output = x + self.ff(x_norm)
+
+    return output if kv is None else (output, kv)
+
+
+class CausalSelfAttention(nn.Module):
+
+  def __init__(
+      self,
+      batch_size: int,
+      dim: int,
+      config: cfg.AttentionConfig,
+      enable_hlfb: bool,
+  ) -> None:
+    """Initialize an instance of CausalSelfAttention.
+
+    Args:
+      batch_size (int): batch size of the input tensor.
+      dim (int): causal attention's input/output dimmension.
+      config (cfg.AttentionConfig): attention specific configurations.
+      enable_hlfb (bool): whether hlfb is enabled or not.
+    """
+    super().__init__()
+    self.kv_cache = None
+    self.batch_size = batch_size
+    qkv_shape = (
+        config.num_heads + 2 * config.num_query_groups
+    ) * config.head_dim
+    output_shape = config.num_heads * config.head_dim
+    # Key, query, value projections for all heads.
+    self.qkv_projection = nn.Linear(dim, qkv_shape, bias=config.qkv_use_bias)
+    self.output_projection = nn.Linear(
+        output_shape, dim, bias=config.output_proj_use_bias
+    )
+    self.query_norm = builder.build_norm(
+        config.head_dim, config.query_norm_config
+    )
+    self.key_norm = builder.build_norm(config.head_dim, config.key_norm_config)
+    self.config = config
+    self.enable_hlfb = enable_hlfb
+    self.sdpa_func = sdpa.scaled_dot_product_attention
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+      mask: Optional[torch.Tensor] = None,
+      input_pos: Optional[torch.Tensor] = None,
+      kv_cache: Optional[kv_utils.KVCacheEntryBase] = None,
+      lora: Optional[lora_utils.LoRAEntry] = None,
+  ) -> Union[torch.Tensor, Tuple[torch.Tensor, kv_utils.KVCacheEntryBase]]:
+    """Forward function of the CausalSelfAttention layer, which can support
+
+       MQA, GQA and MHA.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+      kv_cache (KVCacheEntryBase): the KV cache entry corresponding to this
+        module.
+      lora (LoRAEntry): the optional lora entry.
+
+    Returns:
+      output activation from this self attention layer, and the updated
+        KV Cach Entry (if passed in).
+    """
+    # Batch size, sequence length, embedding dimensionality.
+    B, T, E = x.size()
+    assert B == self.batch_size, (
+        "batch size of input tensor must match with the batch size specified in"
+        " the model configuration."
+    )
+
+    qkv = self.qkv_projection(x)
+
+    # Assemble into a number of query groups to support MHA, MQA and GQA.
+    q_per_kv = self.config.num_heads // self.config.num_query_groups
+    # Each group has >=1 queries, 1 key, and 1 value.
+    if self.config.qkv_transpose_before_split:
+      qkv = qkv.view(B, T, -1, self.config.head_dim)
+      q, k, v = qkv.split(
+          (
+              q_per_kv * self.config.num_query_groups,
+              self.config.num_query_groups,
+              self.config.num_query_groups,
+          ),
+          dim=-2,
+      )
+    else:
+      qkv = qkv.view(B, T, self.config.num_query_groups, -1)
+      q, k, v = qkv.split(
+          (
+              q_per_kv * self.config.head_dim,
+              self.config.head_dim,
+              self.config.head_dim,
+          ),
+          dim=-1,
+      )
+
+    if lora is not None:
+      q += lora_utils.apply_lora(x, lora.attention.query, shape=q.shape)
+      k += lora_utils.apply_lora(x, lora.attention.key, shape=k.shape)
+      v += lora_utils.apply_lora(x, lora.attention.value, shape=v.shape)
+
+    q = self.query_norm(q)
+    k = self.key_norm(k)
+
+    q = q.reshape(B, T, -1, self.config.head_dim)
+    k = k.reshape(B, T, -1, self.config.head_dim)
+    v = v.reshape(B, T, -1, self.config.head_dim)
+
+    if rope is not None:
+      # Compute rotary positional embedding for query and key.
+      n_elem = int(self.config.rotary_percentage * self.config.head_dim)
+      cos, sin = rope
+      q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
+
+    # Transpose k/v to specific layout for GPU implementation.
+    b, _, n, h = q.shape
+    g = n // self.config.num_query_groups
+    # btnh -> bnth -> b(kg)th -> 1(bk)(gt)h
+    q = q.permute(0, 2, 1, 3).reshape(
+        1, b * self.config.num_query_groups, g * T, h
+    )
+
+    k = k.permute(0, 2, 1, 3).reshape(
+        1, -1, T, self.config.head_dim
+    )  # 1, bk, s, h
+    v = v.permute(0, 2, 3, 1).reshape(
+        1, -1, self.config.head_dim, T
+    )  # 1, bk, h, s
+
+    if kv_cache is not None:
+      kv_cache = kv_utils.update(kv_cache, input_pos, k, v)
+      k, v = kv_cache.k_cache, kv_cache.v_cache
+
+    sdpa_out = self.sdpa_func(
+        kv_cache,
+        q,
+        k,
+        v,
+        self.config.head_dim,
+        mask=mask,
+        softcap=self.config.logit_softcap,
+    )  # 1, bk, gt, h
+    sdpa_out = (
+        sdpa_out.reshape(B, -1, T, h).permute(0, 2, 1, 3).reshape(B, T, -1)
+    )
+
+    # Compute the output projection.
+    y = self.output_projection(sdpa_out)
+    if lora is not None:
+      y += lora_utils.apply_lora(sdpa_out, lora.attention.output)
+
+    return y if kv_cache is None else (y, kv_cache)

ai_edge_torch/generative/layers/experimental/kv_cache.py

@@ -0,0 +1,314 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Utility functions for KV Cache.
+
+This is an experimental implementation and is subject to change at any time.
+"""
+
+import dataclasses
+from typing import List, Tuple
+
+from ai_edge_torch import hlfb
+from ai_edge_torch.generative.layers import model_config
+from ai_edge_torch.generative.layers.experimental import types as types
+from ai_edge_torch.generative.utilities.dynamic_update_slice import dynamic_update_slice
+import torch
+import torch.nn as nn
+import torch.utils._pytree as pytree
+
+BATCH_SIZE = 1
+
+
+@dataclasses.dataclass
+class KVCacheEntryBase:
+  """A single cache entry that includes K and V caches.
+
+  The chaches are built based on the provided config with the shape of
+  (batch_size=1, kv_cache_max, num_query_groups, head_dim).
+  """
+
+  k_cache: torch.Tensor
+  v_cache: torch.Tensor
+
+  @classmethod
+  def _from_model_config(
+      cls,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      k_shape: Tuple,
+      v_shape: Tuple,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheEntryBase":
+    """Build an instance of the class based on model config."""
+    k = torch.zeros(k_shape, dtype=dtype, device=device)
+    v = torch.zeros(v_shape, dtype=dtype, device=device)
+    obj = cls(k_cache=k, v_cache=v)
+    return obj
+
+  @classmethod
+  def from_model_config(
+      cls,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheEntryBase":
+    """Build an instance of the class based on model config."""
+    shape = (BATCH_SIZE, kv_cache_max, config.num_query_groups, config.head_dim)
+    return cls._from_model_config(
+        kv_cache_max, config, shape, shape, dtype, device
+    )
+
+
+@dataclasses.dataclass
+class KVCacheEntryBTNH(KVCacheEntryBase):
+  k_type = types.BTNH()
+  v_type = types.BTNH()
+
+
+@dataclasses.dataclass
+class KVCacheEntryTransposed(KVCacheEntryBase):
+
+  k_type = types.BNTH()
+  v_type = types.BNHT()
+
+  @classmethod
+  def from_model_config(
+      cls,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheEntryBase":
+    """Build an instance of the class based on model config."""
+    num_kv_heads = config.num_query_groups
+    k_shape = (
+        1,
+        BATCH_SIZE * num_kv_heads,
+        kv_cache_max,
+        config.head_dim,
+    )  # 1, bk, s, h
+    v_shape = (
+        1,
+        BATCH_SIZE * num_kv_heads,
+        config.head_dim,
+        kv_cache_max,
+    )  # 1, bk, h, s
+    return cls._from_model_config(
+        kv_cache_max, config, k_shape, v_shape, dtype, device
+    )
+
+
+@dataclasses.dataclass
+class KVCacheBase:
+  """A utility class for holding KV cache entries per layer."""
+
+  caches: Tuple[KVCacheEntryBase, ...]
+
+  @classmethod
+  def _from_model_config(
+      cls,
+      kv_entry_cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBase":
+    caches = [
+        kv_entry_cls.from_model_config(
+            config.kv_cache_max,
+            config.block_config(idx).attn_config,
+            dtype,
+            device,
+        )
+        for idx in range(config.num_layers)
+    ]
+    obj = cls(caches=tuple(caches))
+    return obj
+
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBase":
+    """Build an instance of the class based on model config.
+
+    Args:
+        config (ModelConfig): Model config used for building the cache.
+        dtype (torch.dtype, optional): The data type of the cache tensor.
+          Defaults to torch.float32.
+        device (torch.device, optional): The device placement of the cache
+          tensors. Defaults to None.
+
+    Returns:
+        KVCacheBase: The created cache object.
+    """
+    return cls._from_model_config(
+        KVCacheEntryBase, config=config, dtype=dtype, device=device
+    )
+
+  def flatten(self) -> List[torch.Tensor]:
+    """Flatten the cache entries into a list of tensors with order k_i, v_i."""
+    flattened, _ = _flatten_kvc(self)
+    return flattened
+
+
+@dataclasses.dataclass
+class KVCacheBTNH(KVCacheBase):
+
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBTNH":
+    return cls._from_model_config(
+        KVCacheEntryBTNH, config=config, dtype=dtype, device=device
+    )
+
+
+@dataclasses.dataclass
+class KVCacheTransposed(KVCacheBase):
+
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBTNH":
+    return cls._from_model_config(
+        KVCacheEntryTransposed, config=config, dtype=dtype, device=device
+    )
+
+
+def _flatten_kvc(kvc: KVCacheBase) -> Tuple[List[str], List[str]]:
+  flattened = []
+  flat_names = []
+  none_names = []
+  for i, kv_entry in enumerate(kvc.caches):
+    flattened.append(kv_entry.k_cache)
+    flat_names.append(f"k_{i}")
+    flattened.append(kv_entry.v_cache)
+    flat_names.append(f"v_{i}")
+  return flattened, [flat_names, none_names]
+
+
+def _flatten_kvc_with_keys(kvc: KVCacheBase) -> Tuple[List, List]:
+  flattened, (flat_names, none_names) = _flatten_kvc(kvc)
+  return [
+      (pytree.MappingKey(k), v) for k, v in zip(flat_names, flattened)
+  ], flat_names
+
+
+def _unflatten_kvc(
+    values: List[torch.Tensor], context: Tuple[List, List]
+) -> KVCacheBase:
+  assert len(values) % 2 == 0, "Found odd number of K and V entries."
+  num_layers = len(values) // 2
+  flat_names = context[0]
+  kv_entries = []
+  for i in range(num_layers):
+    k_cache_idx = flat_names.index(f"k_{i}")
+    v_cache_idx = flat_names.index(f"v_{i}")
+    kv_entries.append(
+        KVCacheEntryBase(
+            k_cache=values[k_cache_idx], v_cache=values[v_cache_idx]
+        )
+    )
+  obj = KVCacheBase(tuple(kv_entries))
+  return obj
+
+
+pytree.register_pytree_node(
+    KVCacheTransposed,
+    _flatten_kvc,
+    _unflatten_kvc,
+    flatten_with_keys_fn=_flatten_kvc_with_keys,
+    serialized_type_name="",
+)
+
+pytree.register_pytree_node(
+    KVCacheBase,
+    _flatten_kvc,
+    _unflatten_kvc,
+    flatten_with_keys_fn=_flatten_kvc_with_keys,
+    serialized_type_name="",
+)
+
+
+def update(
+    cache: KVCacheEntryBase,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+    use_dus: bool = True,
+) -> KVCacheEntryBase:
+  """Out of place update of Cache buffer.
+
+  Args:
+      cache (KVCacheEntryBase): The original cache buffer.
+      input_pos (torch.Tensor): The update slice positions.
+      k_slice (torch.Tensor): The K slice to be updated in the new cache.
+      v_slice (torch.Tensor): The V slice to be updated in the new cache.
+
+  Returns:
+      KVCacheEntryBase: The updated KVCacheBase entry based on the passed
+      inputs.
+  """
+  update_kv_cache = _update_kv_impl
+  return update_kv_cache(cache, input_pos, k_slice, v_slice)
+
+
+def _get_slice_indices(
+    positions: torch.Tensor, cache_dim: int, ts_idx: int
+) -> torch.Tensor:
+  """Returns the slice indices."""
+  positions = positions.float()[0].reshape(
+      1,
+  )
+
+  zeros = torch.zeros((1,), dtype=torch.float32)
+  indices = []
+  for i in range(cache_dim):
+    if i == ts_idx:
+      indices.append(positions)
+    else:
+      indices.append(zeros)
+  slice_indices = torch.cat(indices, dim=0)
+  slice_indices = slice_indices.int()
+  return slice_indices
+
+
+def _update_kv_impl(
+    cache: KVCacheEntryTransposed,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntryTransposed:
+  """Update the cache buffer with High Level Function Boundary annotation."""
+  cache_dim = 4
+  k_ts_idx = 2
+  v_ts_idx = 3
+  positions = input_pos.clone()
+  k_slice_indices = _get_slice_indices(positions, cache_dim, k_ts_idx)
+  v_slice_indices = _get_slice_indices(positions, cache_dim, v_ts_idx)
+  k = dynamic_update_slice(cache.k_cache, k_slice, [x for x in k_slice_indices])
+  v = dynamic_update_slice(cache.v_cache, v_slice, [x for x in v_slice_indices])
+  return KVCacheEntryTransposed(k, v)

ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py

@@ -0,0 +1,97 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+# Implements scaled dot product attention.  This is experimental and
+# GPU-specific code.
+
+import math
+from typing import Optional
+
+from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
+from ai_edge_torch.generative.layers.experimental import types
+from ai_edge_torch.generative.utilities import bmm_4d as bmm_lib
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+from multipledispatch import dispatch
+import torch
+import torch.nn.functional as F
+
+
+def scaled_dot_product_attention(
+    kv: kv_utils.KVCacheBase,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    head_size: int,
+    mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+    softcap: Optional[float] = None,
+):
+  if hasattr(kv, "k_type") and hasattr(kv, "v_type"):
+    return _sdpa(
+        kv.k_type,
+        kv.v_type,
+        query=query,
+        key=key,
+        value=value,
+        head_size=head_size,
+        mask=mask,
+        scale=scale,
+        softcap=softcap,
+    )
+  raise ValueError(
+      f"SDPA for K type {type(kv.caches[0].k_type)} and V type"
+      f" {type(kv.caches[0].v_type)} not supported."
+  )
+
+
+@dispatch(types.BNTH, types.BNHT)
+def _sdpa(k_type, v_type, *args, **kwargs):
+  query = kwargs["query"]
+  key = kwargs["key"]
+  value = kwargs["value"]
+  head_size = kwargs["head_size"]
+  mask = kwargs.get("mask", None)
+  scale = kwargs.get("scale", None)
+  softcap = kwargs.get("softcap", None)
+
+  if scale is None:
+    scale = 1.0 / math.sqrt(head_size)
+
+  query = query * scale
+
+  assert mask is not None, "Mask should not be None!"
+  t = mask.shape[2]
+
+  logits = bmm_lib.bmm_4d(query, key)
+
+  _, bk, gt, s = logits.shape
+  g = gt // t
+  logits = logits.reshape((bk, g, t, s))
+  if softcap is not None:
+    logits = torch.tanh(logits / softcap)
+    logits = logits * softcap
+
+  padded_logits = logits + mask
+  padded_logits = padded_logits.reshape(1, bk, gt, s)
+  probs = F.softmax(padded_logits, dim=-1).type_as(key)
+
+  encoded = bmm_lib.bmm_4d(probs, value)
+
+  return encoded  # 1, bk, gt, h
+
+
+@dispatch(object, object)
+def _sdpa(k_type, v_type, *args, **kwargs):
+
+  raise ValueError(f"No implementations for k={k_type} and v={v_type}")