ai-edge-torch-nightly 0.7.0.dev20251012__py3-none-any.whl → 0.7.0.dev20251014__py3-none-any.whl

ai_edge_torch/generative/layers/attention.py

@@ -18,6 +18,7 @@
 import abc
 from typing import Optional, Tuple, Union
 
+from ai_edge_torch.generative.layers import attention_utils
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
@@ -240,13 +241,35 @@ class CausalSelfAttention(CausalSelfAttentionBase):
     k = k.reshape(B, T, -1, self.config.head_dim)
     v = v.reshape(B, T, -1, self.config.head_dim)
 
-    if rope is not None:
+    alibi_bias = None
+    if self.config.use_alibi:
+      k_size = T
+      if mask is not None:
+        k_size = mask.shape[-1]
+      elif input_pos is not None:
+        # If mask is not present, assume current sequence length is key length.
+        k_size = input_pos[-1].item() + 1
+      alibi_bias = attention_utils.build_alibi_bias(
+          n_heads=self.config.num_heads,
+          k_size=k_size,
+          dtype=x.dtype,
+          device=x.device,
+      )
+    elif rope is not None:
       # Compute rotary positional embedding for query and key.
       cos, sin = rope
       q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
 
     sdpa_out, kv_cache = sdpa_with_kv_update.sdpa_with_kv_update(
-        q, k, v, kv_cache, input_pos, mask, self.config, self.enable_hlfb
+        q,
+        k,
+        v,
+        kv_cache,
+        input_pos,
+        mask,
+        self.config,
+        self.enable_hlfb,
+        alibi_bias=alibi_bias,
     )
 
     # Compute the output projection.

ai_edge_torch/generative/layers/attention_test.py

