ai-edge-torch-nightly 0.4.0.dev20250305__py3-none-any.whl → 0.4.0.dev20250306__py3-none-any.whl

ai_edge_torch/generative/layers/experimental/kv_cache.py

@@ -19,8 +19,8 @@ This is an experimental implementation and is subject to change at any time.
 """
 
 import dataclasses
-from typing import List, Tuple
-
+import functools
+from typing import Any, List, Tuple, Type
 from ai_edge_torch.generative.layers import model_config
 from ai_edge_torch.generative.layers.experimental import types
 from ai_edge_torch.generative.utilities import dynamic_update_slice as dus_utils
@@ -46,7 +46,7 @@ class KVCacheEntryBase:
       v_shape: Tuple[int, ...],
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
-  ) -> "KVCacheEntryBase":
+  ):
     """Build an instance of the class based on model config."""
     k = torch.zeros(k_shape, dtype=dtype, device=device)
     v = torch.zeros(v_shape, dtype=dtype, device=device)
@@ -61,7 +61,7 @@ class KVCacheEntryBase:
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
       batch_size: int = 1,
-  ) -> "KVCacheEntryBase":
+  ):
     """Build an instance of the class based on model config."""
     shape = (batch_size, kv_cache_max, config.num_query_groups, config.head_dim)
     return cls._from_model_config(shape, shape, dtype, device)
@@ -87,7 +87,7 @@ class KVCacheEntryTransposed(KVCacheEntryBase):
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
       batch_size: int = 1,
-  ) -> "KVCacheEntryBase":
+  ):
     """Build an instance of the class based on model config."""
     k_shape = (
         batch_size,
@@ -104,6 +104,35 @@ class KVCacheEntryTransposed(KVCacheEntryBase):
     return cls._from_model_config(k_shape, v_shape, dtype, device)
 
 
+def _flatten_kv_entry(
+    kv_e: KVCacheEntryBase,
+) -> Tuple[List[torch.Tensor], Any]:
+  return ([kv_e.k_cache, kv_e.v_cache], None)
+
+
+def _unflatten_kv_entry(
+    kv_entry_ty: Type[KVCacheEntryBase],
+    values: List[torch.Tensor],
+    unused_context: Any,
+) -> KVCacheEntryBase:
+  return kv_entry_ty(*values)
+
+
+pytree.register_pytree_node(
+    KVCacheEntryTransposed,
+    _flatten_kv_entry,
+    functools.partial(_unflatten_kv_entry, KVCacheEntryTransposed),
+    serialized_type_name="",
+)
+
+pytree.register_pytree_node(
+    KVCacheEntryBase,
+    _flatten_kv_entry,
+    functools.partial(_unflatten_kv_entry, KVCacheEntryBase),
+    serialized_type_name="",
+)
+
+
 @dataclasses.dataclass
 class KVCacheBase:
   """A utility class for holding KV cache entries per layer."""
@@ -118,7 +147,7 @@ class KVCacheBase:
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
       batch_size: int = 1,
-  ) -> "KVCacheBase":
+  ):
     caches = [
         kv_entry_cls.from_model_config(
             config.kv_cache_max,
@@ -139,7 +168,7 @@ class KVCacheBase:
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
       batch_size: int = 1,
-  ) -> "KVCacheBase":
+  ):
     """Build an instance of the class based on model config.
 
     Args:
@@ -179,7 +208,7 @@ class KVCacheBTNH(KVCacheBase):
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
       batch_size: int = 1,
