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

ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py

@@ -99,6 +99,36 @@ def _aten_hardswish(gm: GraphModule, node: Node):
   node.target = hardswish
 
 
+@_register_composite_builder(torch.ops.aten.gelu.default)
+def _aten_gelu(gm: GraphModule, node: Node):
+  op = node.target
+  args_mapper = TorchOpArgumentsMapper(op)
+
+  def gelu(*args, **kwargs):
+    nonlocal op, args_mapper
+
+    full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
+
+    # TFLite supports exact and tanh approximate.
+    if full_kwargs["approximate"] != "none" and full_kwargs["approximate"] != "tanh":
+      return op(*args, **kwargs)
+
+    builder = StableHLOCompositeBuilder(
+        "aten.gelu.default",
+        attr=_tree_map_to_composite_attr_values(
+            {
+                "approximate": full_kwargs["approximate"],
+            }
+        ),
+    )
+    full_kwargs["self"] = builder.mark_inputs(full_kwargs["self"])
+    output = op(full_kwargs["self"])
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = gelu
+
+
 @_register_composite_builder(torch.ops.aten.avg_pool2d.default)
 def _aten_avg_pool2d(gm: GraphModule, node: Node):
   op = node.target

ai_edge_torch/convert/test/test_convert_composites.py

@@ -169,6 +169,24 @@ class TestConvertComposites(unittest.TestCase):
         model_coverage.compare_tflite_torch(edge_model, torch_module, tracing_args)
     )
 
+  def test_convert_gelu(self):
+    """Tests conversion of a GELU module."""
+
+    args = (torch.randn((5, 10)),)
+    torch_module = torch.nn.GELU().eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+
+  def test_convert_gelu_approximate(self):
+    """Tests conversion of an Approximate GELU module."""
+
+    args = (torch.randn((5, 10)),)
+    torch_module = torch.nn.GELU('tanh').eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+
 
 if __name__ == '__main__':
   unittest.main()

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -15,65 +15,99 @@
 
 import torch
 from torch import nn
-from torch._prims_common import mask_tensor
-from torch._prims_common.wrappers import out_wrapper
 
-from ai_edge_torch.generative.examples.stable_diffusion.attention import SelfAttention  # NOQA
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attention_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="layers.{}.linear_1",
+    ff_down_proj="layers.{}.linear_2",
+    ff_gate_proj="layers.{}.linear_1",
+    attn_fused_qkv_proj="layers.{}.attention.in_proj",
+    attn_output_proj="layers.{}.attention.out_proj",
+    pre_attn_norm="layers.{}.layernorm_1",
+    pre_ff_norm="layers.{}.layernorm_2",
+    embedding="embedding.token_embedding",
+    embedding_position="embedding.position_value",
+    final_norm="layernorm",
+    lm_head=None,
+)
 
 
-class CLIPEmbedding(nn.Module):
-
-  def __init__(self, n_vocab: int, n_embd: int, n_token: int):
-    super().__init__()
-    self.token_embedding = nn.Embedding(n_vocab, n_embd)
-    self.position_value = nn.Parameter(torch.zeros((n_token, n_embd)))
-
-  def forward(self, tokens):
-    x = self.token_embedding(tokens)
-    x += self.position_value
-    return x
-
-
-class CLIPLayer(nn.Module):
+class CLIP(nn.Module):
+  """CLIP text encoder
+  For details, see https://arxiv.org/abs/2103.00020
+  """
 
-  def __init__(self, n_head: int, n_embd: int):
+  def __init__(self, config: cfg.ModelConfig):
     super().__init__()
-    self.layernorm_1 = nn.LayerNorm(n_embd)
-    self.attention = SelfAttention(n_head, n_embd)
-    self.layernorm_2 = nn.LayerNorm(n_embd)
-    self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
-    self.linear_2 = nn.Linear(4 * n_embd, n_embd)
-
-  def forward(self, x):
-    residue = x
-    x = self.layernorm_1(x)
-    x = self.attention(x, causal_mask=True)
-    x += residue
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.tok_embedding_position = nn.Parameter(
+        torch.zeros((config.max_seq_len, config.embedding_dim))
+    )
 
-    residue = x
-    x = self.layernorm_2(x)
-    x = self.linear_1(x)
-    x = x * torch.sigmoid(1.702 * x)  # QuickGELU activation function
-    x = self.linear_2(x)
-    x += residue
+    self.config = config
+    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)
 
