ai-edge-torch-nightly 0.3.0.dev20250124__py3-none-any.whl → 0.3.0.dev20250125__py3-none-any.whl

ai_edge_torch/generative/examples/deepseek/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 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.
+# ==============================================================================

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -0,0 +1,80 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+"""Example of converting DeepSeek R1 distilled models to tflite model."""
+
+import os
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.deepseek import deepseek
+from ai_edge_torch.generative.utilities import converter
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/deepseek'),
+    'The path to the model checkpoint, or directory holding the checkpoint.',
+)
+_OUTPUT_PATH = flags.DEFINE_string(
+    'output_path',
+    '/tmp/',
+    'The path to export the tflite model.',
+)
+_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
+    'output_name_prefix',
+    'deepseek',
+    'The prefix of the output tflite model name.',
+)
+_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
+    'prefill_seq_lens',
+    (8, 64, 128, 256, 512, 1024),
+    'List of the maximum sizes of prefill input tensors.',
+)
+_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
+    'kv_cache_max_len',
+    1280,
+    'The maximum size of KV cache buffer, including both prefill and decode.',
+)
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    True,
+    'Whether the model should be quantized.',
+)
+_LORA_RANKS = flags.DEFINE_multi_integer(
+    'lora_ranks',
+    None,
+    'If set, the model will be converted with the provided list of LoRA ranks.',
+)
+
+
+def main(_):
+  pytorch_model = deepseek.build_model(
+      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+  )
+  converter.convert_to_tflite(
+      pytorch_model,
+      output_path=_OUTPUT_PATH.value,
+      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
+      prefill_seq_len=_PREFILL_SEQ_LENS.value,
+      quantize=_QUANTIZE.value,
+      lora_ranks=_LORA_RANKS.value,
+      export_config=ExportConfig(),
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

ai_edge_torch/generative/examples/deepseek/deepseek.py

@@ -0,0 +1,92 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+"""Example of building DeepSeek R1 distilled models."""
+
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import model_builder
+from torch import nn
+
+TENSOR_NAMES = model_builder.TENSOR_NAMES_WITH_SEPARATE_LM_HEAD
+
+
+class DeepSeekDistillQwen(model_builder.DecoderOnlyModel):
+  """A DeepSeek distilled model based on Qwen."""
+  pass
+
+
+def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for a Qwen 2.5 3B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for a SmolLM model.
+  """
+  attn_config = cfg.AttentionConfig(
+      num_heads=12,
+      head_dim=128,
+      num_query_groups=2,
+      rotary_base=10000,
+      rotary_percentage=1.0,
+      qkv_use_bias=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=8960,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-06,
+  )
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=151936,
+      num_layers=28,
+      max_seq_len=4096,
+      embedding_dim=1536,
+      kv_cache_max_len=kv_cache_max_len,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      lm_head_share_weight_with_embedding=False,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
+  config = get_model_config(**kwargs)
+  config.vocab_size = 128
+  config.num_layers = 2
+  # DeepSeek-R1-Distill-Qwen has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 64
+  return config
+
+
+def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_model_config(**kwargs),
+      tensor_names=TENSOR_NAMES,
+      model_class=DeepSeekDistillQwen,
+  )

ai_edge_torch/generative/examples/deepseek/verify.py

@@ -0,0 +1,70 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+"""Verifies the reauthored DeepSeek R1 distilled 1.5B model."""
+
+import logging
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.deepseek import deepseek
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+def main(_):
+  checkpoint = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = deepseek.build_model(reauthored_checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -85,6 +85,7 @@ def convert_stable_diffusion_to_tflite(
       clip.TENSOR_NAMES,
   )
   loader.load(clip_model, strict=False)
+  clip_model.eval()
 
   diffusion_model = diffusion.Diffusion(
       diffusion.get_model_config(batch_size=2, device_type=_DEVICE_TYPE.value)
@@ -93,6 +94,7 @@ def convert_stable_diffusion_to_tflite(
       diffusion_ckpt_path, diffusion.TENSOR_NAMES
   )
   diffusion_loader.load(diffusion_model, strict=False)
+  diffusion_model.eval()
 
   decoder_model = decoder.Decoder(
       decoder.get_model_config(device_type=_DEVICE_TYPE.value)
@@ -101,6 +103,7 @@ def convert_stable_diffusion_to_tflite(
       decoder_ckpt_path, decoder.TENSOR_NAMES
   )
   decoder_loader.load(decoder_model, strict=False)
+  decoder_model.eval()
 
   # TODO(yichunk): enable image encoder conversion
   # if encoder_ckpt_path is not None:

ai_edge_torch/generative/layers/experimental/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 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.
+#

ai_edge_torch/generative/layers/experimental/attention.py

@@ -0,0 +1,269 @@
+# 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 building blocks for a GPU-specific Attention layer.
+
+This is a temporary implemenation for the GPU. It is subject to change/removal
+at any time.
+"""
+
+from typing import Optional, Tuple, Union
+
+from ai_edge_torch.generative.layers import builder
+from ai_edge_torch.generative.layers import lora as lora_utils
+from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
+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 ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+import torch
+from torch import nn
+
+
+class TransformerBlock(nn.Module):
+
+  def __init__(
+      self,
+      config: cfg.TransformerBlockConfig,
+      model_config: cfg.ModelConfig,
+  ) -> None:
+    """Initialize an instance of the TransformerBlock.
+
+    Args:
+      config (cfg.TransformerBlockConfig): the configuration object for this
+        transformer block.
+      model_config (cfg.ModelConfig): the configuration object for the model
+        this transformer block belongs to.
+    """
+    super().__init__()
+    self.pre_atten_norm = builder.build_norm(
+        model_config.embedding_dim,
+        config.pre_attention_norm_config,
+    )
+    self.atten_func = CausalSelfAttention(
+        model_config.batch_size,
+        model_config.embedding_dim,
+        config.attn_config,
+        model_config.enable_hlfb,
+    )
+    self.post_atten_norm = builder.build_norm(
+        model_config.embedding_dim,
+        config.post_attention_norm_config,
+    )
+    self.ff = builder.build_ff(model_config.embedding_dim, config.ff_config)
+    self.config = config
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+      mask: Optional[torch.Tensor] = None,
+      input_pos: Optional[torch.Tensor] = None,
+      kv_cache: kv_utils.KVCacheEntryBase = None,
+      lora: Optional[lora_utils.LoRAEntry] = None,
+  ) -> Union[torch.Tensor, Tuple[torch.Tensor, kv_utils.KVCacheEntryBase]]:
+    """Forward function of the TransformerBlock.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+      kv_cache (KVCacheEntryBase): the optional kv cache entry.
+      lora (LoRAEntry): the optional lora entry.
+
+    Returns:
+      output activation from this transformer block, and updated kv cache (if
+      passed in).
+    """
+    kv = None
+    if self.config.parallel_residual:
+      x_norm = self.pre_atten_norm(x)
+      atten_func_out = self.atten_func(
+          x_norm, rope, mask, input_pos, kv_cache, lora
+      )
+      if kv_cache is None:
+        attn_out = atten_func_out
+      else:
+        attn_out, kv = atten_func_out
+      ff_out = self.ff(x_norm)
+      output = x + attn_out + ff_out
+    else:
+      x_norm = self.pre_atten_norm(x)
+      atten_func_out = self.atten_func(
+          x_norm, rope, mask, input_pos, kv_cache, lora
+      )
+      if kv_cache is None:
+        attn_out = atten_func_out
+      else:
+        attn_out, kv = atten_func_out
+      x = x + attn_out
+      x_norm = self.post_atten_norm(x)
+      output = x + self.ff(x_norm)
+
+    return output if kv is None else (output, kv)
+
+
+class CausalSelfAttention(nn.Module):
+
+  def __init__(
+      self,
+      batch_size: int,
+      dim: int,
+      config: cfg.AttentionConfig,
+      enable_hlfb: bool,
+  ) -> None:
+    """Initialize an instance of CausalSelfAttention.
+
+    Args:
+      batch_size (int): batch size of the input tensor.
+      dim (int): causal attention's input/output dimmension.
+      config (cfg.AttentionConfig): attention specific configurations.
+      enable_hlfb (bool): whether hlfb is enabled or not.
+    """
+    super().__init__()
+    self.kv_cache = None
+    self.batch_size = batch_size
+    qkv_shape = (
+        config.num_heads + 2 * config.num_query_groups
+    ) * config.head_dim
+    output_shape = config.num_heads * config.head_dim
+    # Key, query, value projections for all heads.
+    self.qkv_projection = nn.Linear(dim, qkv_shape, bias=config.qkv_use_bias)
+    self.output_projection = nn.Linear(
+        output_shape, dim, bias=config.output_proj_use_bias
+    )
+    self.query_norm = builder.build_norm(
+        config.head_dim, config.query_norm_config
+    )
+    self.key_norm = builder.build_norm(config.head_dim, config.key_norm_config)
+    self.config = config
+    self.enable_hlfb = enable_hlfb
+    self.sdpa_func = sdpa.scaled_dot_product_attention
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+      mask: Optional[torch.Tensor] = None,
+      input_pos: Optional[torch.Tensor] = None,
+      kv_cache: Optional[kv_utils.KVCacheEntryBase] = None,
+      lora: Optional[lora_utils.LoRAEntry] = None,
+  ) -> Union[torch.Tensor, Tuple[torch.Tensor, kv_utils.KVCacheEntryBase]]:
+    """Forward function of the CausalSelfAttention layer, which can support
+
+       MQA, GQA and MHA.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+      kv_cache (KVCacheEntryBase): the KV cache entry corresponding to this
+        module.
+      lora (LoRAEntry): the optional lora entry.
+
+    Returns:
+      output activation from this self attention layer, and the updated
+        KV Cach Entry (if passed in).
+    """
+    # Batch size, sequence length, embedding dimensionality.
+    B, T, E = x.size()
+    assert B == self.batch_size, (
+        "batch size of input tensor must match with the batch size specified in"
+        " the model configuration."
+    )
+
+    qkv = self.qkv_projection(x)
+
+    # 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
+    # Each group has >=1 queries, 1 key, and 1 value.
+    if self.config.qkv_transpose_before_split:
+      qkv = qkv.view(B, T, -1, self.config.head_dim)
+      q, k, v = qkv.split(
+          (
+              q_per_kv * self.config.num_query_groups,
+              self.config.num_query_groups,
+              self.config.num_query_groups,
+          ),
+          dim=-2,
+      )
+    else:
+      qkv = qkv.view(B, T, self.config.num_query_groups, -1)
+      q, k, v = qkv.split(
+          (
+              q_per_kv * self.config.head_dim,
+              self.config.head_dim,
+              self.config.head_dim,
+          ),
+          dim=-1,
+      )
+
+    if lora is not None:
+      q += lora_utils.apply_lora(x, lora.attention.query, shape=q.shape)
+      k += lora_utils.apply_lora(x, lora.attention.key, shape=k.shape)
+      v += lora_utils.apply_lora(x, lora.attention.value, shape=v.shape)
+
+    q = self.query_norm(q)
+    k = self.key_norm(k)
+
+    q = q.reshape(B, T, -1, self.config.head_dim)
+    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:
+      # Compute rotary positional embedding for query and key.
+      n_elem = int(self.config.rotary_percentage * self.config.head_dim)
+      cos, sin = rope
+      q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
+
+    # Transpose k/v to specific layout for GPU implementation.
+    b, _, n, h = q.shape
+    g = n // self.config.num_query_groups
+    # btnh -> bnth -> b(kg)th -> 1(bk)(gt)h
+    q = q.permute(0, 2, 1, 3).reshape(
+        1, b * self.config.num_query_groups, g * T, h
+    )
+
+    k = k.permute(0, 2, 1, 3).reshape(
+        1, -1, T, self.config.head_dim
+    )  # 1, bk, s, h
+    v = v.permute(0, 2, 3, 1).reshape(
+        1, -1, self.config.head_dim, T
+    )  # 1, bk, h, s
+
+    if kv_cache is not None:
+      kv_cache = kv_utils.update(kv_cache, input_pos, k, v)
+      k, v = kv_cache.k_cache, kv_cache.v_cache
+
+    sdpa_out = self.sdpa_func(
+        kv_cache,
+        q,
+        k,
+        v,
+        self.config.head_dim,
+        mask=mask,
+        softcap=self.config.logit_softcap,
+    )  # 1, bk, gt, h
+    sdpa_out = (
+        sdpa_out.reshape(B, -1, T, h).permute(0, 2, 1, 3).reshape(B, T, -1)
+    )
+
+    # Compute the output projection.
+    y = self.output_projection(sdpa_out)
+    if lora is not None:
+      y += lora_utils.apply_lora(sdpa_out, lora.attention.output)
+
+    return y if kv_cache is None else (y, kv_cache)

