ai-edge-torch-nightly 0.5.0.dev20250408__py3-none-any.whl → 0.5.0.dev20250410__py3-none-any.whl

ai_edge_torch/generative/layers/kv_cache.py

@@ -16,24 +16,58 @@
 """Utility functions for externalized KV Cache."""
 
 import dataclasses
-from typing import List, Tuple
+from typing import Any, List, Tuple
 
 from ai_edge_torch.generative.custom_ops.dynamic_update_slice import dynamic_update_slice
 from ai_edge_torch.generative.layers import model_config
+from ai_edge_torch.generative.layers.experimental import types
 import torch
 import torch.utils._pytree as pytree
 
 
+KVLayout = Tuple[types.TensorDimensionMeta, types.TensorDimensionMeta]
+
+# Define common layouts for KV Cache.
+KV_LAYOUT_DEFAULT = (types.BTNH, types.BTNH)
+KV_LAYOUT_TRANSPOSED = (types.BNTH, types.BNHT)
+
+
 @dataclasses.dataclass
 class KVCacheEntry:
   """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).
+  The cache layout can be customized based on different use cases.
   """
 
   k_cache: torch.Tensor
   v_cache: torch.Tensor
+  kv_layout: KVLayout = KV_LAYOUT_DEFAULT
+
+  @classmethod
+  def construct_kv_shape_from_layout(
+      cls,
+      shape_spec: types.TensorDimensionMeta,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      batch_size: int,
+  ) -> List[int]:
+    """Constructs the shape of the key or value cache entry based on
+
+    the specified layout.
+    """
+    output_shape = []
+    for dim_spec in shape_spec:
+      if dim_spec is types.TensorDims.BATCH:
+        output_shape.append(batch_size)
+      elif dim_spec is types.TensorDims.SEQUENCE:
+        output_shape.append(kv_cache_max)
+      elif dim_spec is types.TensorDims.NUM_HEADS:
+        output_shape.append(config.num_query_groups)
+      elif dim_spec is types.TensorDims.HEAD_DIM:
+        output_shape.append(config.head_dim)
+      else:
+        raise ValueError(f"Unsupported dimension spec: {dim_spec}")
+    return output_shape
 
   @classmethod
   def from_model_config(
@@ -41,14 +75,20 @@ class KVCacheEntry:
       kv_cache_max: int,
       config: model_config.AttentionConfig,
       dtype: torch.dtype = torch.float32,
-      device: torch.device = None,
+      device: torch.device | None = None,
       batch_size: int = 1,
+      kv_layout: KVLayout = KV_LAYOUT_DEFAULT,
   ) -> "KVCacheEntry":
     """Build an instance of the class based on model config."""
-    shape = (batch_size, kv_cache_max, config.num_query_groups, config.head_dim)
-    k = torch.zeros(shape, dtype=dtype, device=device)
-    v = torch.zeros(shape, dtype=dtype, device=device)
-    obj = cls(k_cache=k, v_cache=v)
+    k_shape = cls.construct_kv_shape_from_layout(
+        kv_layout[0], kv_cache_max, config, batch_size
+    )
+    v_shape = cls.construct_kv_shape_from_layout(
+        kv_layout[1], kv_cache_max, config, batch_size
+    )
+    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, kv_layout=kv_layout)
     return obj
 
 
@@ -63,8 +103,9 @@ class KVCache:
       cls,
       config: model_config.ModelConfig,
       dtype: torch.dtype = torch.float32,
-      device: torch.device = None,
+      device: torch.device | None = None,
       batch_size: int = 1,