-    return x
-
-
-class CLIP(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.embedding = CLIPEmbedding(49408, 768, 77)
-    self.layers = nn.ModuleList([CLIPLayer(12, 768) for i in range(12)])
-    self.layernorm = nn.LayerNorm(768)
+    self.mask_cache = attention_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32
+    )
 
   @torch.inference_mode
   def forward(self, tokens: torch.LongTensor) -> torch.FloatTensor:
     tokens = tokens.type(torch.long)
 
-    state = self.embedding(tokens)
-    for layer in self.layers:
-      state = layer(state)
-    output = self.layernorm(state)
+    state = self.tok_embedding(tokens) + self.tok_embedding_position
+    for layer in self.transformer_blocks:
+      state = layer(state, mask=self.mask_cache)
+    output = self.final_norm(state)
     return output
+
+
+def get_model_config() -> cfg.ModelConfig:
+  max_seq_len = 77
+  vocab_size = 49408
+  num_layers = 12
+  num_heads = 12
+  num_query_groups = 12
+  embedding_dim = 768
+
+  attn_config = cfg.AttentionConfig(
+      num_heads=num_heads,
+      num_query_groups=num_query_groups,
+      rotary_percentage=0.0,
+      qkv_use_bias=True,
+      qkv_transpose_before_split=True,
+      output_proj_use_bias=True,
+      enable_kv_cache=False,
+  )
+
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=cfg.ActivationType.GELU_QUICK,
+      intermediate_size=embedding_dim * 4,
+      use_bias=True,
+  )
+
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
+
+  config = cfg.ModelConfig(
+      vocab_size=vocab_size,
+      num_layers=num_layers,
+      max_seq_len=max_seq_len,
+      embedding_dim=embedding_dim,
+      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

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -19,11 +19,12 @@ from pathlib import Path
 import torch
 
 import ai_edge_torch
-from ai_edge_torch.generative.examples.stable_diffusion.clip import CLIP
+import ai_edge_torch.generative.examples.stable_diffusion.clip as clip
 from ai_edge_torch.generative.examples.stable_diffusion.decoder import Decoder
 from ai_edge_torch.generative.examples.stable_diffusion.diffusion import Diffusion  # NOQA
 from ai_edge_torch.generative.examples.stable_diffusion.encoder import Encoder
 import ai_edge_torch.generative.examples.stable_diffusion.util as util
+import ai_edge_torch.generative.utilities.loader as loading_utils
 
 
 @torch.inference_mode
@@ -36,8 +37,9 @@ def convert_stable_diffusion_to_tflite(
     image_width: int = 512,
 ):
 
-  clip = CLIP()
-  clip.load_state_dict(torch.load(clip_ckpt_path))
+  clip_model = clip.CLIP(clip.get_model_config())
+  loader = loading_utils.ModelLoader(clip_ckpt_path, clip.TENSOR_NAMES)
+  loader.load(clip_model, strict=False)
 
   encoder = Encoder()
   encoder.load_state_dict(torch.load(encoder_ckpt_path))
@@ -59,13 +61,13 @@ def convert_stable_diffusion_to_tflite(
   )
 
   input_latents = encoder(input_image, noise)
-  context_cond = clip(prompt_tokens)
+  context_cond = clip_model(prompt_tokens)
   context_uncond = torch.zeros_like(context_cond)
   context = torch.cat([context_cond, context_uncond], axis=0)
   time_embedding = util.get_time_embedding(timestamp)
 
   # CLIP text encoder
-  ai_edge_torch.signature('encode', clip, (prompt_tokens,)).convert().export(
+  ai_edge_torch.signature('encode', clip_model, (prompt_tokens,)).convert().export(
       '/tmp/stable_diffusion/clip.tflite'
   )
 

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -202,11 +202,6 @@ class UNet(nn.Module):
 
     x = self.bottleneck(x, context, time)
 
-    # print('x shape:')
-    # print(list(x.shape))
-    # print('time shape:')
-    # print(list(time.shape))
-
     for layers in self.decoders:
       x = torch.cat((x, skip_connections.pop()), dim=1)
       x = layers(x, context, time)
@@ -214,199 +209,6 @@ class UNet(nn.Module):
     return x
 
 
-# The encoder component.
-class UNetEncoder(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.time_embedding = TimeEmbedding(320)
-    self.encoders = nn.ModuleList(
-        [
-            SwitchSequential(nn.Conv2d(4, 320, kernel_size=3, padding=1)),
-            SwitchSequential(ResidualBlock(320, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(320, 320), AttentionBlock(8, 40)),
-            SwitchSequential(nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(ResidualBlock(320, 640), AttentionBlock(8, 80)),
-            SwitchSequential(ResidualBlock(640, 640), AttentionBlock(8, 80)),
-            SwitchSequential(nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(ResidualBlock(640, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(ResidualBlock(1280, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(nn.Conv2d(1280, 1280, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(ResidualBlock(1280, 1280)),
-            SwitchSequential(ResidualBlock(1280, 1280)),
-        ]
-    )
-
-  def forward(self, x, context, time):
-    time_embedding = self.time_embedding(time)
-    skip_connections = []
-    for layers in self.encoders:
-      x = layers(x, context, time_embedding)
-      skip_connections.append(x)
-
-    return x, skip_connections, time_embedding
-
-
-class UNetBottleNeck(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.bottleneck = SwitchSequential(
-        ResidualBlock(1280, 1280),
-        AttentionBlock(8, 160),
-        ResidualBlock(1280, 1280),
-    )
-
-  def forward(self, x, context, time):
-    x = self.bottleneck(x, context, time)
-    # print('shape')
-    # print(list(x.shape))
-    return x
-
-
-# Unet decoder.
-class UNetDecoder1(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), Upsample(1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-        ]
-    )
-
-  def forward(self, x, context, time, s9, s10, s11, s12):
-    x = torch.cat((x, s12), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s11), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s10), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s9), dim=1)
-    x = self.decoders[3](x, context, time)
-
-    return x
-
-
-class UNetDecoder2(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(
-                ResidualBlock(1920, 1280), AttentionBlock(8, 160), Upsample(1280)
-            ),
-            SwitchSequential(ResidualBlock(1920, 640), AttentionBlock(8, 80)),
-            SwitchSequential(ResidualBlock(1280, 640), AttentionBlock(8, 80)),
-        ]
-    )
-
-  def forward(self, x, context, time, s5, s6, s7, s8):
-    x = torch.cat((x, s8), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s7), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s6), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s5), dim=1)
-    x = self.decoders[3](x, context, time)
-    return x
-
-
-class UNetDecoder3(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(
-                ResidualBlock(960, 640), AttentionBlock(8, 80), Upsample(640)
-            ),
-            SwitchSequential(ResidualBlock(960, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-        ]
-    )
-    self.final = FinalLayer(320, 4)
-
-  def forward(self, x, context, time, s1, s2, s3, s4):
-    x = torch.cat((x, s4), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s3), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s2), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s1), dim=1)
-    x = self.decoders[3](x, context, time)
-
-    x = self.final(x)
-    return x
-
-
-class UNetDecoder(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), Upsample(1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(
-                ResidualBlock(1920, 1280), AttentionBlock(8, 160), Upsample(1280)
-            ),
-            SwitchSequential(ResidualBlock(1920, 640), AttentionBlock(8, 80)),
-            SwitchSequential(ResidualBlock(1280, 640), AttentionBlock(8, 80)),
-            SwitchSequential(
-                ResidualBlock(960, 640), AttentionBlock(8, 80), Upsample(640)
-            ),
-            SwitchSequential(ResidualBlock(960, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-        ]
-    )
-    self.final = FinalLayer(320, 4)
-
-  def forward(
-      self, x, context, time, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12
-  ):
-    x = torch.cat((x, s12), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s11), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s10), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s9), dim=1)
-    x = self.decoders[3](x, context, time)
-    x = torch.cat((x, s8), dim=1)
-    x = self.decoders[4](x, context, time)
-    x = torch.cat((x, s7), dim=1)
-    x = self.decoders[5](x, context, time)
-    x = torch.cat((x, s6), dim=1)
-    x = self.decoders[6](x, context, time)
-    x = torch.cat((x, s5), dim=1)
-    x = self.decoders[7](x, context, time)
-    x = torch.cat((x, s4), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s3), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s2), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s1), dim=1)
-    x = self.decoders[3](x, context, time)
-
-    x = self.final(x)
-
-    return x
-
-
 class FinalLayer(nn.Module):
 
   def __init__(self, in_channels, out_channels):
