ai-edge-torch-nightly 0.3.0.dev20240916__py3-none-any.whl → 0.3.0.dev20240919__py3-none-any.whl

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -18,69 +18,49 @@
 import os
 import pathlib
 
-import ai_edge_torch
+from absl import app
+from absl import flags
 from ai_edge_torch.generative.examples.smollm import smollm
-from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.quantize import quant_recipes
-import torch
+from ai_edge_torch.generative.utilities import converter
 
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/smollm'),
+    'The path to the model checkpoint, or directory holding the checkpoint.',
+)
+_TFLITE_PATH = flags.DEFINE_string(
+    'tflite_path',
+    '/tmp/smollm_q8_seq512_ekv1024.tflite',
+    'The tflite file path to export.',
+)
+_PREFILL_SEQ_LEN = flags.DEFINE_integer(
+    'prefill_seq_len',
+    512,
+    'The maximum size of prefill input tensor.',
+)
+_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
+    'kv_cache_max_len',
+    1024,
+    'The maximum size of KV cache buffer, including both prefill and decode.',
+)
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    True,
+    'Whether the model should be quantized.',
+)
 
-def convert_smollm_to_tflite(
-    checkpoint_path: str,
-    prefill_seq_len: int = 512,
-    kv_cache_max_len: int = 1024,
-    quantize: bool = True,
-):
-  """Converts SmolLM model to multi-signature tflite model.
 
-  Args:
-      checkpoint_path (str): The filepath to the model checkpoint, or directory
-        holding the checkpoint.
-      prefill_seq_len (int, optional): The maximum size of prefill input tensor.
-        Defaults to 512.
-      kv_cache_max_len (int, optional): The maximum size of KV cache buffer,
-        including both prefill and decode. Defaults to 1024.
-      quantize (bool, optional): Whether the model should be quanized. Defaults
-        to True.
-  """
+def main(_):
   pytorch_model = smollm.build_model(
-      checkpoint_path, kv_cache_max_len=kv_cache_max_len
+      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
-  # Tensors used to trace the model graph during conversion.
-  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.int)
-  prefill_input_pos = torch.arange(0, prefill_seq_len, dtype=torch.int)
-  decode_token = torch.tensor([[0]], dtype=torch.int)
-  decode_input_pos = torch.tensor([0], dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(pytorch_model.config)
-
-  quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
-  edge_model = (
-      ai_edge_torch.signature(
-          'prefill',
-          pytorch_model,
-          sample_kwargs={
-              'tokens': prefill_tokens,
-              'input_pos': prefill_input_pos,
-              'kv_cache': kv,
-          },
-      )
-      .signature(
-          'decode',
-          pytorch_model,
-          sample_kwargs={
-              'tokens': decode_token,
-              'input_pos': decode_input_pos,
-              'kv_cache': kv,
-          },
-      )
-      .convert(quant_config=quant_config)
-  )
-  quant_suffix = 'q8' if quantize else 'f32'
-  edge_model.export(
-      f'/tmp/smollm_{quant_suffix}_seq{prefill_seq_len}_ekv{kv_cache_max_len}.tflite'
+  converter.convert_to_tflite(
+      pytorch_model,
+      tflite_path=_TFLITE_PATH.value,
+      prefill_seq_len=_PREFILL_SEQ_LEN.value,
+      quantize=_QUANTIZE.value,
   )
 
 
 if __name__ == '__main__':
-  path = os.path.join(pathlib.Path.home(), 'Downloads/llm_data/smollm')
-  convert_smollm_to_tflite(path)
+  app.run(main)

ai_edge_torch/generative/examples/smollm/smollm.py

@@ -16,15 +16,10 @@
 """Example of building a SmolLM model."""
 
 import copy
-import os
-import pathlib
 
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
-from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
-import numpy as np
-import torch
 from torch import nn
 
 TENSOR_NAMES = copy.copy(tiny_llama.TENSOR_NAMES)
@@ -104,28 +99,3 @@ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model, strict=False)
   model.eval()
   return model
-
-
-def define_and_run(checkpoint_path: str) -> None:
-  """Instantiates and runs a SmolLM model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  smollm_goldens = torch.load(current_dir / "smollm_lm_logits.pt")
-  kv_cache_max_len = 1024
-  model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  assert torch.allclose(
-      smollm_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-05
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/smollm"
-  )
-  define_and_run(input_checkpoint_path)

ai_edge_torch/generative/examples/smollm/verify.py

@@ -0,0 +1,59 @@
+# 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.
+# ==============================================================================
+
+"""Verifies the reauthored SmolLM-135M model."""
+
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.smollm import smollm
+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.",
+)
+
+
+def main(_):
+  checkpoint = "HuggingFaceTB/SmolLM-135M"
+  verifier.log_msg("Loading the original model from", 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
+  verifier.log_msg("Building the reauthored model from", reauthored_checkpoint)
+  reauthored_model = smollm.build_model(reauthored_checkpoint)
+
+  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=original_model,
+      reauthored_model=reauthored_model,
+      tokenizer=tokenizer,
+      prompts=_PROMPTS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -336,6 +336,8 @@ class Diffusion(nn.Module):
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
                             cross_dim=config.transformer_cross_attention_dim,
+                            hidden_dim=output_channel,
+                            output_dim=output_channel,
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=build_attention_config(
@@ -406,6 +408,8 @@ class Diffusion(nn.Module):
                 cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                     query_dim=mid_block_channels,
                     cross_dim=config.transformer_cross_attention_dim,
+                    hidden_dim=mid_block_channels,
+                    output_dim=mid_block_channels,
                     attention_batch_size=config.transformer_batch_size,
                     normalization_config=config.transformer_norm_config,
                     attention_config=build_attention_config(
@@ -477,6 +481,8 @@ class Diffusion(nn.Module):
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
                             cross_dim=config.transformer_cross_attention_dim,
+                            hidden_dim=output_channel,
+                            output_dim=output_channel,
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=build_attention_config(

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -18,69 +18,49 @@
 import os
 import pathlib
 
-import ai_edge_torch
+from absl import app
+from absl import flags
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
-from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.quantize import quant_recipes
-import torch
+from ai_edge_torch.generative.utilities import converter
 
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/tiny_llama'),
+    'The path to the model checkpoint, or directory holding the checkpoint.',
+)
+_TFLITE_PATH = flags.DEFINE_string(
+    'tflite_path',
+    '/tmp/tiny_llama_q8_seq512_ekv1024.tflite',
+    'The tflite file path to export.',
+)
+_PREFILL_SEQ_LEN = flags.DEFINE_integer(
+    'prefill_seq_len',
+    512,
+    'The maximum size of prefill input tensor.',
+)
+_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
+    'kv_cache_max_len',
+    1024,
+    'The maximum size of KV cache buffer, including both prefill and decode.',
+)
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    True,
+    'Whether the model should be quantized.',
+)
 
-def convert_tiny_llama_to_tflite(
-    checkpoint_path: str,
-    prefill_seq_len: int = 512,
-    kv_cache_max_len: int = 1024,
-    quantize: bool = True,
-):
-  """Converts TinyLlama model to multi-signature tflite model.
 
-  Args:
-      checkpoint_path (str): The filepath to the model checkpoint, or directory
-        holding the checkpoint.
-      prefill_seq_len (int, optional): The maximum size of prefill input tensor.
-        Defaults to 512.
-      kv_cache_max_len (int, optional): The maximum size of KV cache buffer,
-        including both prefill and decode. Defaults to 1024.
-      quantize (bool, optional): Whether the model should be quanized. Defaults
-        to True.
-  """
+def main(_):
   pytorch_model = tiny_llama.build_model(
-      checkpoint_path, kv_cache_max_len=kv_cache_max_len
+      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
-  # Tensors used to trace the model graph during conversion.
-  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.int)
-  prefill_input_pos = torch.arange(0, prefill_seq_len, dtype=torch.int)
-  decode_token = torch.tensor([[0]], dtype=torch.int)
-  decode_input_pos = torch.tensor([0], dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(pytorch_model.config)
-
-  quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
-  edge_model = (
-      ai_edge_torch.signature(
-          'prefill',
-          pytorch_model,
-          sample_kwargs={
-              'tokens': prefill_tokens,
-              'input_pos': prefill_input_pos,
-              'kv_cache': kv,
-          },
-      )
-      .signature(
-          'decode',
-          pytorch_model,
-          sample_kwargs={
-              'tokens': decode_token,
-              'input_pos': decode_input_pos,
-              'kv_cache': kv,
-          },
-      )
-      .convert(quant_config=quant_config)
-  )
-  quant_suffix = 'q8' if quantize else 'f32'
-  edge_model.export(
-      f'/tmp/tiny_llama_{quant_suffix}_seq{prefill_seq_len}_ekv{kv_cache_max_len}.tflite'
+  converter.convert_to_tflite(
+      pytorch_model,
+      tflite_path=_TFLITE_PATH.value,
+      prefill_seq_len=_PREFILL_SEQ_LEN.value,
+      quantize=_QUANTIZE.value,
   )
 
 
 if __name__ == '__main__':
-  path = os.path.join(pathlib.Path.home(), 'Downloads/llm_data/tiny_llama')
-  convert_tiny_llama_to_tflite(path)
+  app.run(main)

ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py

@@ -15,16 +15,12 @@
 
 """Example of building a TinyLlama model."""
 
-import os
-import pathlib
-
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
-import numpy as np
 import torch
 from torch import nn
 
@@ -179,28 +175,3 @@ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model)
   model.eval()
   return model
-
-
-def define_and_run(checkpoint_path: str) -> None:
-  """Instantiates and runs a TinyLlama model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  tiny_llama_goldens = torch.load(current_dir / "tiny_llama_lm_logits.pt")
-  kv_cache_max_len = 1024
-  model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  assert torch.allclose(
-      tiny_llama_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-02
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/tiny_llama"
-  )
-  define_and_run(input_checkpoint_path)

ai_edge_torch/generative/examples/tiny_llama/verify.py

@@ -0,0 +1,61 @@
+# 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.
+# ==============================================================================
+
+"""Verifies the reauthored TinyLlama-1.1B model."""
+
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "Show me the program to add 2 and 3.",
+    "The input prompts to generate answers.",
+)
+
+
+def main(_):
+  checkpoint = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+  verifier.log_msg("Loading the original model from", checkpoint)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(
+      checkpoint, trust_remote_code=True
+  )
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  verifier.log_msg("Building the reauthored model from", reauthored_checkpoint)
+  reauthored_model = tiny_llama.build_model(reauthored_checkpoint)
+
+  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=original_model,
+      reauthored_model=reauthored_model,
+      tokenizer=tokenizer,
+      prompts=_PROMPTS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/layers/attention.py