-  ) -> "KVCacheBTNH":
+  ):
     return cls._from_model_config(
         KVCacheEntryBTNH,
         config=config,
@@ -199,7 +228,7 @@ class KVCacheTransposed(KVCacheBase):
       dtype: torch.dtype = torch.float32,
       device: torch.device = None,
       batch_size: int = 1,
-  ) -> "KVCacheBTNH":
+  ):
     return cls._from_model_config(
         KVCacheEntryTransposed,
         config=config,
@@ -229,7 +258,10 @@ def _flatten_kvc_with_keys(kvc: KVCacheBase) -> Tuple[List, List]:
 
 
 def _unflatten_kvc(
-    values: List[torch.Tensor], context: Tuple[List, List]
+    kv_ty: Type[KVCacheBase],
+    kv_entry_type: Type[KVCacheEntryBase],
+    values: List[torch.Tensor],
+    context: Tuple[List, List],
 ) -> KVCacheBase:
   assert len(values) % 2 == 0, "Found odd number of K and V entries."
   num_layers = len(values) // 2
@@ -239,18 +271,18 @@ def _unflatten_kvc(
     k_cache_idx = flat_names.index(f"k_{i}")
     v_cache_idx = flat_names.index(f"v_{i}")
     kv_entries.append(
-        KVCacheEntryBase(
-            k_cache=values[k_cache_idx], v_cache=values[v_cache_idx]
-        )
+        kv_entry_type(k_cache=values[k_cache_idx], v_cache=values[v_cache_idx])
     )
-  obj = KVCacheBase(tuple(kv_entries))
+  obj = kv_ty(tuple(kv_entries))
   return obj
 
 
 pytree.register_pytree_node(
     KVCacheTransposed,
     _flatten_kvc,
-    _unflatten_kvc,
+    functools.partial(
+        _unflatten_kvc, KVCacheTransposed, KVCacheEntryTransposed
+    ),
     flatten_with_keys_fn=_flatten_kvc_with_keys,
     serialized_type_name="",
 )
@@ -258,7 +290,7 @@ pytree.register_pytree_node(
 pytree.register_pytree_node(
     KVCacheBase,
     _flatten_kvc,
-    _unflatten_kvc,
+    functools.partial(_unflatten_kvc, KVCacheBase, KVCacheEntryBase),
     flatten_with_keys_fn=_flatten_kvc_with_keys,
     serialized_type_name="",
 )

ai_edge_torch/generative/test/test_kv_cache.py

@@ -16,8 +16,10 @@
 """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
 
 from absl.testing import absltest as googletest
 
@@ -115,6 +117,66 @@ 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):
+    NUM_LAYERS = 4
+    config = self._get_test_config(
+        num_layers=NUM_LAYERS,
+        head_dim=2,
+        num_query_groups=1,
+        kv_cache_max_len=4,
+    )
+    kv = kv_utils_experimental.KVCacheBase.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):
+    NUM_LAYERS = 4
+    config = self._get_test_config(
+        num_layers=NUM_LAYERS,
+        head_dim=2,
+        num_query_groups=1,
+        kv_cache_max_len=4,
+    )
+    kv = kv_utils_experimental.KVCacheTransposed.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_entry_base(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
+    )
+    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)
+
+  def test_pytree_roundtrip_experimental_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
+    )
+    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
+    )
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250305"
+__version__ = "0.4.0.dev20250306"

{ai_edge_torch_nightly-0.4.0.dev20250305.dist-info → ai_edge_torch_nightly-0.4.0.dev20250306.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250305
+Version: 0.4.0.dev20250306
 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.4.0.dev20250305.dist-info → ai_edge_torch_nightly-0.4.0.dev20250306.dist-info}/RECORD

@@ -2,7 +2,7 @@ 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=KSuxuG4iAZMdD9Pi3Eg36fWk8j5YbKwBegeGyh08BIg,706
+ai_edge_torch/version.py,sha256=3TrrWqlr-XarP5R47N_A6I7W4epX_4Iuuv5YVSKn-rQ,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_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -151,7 +151,7 @@ ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIr
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
 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=VN4gn4ylaVOwaTR5EXKv0YTVgpQ850bmjGLCgCCI1ps,9267
+ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=0H-Rqtm6ArMxchHSv3eeX8W3AryoF73EFEpGNfjciK8,9996
 ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=1vMh1L3uYX4ptKQMWcAjxkL1v2-g0jmOiuai8ydp0dc,2879
 ai_edge_torch/generative/layers/experimental/types.py,sha256=bPPxw6TOCZVWdeDP3vCbOnjNP5-bdUMmfsfO-EtdazQ,2847
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -167,7 +167,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=gxG78CcTpXF3iLzDR15Rlz1ey1tNTlSdkp6TeYEijp0,3301
-ai_edge_torch/generative/test/test_kv_cache.py,sha256=2AulHBS3hC4b_68PNNBkRVOrypy4IM5YjC4p-6dgCMM,3793
+ai_edge_torch/generative/test/test_kv_cache.py,sha256=MBPS-0bDXB0tQSKHa1XwDQeVIfabRbc8JQA99h9fzlQ,5961
 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
@@ -233,8 +233,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 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=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.4.0.dev20250305.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250305.dist-info/METADATA,sha256=PFPy_Qd9oHyUNQFnd04iTXM5UKD0R_azKA3QLYAWr8o,1966
-ai_edge_torch_nightly-0.4.0.dev20250305.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250305.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250305.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250306.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250306.dist-info/METADATA,sha256=AGXvfH7AuPCCnrW0vgLAcW2e2jUQLqx1berQqniMFc0,1966
+ai_edge_torch_nightly-0.4.0.dev20250306.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250306.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250306.dist-info/RECORD,,