+      kv_layout: KVLayout = KV_LAYOUT_DEFAULT,
   ) -> "KVCache":
     """Build an instance of the class based on model config.
 
@@ -89,6 +130,7 @@ class KVCache:
             dtype,
             device,
             batch_size,
+            kv_layout,
         )
         for idx in range(config.num_layers)
     ]
@@ -104,7 +146,7 @@ class KVCache:
 def _flatten_kvc(kvc: KVCache) -> Tuple[List[str], List[str]]:
   flattened = []
   flat_names = []
-  none_names = []
+  none_names = [kvc.caches[0].kv_layout]
   for i, kv_entry in enumerate(kvc.caches):
     flattened.append(kv_entry.k_cache)
     flat_names.append(f"k_{i}")
@@ -121,22 +163,48 @@ def _flatten_kvc_with_keys(kvc: KVCache) -> Tuple[List, List]:
 
 
 def _unflatten_kvc(
-    values: List[torch.Tensor], context: Tuple[List, List]
+    values: List[torch.Tensor],
+    context: Tuple[List, List],
 ) -> KVCache:
   assert len(values) % 2 == 0, "Found odd number of K and V entries."
   num_layers = len(values) // 2
   flat_names = context[0]
+  kv_layout = context[1][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(
-        KVCacheEntry(k_cache=values[k_cache_idx], v_cache=values[v_cache_idx])
+        KVCacheEntry(
+            k_cache=values[k_cache_idx],
+            v_cache=values[v_cache_idx],
+            kv_layout=kv_layout,
+        )
     )
   obj = KVCache(tuple(kv_entries))
   return obj
 
 
+def _flatten_kv_entry(
+    kv_e: KVCacheEntry,
+) -> Tuple[List[torch.Tensor], Any]:
+  return ([kv_e.k_cache, kv_e.v_cache], kv_e.kv_layout)
+
+
+def _unflatten_kv_entry(
+    values: List[torch.Tensor],
+    context: Any,
+) -> KVCacheEntry:
+  return KVCacheEntry(*values, kv_layout=context)
+
+
+pytree.register_pytree_node(
+    KVCacheEntry,
+    _flatten_kv_entry,
+    _unflatten_kv_entry,
+    serialized_type_name="",
+)
+
 pytree.register_pytree_node(
     KVCache,
     _flatten_kvc,
@@ -145,7 +213,6 @@ pytree.register_pytree_node(
     serialized_type_name="",
 )
 
-
 def update(
     cache: KVCacheEntry,
     input_pos: torch.Tensor,
@@ -204,5 +271,5 @@ def _update_kv_impl(
   k = dynamic_update_slice(cache.k_cache, k_slice, k_slice_indices)
   v = dynamic_update_slice(cache.v_cache, v_slice, v_slice_indices)
 
-  updated_cache = KVCacheEntry(k, v)
+  updated_cache = KVCacheEntry(k, v, cache.kv_layout)
   return updated_cache

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -0,0 +1,124 @@
+# 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 utility functions for data loading etc.
+from dataclasses import dataclass
+from typing import Tuple
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa_default
+from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
+from ai_edge_torch.generative.layers.experimental import scaled_dot_product_attention as sdpa
+from ai_edge_torch.generative.layers.experimental import types
+import ai_edge_torch.generative.layers.model_config as cfg
+from multipledispatch import dispatch
+import torch
+
+
+def sdpa_with_kv_update(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv: kv_utils.KVCacheEntry,
+    input_pos: torch.Tensor,
+    mask: torch.Tensor,
+    config: cfg.AttentionConfig,
+) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
+  return sdpa_with_kv_update_impl(
+      kv.kv_layout[0](),  # key layout
+      kv.kv_layout[1](),  # value layout
+      query=query,
+      key=key,
+      value=value,
+      kv=kv,
+      input_pos=input_pos,
+      mask=mask,
+      config=config,
+  )
+
+
+@dispatch(types.BNTH, types.BNHT)
+def sdpa_with_kv_update_impl(
+    k_type, v_type, *args, **kwargs
+) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
+  query = kwargs["query"]
+  key = kwargs["key"]
+  value = kwargs["value"]
+  kv = kwargs["kv"]
+  input_pos = kwargs["input_pos"]
+  mask = kwargs["mask"]
+  config = kwargs["config"]
+
+  # Transpose k/v to specific layout for GPU implementation.
+  b, seq_len, n, h = query.shape
+  g = n // config.num_query_groups
+  # btnh -> bnth -> b(kg)th -> 1(bk)(gt)h
+  query = query.permute(0, 2, 1, 3).reshape(
+      1, b * config.num_query_groups, g * seq_len, h
+  )
+
+  key = key.permute(0, 2, 1, 3).reshape(
+      1, -1, seq_len, config.head_dim
+  )  # 1, bk, s, h
+  value = value.permute(0, 2, 3, 1).reshape(
+      1, -1, config.head_dim, seq_len
+  )  # 1, bk, h, s
+
+  if kv is not None:
+    kv = kv_utils_experimental.update(kv, input_pos, key, value)
+    key, value = kv.k_cache, kv.v_cache
+
+  sdpa_out = sdpa.scaled_dot_product_attention(
+      kv,
+      query,
+      key,
+      value,
+      config.head_dim,
+      mask=mask,
+      softcap=config.logit_softcap,
+  )  # 1, bk, gt, h
+  sdpa_out = (
+      sdpa_out.reshape(b, -1, seq_len, h)
+      .permute(0, 2, 1, 3)
+      .reshape(b, seq_len, -1)
+  )
+  return sdpa_out, kv
+
+
+@dispatch(object, object)
+def sdpa_with_kv_update_impl(
+    k_type, v_type, *args, **kwargs
+) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
+  query = kwargs["query"]
+  key = kwargs["key"]
+  value = kwargs["value"]
+  kv = kwargs["kv"]
+  input_pos = kwargs["input_pos"]
+  mask = kwargs["mask"]
+  config = kwargs["config"]
+
+  b, seq_len, _, _ = query.shape
+  if kv is not None:
+    kv = kv_utils.update(kv, input_pos, key, value)
+    key, value = kv.k_cache, kv.v_cache
+
+  sdpa_out = sdpa_default.scaled_dot_product_attention(
+      query,
+      key,
+      value,
+      config.head_dim,
+      mask=mask,
+      softcap=config.logit_softcap,
+  )
+  sdpa_out = sdpa_out.reshape(b, seq_len, -1)
+  return sdpa_out, kv

ai_edge_torch/generative/test/test_kv_cache.py

@@ -16,7 +16,6 @@
 """A suite of tests to validate KV Cache layer."""
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
 import ai_edge_torch.generative.layers.model_config as cfg
 import torch
 import torch.utils._pytree as pytree
@@ -117,7 +116,7 @@ class TestKVLayers(googletest.TestCase):
     self.assertEqual(input_specs[0].arg.name, "kv_k_0")
     self.assertEqual(input_specs[1].arg.name, "kv_v_0")
 
-  def test_pytree_roundtrip_experimental_kv_cache_base(self):
+  def test_pytree_roundtrip_kv_cache(self):
     NUM_LAYERS = 4
     config = self._get_test_config(
         num_layers=NUM_LAYERS,
@@ -125,15 +124,13 @@ class TestKVLayers(googletest.TestCase):
         num_query_groups=1,
         kv_cache_max_len=4,
     )
-    kv = kv_utils_experimental.KVCacheBase.from_model_config(
-        config, batch_size=1
-    )
+    kv = kv_utils.KVCache.from_model_config(config, batch_size=1)
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, NUM_LAYERS * 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
     self.assertEqual(kv, kv_unflat)
 
-  def test_pytree_roundtrip_experimental_kv_cache_derived(self):
+  def test_pytree_roundtrip_kv_cache_derived(self):
     NUM_LAYERS = 4
     config = self._get_test_config(
         num_layers=NUM_LAYERS,
@@ -141,41 +138,37 @@ class TestKVLayers(googletest.TestCase):
         num_query_groups=1,
         kv_cache_max_len=4,
     )
-    kv = kv_utils_experimental.KVCacheTransposed.from_model_config(
-        config, batch_size=1
+    kv = kv_utils.KVCache.from_model_config(
+        config, batch_size=1, kv_layout=kv_utils.KV_LAYOUT_TRANSPOSED
     )
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, NUM_LAYERS * 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
     self.assertEqual(kv, kv_unflat)
 
-  def test_pytree_roundtrip_experimental_kv_entry_base(self):
+  def test_pytree_roundtrip_kv_entry(self):
     attn_config = cfg.AttentionConfig(
         num_heads=1, head_dim=1, num_query_groups=1
     )
-    kv = kv_utils_experimental.KVCacheEntryBase.from_model_config(
-        32, attn_config
-    )
+    kv = kv_utils.KVCacheEntry.from_model_config(32, attn_config)
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
     self.assertEqual(kv, kv_unflat)
-    self.assertIsInstance(kv_unflat, kv_utils_experimental.KVCacheEntryBase)
+    self.assertIsInstance(kv_unflat, kv_utils.KVCacheEntry)
 
-  def test_pytree_roundtrip_experimental_kv_entry_derived(self):
+  def test_pytree_roundtrip_kv_entry_derived(self):
     attn_config = cfg.AttentionConfig(
         num_heads=1, head_dim=1, num_query_groups=1
     )
-    kv = kv_utils_experimental.KVCacheEntryTransposed.from_model_config(
-        32, attn_config
+    kv = kv_utils.KVCacheEntry.from_model_config(
+        32, attn_config, kv_layout=kv_utils.KV_LAYOUT_TRANSPOSED
     )
     flat, treespec = pytree.tree_flatten(kv)
     self.assertLen(flat, 2)
     kv_unflat = pytree.tree_unflatten(flat, treespec)
     self.assertEqual(kv, kv_unflat)
-    self.assertIsInstance(
-        kv_unflat, kv_utils_experimental.KVCacheEntryTransposed
-    )
+    self.assertIsInstance(kv_unflat, kv_utils.KVCacheEntry)
 
 
 if __name__ == "__main__":

ai_edge_torch/generative/utilities/converter.py

@@ -20,6 +20,7 @@ import pathlib
 from typing import Optional, Union
 from absl import flags
 from ai_edge_torch._convert import converter as converter_utils
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.quantize import quant_recipes
@@ -218,9 +219,13 @@ def _export_helper(
       [[0] for _ in range(export_config.decode_batch_size)], dtype=torch.int
   )
   decode_input_pos = torch.tensor([0], dtype=torch.int)
-  prefill_kv = export_config.kvcache_cls.from_model_config(config)
-  decode_kv = export_config.kvcache_cls.from_model_config(
-      config, batch_size=export_config.decode_batch_size
+  prefill_kv = kv_utils.KVCache.from_model_config(
+      config, kv_layout=export_config.kvcache_layout
+  )
+  decode_kv = kv_utils.KVCache.from_model_config(
+      config,
+      batch_size=export_config.decode_batch_size,
+      kv_layout=export_config.kvcache_layout,
   )
 
   quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None

ai_edge_torch/generative/utilities/export_config.py

@@ -32,7 +32,9 @@ class ExportConfig:
   # Attention masks given as inputs to the model.
   prefill_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
   decode_mask: Optional[torch.Tensor | List[torch.Tensor]] = None
-  # The KV Cache class for K and V buffers in attention.
+  # The KV Cache layout for K and V buffers in attention.
+  kvcache_layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT
+  # TODO(b/409373223): The KV Cache class for K and V buffers in attention.
   kvcache_cls: type = kv_utils.KVCache
   # The batch size of the decode signature.
   decode_batch_size: int = 1

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -52,6 +52,7 @@ def torch_dtype_to_tf(dtype):
       torch.int32: tf.int32,
       torch.int16: tf.int16,
       torch.bool: tf.bool,
+      torch.bfloat16: tf.bfloat16,
   }.get(dtype)
 
 

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -301,3 +301,22 @@ def _aten_slice_scatter(lctx, self, src, dim=0, start=None, end=None, step=1):
   )
   out = stablehlo.select(pred, self, src)
   return out
+
+
+# Schema:
+#   - aten::_to_copy(Tensor self, *, ScalarType? dtype=None,
+#       Layout? layout=None, Device? device=None, bool? pin_memory=None,
+#       bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor
+@lower(torch.ops.aten._to_copy.default)
+def _aten_to_copy(
+    lctx, x: ir.Value, dtype: torch.dtype | None = None, **kwargs
+):
+  if not dtype:
+    return x
+
+  return stablehlo.convert(
+      ir.RankedTensorType.get(
+          x.type.shape, utils.torch_dtype_to_ir_element_type(dtype)
+      ),
+      x,
+  )

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -74,7 +74,6 @@ lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit)
 lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit_no_training)
 lower_by_torch_xla2(torch.ops.aten._pdist_forward)
 lower_by_torch_xla2(torch.ops.aten._softmax)
-lower_by_torch_xla2(torch.ops.aten._to_copy)
 lower_by_torch_xla2(torch.ops.aten._unsafe_index)
 lower_by_torch_xla2(torch.ops.aten._unsafe_view)
 lower_by_torch_xla2(torch.ops.aten.acos)

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -37,6 +37,7 @@ def torch_dtype_to_ir_element_type(dtype) -> ir.Type:
       torch.int16: functools.partial(ir.IntegerType.get_signless, 16),
       torch.int8: functools.partial(ir.IntegerType.get_signless, 8),
       torch.bool: functools.partial(ir.IntegerType.get_signless, 1),
+      torch.bfloat16: ir.BF16Type.get,
   }[dtype]
   return ty_get()
 

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250408"
+__version__ = "0.5.0.dev20250410"

{ai_edge_torch_nightly-0.5.0.dev20250408.dist-info → ai_edge_torch_nightly-0.5.0.dev20250410.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250408
+Version: 0.5.0.dev20250410
 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
@@ -25,7 +25,9 @@ License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: safetensors
-Requires-Dist: multipledispatchtransformerskagglehub
+Requires-Dist: multipledispatch
+Requires-Dist: transformers
+Requires-Dist: kagglehub
 Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
 Requires-Dist: tf-nightly>=2.19.0.dev20250101

{ai_edge_torch_nightly-0.5.0.dev20250408.dist-info → ai_edge_torch_nightly-0.5.0.dev20250410.dist-info}/RECORD

@@ -2,16 +2,17 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=Lf4c2aVfixNX2KTgdqQTLOGBdi0vVxNOkJuNt4SvQ8c,706
+ai_edge_torch/version.py,sha256=dQvyVQmvNYF8n8HwlkY-9fdSo-n3_bdLO9EAZpJnC8s,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
+ai_edge_torch/_convert/conversion.py,sha256=GPDsXhfECjDzOut4vh_d9qWcyfpxobFMBTsC7MyJbM0,5557
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
-ai_edge_torch/_convert/fx_passes/__init__.py,sha256=dG4WIICk0FqCH9euvbYHHsybRN7B1cYcuxN_OYxmjWo,1263
+ai_edge_torch/_convert/fx_passes/__init__.py,sha256=6LtGzzqT2IXprfI_vPYKhE7IuN5XmPG0xy-v0UtZ9yk,1361
 ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=a1KhqLetFb_efRHjX4T-zH0vF-U37Ha5I1CPIAsIluE,9211
 ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=3JyjiHpn17Zhfq3yGQXK5LMH71DQPXHb_4GOkP9uAjY,4251
+ai_edge_torch/_convert/fx_passes/cast_inputs_bf16_to_f32_pass.py,sha256=90YxLVAAkiA3qKr4Um__JmPeC1bTeA2PxBCj0GETq1Q,1748
 ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=Z6E3U7SYZvMl3Ivpqa3burVOLKFndEZuNmWKNxjq2mM,2386
 ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=HCOkj0k3NhaYbtfjE8HDXVmYhZ9fL5V_u6VunVh9mN4,2116
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=UKC-wM93-oe8spxyFqgybJ0TwnSRw8f-SOA2glCh2FA,890
@@ -26,7 +27,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=o6tuJkD-ESaQxLxJpN104qpchm3LCtPmHinzQxe6PSg,17226
+ai_edge_torch/_convert/test/test_convert.py,sha256=6vQa0UJn2L3qxR967_-vkfLrO7JdrLLBk4BfguOtHRI,17874
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -65,12 +66,12 @@ ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSd
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=IoBhEMwH07-tFm5-U6F2hpCsI8xynglhq1x9tIOdaPQ,1322
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=tR8RflXocDZqvuStyw9aFlzuiTllEC8rNnjrxms6_Is,5727
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=6Dkxi7Vs8xBaqMif00ATQSr_hTPhYXMdDqHwzOsAzq8,2952
-ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=__kpzArZ0mLfX7IzpHPmYFuhKTP9uI_9Lrzk_EfFDlE,15701
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=szssSBrIUYdNIoU7LHdAq7wCqgjaY6qbV8yvTgg796Q,2945
+ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=n6ZQfqNEHuOhY7Pu21bb8Eax8yn2Sx5osTKJKmhonXY,15659
 ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=5PEt0aWJ5wkUBvMoWFOJ-C48ZhG7uCVb8PCKQtZ8Fvw,6485
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=uRoLoBWzFtQz5wFZfPCxbkvZsgPAqSkUUsV3977GbYc,5184
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
-ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=u30qiZu3HJCTt5noWqtf9PgGLKQ87ke4Zpa4cpG6-As,8883
+ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=nEv0qQ0l6gSXKxP5mNwkd2lRGxpFfD4e7FNV3V76zhw,8915
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=A4uLUdqvU1NKo3seqZlWSS3fqYahnEKqNBQBJO6yXvE,1762
 ai_edge_torch/generative/examples/llama/llama.py,sha256=UKvMO85_5z1vEY5MVu6QBW_vpQYA8LWHbJI4Yx6BrCc,6592
@@ -153,17 +154,18 @@ ai_edge_torch/generative/layers/attention.py,sha256=wLZ1jgUlcODBWgK3hnnhclHuuQDq
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
-ai_edge_torch/generative/layers/kv_cache.py,sha256=zjdovWqgEKtx7cvbA0apOwXaNft5AXxNTbJhBT4CXyg,6541
+ai_edge_torch/generative/layers/kv_cache.py,sha256=9kkFpB9msgUDStFxEyQYYsavKPP4Dgqb_NFcd4hA4aU,8502
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
 ai_edge_torch/generative/layers/model_config.py,sha256=nLXvTkDAIHJQ0PTaWODF8oxJQoJ-K8D10cKR9229SAw,8355
 ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
+ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=TcwiI1IHhcYUrTx0kpSPAJMxFfjFcDwAHHULfZm67U4,3785
 ai_edge_torch/generative/layers/experimental/__init__.py,sha256=nz-K0h8DfiATHzR6s1_bCw2akUmHWffU1bDRSkIzSqI,592
-ai_edge_torch/generative/layers/experimental/attention.py,sha256=95djjlJItDVuSNE3BL0b6u3lQoIhmmdvaik7qBBvQA0,8909
-ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=uXUxiQjPndXYZVGKgm9FxzHgQDal8GdY7cUZDpc_Sno,9997
-ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=YFW0iGcZjTuej6VFIkwdSY28fIQi_KTAVdT8gWNmq7o,2880
-ai_edge_torch/generative/layers/experimental/types.py,sha256=bPPxw6TOCZVWdeDP3vCbOnjNP5-bdUMmfsfO-EtdazQ,2847
+ai_edge_torch/generative/layers/experimental/attention.py,sha256=XYbo1KlmiMEuwArye0Ul86jEsdxLr1RG-usRpidZiT8,8001
+ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=zgpFVftOfllvjh9-UEBSvUbm152SnQETn29rUMMMvAM,2978
+ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=8M6tC5kIUus-wbMEKDSMbCLnsobs6rgbujycsmhYa5g,2807
+ai_edge_torch/generative/layers/experimental/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=ZteHZXK6HKyxYji49DQ46sA9aIy7U3Jnz0HZp6hfevY,28996
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
@@ -177,7 +179,7 @@ ai_edge_torch/generative/quantize/quant_recipes.py,sha256=0Kvr_o7pbMnE8VMe6Ml0FB
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVufetPvkDoD7rUowIOSKV9oV5Kk,1418
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_custom_dus.py,sha256=MjIhTvkTko872M35XMciobvICcDWTcIDJ3rociko-wM,3267
-ai_edge_torch/generative/test/test_kv_cache.py,sha256=MBPS-0bDXB0tQSKHa1XwDQeVIfabRbc8JQA99h9fzlQ,5961
+ai_edge_torch/generative/test/test_kv_cache.py,sha256=1sXN2RPntq0PP3IEy0NkvIbzQ0Y8JhPIwRSFwO9JLlE,5728
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_lora.py,sha256=6QIM6RLTc2HrodGpp_aS3OxM9Rco2KAzEnYgotkg41M,5310
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=jfqkECCX7XKHeBAuDXrkwQJf0vM72eG3LMc5rluha84,6191
@@ -185,8 +187,8 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=-v2Vj7Qdd3Gy
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=ycXDcd3ZE-EdjksDjHi4ru3JpfhtrfOompg_990qvWI,9607
-ai_edge_torch/generative/utilities/export_config.py,sha256=-UuukWqUUj8RM8lTtMCa_PD6SqCZv97i4BMiJA2zBPg,1491
+ai_edge_torch/generative/utilities/converter.py,sha256=87Tzj-gLydx8_xnHxKlCbMmM1XHShstpKi8RH3xY7Xw,9757
+ai_edge_torch/generative/utilities/export_config.py,sha256=8-795nyd3M34LkGhgW7hwHlJyTc2Oz1iipHK8yBhdFs,1633
 ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
 ai_edge_torch/generative/utilities/model_builder.py,sha256=ZYX1TxpFdj573du2QCyHJlFjx4q1m12R74fp4Gwl92A,6343
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
@@ -203,7 +205,7 @@ ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=5kmOJWCc7sU1Hrqr1y17BtShUrss
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
 ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5SIrG0,3276
 ai_edge_torch/lowertools/common_utils.py,sha256=4HQtquPZ6oiId8vR_1ykW_uK4ELnyo5zo3MlX1QYW4c,4513
-ai_edge_torch/lowertools/odml_torch_utils.py,sha256=dxg2pBuVhSZeY2Ouc0F6nHiZilWZmpSPA7I8kGqSkVI,8282
+ai_edge_torch/lowertools/odml_torch_utils.py,sha256=QRuS7S5lULRWEh3J1sWIsnKh-rbX7rd9tt6JJHbMPfo,8317
 ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUGdSY1ieZjw,1949
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
 ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
@@ -223,17 +225,17 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNi
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
 ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=uJ-niilt1c-D6QJzLwgvCUf62le_JsxQTlqj_iP_Ps0,1009
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=Jq8_yAxC7ilzd6tOaRyBsOUEeenFF_EAC5haacZT4Pg,10247
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=4syWstepGiw3IKa8O7lciXywY7RFJ7OCWFMU1Lg3h-s,10777
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=VhmeGFnB5hrUsALiVWV96JJOqPDrTIWouHjTvLuT5eU,2477
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=h-YHW7xmvt9dpea-7Zj82HW7h5TKzW6GBEE13dIJQ40,11518
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=JRGLXW8EQ1L-vdiVTkD1kb4AnTU05eRwZ7Ke010hZmg,11473
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
-ai_edge_torch/odml_torch/lowerings/utils.py,sha256=-TzK1igPgR38oZkU1iPh-DZhlKVwuBtGWVC-y81PXzY,8935
+ai_edge_torch/odml_torch/lowerings/utils.py,sha256=uJaFbbgvYMI4-VFpFcMpaObNfBQl6nV0x8Yo8LaSAOE,8974
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
 ai_edge_torch/quantize/pt2e_quantizer.py,sha256=CKIEhs9jCcna64qj1jFH9zEbMbRdyeGV_TmSqEBPjes,15741
@@ -243,8 +245,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/METADATA,sha256=-Bw-LUn9l-B66aMZiFiUiYBifr1B6Fr86LU8KXtBieo,2019
-ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/METADATA,sha256=8m6hxUmTT0arSNdEuo-mOyg1w9T3ekAvedvY2T6Opgw,2051
+ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/RECORD,,