ai-edge-torch-nightly 0.2.0.dev20240604__py3-none-any.whl → 0.2.0.dev20240606__py3-none-any.whl

ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py

@@ -0,0 +1,161 @@
+# 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.
+# ==============================================================================
+# A toy example which has basic transformer block (w/ externalized KV-Cache).
+
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch_xla
+
+import ai_edge_torch
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+from ai_edge_torch.generative.layers.experimental.attention import TransformerBlock  # NOQA
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+
+
+class ToyModelWithExternalKV(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.transformer_blocks = nn.ModuleList(
+        TransformerBlock(config) for _ in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        base=10_000,
+        condense_ratio=1,
+        dtype=torch.float32,
+        device=torch.device('cpu'),
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+    )
+    self.config = config
+
+  def forward(
+      self,
+      idx: torch.Tensor,
+      input_pos: torch.Tensor,
+      k_caches: torch.Tensor,
+      v_caches: torch.Tensor,
+  ) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.max_seq_len]
+
+    for i, block in enumerate(self.transformer_blocks):
+      input_k, input_v = k_caches[i], v_caches[i]
+      x, (updated_k, updated_v) = block(
+          x, (cos, sin), mask, input_pos, (input_k, input_v)
+      )
+      k_caches[i], v_caches[i] = updated_k, updated_v
+
+    x = self.final_norm(x)
+    return self.lm_head(x), k_caches, v_caches
+
+
+def _export_stablehlo_mlir(model, args):
+  ep = torch.export.export(model, args)
+  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+  return stablehlo_gm.get_stablehlo_text()
+
+
+def get_model_config() -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32, num_query_groups=4, rotary_percentage=1.0
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationType.SILU,
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  config = cfg.ModelConfig(
+      vocab_size=150,
+      num_layers=2,
+      max_seq_len=100,
+      embedding_dim=128,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      pre_ff_norm_config=norm_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.unsqueeze(torch.arange(0, 100), 0)
+  input_pos = torch.arange(0, 100)
+  return idx, input_pos
+
+
+def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.tensor([[1]], dtype=torch.long)
+  input_pos = torch.tensor([10])
+  return idx, input_pos
+
+
+def define_and_run() -> None:
+  dump_mlir = False
+
+  config = get_model_config()
+  model = ToyModelWithExternalKV(config)
+  print('running an inference')
+  k_caches = torch.zeros((2, 1, 100, 4, 4), dtype=torch.float32)
+  v_caches = torch.zeros((2, 1, 100, 4, 4), dtype=torch.float32)
+
+  idx, input_pos = get_sample_prefill_inputs()
+  decode_idx, decode_input_pos = get_sample_decode_inputs()
+  print(model.forward(idx, input_pos, k_caches, v_caches))
+
+  if dump_mlir:
+    mlir_text = _export_stablehlo_mlir(model, (idx, input_pos, k_caches, v_caches))
+    with open('/tmp/toy_model_with_external_kv.stablehlo.mlir', 'w') as f:
+      f.write(mlir_text)
+
+  # Convert model to tflite with 2 signatures (prefill + decode).
+  # TODO(b/344014416): currently conversion will fail, because we generate int64 index
+  # in dynamic update slice op.
+  print('converting toy model to tflite with 2 signatures (prefill + decode)')
+  edge_model = (
+      ai_edge_torch.signature('prefill', model, (idx, input_pos, k_caches, v_caches))
+      .signature('decode', model, (decode_idx, decode_input_pos, k_caches, v_caches))
+      .convert()
+  )
+  edge_model.export('/tmp/toy_external_kv_cache.tflite')
+
+
+if __name__ == '__main__':
+  with torch.inference_mode():
+    define_and_run()

ai_edge_torch/generative/layers/attention.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 # Common building blocks for Attention layer.
 
-import math
 from typing import Optional, Tuple
 
 import torch
@@ -25,101 +24,8 @@ import ai_edge_torch.generative.layers.builder as builder
 from ai_edge_torch.generative.layers.kv_cache import KVCache
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
-
-
-def scaled_dot_product_attention(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    head_size: int,
-    mask: Optional[torch.Tensor] = None,
-    scale: Optional[float] = None,
-):
-  """Scaled dot product attention.
-
-  Args:
-    q (torch.Tensor): Query tensor, with shape [B, T, N, H].
-    k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
-    v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
-    head_size (int): head dimension.
-    mask (torch.Tensor): the optional mask tensor.
-
-  Returns:
-    The output tensor of scaled_dot_product_attention.
-  """
-
-  if scale is None:
-    scale = 1.0 / math.sqrt(head_size)
-
-  q = q.transpose(1, 2)
-  k = k.transpose(1, 2)
-  v = v.transpose(1, 2)
-  if q.size() != k.size():
-    # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
-    k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
-    v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
-  y = F.scaled_dot_product_attention(
-      q,
-      k,
-      v,
-      attn_mask=mask,
-      dropout_p=0.0,
-      is_causal=mask is None,
-      scale=scale,
-  )
-  return y.transpose(1, 2)
-
-
-def scaled_dot_product_attention_with_hlfb(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    head_size: int,
-    mask: Optional[torch.Tensor] = None,
-    scale: Optional[float] = None,
-):
-  """Scaled dot product attention with high-level function boundary enabled.
-
-  Args:
-    q (torch.Tensor): Query tensor, with shape [B, T, N, H].
-    k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
-    v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
-    head_size (int): head dimension.
-    mask (torch.Tensor): the optional mask tensor.
-
-  Returns:
-    The output tensor of scaled_dot_product_attention.
-  """
-
-  if scale is None:
-    scale = 1.0 / math.sqrt(head_size)
-
-  builder = StableHLOCompositeBuilder(
-      name="odml.scaled_dot_product_attention", attr={"scale": scale}
-  )
-  q, k, v, mask = builder.mark_inputs(q, k, v, mask)
-
-  q = q.transpose(1, 2)
-  k = k.transpose(1, 2)
-  v = v.transpose(1, 2)
-  if q.size() != k.size():
-    # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
-    k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
-    v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
-  y = F.scaled_dot_product_attention(
-      q,
-      k,
-      v,
-      attn_mask=mask,
-      dropout_p=0.0,
-      is_causal=mask is None,
-      scale=scale,
-  )
-
-  result = y.transpose(1, 2)
-  result = builder.mark_outputs(result)
-  return result
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention  # NOQA
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention_with_hlfb  # NOQA
 
 
 class TransformerBlock(nn.Module):
@@ -151,7 +57,7 @@ class TransformerBlock(nn.Module):
   def forward(
       self,
       x: torch.Tensor,
-      rope: Tuple[torch.Tensor, torch.Tensor],
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
       mask: Optional[torch.Tensor] = None,
       input_pos: Optional[torch.Tensor] = None,
   ) -> torch.Tensor:
@@ -182,7 +88,6 @@ class TransformerBlock(nn.Module):
     return output
 
 
-# CausalSelfAttention which can support MHQ, MQA or GQA.
 class CausalSelfAttention(nn.Module):
 
   def __init__(
@@ -229,11 +134,12 @@ class CausalSelfAttention(nn.Module):
   def forward(
       self,
       x: torch.Tensor,
-      rope: Tuple[torch.Tensor, torch.Tensor],
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
       mask: Optional[torch.Tensor] = None,
       input_pos: Optional[torch.Tensor] = None,
   ) -> torch.Tensor:
-    """Forward function of the CausalSelfAttention layer.
+    """Forward function of the CausalSelfAttention layer, which can support
+       MQA, GQA and MHA.
 
     Args:
       x (torch.Tensor): the input tensor.
@@ -253,28 +159,35 @@ class CausalSelfAttention(nn.Module):
     # 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
     total_qkv = q_per_kv + 2  # Each group has >=1 queries, 1 key, and 1 value.
-    qkv = qkv.view(
-        B, T, self.config.num_query_groups, total_qkv, self.head_dim
-    )  # (B, T, num_query_groups, total_qkv, head_dim)
+    if self.config.qkv_transpose_before_split:
+      qkv = qkv.view(
+          B, T, total_qkv, self.config.num_query_groups, self.head_dim
+      )  # (B, T, total_qkv, num_query_groups, head_dim)
+      qkv_axis = -3
+    else:
+      qkv = qkv.view(
+          B, T, self.config.num_query_groups, total_qkv, self.head_dim
+      )  # (B, T, num_query_groups, total_qkv, head_dim)
+      qkv_axis = -2
 
     # Split batched computation into three.
-    q, k, v = qkv.split((q_per_kv, 1, 1), dim=-2)
-
+    q, k, v = qkv.split((q_per_kv, 1, 1), dim=qkv_axis)
     q = q.reshape(B, T, -1, self.head_dim)
     k = k.reshape(B, T, -1, self.head_dim)
     v = v.reshape(B, T, -1, self.head_dim)
 
     # Compute rotary positional embedding for query and key.
     n_elem = int(self.config.rotary_percentage * self.head_dim)
-    cos, sin = rope
-    q_roped = rotary_pos_emb.apply_rope(
-        q[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
-    )
-    k_roped = rotary_pos_emb.apply_rope(
-        k[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
-    )
-    q = torch.cat((q_roped, q[..., n_elem:]), dim=-1)
-    k = torch.cat((k_roped, k[..., n_elem:]), dim=-1)
+    if n_elem > 0:
+      cos, sin = rope
+      q_roped = rotary_pos_emb.apply_rope(
+          q[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+      )
+      k_roped = rotary_pos_emb.apply_rope(
+          k[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+      )
+      q = torch.cat((q_roped, q[..., n_elem:]), dim=-1)
+      k = torch.cat((k_roped, k[..., n_elem:]), dim=-1)
 
     if self.kv_cache is not None:
       # TODO(haoliang): Handle when execeeding max sequence length.

ai_edge_torch/generative/layers/builder.py

@@ -97,6 +97,10 @@ def _get_activation(type_: cfg.ActivationType):
     return F.gelu
   elif type_ == cfg.ActivationType.GELU_TANH:
     return lambda x: F.gelu(x, approximate="tanh")
+  elif type_ == cfg.ActivationType.GELU_QUICK:
+    # GELU approximation that is fast but somewhat inaccurate.
+    # See: https://github.com/hendrycks/GELUs
+    return lambda x: x * F.sigmoid(1.702 * x)
   elif type_ == cfg.ActivationType.RELU:
     return F.relu
   else:

ai_edge_torch/generative/layers/model_config.py

@@ -27,6 +27,7 @@ class ActivationType(enum.Enum):
   SILU = enum.auto()
   GELU = enum.auto()
   GELU_TANH = enum.auto()
+  GELU_QUICK = enum.auto()
   RELU = enum.auto()
 
 
@@ -46,7 +47,7 @@ class FeedForwardType(enum.Enum):
 
   # `output = linear(act(linear(x)))`.
   SEQUENTIAL = enum.auto()
-  # `output = linear(act(linear(x)) * lienar(x))`.
+  # `output = linear_2(act(linear_1(x)) * lienar_3(x))`.
   GATED = enum.auto()
 
 
@@ -60,6 +61,9 @@ class AttentionConfig:
   num_query_groups: Optional[int]
   # Percentage of Rotary Positional Embedding added Q and K projections.
   rotary_percentage: Optional[float] = None
+  # Whether to transpose the query groups of qkv bundled tensor before
+  # splitting into separated tensors.
+  qkv_transpose_before_split: bool = False
   # Whether to use bias with Query, Key, and Value projection.
   qkv_use_bias: bool = False
   # Whether to use bias with attention output projection.

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -0,0 +1,117 @@
+# 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.
+# ==============================================================================
+# Implements scaled dot product attention.
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+
+def scaled_dot_product_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    head_size: int,
+    mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+):
+  """Scaled dot product attention.
+
+  Args:
+    q (torch.Tensor): Query tensor, with shape [B, T, N, H].
+    k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
+    v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
+    head_size (int): head dimension.
+    mask (torch.Tensor): the optional mask tensor.
+
+  Returns:
+    The output tensor of scaled_dot_product_attention.
+  """
+
+  if scale is None:
+    scale = 1.0 / math.sqrt(head_size)
+
+  q = q.transpose(1, 2)
+  k = k.transpose(1, 2)
+  v = v.transpose(1, 2)
+  if q.size() != k.size():
+    # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
+    k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
+    v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
+  y = F.scaled_dot_product_attention(
+      q,
+      k,
+      v,
+      attn_mask=mask,
+      dropout_p=0.0,
+      is_causal=mask is None,
+      scale=scale,
+  )
+  return y.transpose(1, 2)
+
+
+def scaled_dot_product_attention_with_hlfb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    head_size: int,
+    mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+):
+  """Scaled dot product attention with high-level function boundary enabled.
+
+  Args:
+    q (torch.Tensor): Query tensor, with shape [B, T, N, H].
+    k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
+    v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
+    head_size (int): head dimension.
+    mask (torch.Tensor): the optional mask tensor.
+
+  Returns:
+    The output tensor of scaled_dot_product_attention.
+  """
+
+  if scale is None:
+    scale = 1.0 / math.sqrt(head_size)
+
+  builder = StableHLOCompositeBuilder(
+      name="odml.scaled_dot_product_attention", attr={"scale": scale}
+  )
+  q, k, v, mask = builder.mark_inputs(q, k, v, mask)
+
+  q = q.transpose(1, 2)
+  k = k.transpose(1, 2)
+  v = v.transpose(1, 2)
+  if q.size() != k.size():
+    # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
+    k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
+    v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
+  y = F.scaled_dot_product_attention(
+      q,
+      k,
+      v,
+      attn_mask=mask,
+      dropout_p=0.0,
+      is_causal=mask is None,
+      scale=scale,
+  )
+
+  result = y.transpose(1, 2)
+  result = builder.mark_outputs(result)
+  return result