@@ -298,6 +298,8 @@ class CrossAttention(nn.Module):
       batch_size: int,
       query_dim: int,
       cross_dim: int,
+      hidden_dim: int,
+      output_dim: int,
       config: cfg.AttentionConfig,
       enable_hlfb: bool,
   ):
@@ -307,6 +309,8 @@ class CrossAttention(nn.Module):
       batch_size (int): batch size of the input tensor.
       query_dim (int): query tensor's dimension.
       cross_dim (int): cross attention's dimensions, for key and value tensors.
+      hidden_dim (int): hidden dimension that q, k, v tensors project to.
+      output_dim (int): output tensor's dimension.
       config (cfg.AttentionConfig): attention specific configurations.
       enable_hlfb (bool): whether hlfb is enabled or not.
     """
@@ -314,16 +318,16 @@ class CrossAttention(nn.Module):
     self.config = config
     self.n_heads = config.num_heads
     self.q_projection = nn.Linear(
-        query_dim, query_dim, bias=config.qkv_use_bias
+        query_dim, hidden_dim, bias=config.qkv_use_bias
     )
     self.k_projection = nn.Linear(
-        cross_dim, query_dim, bias=config.qkv_use_bias
+        cross_dim, hidden_dim, bias=config.qkv_use_bias
     )
     self.v_projection = nn.Linear(
-        cross_dim, query_dim, bias=config.qkv_use_bias
+        cross_dim, hidden_dim, bias=config.qkv_use_bias
     )
     self.output_projection = nn.Linear(
-        query_dim, query_dim, bias=config.output_proj_use_bias
+        hidden_dim, output_dim, bias=config.output_proj_use_bias
     )
 
     self.sdpa_func = (

ai_edge_torch/generative/layers/builder.py

@@ -75,7 +75,9 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
         zero_centered_gamma=config.zero_centered,
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
-    return normalization.LayerNorm(dim, config.epsilon, config.enable_hlfb)
+    return normalization.LayerNorm(
+        dim, config.epsilon, config.enable_hlfb, config.use_input_shape
+    )
   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

@@ -69,6 +69,9 @@ class NormalizationConfig:
   zero_centered: bool = False
   # Number of groups used in group normalization.
   group_num: Optional[float] = None
+  # Whether to use the input shape to determine the dimension of normalization
+  # when type is LAYER_NORM.
+  use_input_shape: bool = True
 
 
 @dataclass

ai_edge_torch/generative/layers/normalization.py

@@ -78,7 +78,7 @@ class GroupNorm(torch.nn.Module):
       group_num (int): Number of groups to separate the channels into.
       dim (int): Dimension of the input tensor.
       eps (float): A small float value to ensure numerical stability (default:
-        1e-6).
+        1e-5).
       enable_hlfb (bool): Whether to convert this normalization into a single
         op.
     """