@@ -432,68 +234,6 @@ class Diffusion(nn.Module):
   @torch.inference_mode
   def forward(self, latent, context, time):
     time = self.time_embedding(time)
-    # print('time:')
-    # print(list(time.shape))
     output = self.unet(latent, context, time)
     output = self.final(output)
     return output
-
-
-# Calling code as if Diffusion is splitted into two parts.
-class DiffusionSplitted(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.unet_encoder = UNetEncoder()
-    self.bottleneck = UNetBottleNeck()
-    self.unet_decoder1 = UNetDecoder1()
-    self.unet_decoder2 = UNetDecoder2()
-    self.unet_decoder3 = UNetDecoder3()
-
-  def get_skip_connections(self, latent, context, time):
-    _, skip_connections, _ = self.unet_encoder(latent, context, time)
-    return skip_connections
-
-  def forward(self, latent, context, time):
-    output, skip_connections, time = self.unet_encoder(latent, context, time)
-    # print("output shape of unet encoder...")
-    # print(list(output.shape))
-    # print("output shape of time...")
-    # print(list(time.shape))
-    output = self.bottleneck(output, context, time)
-    # print("output shape of bn")
-    # print(list(output.shape))
-    output = self.unet_decoder1(
-        output,
-        context,
-        time,
-        skip_connections[8],
-        skip_connections[9],
-        skip_connections[10],
-        skip_connections[11],
-    )
-    # print("output shape of d1:")
-    # print(list(output.shape))
-
-    output = self.unet_decoder2(
-        output,
-        context,
-        time,
-        skip_connections[4],
-        skip_connections[5],
-        skip_connections[6],
-        skip_connections[7],
-    )
-
-    # print("output shape of d2:")
-    # print(list(output.shape))
-    output = self.unet_decoder3(
-        output,
-        context,
-        time,
-        skip_connections[0],
-        skip_connections[1],
-        skip_connections[2],
-        skip_connections[3],
-    )
-    return output

ai_edge_torch/generative/examples/t5/t5_attention.py

@@ -20,11 +20,11 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 
-from ai_edge_torch.generative.layers.attention import scaled_dot_product_attention  # NOQA
-from ai_edge_torch.generative.layers.attention import scaled_dot_product_attention_with_hlfb  # NOQA
 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
+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 EncoderDecoderBlock(nn.Module):

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

ai_edge_torch/generative/utilities/loader.py

@@ -69,10 +69,16 @@ def load_pytorch_statedict(full_path: str):
   Raises:
     ValueError: If no tensors are loaded from the provided directory or file.
   """
-  pattern = os.path.join(full_path, "*.bin") if os.path.isdir(full_path) else full_path
   files = []
-  for file in glob.glob(pattern):
-    files.append(file)
+  patterns = []
+  if os.path.isdir(full_path):
+    patterns.append(os.path.join(full_path, "*.bin"))
+    patterns.append(os.path.join(full_path, "*.pt"))
+  else:
+    patterns.append(full_path)
+  for pattern in patterns:
+    for file in glob.glob(pattern):
+      files.append(file)
 
   tensors = {}
   for file in files:
@@ -93,18 +99,20 @@ class ModelLoader:
 
   @dataclass
   class TensorNames:
-    attn_query_proj: str
-    attn_key_proj: str
-    attn_value_proj: str
-    attn_output_proj: str
-
-    ff_up_proj: str
-    ff_down_proj: str
+    attn_query_proj: str = None
+    attn_key_proj: str = None
+    attn_value_proj: str = None
+    attn_fused_qkv_proj: str = None
+    attn_output_proj: str = None
+
+    ff_up_proj: str = None
+    ff_down_proj: str = None
     ff_gate_proj: str = None
 
     pre_attn_norm: str = None
     pre_ff_norm: str = None
     embedding: str = None
+    embedding_position: str = None
     final_norm: str = None
     lm_head: str = None
 
@@ -129,6 +137,10 @@ class ModelLoader:
         strict (bool, optional): Whether the converted keys are strictly
           matched. Defaults to True.
 
+    Returns:
+        missing_keys (List[str]): a list of str containing the missing keys
+        unexpected_keys (List[str]): a list of str containing the unexpected keys
+
     Raises:
         ValueError: If conversion results in unmapped tensors and strict mode is
           enabled.
@@ -139,6 +151,10 @@ class ModelLoader:
       converted_state["tok_embedding.weight"] = state.pop(
           f"{self._names.embedding}.weight"
       )
+      if self._names.embedding_position is not None:
+        converted_state["tok_embedding_position"] = state.pop(
+            f"{self._names.embedding_position}"
+        )
     if self._names.lm_head is not None:
       converted_state["lm_head.weight"] = state.pop(f"{self._names.lm_head}.weight")
       if model.config.lm_head_use_bias:
@@ -158,7 +174,7 @@ class ModelLoader:
       raise ValueError(
           f"Failed to map all tensor. Remaing tensor are: {list(state.keys())}"
       )
-    model.load_state_dict(converted_state, strict=strict)
+    return model.load_state_dict(converted_state, strict=strict)
 
   def _get_loader(self) -> Callable[[str], Dict[str, torch.Tensor]]:
     """A best effort method for finding appropriate state loader.
@@ -172,13 +188,15 @@ class ModelLoader:
     if os.path.isdir(self._file_name):
       if glob.glob(os.path.join(self._file_name, "*.safetensors")):
         return load_safetensors
-      if glob.glob(os.path.join(self._file_name, "*.bin")):
+      if glob.glob(os.path.join(self._file_name, "*.bin")) or glob.glob(
+          os.path.join(self._file_name, "*.pt")
+      ):
         return load_pytorch_statedict
 
     if self._file_name.endswith(".safetensors"):
       return load_safetensors
 
-    if self._file_name.endswith(".bin"):
+    if self._file_name.endswith(".bin") or self._file_name.endswith(".pt"):
       return load_pytorch_statedict
 
     raise ValueError(f"File format not supported.")
@@ -225,22 +243,33 @@ class ModelLoader:
       converted_state: Dict[str, torch.Tensor],
   ):
     prefix = f"transformer_blocks.{idx}"
