ai-edge-torch-nightly 0.5.0.dev20250424__py3-none-any.whl → 0.5.0.dev20250426__py3-none-any.whl

ai_edge_torch/generative/examples/qwen/qwen.py

@@ -53,6 +53,7 @@ def get_3b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   norm_config = cfg.NormalizationConfig(
       type=cfg.NormalizationType.RMS_NORM,
       epsilon=1e-06,
+      enable_hlfb=True,
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -35,6 +35,10 @@ def main(_):
   pytorch_model = smollm.build_model(
       flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
+
+  export_config = export_cfg.get_from_flags()
+  export_config.decode_batch_size = _DECODE_BATCH_SIZE.value
+
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
@@ -42,9 +46,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_cfg.ExportConfig(
-          decode_batch_size=_DECODE_BATCH_SIZE.value
-      ),
+      export_config=export_config,
   )
 
 

ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py

@@ -34,6 +34,9 @@ def main(_):
       flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
 
+  export_config = export_cfg.get_from_flags()
+  export_config.decode_batch_size = _DECODE_BATCH_SIZE.value
+
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
@@ -41,9 +44,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_cfg.ExportConfig(
-          decode_batch_size=_DECODE_BATCH_SIZE.value
-      ),
+      export_config=export_config,
   )
 
 

ai_edge_torch/generative/examples/smollm/smollm.py

@@ -49,7 +49,9 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=1536,
   )
-  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
+  )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -21,7 +21,6 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("tiny_llama")
-ExportConfig = export_config.ExportConfig
 
 
 def main(_):
@@ -35,7 +34,7 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=ExportConfig(),
+      export_config=export_config.get_from_flags(),
   )
 
 

ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py

@@ -49,7 +49,9 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=5632,
   )
-  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
+  )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,

ai_edge_torch/generative/layers/kv_cache.py

@@ -51,10 +51,7 @@ class KVCacheEntry:
       config: model_config.AttentionConfig,
       batch_size: int,
   ) -> List[int]:
-    """Constructs the shape of the key or value cache entry based on
-
-    the specified layout.
-    """
+    """Construct the shape of KV cache entry based on the specified layout."""
     output_shape = []
     for dim_spec in shape_spec:
       if dim_spec is types.TensorDims.BATCH:
@@ -213,6 +210,7 @@ pytree.register_pytree_node(
     serialized_type_name="",
 )
 