@@ -112,7 +112,13 @@ class GroupNorm(torch.nn.Module):
 
 class LayerNorm(torch.nn.Module):
 
-  def __init__(self, dim: int, eps: float = 1e-5, enable_hlfb: bool = False):
+  def __init__(
+      self,
+      dim: int,
+      eps: float = 1e-5,
+      enable_hlfb: bool = False,
+      use_input_shape: bool = True,
+  ):
     """Initialize the LayerNorm layer.
 
     Args:
@@ -121,9 +127,12 @@ class LayerNorm(torch.nn.Module):
         1e-6).
       enable_hlfb (bool): Whether to convert this normalization into a single
         op.
+      use_input_shape (bool): Whether to use the input shape to determine the
+        dimension of normalization (default: True).
     """
     super().__init__()
     self.enable_hlfb = enable_hlfb
+    self.use_input_shape = use_input_shape
     self.eps = eps
     self.weight = torch.nn.Parameter(torch.ones(dim))
     self.bias = torch.nn.Parameter(torch.ones(dim))
@@ -139,19 +148,18 @@ class LayerNorm(torch.nn.Module):
     """
     if self.enable_hlfb:
       return layer_norm_with_hlfb(
-          x,
-          self.weight,
-          self.bias,
-          self.eps,
+          x, self.weight, self.bias, self.eps, self.use_input_shape
       )
+
+    if self.use_input_shape:
+      normalized_shape = x.shape
+      weight = self.weight.broadcast_to(x.shape)
+      bias = self.bias.broadcast_to(x.shape)
     else:
-      return F.layer_norm(
-          x,
-          x.shape,
-          self.weight.broadcast_to(x.shape),
-          self.bias.broadcast_to(x.shape),
-          self.eps,
-      )
+      normalized_shape = self.weight.shape
+      weight = self.weight
+      bias = self.bias
+    return F.layer_norm(x, normalized_shape, weight, bias, self.eps)
 
 
 def group_norm_with_hlfb(
@@ -193,6 +201,7 @@ def layer_norm_with_hlfb(
     w: torch.Tensor,
     b: torch.Tensor,
     eps: float,
+    use_input_shape: bool,
 ):
   """Layer Normalization with high-level function boundary enabled.
 