-    q_name = self._names.attn_query_proj.format(idx)
-    k_name = self._names.attn_key_proj.format(idx)
-    v_name = self._names.attn_value_proj.format(idx)
-    converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = self._fuse_qkv(
-        config,
-        state.pop(f"{q_name}.weight"),
-        state.pop(f"{k_name}.weight"),
-        state.pop(f"{v_name}.weight"),
-    )
-    if config.attn_config.qkv_use_bias:
-      converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = self._fuse_qkv(
+    if self._names.attn_fused_qkv_proj:
+      fused_qkv_name = self._names.attn_fused_qkv_proj.format(idx)
+      converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = state.pop(
+          f"{fused_qkv_name}.weight"
+      )
+    else:
+      q_name = self._names.attn_query_proj.format(idx)
+      k_name = self._names.attn_key_proj.format(idx)
+      v_name = self._names.attn_value_proj.format(idx)
+      converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = self._fuse_qkv(
           config,
-          state.pop(f"{q_name}.bias"),
-          state.pop(f"{k_name}.bias"),
-          state.pop(f"{v_name}.bias"),
+          state.pop(f"{q_name}.weight"),
+          state.pop(f"{k_name}.weight"),
+          state.pop(f"{v_name}.weight"),
       )
+    if config.attn_config.qkv_use_bias:
+      if self._names.attn_fused_qkv_proj:
+        converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = state.pop(
+            f"{fused_qkv_name}.bias"
+        )
+      else:
+        converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = self._fuse_qkv(
+            config,
+            state.pop(f"{q_name}.bias"),
+            state.pop(f"{k_name}.bias"),
+            state.pop(f"{v_name}.bias"),
+        )
 
     o_name = self._names.attn_output_proj.format(idx)
     converted_state[f"{prefix}.atten_func.output_projection.weight"] = state.pop(

{ai_edge_torch_nightly-0.2.0.dev20240604.dist-info → ai_edge_torch_nightly-0.2.0.dev20240605.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240604
+Version: 0.2.0.dev20240605
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.2.0.dev20240604.dist-info → ai_edge_torch_nightly-0.2.0.dev20240605.dist-info}/RECORD

@@ -6,7 +6,7 @@ ai_edge_torch/convert/conversion_utils.py,sha256=NpVm3Ms81_cIW5IYgGsr0BVganJJgBK
 ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=Ll2nNwufjcV5nSruQPXiloq7F1E7pWJ2T5clXmy1lk8,2825
 ai_edge_torch/convert/fx_passes/_pass_base.py,sha256=ijVyDclPnd6a0DWWUJkwR4igj6f82S-cE1-83QGPvgw,1652
-ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=wHVWNNMu5h_ya6GnnJn0cNif9xmdSqr8Vm-R7lllxZM,6213
+ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=2yqUwJJ2R233_X9FNMOP9oYRTTzH34TR_BIUj-wfnKw,7080
 ai_edge_torch/convert/fx_passes/build_upsample_bilinear2d_composite_pass.py,sha256=76XYoIlFDgrzp5QemoaEalPFcEbfszkEH_PLvO1ASCk,2607
 ai_edge_torch/convert/fx_passes/canonicalize_pass.py,sha256=UX6dJsxCqSkftXXvNBV-i7Bjk6H7qTyqzUnE640Itfg,1673
 ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=aRT8hTS3n9ie28lgu6mygtFO6Ypwu0qjNb0c81v9HLs,2448
@@ -22,7 +22,7 @@ ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partition
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=FlNKt2EhIKnlVEeUWTiv5sz446YKU6Yy1H0Gd6VRgkU,6432
 ai_edge_torch/convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/convert/test/test_convert.py,sha256=2qPmmGqnfV_o1gfsSdjGq3-JR1b323ligiy5MdAv9NA,8021
-ai_edge_torch/convert/test/test_convert_composites.py,sha256=SrVn_cEMtQhYYCMOUKK0K7M57MQNQX-lOUwieln0HGA,6616
+ai_edge_torch/convert/test/test_convert_composites.py,sha256=_Ojc-H6GOS5s8ek3_8eRBL_AiCs-k3srziPJ2R4Ulrg,7255
 ai_edge_torch/convert/test/test_convert_multisig.py,sha256=kMaGnHe9ylfyU68qCifYcaGwJqyejKz--QQt9jS2oUA,4537
 ai_edge_torch/debug/__init__.py,sha256=TKvmnjVk3asvYcVh6C-LPr6srgAF_nppSAupWEXqwPY,707
 ai_edge_torch/debug/culprit.py,sha256=vklaxBUfINdo44OsH7csILK70N41gEThCGchGEfbTZw,12789
@@ -40,10 +40,10 @@ ai_edge_torch/generative/examples/phi2/convert_to_tflite.py,sha256=6nOuwx9q3AUlY
 ai_edge_torch/generative/examples/phi2/phi2.py,sha256=VvigzPQ_LJHeADTsMliwFwPe2BcnOhFgKDqr_WZ2JQ8,5540
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=Lo4Dq7a3Kg-lyH56iqGtqCo5UaClQHRCTDdNagXGTo8,3535
-ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=KR1Ci4rlJeeGfsFRliCxUve9K7RTJLZfTRMgFtfQ4MU,2434
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=6REAYy1Bv-Iv5zcmA_m_W6fH6jt5a3IS6Vge18jS_Wo,3633
+ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=yUCJemEh4n8ez-yLgVU0HZAki-PZ9nY04DFjgpx9PUc,3698
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=MI73RjOeD4Kh7AL0j5_QXiZq-rl_qCdibSE6eCQCyeY,3804
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=AgVAdUbSkHXONVUjAyBQEXhIUUlinf9kNljcBpWnj3A,3276
-ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=nq94VpQ103eOimnmdyg7u3Xk1LH1IxGlmIbr2AttRIk,16224
+ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=TfbfsmuKoGsBENF9fYIAN_SMEQNhj-kjNdqQXFJGxpg,7784
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=L6hLaMQGb8-_BwSvTLIuDnZwfTqn0K4swBUjfPnYWZo,2341
 ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=FCbnwlkpYYb-tF7KscbSYjNEdg7XnuLju1cDuIRoQv8,8277
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=r9RqbyNvuvXOGu3ojtl7ZmbC7o4Pt8aUKAhN1yCdtEc,3397
@@ -56,22 +56,24 @@ ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py,sha256=5i
 ai_edge_torch/generative/examples/t5/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/t5/convert_to_tflite.py,sha256=bWtwtUacvJOEDUpuYvLTgkP7oTkXKJA-Tf4FPxlD1Cw,4536
 ai_edge_torch/generative/examples/t5/t5.py,sha256=q2gG5RRo7RgNzvHXYC0Juh6Tgt5d_RTMSWFaYvOKiZU,21065
-ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=anR99IrzR21x6yswFHYG5QQtPDZ7rVicf6STfMp54fU,8998
+ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=rRgwCEdVtzcJEaGbbBjw8HxCxrCX3pXA5nelawdYiME,9036
 ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=EV07_MEG3fv9g0ZGu9gbBd5BjjrGkxCT1pv7dvhz4TI,3791
+ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py,sha256=rzL5h7Z5DIEgfpc1pWgYHdKt2aR8ha_CUqTKQBSPBaU,5521
 ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=MUr6fSj2hBuYSlNbZtrBBpzqB_0WY-l_xYcd_TFFUjY,4831
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=E4I5OlC4zyl5cxiiu7uTED-zcwYRu210lP1zuT3xLBE,2566
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=hVGpuI8gpj4Rn9k4otsRE22MSLFHBDlUOgioY6Ru6VI,5629
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=PxixRZb00v5BQkWbDwaJgke4Rd5LwzdWe0zH9SG4Tj0,9127
+ai_edge_torch/generative/layers/attention.py,sha256=zNIBXxCOA5Mz_F_dfBbKpIovhtcB6q5a-i8oAxls1d0,7071
 ai_edge_torch/generative/layers/attention_utils.py,sha256=hXhuyKblPPxKIRzlAf1YNlwHgpbj-6nReRLhRHELx5k,6350
-ai_edge_torch/generative/layers/builder.py,sha256=ZSBVLv5EOtCkSW_Z8C2Hd7jN52nIAA2as1-qpmHGbCg,3201
+ai_edge_torch/generative/layers/builder.py,sha256=WLTeDId9t3Xwt0h1zxzqoYyFvfrNzPKLskcl39q8Aqw,3403
 ai_edge_torch/generative/layers/feed_forward.py,sha256=4j2QaSCw59Jkk_ixKDpKEj7FLRauzuExTiSNRzAjAhE,2820
 ai_edge_torch/generative/layers/kv_cache.py,sha256=4uiZLO3om5G3--kT04Jt0esEYznbkJ7QLzSHfb8mjc4,3090
-ai_edge_torch/generative/layers/model_config.py,sha256=KpJRIHV5BJH8QOa7h6LXLZyC7UDWgbCEsw0CvArz49Q,4064
+ai_edge_torch/generative/layers/model_config.py,sha256=2zT9nyoyuuyk5ziiww0VSJ6_JO7pDf7uOYbO9O3OQc4,4249
 ai_edge_torch/generative/layers/normalization.py,sha256=M27eW3TcNK20oaXClXtfnu0lLWrAGrSKSsbegRWnj3c,1867
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=12SsCuoRuLNCwnFGe_pHDOZEBwBcqXs87Aj0PaWWw4E,1383
+ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=dYafGC205QE5CLIbBTCI-7eVvEGZEHzs1toPEhemeDs,3391
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/quantize/example.py,sha256=t-YwyKSPAG-OZC1DfH-0vfie2RHHpTSQjxUY-tmhu5g,1543
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=ffBALrrbrfiG_mrOr-f3B1Gc6PlAma9gtvVnfP7SDzI,1862
@@ -84,7 +86,7 @@ ai_edge_torch/generative/test/loader_test.py,sha256=N88CbrLW7Q2x1EyurwdXQ6YjsA-y
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=1NfZxKo9Gx6CmVfd86K1FkmsNQnjzIV1ojBS85UGvT0,6500
 ai_edge_torch/generative/test/test_quantize.py,sha256=f70sH1ZFzdCwYj0MG-eg54WOC4LasR0D8CTUYpjxZYM,3728
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/loader.py,sha256=c-ZOIDBVnat_5l2W5sWU7HQm7CL-wducS8poSu5PlUg,10107
+ai_edge_torch/generative/utilities/loader.py,sha256=r-_hSanSjLZ_YXFpZUb0Up94u5F8JHp70Vf2nlONPSg,11269
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=guDTv-12UUvJGl4eDvvZX3t4rRKewfXO8SpcYXM6gbc,16156
 ai_edge_torch/hlfb/__init__.py,sha256=rrje8a2iuKboBoV96bVq7nlS9HsnuEMbHE5JiWmCxFA,752
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=2VXnHcGf23VOuP-1GriGIpuL98leBB8twp_qaScMnmc,4799
@@ -100,8 +102,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=ExThdTXqnWmGC3-F6sdXbXr8nYzkEe_qCz
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/METADATA,sha256=0Hjm-wBxxZRbVdxcFhmOKVwsU52_l5F6IblrflOWMWk,1748
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240604.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/METADATA,sha256=GJzwmKkM4T0H-vTvMyoxiD80WfppEpE_sd2Ip4aSbgM,1748
+ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240605.dist-info/RECORD,,