+
 def update(
     cache: KVCacheEntry,
     input_pos: torch.Tensor,

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -17,6 +17,8 @@
 import math
 from typing import Optional
 
+from ai_edge_torch.generative.custom_ops import bmm_4d as bmm_lib
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 import torch.nn.functional as F
@@ -142,3 +144,52 @@ def scaled_dot_product_attention_with_hlfb(
   result = y.transpose(1, 2)
   result = builder.mark_outputs(result)
   return result
+
+
+def scaled_dot_product_attention_transposed(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    head_size: int,
+    mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+    softcap: Optional[float] = None,
+):
+  """Scaled dot product attention with transposed key and value.
+
+  Args:
+    query: Query tensor, with shape [B, T, N, H].
+    key: Key tensor, with shape [B, T, KV_LEN, H].
+    value: Value tensor, with shape [B, T, KV_LEN, H].
+    head_size (int): head dimension.
+    mask (torch.Tensor): the optional mask tensor.
+    scale (float): the optional scale factor.
+    softcap (float): the optional softcap for the logits.
+
+  Returns:
+    The output tensor of scaled_dot_product_attention_transposed.
+  """
+
+  if scale is None:
+    scale = 1.0 / math.sqrt(head_size)
+
+  query = query * scale
+
+  assert mask is not None, "Mask should not be None!"
+  t = mask.shape[2]
+
+  logits = bmm_lib.bmm_4d(query, key)
+
+  _, bk, gt, s = logits.shape
+  g = gt // t
+  logits = logits.reshape((bk, g, t, s))
+  if softcap is not None:
+    logits = torch.tanh(logits / softcap)
+    logits = logits * softcap
+
+  padded_logits = logits + mask
+  padded_logits = padded_logits.reshape(1, bk, gt, s)
+  probs = F.softmax(padded_logits, dim=-1).type_as(key)
+  encoded = bmm_lib.bmm_4d(probs, value)
+
+  return encoded  # 1, bk, gt, h

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -18,9 +18,8 @@
 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 import scaled_dot_product_attention as sdpa
 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
 import ai_edge_torch.generative.layers.model_config as cfg
 import torch
 
@@ -72,8 +71,7 @@ def _sdpa_with_kv_update_transposed(
   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,
+  sdpa_out = sdpa.scaled_dot_product_attention_transposed(
       query,
       key,
       value,
@@ -105,9 +103,9 @@ def _sdpa_with_kv_update_default(
     key, value = kv.k_cache, kv.v_cache
 
   if enable_hlfb:
-    sdpa_func = sdpa_default.scaled_dot_product_attention_with_hlfb
+    sdpa_func = sdpa.scaled_dot_product_attention_with_hlfb
   else:
-    sdpa_func = sdpa_default.scaled_dot_product_attention
+    sdpa_func = sdpa.scaled_dot_product_attention
   sdpa_out = sdpa_func(
       query,
       key,

ai_edge_torch/generative/test/test_model_conversion.py

@@ -32,10 +32,8 @@ class TestModelConversion(googletest.TestCase):
 
   def setUp(self):
     super().setUp()
-    # Builder function for an Interpreter that supports custom ops.
     self._interpreter_builder = (
-        lambda tflite_model: lambda: interpreter.InterpreterWithCustomOps(
-            custom_op_registerers=["GenAIOpsRegisterer"],
+        lambda tflite_model: lambda: interpreter.Interpreter(
             model_content=tflite_model,
             experimental_default_delegate_latest_features=True,
         )
@@ -85,44 +83,24 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_toy_model_with_kv_cache(self):
     self._test_model_with_kv_cache(enable_hlfb=False)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_toy_model_with_kv_cache_with_hlfb(self):
     self._test_model_with_kv_cache(enable_hlfb=True)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_toy_model_with_kv_cache_transposed(self):
     self._test_model_with_kv_cache(kv_layout=kv_cache.KV_LAYOUT_TRANSPOSED)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_toy_model_has_dus_op(self):
     """Tests that the model has the dynamic update slice op."""
     _, edge_model, _ = self._get_params(
         enable_hlfb=True, kv_layout=kv_cache.KV_LAYOUT_DEFAULT
     )
-    interpreter_ = interpreter.InterpreterWithCustomOps(
-        custom_op_registerers=["GenAIOpsRegisterer"],
-        model_content=edge_model.tflite_model(),
-        experimental_default_delegate_latest_features=True,
-    )
+    interpreter = self._interpreter_builder(edge_model.tflite_model())()
 
     # pylint: disable=protected-access
-    op_names = [op["op_name"] for op in interpreter_._get_ops_details()]
+    op_names = [op["op_name"] for op in interpreter._get_ops_details()]
     self.assertIn("DYNAMIC_UPDATE_SLICE", op_names)
 
   def _test_multisig_model(
@@ -197,19 +175,11 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_tiny_llama_multisig(self):
     config = tiny_llama.get_fake_model_config()
     pytorch_model = tiny_llama.TinyLlama(config).eval()
     self._test_multisig_model(config, pytorch_model, atol=1e-5, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_tiny_llama_multisig_kv_layout_transposed(self):
     config = tiny_llama.get_fake_model_config()
     pytorch_model = tiny_llama.TinyLlama(config).eval()

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -20,6 +20,7 @@ from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
 from ai_edge_torch.generative.examples.deepseek import deepseek
 from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.examples.gemma import gemma2
+from ai_edge_torch.generative.examples.hammer import hammer
 from ai_edge_torch.generative.examples.llama import llama
 from ai_edge_torch.generative.examples.openelm import openelm
 from ai_edge_torch.generative.examples.paligemma import decoder
@@ -48,10 +49,8 @@ class TestModelConversion(googletest.TestCase):
 
   def setUp(self):
     super().setUp()
-    # Builder function for an Interpreter that supports custom ops.
     self._interpreter_builder = (
-        lambda tflite_model: lambda: interpreter.InterpreterWithCustomOps(
-            custom_op_registerers=["GenAIOpsRegisterer"],
+        lambda tflite_model: lambda: interpreter.Interpreter(
             model_content=tflite_model,
             experimental_default_delegate_latest_features=True,
         )
@@ -94,110 +93,68 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_gemma1(self):
     config = gemma1.get_fake_model_config()
     pytorch_model = gemma1.Gemma1(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_gemma2(self):
     config = gemma2.get_fake_model_config()
     pytorch_model = gemma2.Gemma2(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_llama(self):
     config = llama.get_fake_model_config()
     pytorch_model = llama.Llama(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_phi2(self):
     config = phi2.get_fake_model_config()
     pytorch_model = phi2.Phi2(config).eval()
     # Phi-2 logits are very big, so we need a larger absolute tolerance.
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_phi3(self):
     config = phi3.get_fake_model_config()
     pytorch_model = phi3.Phi3_5Mini(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-5, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_phi4(self):
     config = phi4.get_fake_model_config()
     pytorch_model = phi4.Phi4Mini(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_smollm(self):
     config = smollm.get_fake_model_config()
     pytorch_model = smollm.SmolLM(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_smollm2(self):
     config = smollm.get_fake_model_config_v2()
     pytorch_model = smollm.SmolLM2(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_openelm(self):
     config = openelm.get_fake_model_config()
     pytorch_model = openelm.OpenELM(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-4, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_qwen(self):
     config = qwen.get_fake_model_config()
     pytorch_model = qwen.Qwen(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-3, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_deepseek(self):
     config = deepseek.get_fake_model_config()
     pytorch_model = deepseek.DeepSeekDistillQwen(config).eval()
     self._test_model(config, pytorch_model, "prefill", atol=1e-5, rtol=1e-5)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
+  def test_hammer(self):
+    config = hammer.get_fake_model_config()
+    pytorch_model = hammer.Hammer(config).eval()
+    self._test_model(config, pytorch_model, "prefill", atol=1e-5, rtol=1e-5)
+
+
   def test_amd_llama_135m(self):
     config = amd_llama_135m.get_fake_model_config()
     pytorch_model = amd_llama_135m.AmdLlama(config).eval()
@@ -246,19 +203,11 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_paligemma1(self):
     self._test_paligemma_model(
         decoder.Decoder, decoder.get_fake_decoder_config, atol=1e-3, rtol=1e-5
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_paligemma2(self):
     self._test_paligemma_model(
         decoder2.Decoder2,
@@ -267,10 +216,6 @@ class TestModelConversion(googletest.TestCase):
         rtol=1e-5,
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_qwen_vl_model(self):
     config = qwen_vl.get_fake_model_config()
     pytorch_model = qwen_vl.QwenVL(config).eval()
@@ -316,10 +261,7 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
+  @googletest.skipIf(ai_edge_torch.config.in_oss, reason="flaky")
   def test_stable_diffusion_clip(self):
     config = sd_clip.get_fake_model_config()
     prompt_tokens = torch.from_numpy(
@@ -348,10 +290,7 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
+  @googletest.skipIf(ai_edge_torch.config.in_oss, reason="b/413106901")
   def test_stable_diffusion_diffusion(self):
     config = sd_diffusion.get_fake_model_config(2)
     # Reduce stddev(scale) of input values to avoid too big output logits which
@@ -390,10 +329,6 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_stable_diffusion_decoder(self):
     config = sd_decoder.get_fake_model_config()
     # Reduce stddev(scale) of input values to avoid too big output logits which

ai_edge_torch/generative/utilities/converter.py

@@ -81,7 +81,17 @@ def define_conversion_flags(model_name: str):
       'If set, the model will be converted with the provided list of LoRA'
       ' ranks.',
   )
-
+  flags.DEFINE_bool(
+      'mask_as_input',
+      False,
+      'If true, the mask will be passed in as input. Otherwise, mask will be '
+      'built by the model internally.',
+  )
+  flags.DEFINE_bool(
+      'transpose_kv_cache',
+      False,
+      'If true, the model will be converted with transposed KV cache.',
+  )
   return flags
 
 

ai_edge_torch/generative/utilities/export_config.py

@@ -14,8 +14,11 @@
 # ==============================================================================
 
 """Config for customizing model export process."""
+
 import dataclasses
 from typing import List, Optional
+
+from absl import flags
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import torch
 
@@ -38,3 +41,30 @@ class ExportConfig:
   kvcache_cls: type = kv_utils.KVCache
   # The batch size of the decode signature.
   decode_batch_size: int = 1
+
+
+def _build_mask(mask_len, kv_cache_max_len) -> torch.Tensor:
+  if isinstance(mask_len, list):
+    return [_build_mask(i, kv_cache_max_len) for i in mask_len]
+
+  mask = torch.full(
+      (mask_len, kv_cache_max_len), float('-inf'), dtype=torch.float32
+  )
+  mask = torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+  return mask
+
+
+def get_from_flags() -> ExportConfig:
+  """Builds an export config according to the commandline flags."""
+  export_config = ExportConfig()
+
+  if flags.FLAGS.mask_as_input:
+    export_config.prefill_mask = _build_mask(
+        flags.FLAGS.prefill_seq_lens, flags.FLAGS.kv_cache_max_len
+    )
+    export_config.decode_mask = _build_mask(1, flags.FLAGS.kv_cache_max_len)
+
+  if flags.FLAGS.transpose_kv_cache:
+    export_config.kvcache_layout = kv_utils.KV_LAYOUT_TRANSPOSED
+
+  return export_config

ai_edge_torch/model.py

@@ -22,6 +22,7 @@ from __future__ import annotations
 
 import abc
 import re
+import os
 from typing import Callable
 
 import numpy.typing as npt
@@ -154,6 +155,7 @@ class TfLiteModel(Model):
     Args:
       path: The path to file to which the model is serialized.
     """
+    os.makedirs(os.path.dirname(path), exist_ok=True)
     with open(path, 'wb') as file_handle:
       file_handle.write(self._tflite_model)
 

ai_edge_torch/odml_torch/lowerings/_decomp_registry.py

@@ -34,6 +34,8 @@ fx_infra.decomp.update_pre_lower_decomp(
         torch.ops.aten.replication_pad1d,
         torch.ops.aten.replication_pad2d,
         torch.ops.aten.replication_pad3d,
+        torch.ops.aten.upsample_bilinear2d.vec,
+        torch.ops.aten.upsample_nearest2d.vec,
         torch.ops.aten.addmm,
     ])
 )

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.5.0.dev20250424.dist-info → ai_edge_torch_nightly-0.5.0.dev20250426.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250424
+Version: 0.5.0.dev20250426
 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