@@ -201,18 +210,20 @@ def layer_norm_with_hlfb(
     w (torch.Tensor): The weight tensor for the normalization.
     b (torch.Tensor): The bias tensor for the normalization.
     eps (float): A small float value to ensure numerical stability.
+    use_input_shape (bool): Whether to use the input shape to determine the
+      dimension of normalization.
 
   Returns:
     The output tensor of Layer Normalization.
   """
   builder = StableHLOCompositeBuilder(name="odml.layer_norm", attr={"eps": eps})
   x, w, b = builder.mark_inputs(x, w, b)
-  y = F.layer_norm(
-      x,
-      x.shape,
-      weight=w.broadcast_to(x.shape),
-      bias=b.broadcast_to(x.shape),
-      eps=eps,
-  )
+  if use_input_shape:
+    normalized_shape = x.shape
+    w = w.broadcast_to(x.shape)
+    b = b.broadcast_to(x.shape)
+  else:
+    normalized_shape = w.shape
+  y = F.layer_norm(x, normalized_shape, w, b, eps=eps)
   y = builder.mark_outputs(y)
   return y

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -119,15 +119,25 @@ def scaled_dot_product_attention_with_hlfb(
     # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
     k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
     v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
-  y = F.scaled_dot_product_attention(
-      q,
-      k,
-      v,
-      attn_mask=mask,
-      dropout_p=0.0,
-      is_causal=mask is None,
-      scale=scale,
-  )
+  if softcap is None:
+    y = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=mask,
+        dropout_p=0.0,
+        is_causal=mask is None,
+        scale=scale,
+    )
+  else:
+    q.mul_(scale)
+    scores = q @ k.transpose(-1, -2)
+    scores = scores / softcap
+    scores = torch.tanh(scores)
+    scores = scores * softcap
+    scores = scores + mask
+    out = F.softmax(scores.float(), dim=-1).type_as(q)
+    y = torch.matmul(out, v)
 
   result = y.transpose(1, 2)
   result = builder.mark_outputs(result)