Introduce sdpa_with_kv_update function. It will perform different cache update, pre/post sdpa logic based on kv cache layout.

PiperOrigin-RevId: 745609775

Author: haozha111

ai_edge_torch/generative/layers/experimental/attention.py

@@ -24,8 +24,7 @@
 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
-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
+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
@@ -147,7 +146,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
 
   def forward(
       self,
@@ -221,36 +219,8 @@ def forward(
       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_experimental.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)
+    sdpa_out, kv_cache = sdpa_with_kv_update.sdpa_with_kv_update(
+        q, k, v, kv_cache, input_pos, mask, self.config
     )
 
     # Compute the output projection.

ai_edge_torch/generative/layers/experimental/kv_cache.py

@@ -44,7 +44,8 @@ def update(
   assert (
       cache.kv_layout == kv_utils.KV_LAYOUT_TRANSPOSED
   ), "KV entry must have transposed layout."
-  return _update_kv_impl_transposed(cache, input_pos, k_slice, v_slice)
+  update_kv_cache = _update_kv_impl_transposed
+  return update_kv_cache(cache, input_pos, k_slice, v_slice)
 
 
 def _get_slice_indices(

ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py

@@ -82,7 +82,6 @@ def _sdpa(k_type, v_type, *args, **kwargs):
   padded_logits = logits + mask
   padded_logits = padded_logits.reshape(1, bk, gt, s)
   probs = F.softmax(padded_logits, dim=-1).type_as(key)
-
   encoded = bmm_lib.bmm_4d(probs, value)
 
   return encoded  # 1, bk, gt, h

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -0,0 +1,124 @@
+# 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 utility functions for data loading etc.
+from dataclasses import dataclass
+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
+from ai_edge_torch.generative.layers.experimental import types
+import ai_edge_torch.generative.layers.model_config as cfg
+from multipledispatch import dispatch
+import torch
+
+
+def sdpa_with_kv_update(
+    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]:
+  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,
+  )
+
+
+@dispatch(types.BNTH, types.BNHT)
+def sdpa_with_kv_update_impl(
+    k_type, v_type, *args, **kwargs
+) -> 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
+  # btnh -> bnth -> b(kg)th -> 1(bk)(gt)h
+  query = query.permute(0, 2, 1, 3).reshape(
+      1, b * config.num_query_groups, g * seq_len, h
+  )
+
+  key = key.permute(0, 2, 1, 3).reshape(
+      1, -1, seq_len, config.head_dim
+  )  # 1, bk, s, h
+  value = value.permute(0, 2, 3, 1).reshape(
+      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
+
+  sdpa_out = sdpa.scaled_dot_product_attention(
+      kv,
+      query,
+      key,
+      value,
+      config.head_dim,
+      mask=mask,
+      softcap=config.logit_softcap,
+  )  # 1, bk, gt, h
+  sdpa_out = (
+      sdpa_out.reshape(b, -1, seq_len, h)
+      .permute(0, 2, 1, 3)
+      .reshape(b, seq_len, -1)
+  )
+  return sdpa_out, kv
+
+
+@dispatch(object, object)
+def sdpa_with_kv_update_impl(
+    k_type, v_type, *args, **kwargs
+) -> 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(
+      query,
+      key,
+      value,
+      config.head_dim,
+      mask=mask,
+      softcap=config.logit_softcap,
+  )
+  sdpa_out = sdpa_out.reshape(b, seq_len, -1)
+  return sdpa_out, kv

ai_edge_torch/generative/utilities/experimental/verifier.py

@@ -19,7 +19,7 @@
 from typing import Any, List, Optional
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 import torch
 
 ExportConfig = export_config.ExportConfig