Call sdpa_with_kv_update from non-experimental attention

- Replace multidispatch with an explicit if clause for strong-typed function calls.
- Add unittests for transposed KV cache.

PiperOrigin-RevId: 750707606

Author: protobird-git

ai_edge_torch/generative/layers/attention.py

@@ -21,6 +21,7 @@
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
 from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa
+from ai_edge_torch.generative.layers import sdpa_with_kv_update
 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
@@ -142,11 +143,6 @@ def __init__(
     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_with_hlfb
-        if enable_hlfb
-        else sdpa.scaled_dot_product_attention
-    )
 
   def forward(
       self,
@@ -174,7 +170,7 @@ def forward(
         KV Cach Entry (if passed in).
     """
     # Batch size, sequence length, embedding dimensionality.
-    B, T, E = x.size()
+    B, T, _ = x.size()
     qkv = self.qkv_projection(x)
 
     # Assemble into a number of query groups to support MHA, MQA and GQA.
@@ -218,19 +214,9 @@ def forward(
       cos, sin = rope
       q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
 
-    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(
-        q,
-        k,
-        v,
-        self.config.head_dim,
-        mask=mask,
-        softcap=self.config.logit_softcap,
+    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
     )
-    sdpa_out = sdpa_out.reshape(B, T, -1)
 
     # Compute the output projection.
     y = self.output_projection(sdpa_out)

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Common utility functions for data loading etc.
-from dataclasses import dataclass
+
+"""Common utility functions for data loading etc."""
+
 from typing import Tuple
+
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa_default
 from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
 from ai_edge_torch.generative.layers.experimental import scaled_dot_product_attention as sdpa
 import ai_edge_torch.generative.layers.model_config as cfg
-from ai_edge_torch.generative.utilities import types
-from multipledispatch import dispatch
 import torch
 
 
@@ -33,32 +33,27 @@ def sdpa_with_kv_update(
     input_pos: torch.Tensor,
     mask: torch.Tensor,
     config: cfg.AttentionConfig,
+    enable_hlfb: bool,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
-  return sdpa_with_kv_update_impl(
-      kv.kv_layout[0](),  # key layout
-      kv.kv_layout[1](),  # value layout
-      query=query,
-      key=key,
-      value=value,
-      kv=kv,
-      input_pos=input_pos,
-      mask=mask,
-      config=config,
+  """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
+    )
+  return _sdpa_with_kv_update_default(
+      query, key, value, kv, input_pos, mask, config, enable_hlfb
   )
 
 
-@dispatch(types.BNTH, types.BNHT)
-def sdpa_with_kv_update_impl(
-    k_type, v_type, *args, **kwargs
+def _sdpa_with_kv_update_transposed(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv: kv_utils.KVCacheEntry,
+    input_pos: torch.Tensor,
+    mask: torch.Tensor,
+    config: cfg.AttentionConfig,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
-  query = kwargs["query"]
-  key = kwargs["key"]
-  value = kwargs["value"]
-  kv = kwargs["kv"]
-  input_pos = kwargs["input_pos"]
-  mask = kwargs["mask"]
-  config = kwargs["config"]
-
   # Transpose k/v to specific layout for GPU implementation.
   b, seq_len, n, h = query.shape
   g = n // config.num_query_groups
@@ -74,9 +69,8 @@ def sdpa_with_kv_update_impl(
       1, -1, config.head_dim, seq_len
   )  # 1, bk, h, s
 
-  if kv is not None:
-    kv = kv_utils_experimental.update(kv, input_pos, key, value)
-    key, value = kv.k_cache, kv.v_cache
+  kv = kv_utils_experimental.update(kv, input_pos, key, value)
+  key, value = kv.k_cache, kv.v_cache
 
   sdpa_out = sdpa.scaled_dot_product_attention(
       kv,
@@ -95,24 +89,26 @@ def sdpa_with_kv_update_impl(
   return sdpa_out, kv
 
 
-@dispatch(object, object)
-def sdpa_with_kv_update_impl(
-    k_type, v_type, *args, **kwargs
+def _sdpa_with_kv_update_default(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv: kv_utils.KVCacheEntry,
+    input_pos: torch.Tensor,
+    mask: torch.Tensor,
+    config: cfg.AttentionConfig,
+    enable_hlfb: bool,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
-  query = kwargs["query"]
-  key = kwargs["key"]
-  value = kwargs["value"]
-  kv = kwargs["kv"]
-  input_pos = kwargs["input_pos"]
-  mask = kwargs["mask"]
-  config = kwargs["config"]
-
   b, seq_len, _, _ = query.shape
   if kv is not None:
     kv = kv_utils.update(kv, input_pos, key, value)
     key, value = kv.k_cache, kv.v_cache
 
-  sdpa_out = sdpa_default.scaled_dot_product_attention(
+  if enable_hlfb:
+    sdpa_func = sdpa_default.scaled_dot_product_attention_with_hlfb
+  else:
+    sdpa_func = sdpa_default.scaled_dot_product_attention
+  sdpa_out = sdpa_func(
       query,
       key,
       value,

ai_edge_torch/generative/test/test_model_conversion.py

@@ -41,15 +41,15 @@ def setUp(self):
         )
     )
 
-  def _get_params(self, enable_hlfb: bool):
+  def _get_params(self, enable_hlfb: bool, kv_layout: kv_cache.KVLayout):
     """Returns a model, edge model and the kwargs to use for testing."""
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = enable_hlfb
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKVCache(config).eval()
     tokens, input_pos = torch.tensor([[1]], dtype=torch.int), torch.tensor(
         [10], dtype=torch.int
     )
-    kv = kv_cache.KVCache.from_model_config(config)
+    kv = kv_cache.KVCache.from_model_config(config, kv_layout=kv_layout)
     kwargs = {
         "tokens": tokens,
         "input_pos": input_pos,
@@ -65,8 +65,12 @@ def _get_params(self, enable_hlfb: bool):
     )
     return pytorch_model, edge_model, kwargs
 
-  def _test_model_with_kv_cache(self, enable_hlfb: bool):
-    pytorch_model, edge_model, kwargs = self._get_params(enable_hlfb)
+  def _test_model_with_kv_cache(
+      self,
+      enable_hlfb: bool = False,
+      kv_layout: kv_cache.KVLayout = kv_cache.KV_LAYOUT_DEFAULT,
+  ):
+    pytorch_model, edge_model, kwargs = self._get_params(enable_hlfb, kv_layout)
 
     self.assertTrue(
         test_utils.compare_tflite_torch(
@@ -95,13 +99,22 @@ def test_toy_model_with_kv_cache(self):
   def test_toy_model_with_kv_cache_with_hlfb(self):
     self._test_model_with_kv_cache(enable_hlfb=True)
 
+  @googletest.skipIf(
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
+  )
+  def test_toy_model_with_kv_cache_transposed(self):
+    self._test_model_with_kv_cache(kv_layout=kv_cache.KV_LAYOUT_TRANSPOSED)
+
   @googletest.skipIf(
       ai_edge_torch.config.in_oss,
       reason="tests with custom ops are not supported in oss",
   )
   def test_toy_model_has_dus_op(self):
     """Tests that the model has the dynamic update slice op."""
-    _, edge_model, _ = self._get_params(enable_hlfb=True)
+    _, edge_model, _ = self._get_params(
+        enable_hlfb=True, kv_layout=kv_cache.KV_LAYOUT_DEFAULT
+    )
     interpreter_ = interpreter.InterpreterWithCustomOps(
         custom_op_registerers=["GenAIOpsRegisterer"],
         model_content=edge_model.tflite_model(),
@@ -112,7 +125,14 @@ def test_toy_model_has_dus_op(self):
     op_names = [op["op_name"] for op in interpreter_._get_ops_details()]
     self.assertIn("DYNAMIC_UPDATE_SLICE", op_names)
 
-  def _test_multisig_model(self, config, pytorch_model, atol, rtol):
+  def _test_multisig_model(
+      self,
+      config,
+      pytorch_model,
+      atol,
+      rtol,
+      kv_layout=kv_cache.KV_LAYOUT_DEFAULT,
+  ):
     # prefill
     seq_len = 10
     prefill_tokens = torch.zeros((1, seq_len), dtype=torch.int, device="cpu")
@@ -124,7 +144,7 @@ def _test_multisig_model(self, config, pytorch_model, atol, rtol):
     decode_token = torch.tensor([[1]], dtype=torch.int)
     decode_input_pos = torch.tensor([5], dtype=torch.int)
 
-    kv = kv_cache.KVCache.from_model_config(config)
+    kv = kv_cache.KVCache.from_model_config(config, kv_layout=kv_layout)
 
     edge_model = (
         ai_edge_torch.signature(
@@ -160,7 +180,7 @@ def _test_multisig_model(self, config, pytorch_model, atol, rtol):
             kv,
             signature_name="prefill",
             atol=atol,
-            rtol=atol,
+            rtol=rtol,
         )
     )
 
@@ -173,7 +193,7 @@ def _test_multisig_model(self, config, pytorch_model, atol, rtol):
             kv,
             signature_name="decode",
             atol=atol,
-            rtol=atol,
+            rtol=rtol,
         )
     )
 
@@ -186,6 +206,21 @@ def test_tiny_llama_multisig(self):
     pytorch_model = tiny_llama.TinyLlama(config).eval()
     self._test_multisig_model(config, pytorch_model, atol=1e-5, rtol=1e-5)
 
+  @googletest.skipIf(
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
+  )
+  def test_tiny_llama_multisig_kv_layout_transposed(self):
+    config = tiny_llama.get_fake_model_config()
+    pytorch_model = tiny_llama.TinyLlama(config).eval()
+    self._test_multisig_model(
+        config,
+        pytorch_model,
+        atol=1e-5,
+        rtol=1e-5,
+        kv_layout=kv_cache.KV_LAYOUT_TRANSPOSED,
+    )
+
 
 if __name__ == "__main__":
   googletest.main()