@@ -27,16 +27,27 @@ class AttentionTest(parameterized.TestCase):
       dict(
           testcase_name="local_causal_self_attention",
           attn_type=cfg.AttentionType.LOCAL_SLIDING,
+          use_alibi=False,
           expected_shape=(1, 10, 16),
       ),
       dict(
           testcase_name="global_causal_self_attention",
           attn_type=cfg.AttentionType.GLOBAL,
+          use_alibi=False,
+          expected_shape=(1, 10, 16),
+      ),
+      dict(
+          testcase_name="alibi_attention",
+          attn_type=cfg.AttentionType.GLOBAL,
+          use_alibi=True,
           expected_shape=(1, 10, 16),
       ),
   )
   def test_causal_self_attention(
-      self, attn_type: cfg.AttentionType, expected_shape: tuple[int, ...]
+      self,
+      attn_type: cfg.AttentionType,
+      use_alibi: bool,
+      expected_shape: tuple[int, ...],
   ):
     norm_config = cfg.NormalizationConfig(
         type=cfg.NormalizationType.RMS_NORM,
@@ -56,6 +67,7 @@ class AttentionTest(parameterized.TestCase):
         logit_softcap=None,
         sliding_window_size=16,
         attn_type=attn_type,
+        use_alibi=use_alibi,
     )
     self_atten = attention.CausalSelfAttention(
         dim=16,

ai_edge_torch/generative/layers/attention_utils.py

@@ -15,11 +15,72 @@
 # Common utility functions used with attention module.
 
 import math
-from typing import Tuple
+from typing import List, Tuple
 
 import torch
 
 
+def _get_alibi_slopes(n_heads: int) -> List[float]:
+  """Returns slopes for ALiBi implementation.
+
+  The slopes are taken from the ALiBi paper
+  [https://arxiv.org/abs/2108.12409].
+  The slopes are later used to calculate the bias which is added to the
+  attention scores.
+
+  Args:
+      n_heads (int): The number of attention heads.
+  """
+
+  def get_slopes_power_of_2(n):
+    start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+    return [start**i for i in range(1, n + 1)]
+
+  if math.log2(n_heads).is_integer():
+    return get_slopes_power_of_2(n_heads)
+  else:
+    closest_power_of_2 = 2 ** math.floor(math.log2(n_heads))
+    return (
+        get_slopes_power_of_2(closest_power_of_2)
+        + _get_alibi_slopes(2 * closest_power_of_2)[0::2][
+            : n_heads - closest_power_of_2
+        ]
+    )
+
+
+def build_alibi_bias(
+    n_heads: int,
+    k_size: int,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device = None,
+) -> torch.Tensor:
+  """Builds ALiBi bias tensor based on key position.
+
+  The bias tensor is added to the attention scores before softmax.
+  Replicates HuggingFace Falcon implementation behavior where bias only depends
+  on key position j, not relative position j-i.
+
+  Args:
+      n_heads (int): The number of attention heads.
+      k_size (int): The key size of the bias tensor.
+      dtype (torch.dtype, optional): Output tensor's data type. Defaults to
+        torch.float32.
+      device (torch.device, optional): Output tensor's data type. Defaults to
+        None in which case "cpu" is used.
+
+  Returns:
+      torch.Tensor: The ALiBi bias tensor of shape (1, n_heads, 1, k_size).
+  """
+  if device is None:
+    device = torch.device('cpu')
+  slopes = torch.tensor(_get_alibi_slopes(n_heads), dtype=dtype, device=device)
+  k_pos = torch.arange(k_size, device=device)
+  # According to HF implementation, bias only depends on key position.
+  # slopes[h] * k_pos[j]
+  alibi_bias = slopes.unsqueeze(-1) * k_pos.unsqueeze(0)  # Shape: H, K
+  return alibi_bias[None, :, None, :].to(dtype)
+
+
 def build_rope_cache(
     size: int,
     dim: int,

ai_edge_torch/generative/layers/attention_utils_test.py

@@ -21,6 +21,26 @@ from absl.testing import absltest as googletest
 
 class AttentionUtilsTest(googletest.TestCase):
 
+  def test_get_alibi_slopes(self):
+    slopes = attention_utils._get_alibi_slopes(1)
+    self.assertSequenceAlmostEqual(slopes, [0.00390625], places=6)
+    slopes = attention_utils._get_alibi_slopes(2)
+    self.assertSequenceAlmostEqual(slopes, [0.0625, 0.00390625], places=6)
+    slopes = attention_utils._get_alibi_slopes(4)
+    self.assertSequenceAlmostEqual(
+        slopes, [0.25, 0.0625, 0.015625, 0.00390625], places=6
+    )
+    slopes = attention_utils._get_alibi_slopes(3)
+    self.assertSequenceAlmostEqual(slopes, [0.0625, 0.00390625, 0.25], places=6)
+
+  def test_build_alibi_bias(self):
+    bias = attention_utils.build_alibi_bias(n_heads=2, k_size=3)
+    self.assertEqual(bias.shape, (1, 2, 1, 3))
+    expected = torch.tensor(
+        [[[[0.0, 0.0625, 0.125]], [[0.0, 0.00390625, 0.0078125]]]]
+    )
+    torch.testing.assert_close(bias, expected)
+
   def test_build_causal_mask_cache(self):
     mask = attention_utils.build_causal_mask_cache(3)
     self.assertEqual(mask.shape, (1, 1, 3, 3))

ai_edge_torch/generative/layers/builder.py

@@ -71,7 +71,7 @@ def build_norm(
   Raises:
     ValueError: If config's `layer_norm_type` is not supported.
   """
-  if config.type == cfg.NormalizationType.NONE:
+  if config is None or config.type == cfg.NormalizationType.NONE:
     return lambda x: x
   elif config.type == cfg.NormalizationType.RMS_NORM:
     return normalization.RMSNorm(
@@ -84,7 +84,9 @@ def build_norm(
         init_fn=init_fn,
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
-    return normalization.LayerNorm(dim, config.epsilon, config.enable_hlfb)
+    return normalization.LayerNorm(
+        dim, config.epsilon, config.use_bias, config.enable_hlfb
+    )
   elif config.type == cfg.NormalizationType.GROUP_NORM:
     return normalization.GroupNorm(
         config.group_num, dim, config.epsilon, config.enable_hlfb

ai_edge_torch/generative/layers/model_config.py

@@ -75,6 +75,8 @@ class NormalizationConfig:
   scale_shift: float = 0.0
   # Number of groups used in group normalization.
   group_num: Optional[float] = None
+  # Whether to use bias in norm.
+  use_bias: bool = True
 
 
 # Exprimental feature and may subject to change.
@@ -108,6 +110,8 @@ class AttentionConfig:
   rotary_base: int = 10_000
   # Percentage of Rotary Positional Embedding added Q and K projections.
   rotary_percentage: Optional[float] = None
+  # Whether to use ALiBi positional encoding.
+  use_alibi: bool = False
   # Whether to transpose the query groups of qkv bundled tensor before
   # splitting into separated tensors.
   qkv_transpose_before_split: bool = False

ai_edge_torch/generative/layers/normalization.py

@@ -148,6 +148,7 @@ class LayerNorm(torch.nn.Module):
       self,
       dim: int,
       eps: float = 1e-5,
+      use_bias: bool = True,
       enable_hlfb: bool = False,
   ):
     """Initialize the LayerNorm layer.
@@ -156,6 +157,7 @@ class LayerNorm(torch.nn.Module):
       dim (int): dimension of the input tensor.
       eps (float): A small float value to ensure numerical stability (default:
         1e-5).
+      use_bias (bool): Whether to use bias in LayerNorm.
       enable_hlfb (bool): Whether to convert this normalization into a single
         op.
     """
@@ -164,7 +166,11 @@ class LayerNorm(torch.nn.Module):
     self.normalized_shape = (dim,)
     self.eps = eps
     self.weight = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
-    self.bias = torch.nn.Parameter(torch.empty(dim), requires_grad=False)
+    self.bias = (
+        torch.nn.Parameter(torch.empty(dim), requires_grad=False)
+        if use_bias
+        else None
+    )
 
   def forward(self, x):
     """Running the forward pass of LayerNorm layer.
@@ -175,7 +181,7 @@ class LayerNorm(torch.nn.Module):
     Returns:
       torch.Tensor: output tensor after applying LayerNorm.
     """
-    if self.enable_hlfb:
+    if self.enable_hlfb and self.bias is not None:
       return layer_norm_with_hlfb(
           x, self.normalized_shape, self.weight, self.bias, self.eps
       )

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -32,6 +32,7 @@ def scaled_dot_product_attention(
     mask: Optional[torch.Tensor] = None,
     scale: Optional[float] = None,
     softcap: Optional[float] = None,
+    alibi_bias: Optional[torch.Tensor] = None,
 ):
   """Scaled dot product attention.
 
@@ -41,14 +42,23 @@ def scaled_dot_product_attention(
     v (torch.Tensor): 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.
+    alibi_bias (torch.Tensor): optional alibi bias tensor.
 
   Returns:
     The output tensor of scaled_dot_product_attention.
   """
-
   if scale is None:
     scale = 1.0 / math.sqrt(head_size)
 
+  if alibi_bias is not None:
+    alibi_bias = alibi_bias * scale
+    if mask is None:
+      mask = alibi_bias
+    else:
+      mask = mask + alibi_bias
+
   q = q.transpose(1, 2)
   k = k.transpose(1, 2)
   v = v.transpose(1, 2)
@@ -72,7 +82,8 @@ def scaled_dot_product_attention(
     scores = scores / softcap
     scores = torch.tanh(scores)
     scores = scores * softcap
-    scores = scores + mask
+    if mask is not None:
+      scores = scores + mask
     out = F.softmax(scores.float(), dim=-1).type_as(q)
     y = torch.matmul(out, v)
 
@@ -87,6 +98,7 @@ def scaled_dot_product_attention_with_hlfb(
     mask: Optional[torch.Tensor] = None,
     scale: Optional[float] = None,
     softcap: Optional[float] = None,
+    alibi_bias: Optional[torch.Tensor] = None,
 ):
   """Scaled dot product attention with high-level function boundary enabled.
 
@@ -96,14 +108,23 @@ def scaled_dot_product_attention_with_hlfb(
     v (torch.Tensor): 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.
+    alibi_bias (torch.Tensor): optional alibi bias tensor.
 
   Returns:
     The output tensor of scaled_dot_product_attention.
   """
-
   if scale is None:
     scale = 1.0 / math.sqrt(head_size)
 
+  if alibi_bias is not None:
+    alibi_bias = alibi_bias * scale
+    if mask is None:
+      mask = alibi_bias
+    else:
+      mask = mask + alibi_bias
+
   attrs = {"scale": scale}
 
   if softcap is not None:
@@ -137,7 +158,8 @@ def scaled_dot_product_attention_with_hlfb(
     scores = scores / softcap
     scores = torch.tanh(scores)
     scores = scores * softcap
-    scores = scores + mask
+    if mask is not None:
+      scores = scores + mask
     out = F.softmax(scores.float(), dim=-1).type_as(q)
     y = torch.matmul(out, v)
 
@@ -154,6 +176,7 @@ def scaled_dot_product_attention_transposed(
     mask: Optional[torch.Tensor] = None,
     scale: Optional[float] = None,
     softcap: Optional[float] = None,
+    alibi_bias: Optional[torch.Tensor] = None,
 ):
   """Scaled dot product attention with transposed key and value.
 
@@ -165,14 +188,21 @@ def scaled_dot_product_attention_transposed(
     mask (torch.Tensor): the optional mask tensor.
     scale (float): the optional scale factor.
     softcap (float): the optional softcap for the logits.
+    alibi_bias (torch.Tensor): optional alibi bias tensor.
 
   Returns:
     The output tensor of scaled_dot_product_attention_transposed.
   """
-
   if scale is None:
     scale = 1.0 / math.sqrt(head_size)
 
+  if alibi_bias is not None:
+    alibi_bias = alibi_bias * scale
+    if mask is None:
+      mask = alibi_bias
+    else:
+      mask = mask + alibi_bias
+
   query = query * scale
 
   assert mask is not None, "Mask should not be None!"

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -15,7 +15,7 @@
 
 """Common utility functions for data loading etc."""
 
-from typing import Tuple
+from typing import Optional, 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
@@ -32,14 +32,15 @@ def sdpa_with_kv_update(
     mask: torch.Tensor,
     config: cfg.AttentionConfig,
     enable_hlfb: bool,
+    alibi_bias: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
   """Wrapper function for scaled dot product attention with KV cache update."""
   if kv is not None and kv.kv_layout == kv_utils.KV_LAYOUT_TRANSPOSED:
     return _sdpa_with_kv_update_transposed(
-        query, key, value, kv, input_pos, mask, config
+        query, key, value, kv, input_pos, mask, config, alibi_bias
     )
   return _sdpa_with_kv_update_default(
-      query, key, value, kv, input_pos, mask, config, enable_hlfb
+      query, key, value, kv, input_pos, mask, config, enable_hlfb, alibi_bias
   )
 
 
@@ -51,6 +52,7 @@ def _sdpa_with_kv_update_transposed(
     input_pos: torch.Tensor,
     mask: torch.Tensor,
     config: cfg.AttentionConfig,
+    alibi_bias: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
   # Transpose k/v to specific layout for GPU implementation.
   b, seq_len, n, h = query.shape
@@ -77,6 +79,7 @@ def _sdpa_with_kv_update_transposed(
       config.head_dim,
       mask=mask,
       softcap=config.logit_softcap,
+      alibi_bias=alibi_bias,
   )  # 1, bk, gt, h
   sdpa_out = (
       sdpa_out.reshape(b, -1, seq_len, h)
@@ -95,6 +98,7 @@ def _sdpa_with_kv_update_default(
     mask: torch.Tensor,
     config: cfg.AttentionConfig,
     enable_hlfb: bool,
+    alibi_bias: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
   b, seq_len, _, _ = query.shape
   if kv is not None:
@@ -112,6 +116,7 @@ def _sdpa_with_kv_update_default(
       config.head_dim,
       mask=mask,
       softcap=config.logit_softcap,
+      alibi_bias=alibi_bias,
   )
   sdpa_out = sdpa_out.reshape(b, seq_len, -1)
   return sdpa_out, kv

ai_edge_torch/generative/utilities/loader.py

@@ -135,7 +135,8 @@ def load_pytorch_statedict(full_path: str):
 
   tensors = {}
   for file in files:
-    this_file_tensors = torch.load(file)
+    map_location = "cpu" if not torch.cuda.is_available() else None
+    this_file_tensors = torch.load(file, map_location=map_location)
     for k in this_file_tensors:
       assert k not in tensors
     tensors.update(this_file_tensors)

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.7.0.dev20251012"
+__version__ = "0.7.0.dev20251014"

{ai_edge_torch_nightly-0.7.0.dev20251012.dist-info → ai_edge_torch_nightly-0.7.0.dev20251014.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.7.0.dev20251012
+Version: 0.7.0.dev20251014
 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.7.0.dev20251012.dist-info → ai_edge_torch_nightly-0.7.0.dev20251014.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=JGvOmXPEMBm7uXPnxRy3Gt3Z4nNCmkWxljpP6yAfgGQ,806
+ai_edge_torch/version.py,sha256=y2xP-8zCqeG6cVDqBVnnBw10Uvwy2VxztZNUfI-7Z-8,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -169,24 +169,24 @@ ai_edge_torch/generative/examples/tiny_llama/verify_util.py,sha256=z6vPBXDWAL6gN
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=PFSMsA1vfBfrV9ssBCkYJNl8Hx_bLdWjN01iyjPM5jE,1094
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=RaXENRRQo1MsLdt3U8h3kYTCmd6imHQ-aCXtmPXCh_o,13911
-ai_edge_torch/generative/layers/attention_test.py,sha256=9v8v96TLyFPdqxEylU1JOAeRFAp2s0YoDHZN83SFJJA,4764
-ai_edge_torch/generative/layers/attention_utils.py,sha256=2qfg7Tzk9ikKph5w3geOHC1I6EyOCdDsWXMr7F7IOZM,7630
-ai_edge_torch/generative/layers/attention_utils_test.py,sha256=22gQ1gcRPkwqFG3_p82GZfRKVE3udEssSy58wNOqv0w,2431
-ai_edge_torch/generative/layers/builder.py,sha256=UiLOyvwd-bc0n5XcbVxi6JCn_qiKSC6zrDKSZT_TSDA,5030
+ai_edge_torch/generative/layers/attention.py,sha256=ZjU3vX-7gOq1KQb3xSZ1NT3xryOTXbYb_vkx_DlcizA,14524
+ai_edge_torch/generative/layers/attention_test.py,sha256=ON9jQRY1r2kFpVq-Qkg6b13Ob95fd4PqHo1hic3RbOQ,5057
+ai_edge_torch/generative/layers/attention_utils.py,sha256=3Ox1XjW_vaqz1-RuVG9RbzRKUqCberFW8P2BQcoNm7A,9659
+ai_edge_torch/generative/layers/attention_utils_test.py,sha256=IHIk39wqaPvxmkZtW27VD3_4xUpyFow_7mScf8OWdqU,3292
+ai_edge_torch/generative/layers/builder.py,sha256=5QL59CbOOW_mk3mlPdcdirGcAxdLee5atbZlnu5Z3ts,5079
 ai_edge_torch/generative/layers/einsum.py,sha256=LH4CNHr-pFfLUuCpwbYL3GpoAMgHJ4nLju3XCqA4VwM,1416
 ai_edge_torch/generative/layers/einsum_test.py,sha256=ltIE773bvvNLv_9aLQxFwe1MgQ762sez0c5E2tejxuA,1079
 ai_edge_torch/generative/layers/feed_forward.py,sha256=_GmtHxwL068l9gh_F_WFcFk7La-Tl5SfoQ9v2hMabZM,5541
 ai_edge_torch/generative/layers/feed_forward_test.py,sha256=Y5l1eC9NgfYixHcfIfE1W4FGh7oC-9UGGyHdKS9tQKc,1880
 ai_edge_torch/generative/layers/kv_cache.py,sha256=A0IFXZ1HD2ZHOWRLfsDO4almgE0KQfjyBOdBFZIGnAs,10893
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
-ai_edge_torch/generative/layers/model_config.py,sha256=HP-vu1UmAiTmdLlTyZGDUF3le0gji8a61mLCy966NZw,10261
-ai_edge_torch/generative/layers/normalization.py,sha256=WAhPcLbcC3SCEa6oIgdsojvN306_S8d90WyMQ7ZVP6I,7269
+ai_edge_torch/generative/layers/model_config.py,sha256=Pvoa766jIf1LWvRwEDMNce43C9NgPOvIpT30VUcnpqA,10390
+ai_edge_torch/generative/layers/normalization.py,sha256=syasVh3dRDVp2Nwhl0x7zucL-chTnCqWgeV1mb87DFY,7435
 ai_edge_torch/generative/layers/normalization_test.py,sha256=zwurZly-TgFxdgVVdpzu9vCpcLbd5RYt_gKg9Lfg1jI,2248
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
-ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=2_AgwENsaOgaxgiSqgoj0V0JzQ09dFtP_nBhX-lJK2g,5648
+ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=1zhOsJpI4CTn78weOs0uRwkRxYu6wGfBvYVFpGFl0qQ,6681
 ai_edge_torch/generative/layers/scaled_dot_product_attention_test.py,sha256=c6JBMQsq9XeMmR1XvGEIidNsoh-YIvichXo2LwVHgr4,3301
-ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=iw7D_46CFe9iRvU0UumbkIoqWQEhDroxm9ABcK-CLlM,3600
+ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=fK_h9M-03ai5dV8ZyQzvB0y84IKlNg9h-4bt9F6bU0g,3833
 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
@@ -211,7 +211,7 @@ ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5l
 ai_edge_torch/generative/utilities/converter.py,sha256=d8pehTq6EzEdVR8ioL2b1ECGTR4G1K1fczc9amu_Oyk,23106
 ai_edge_torch/generative/utilities/export_config.py,sha256=5B15nYyqf96kjjYlHfPctUfsIdsBsh1f8rxKitJpwKQ,2384
 ai_edge_torch/generative/utilities/litertlm_builder.py,sha256=0cNuaqhc7cQcAa4NRalUXyoPQUQC9O3-aHAJEDV1Mps,4265
-ai_edge_torch/generative/utilities/loader.py,sha256=drgKBmNibuc3PCdc0kU0pVcp2Nt1_mjLYh67RyXOn7U,15952
+ai_edge_torch/generative/utilities/loader.py,sha256=QQeEu0cTC7gWnB7RkHonjWLdVGjMbDHd1lfYO_TcJyU,16047
 ai_edge_torch/generative/utilities/model_builder.py,sha256=xBvcTxihB9TN88UtQiXA9sAITQgf-pA77R-VZlLgUeU,6950
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
@@ -270,8 +270,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 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.7.0.dev20251012.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20251012.dist-info/METADATA,sha256=kSPmOUgL8pjpVMiiiRVuBQ-oXp4bDCU5P5N2BtNZ56c,2074
-ai_edge_torch_nightly-0.7.0.dev20251012.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.7.0.dev20251012.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20251012.dist-info/RECORD,,
+ai_edge_torch_nightly-0.7.0.dev20251014.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.7.0.dev20251014.dist-info/METADATA,sha256=xbKvth_nIjCuHT1ELDvsnYl6qB0G0BWm695-i1Wyh_s,2074
+ai_edge_torch_nightly-0.7.0.dev20251014.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.7.0.dev20251014.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.7.0.dev20251014.dist-info/RECORD,,