ai_edge_torch/generative/layers/experimental/kv_cache.py

@@ -0,0 +1,314 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+"""Utility functions for KV Cache.
+
+This is an experimental implementation and is subject to change at any time.
+"""
+
+import dataclasses
+from typing import List, Tuple
+
+from ai_edge_torch import hlfb
+from ai_edge_torch.generative.layers import model_config
+from ai_edge_torch.generative.layers.experimental import types as types
+from ai_edge_torch.generative.utilities.dynamic_update_slice import dynamic_update_slice
+import torch
+import torch.nn as nn
+import torch.utils._pytree as pytree
+
+BATCH_SIZE = 1
+
+
+@dataclasses.dataclass
+class KVCacheEntryBase:
+  """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).
+  """
+
+  k_cache: torch.Tensor
+  v_cache: torch.Tensor
+
+  @classmethod
+  def _from_model_config(
+      cls,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      k_shape: Tuple,
+      v_shape: Tuple,
+      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)
+    obj = cls(k_cache=k, v_cache=v)
+    return obj
+
+  @classmethod
+  def from_model_config(
+      cls,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "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(
+        kv_cache_max, config, shape, shape, dtype, device
+    )
+
+
+@dataclasses.dataclass
+class KVCacheEntryBTNH(KVCacheEntryBase):
+  k_type = types.BTNH()
+  v_type = types.BTNH()
+
+
+@dataclasses.dataclass
+class KVCacheEntryTransposed(KVCacheEntryBase):
+
+  k_type = types.BNTH()
+  v_type = types.BNHT()
+
+  @classmethod
+  def from_model_config(
+      cls,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheEntryBase":
+    """Build an instance of the class based on model config."""
+    num_kv_heads = config.num_query_groups
+    k_shape = (
+        1,
+        BATCH_SIZE * num_kv_heads,
+        kv_cache_max,
+        config.head_dim,
+    )  # 1, bk, s, h
+    v_shape = (
+        1,
+        BATCH_SIZE * num_kv_heads,
+        config.head_dim,
+        kv_cache_max,
+    )  # 1, bk, h, s
+    return cls._from_model_config(
+        kv_cache_max, config, k_shape, v_shape, dtype, device
+    )
+
+
+@dataclasses.dataclass
+class KVCacheBase:
+  """A utility class for holding KV cache entries per layer."""
+
+  caches: Tuple[KVCacheEntryBase, ...]
+
+  @classmethod
+  def _from_model_config(
+      cls,
+      kv_entry_cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBase":
+    caches = [
+        kv_entry_cls.from_model_config(
+            config.kv_cache_max,
+            config.block_config(idx).attn_config,
+            dtype,
+            device,
+        )
+        for idx in range(config.num_layers)
+    ]
+    obj = cls(caches=tuple(caches))
+    return obj
+
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBase":
+    """Build an instance of the class based on model config.
+
+    Args:
+        config (ModelConfig): Model config used for building the cache.
+        dtype (torch.dtype, optional): The data type of the cache tensor.
+          Defaults to torch.float32.
+        device (torch.device, optional): The device placement of the cache
+          tensors. Defaults to None.
+
+    Returns:
+        KVCacheBase: The created cache object.
+    """
+    return cls._from_model_config(
+        KVCacheEntryBase, config=config, dtype=dtype, device=device
+    )
+
+  def flatten(self) -> List[torch.Tensor]:
+    """Flatten the cache entries into a list of tensors with order k_i, v_i."""
+    flattened, _ = _flatten_kvc(self)
+    return flattened
+
+
+@dataclasses.dataclass
+class KVCacheBTNH(KVCacheBase):
+
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBTNH":
+    return cls._from_model_config(
+        KVCacheEntryBTNH, config=config, dtype=dtype, device=device
+    )
+
+
+@dataclasses.dataclass
+class KVCacheTransposed(KVCacheBase):
+
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheBTNH":
+    return cls._from_model_config(
+        KVCacheEntryTransposed, config=config, dtype=dtype, device=device
+    )
+
+
+def _flatten_kvc(kvc: KVCacheBase) -> Tuple[List[str], List[str]]:
+  flattened = []
+  flat_names = []
+  none_names = []
+  for i, kv_entry in enumerate(kvc.caches):
+    flattened.append(kv_entry.k_cache)
+    flat_names.append(f"k_{i}")
+    flattened.append(kv_entry.v_cache)
+    flat_names.append(f"v_{i}")
+  return flattened, [flat_names, none_names]
+
+
+def _flatten_kvc_with_keys(kvc: KVCacheBase) -> Tuple[List, List]:
+  flattened, (flat_names, none_names) = _flatten_kvc(kvc)
+  return [
+      (pytree.MappingKey(k), v) for k, v in zip(flat_names, flattened)
+  ], flat_names
+
+
+def _unflatten_kvc(
+    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
+  flat_names = context[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(
+        KVCacheEntryBase(
+            k_cache=values[k_cache_idx], v_cache=values[v_cache_idx]
+        )
+    )
+  obj = KVCacheBase(tuple(kv_entries))
+  return obj
+
+
+pytree.register_pytree_node(
+    KVCacheTransposed,
+    _flatten_kvc,
+    _unflatten_kvc,
+    flatten_with_keys_fn=_flatten_kvc_with_keys,
+    serialized_type_name="",
+)
+
+pytree.register_pytree_node(
+    KVCacheBase,
+    _flatten_kvc,
+    _unflatten_kvc,
+    flatten_with_keys_fn=_flatten_kvc_with_keys,
+    serialized_type_name="",
+)
+
+
+def update(
+    cache: KVCacheEntryBase,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+    use_dus: bool = True,
+) -> KVCacheEntryBase:
+  """Out of place update of Cache buffer.
+
+  Args:
+      cache (KVCacheEntryBase): The original cache buffer.
+      input_pos (torch.Tensor): The update slice positions.
+      k_slice (torch.Tensor): The K slice to be updated in the new cache.
+      v_slice (torch.Tensor): The V slice to be updated in the new cache.
+
+  Returns:
+      KVCacheEntryBase: The updated KVCacheBase entry based on the passed
+      inputs.
+  """
+  update_kv_cache = _update_kv_impl
+  return update_kv_cache(cache, input_pos, k_slice, v_slice)
+
+
+def _get_slice_indices(
+    positions: torch.Tensor, cache_dim: int, ts_idx: int
+) -> torch.Tensor:
+  """Returns the slice indices."""
+  positions = positions.float()[0].reshape(
+      1,
+  )
+
+  zeros = torch.zeros((1,), dtype=torch.float32)
+  indices = []
+  for i in range(cache_dim):
+    if i == ts_idx:
+      indices.append(positions)
+    else:
+      indices.append(zeros)
+  slice_indices = torch.cat(indices, dim=0)
+  slice_indices = slice_indices.int()
+  return slice_indices
+
+
+def _update_kv_impl(
+    cache: KVCacheEntryTransposed,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntryTransposed:
+  """Update the cache buffer with High Level Function Boundary annotation."""
+  cache_dim = 4
+  k_ts_idx = 2
+  v_ts_idx = 3
+  positions = input_pos.clone()
+  k_slice_indices = _get_slice_indices(positions, cache_dim, k_ts_idx)
+  v_slice_indices = _get_slice_indices(positions, cache_dim, v_ts_idx)
+  k = dynamic_update_slice(cache.k_cache, k_slice, [x for x in k_slice_indices])
+  v = dynamic_update_slice(cache.v_cache, v_slice, [x for x in v_slice_indices])
+  return KVCacheEntryTransposed(k, v)

ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py

@@ -0,0 +1,97 @@
+# Copyright 2025 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.  This is experimental and
+# GPU-specific code.
+
+import math
+from typing import Optional
+
+from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
+from ai_edge_torch.generative.layers.experimental import types
+from ai_edge_torch.generative.utilities import bmm_4d as bmm_lib
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+from multipledispatch import dispatch
+import torch
+import torch.nn.functional as F
+
+
+def scaled_dot_product_attention(
+    kv: kv_utils.KVCacheBase,
+    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,
+):
+  if hasattr(kv, "k_type") and hasattr(kv, "v_type"):
+    return _sdpa(
+        kv.k_type,
+        kv.v_type,
+        query=query,
+        key=key,
+        value=value,
+        head_size=head_size,
+        mask=mask,
+        scale=scale,
+        softcap=softcap,
+    )
+  raise ValueError(
+      f"SDPA for K type {type(kv.caches[0].k_type)} and V type"
+      f" {type(kv.caches[0].v_type)} not supported."
+  )
+
+
+@dispatch(types.BNTH, types.BNHT)
+def _sdpa(k_type, v_type, *args, **kwargs):
+  query = kwargs["query"]
+  key = kwargs["key"]
+  value = kwargs["value"]
+  head_size = kwargs["head_size"]
+  mask = kwargs.get("mask", None)
+  scale = kwargs.get("scale", None)
+  softcap = kwargs.get("softcap", None)
+
+  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
+
+
+@dispatch(object, object)
+def _sdpa(k_type, v_type, *args, **kwargs):
+
+  raise ValueError(f"No implementations for k={k_type} and v={v_type}")

ai_edge_torch/generative/layers/experimental/types.py

@@ -0,0 +1,97 @@
+# Copyright 2025 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 listing of types describes the K and V tensors in KV caches.
+
+import enum
+from enum import Enum, auto
+from typing import Tuple
+from torch import nn
+
+
+@enum.unique
+class TensorDims(Enum):
+  BATCH = enum.auto()
+  SEQUENCE = enum.auto()
+  NUM_HEADS = enum.auto()
+  HEAD_DIM = enum.auto()
+  MODEL_DIM = enum.auto()  # often num_heads * head_dim
+
+
+DIM_TO_LETTER = {
+    TensorDims.BATCH: 'B',
+    TensorDims.SEQUENCE: 'T',
+    TensorDims.NUM_HEADS: 'N',
+    TensorDims.HEAD_DIM: 'H',
+    TensorDims.MODEL_DIM: 'D',
+}
+
+
+class TensorDimensionMeta(type):
+  """Metaclass to create classes representing an order of tensor dimensions."""
+
+  def __new__(cls, name, bases, attrs, dimensions: Tuple[TensorDims]):
+    """Creates a new class with the given name and tensor dimension order.
+
+    Args:
+      name: Name of the new class.
+      bases: Base classes for the new class.
+      attrs: Attributes for the new class.
+      dimensions: A tuple of TensorDims defining the order.
+    """
+
+    attrs['dimensions'] = (
+        dimensions  # Store the dimensions as a class attribute
+    )
+    return super().__new__(cls, name, bases, attrs)
+
+  def __init__(cls, name, bases, attrs, dimensions: Tuple[TensorDims]):
+    super().__init__(name, bases, attrs)
+
+  def __repr__(cls):
+    return f'{cls.__name__}'
+
+
+def create_tensor_dimension_order_class(dims: Tuple[TensorDims]):
+  """Creates a TensorDimensionMeta class with the specified dimensions.
+
+  Args:
+    dimensions: A tuple of TensorDims.
+
+  Returns:
+    A new class representing the tensor dimension order.
+  """
+  name = ''.join(DIM_TO_LETTER[d] for d in dims)
+  # Derive from nn.Module for torch tracing compatiblity.
+  return TensorDimensionMeta(name, (nn.Module,), {}, dimensions=dims)
+
+
+BTNH = create_tensor_dimension_order_class((
+    TensorDims.BATCH,
+    TensorDims.SEQUENCE,
+    TensorDims.NUM_HEADS,
+    TensorDims.HEAD_DIM,
+))
+BNTH = create_tensor_dimension_order_class((
+    TensorDims.BATCH,
+    TensorDims.NUM_HEADS,
+    TensorDims.SEQUENCE,
+    TensorDims.HEAD_DIM,
+))
+BNHT = create_tensor_dimension_order_class((
+    TensorDims.BATCH,
+    TensorDims.NUM_HEADS,
+    TensorDims.HEAD_DIM,
+    TensorDims.SEQUENCE,
+))

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -17,6 +17,7 @@
 
 import ai_edge_torch
 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.llama import llama
@@ -150,16 +151,15 @@ class TestModelConversion(googletest.TestCase):
       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()
@@ -174,6 +174,15 @@ class TestModelConversion(googletest.TestCase):
     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",

ai_edge_torch/generative/utilities/converter.py

@@ -19,7 +19,6 @@ import os
 from typing import Optional, Union
 from ai_edge_torch._convert import converter as converter_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
-import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.quantize import quant_recipes
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
@@ -151,9 +150,21 @@ def _export_helper(
       else None
   )
 
+  if export_config.prefill_mask is None:
+    prefill_masks = None
+  elif isinstance(export_config.prefill_mask, torch.Tensor):
+    prefill_masks = [export_config.prefill_mask]
+  elif isinstance(export_config.prefill_mask, list):
+    prefill_masks = export_config.prefill_mask
+  else:
+    raise ValueError('Prefill masks unrecognized.')
+
+  if prefill_masks:
+    assert len(prefill_masks) == len(prefill_seq_lens)
+
   decode_token = torch.tensor([[0]], dtype=torch.int)
   decode_input_pos = torch.tensor([0], dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(config)
+  kv = export_config.kvcache_cls.from_model_config(config)
 
   quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
 
@@ -174,8 +185,8 @@ def _export_helper(
           'input_pos': prefill_input_pos,
           'kv_cache': kv,
       }
-      if export_config.prefill_mask is not None:
-        sample_kwargs['mask'] = export_config.prefill_mask
+      if prefill_masks is not None:
+        sample_kwargs['mask'] = prefill_masks[i]
 
       if lora is not None:
         prefill_signature_name += f'_lora_r{lora.get_rank()}'

ai_edge_torch/generative/utilities/model_builder.py

@@ -17,7 +17,7 @@
 
 import copy
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
@@ -56,8 +56,10 @@ class ExportConfig:
   # When False, only decode signatures will produce output.
   output_logits_on_prefill: bool = False
   # Attention masks given as inputs to the model.
-  prefill_mask: Optional[torch.Tensor] = None
-  decode_mask: Optional[torch.Tensor] = None
+  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.
+  kvcache_cls: type = kv_utils.KVCache
 
 
 class DecoderOnlyModel(nn.Module):

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20250124"
+__version__ = "0.3.0.dev20250125"

{ai_edge_torch_nightly-0.3.0.dev20250124.dist-info → ai_edge_torch_nightly-0.3.0.dev20250125.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250124
+Version: 0.3.0.dev20250125
 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.3.0.dev20250124.dist-info → ai_edge_torch_nightly-0.3.0.dev20250125.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=xloGd_dX0MD8k-quT07WLlEN1zIGVtCKu6xBSvjofrc,706
+ai_edge_torch/version.py,sha256=yuz53SwRvngiQ41D-VX7MPmVGe-Vi-UR3v12E-o3P4I,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -49,6 +49,10 @@ ai_edge_torch/generative/examples/amd_llama_135m/__init__.py,sha256=hHLluseD2R0H
 ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py,sha256=urNif89PyCXbdXT5spOeDvdM5luJ-a5HaXHM86v4JnU,2766
 ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py,sha256=Oqlg5ZoUuG2aU3067QaPpmEXWOdB8GEq7u_NWoBpoB4,2337
 ai_edge_torch/generative/examples/amd_llama_135m/verify.py,sha256=-9Nb9D818YSJR3olVtBwoLNeMMD5qE58YBnsA67hlHg,2421
+ai_edge_torch/generative/examples/deepseek/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
+ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py,sha256=I5eA-XfFdHjYwDsLIjn23T2e-IgnSCQ129-5DOU8j44,2532
+ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=AOAJ7ltXwY5IbmcCP2nVHW9FmRwexzfNxnoDlR-sW9c,2885
+ai_edge_torch/generative/examples/deepseek/verify.py,sha256=sDYBhmE_CeZw5iLIQ7rJNGLjhcTyKUQGdg7_QQBh9WM,2398
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=8HJi0cutxPstafVNs2LfBKdUzufVucje1Vrfjw_RS_g,2527
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=MX8fZhJJPZ5IoMiNHX0tLkRpHYqVuh4qhW0rkeIfmYw,2529
@@ -97,7 +101,7 @@ ai_edge_torch/generative/examples/smollm/verify.py,sha256=KpYxVz_lv61YWy6HLfwT68
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=5M4auM33SgCTODt0VT8TO-EVILruqGDRiNILBPeB83Y,6072
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=sB_7-PVri8PxKnFG7c8GsTGyrxGEda-oZwGyyScTL3o,5239
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=GtwKAByEk0ENGEWbUmC2mAAPkbLZ3M5xH1HIToyu8QE,5307
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=sQKQ-k6H9kG2brgwLsktjCMeN2h0POyfMP6iNsPNKWc,16271
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=6W58LxmHHkz2ctgpknQkyoDANZAnE9Byp_svfqLpQf0,34793
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=CAPsW84A8f00nS6fLFeh_XUjCPsDCA5UxHOUsMrLfSU,3450
@@ -134,6 +138,11 @@ ai_edge_torch/generative/layers/model_config.py,sha256=ZVRWEGw1BnLbLCuoR71kWGqQt
 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/experimental/__init__.py,sha256=nz-K0h8DfiATHzR6s1_bCw2akUmHWffU1bDRSkIzSqI,592
+ai_edge_torch/generative/layers/experimental/attention.py,sha256=KC1UkIhaPx2DNRfkxCXO7eZZMeNm2UxkjFi-fB8HVhw,9212
+ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=gE_q8YoSzOhGgbSm0K91jXkbFKnFJpuYf-hxMzLNw78,8976
+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
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=9jKzOfiBQ66bp1ZnVIAoREIifVNFx4aTlQeYMAx2_pA,29062
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
@@ -151,15 +160,15 @@ ai_edge_torch/generative/test/test_kv_cache.py,sha256=2AulHBS3hC4b_68PNNBkRVOryp
 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
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=yzMgXkiZxHUF_xz0UR3kD3x74ELsmJetbQnmv7-9gyQ,12473
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=AJs_ARfWUqwuFRwYtQQOLd87CiD4mUDwAhq885cqc4Q,12875
 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/bmm_4d.py,sha256=2BMOYiFVUsl-bjxmLkrX4N7kpO0CnhB7eDYxm_iBCr8,2533
-ai_edge_torch/generative/utilities/converter.py,sha256=QIYxT-zATMzsD3LG-keRkxpJqDKXkbil4Se1KXthWFg,7726
+ai_edge_torch/generative/utilities/converter.py,sha256=6siSpCvH_cLV-eP40lkF_AqjBpYv68xeMRQ722fKgE0,8065
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
-ai_edge_torch/generative/utilities/model_builder.py,sha256=aXigoFEMLAKk7HQuWJM5ILs3igA4z2VH64ZCzCuBhDE,6671
+ai_edge_torch/generative/utilities/model_builder.py,sha256=5WqcxpeTdt51nVoUwt9g5kKB5wQKj2eYbiaz7k6Ofxg,6815
 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
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
@@ -213,8 +222,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.3.0.dev20250124.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/METADATA,sha256=WKVqBXJtXvMv3JfqtYKcl1GFgKrKtSbZ-tJAol5PPHk,1966
-ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250124.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250125.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250125.dist-info/METADATA,sha256=BkUH2iAinJYGmBLTMdeYSpihXAHY_mBOkeprZLPaDGk,1966
+ai_edge_torch_nightly-0.3.0.dev20250125.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250125.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250125.dist-info/RECORD,,