Update on " [ExecuTorch][BE] Split kv cache and SDPA for better code sharing"

Summary:

Why?
We have coupled SDPA with kv cache for a while. Initially this was done
as we implemented sdpa_with_kv_cache custom op to reduce multiple copy
overheads from kv cache update. (This could have been done by having
separate custom kv cache update and custom sdpa op. Recent changes
enabled this.)
As a result of SDPA module owning kv cache, we get a) non-composable
implementation and b) harder to reuse model definition and components
from repos like tune. Output of this is that we have multiple definition
of the same model, llama, lying around in ET, TorchChat and Tune. This
diff and subsequent ones will try to move in the direction where custom
kv cache and custom sdpa become decoupled and composable, making it more
module-swap friendly with tune's model definition.

How.
Earlier PRs decoupled kv cache update from sdpa. So now
1. Decouple SDPA nn.Module from KV cache.
2. Standardize on KVCache and SDPA interface. That is KVCache and SDPA
   both operate on q, k, v in [B, # heads, seq_len, head_dim] formatted
   tensors.
3. 2 will introduce multiple tranposes when KVCache and SDPA are
   replaced by custom modules, but we will write graph pass to undo
   those.

Test Plan:
Existing tests.
Make sure perf doesnt regress

Differential Revision: [D67914054](https://our.internmc.facebook.com/intern/diff/D67914054)

[ghstack-poisoned]

Author: kimishpatel

examples/models/llama/llama_transformer.py

@@ -286,13 +286,15 @@ def update(
 class SDPA(nn.Module):
     def __init__(
         self,
+        kv_cache: KVCache,
         dim: int,
         head_dim: int,
         n_rep: int,
         max_seq_len: int,
         enable_dynamic_shape: bool,
     ):
         super().__init__()
+        self.kv_cache = kv_cache
         self.dim = dim
         self.head_dim = head_dim
         self.n_rep = n_rep
@@ -371,6 +373,7 @@ def __init__(self, args: ModelArgs, layer_id: int, rope: Rope):
                 args.enable_dynamic_shape,
             )
             self.SDPA = SDPA(
+                kv_cache=self.kv_cache,
                 dim=self.n_local_heads * self.head_dim,
                 head_dim=self.head_dim,
                 n_rep=self.n_rep,
@@ -403,7 +406,7 @@ def forward(
 
         if self.use_kv_cache:
             assert input_pos is not None
-            k, v = self.kv_cache.update(input_pos, k, v)
+            k, v = self.SDPA.kv_cache.update(input_pos, k, v)
             output = self.SDPA(input_pos, q, k, v, bsz, seqlen, self.mask)
             return self.wo(output)
 

examples/models/llama/source_transformation/sdpa.py

@@ -19,9 +19,11 @@
 class SDPACustom(torch.nn.Module):
     def __init__(
         self,
+        kv_cache: KVCache,
         dim: int,
     ):
         super().__init__()
+        self.kv_cache = kv_cache
         self.dim = dim
 
     def forward(
@@ -63,7 +65,7 @@ def _replace_sdpa_with_custom_op(module: torch.nn.Module):
             setattr(
                 module,
                 name,
-                SDPACustom(child.dim),
+                SDPACustom(child.kv_cache, child.dim),
             )
         else:
             _replace_sdpa_with_custom_op(child)
@@ -79,11 +81,13 @@ def replace_sdpa_with_custom_op(module: torch.nn.Module) -> torch.nn.Module:
 class SDPASimple(torch.nn.Module):
     def __init__(
         self,
+        kv_cache: KVCache,
         dim: int,
         head_dim: int,
         n_rep: int,
     ):
         super().__init__()
+        self.kv_cache = kv_cache
         self.dim = dim
         self.head_dim = head_dim
         self.n_rep = n_rep
@@ -131,10 +135,12 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
 class SDPAFlex(torch.nn.Module):
     def __init__(
         self,
+        kv_cache: KVCache,
         dim: int,
         n_rep: int,
     ):
         super().__init__()
+        self.kv_cache = kv_cache
         self.dim = dim
         self.n_rep = n_rep
 
@@ -171,7 +177,7 @@ def replace_sdpa_with_simple_sdpa(module: torch.nn.Module):
             setattr(
                 module,
                 name,
-                SDPASimple(child.dim, child.head_dim, child.n_rep),
+                SDPASimple(child.kv_cache, child.dim, child.head_dim, child.n_rep),
             )
         else:
             replace_sdpa_with_simple_sdpa(child)
@@ -184,7 +190,7 @@ def replace_sdpa_with_flex_sdpa(module: torch.nn.Module):
             setattr(
                 module,
                 name,
-                SDPAFlex(child.dim, child.n_rep),
+                SDPAFlex(child.kv_cache, child.dim, child.n_rep),
             )
         else:
             replace_sdpa_with_flex_sdpa(child)
@@ -216,11 +222,13 @@ class SDPACoreML(torch.nn.Module):
 
     def __init__(
         self,
+        kv_cache: KVCache,
         dim: int,
         head_dim: int,
         n_rep: int,
     ):
         super().__init__()
+        self.kv_cache = kv_cache
         self.dim = dim
         self.head_dim = head_dim
         self.n_rep = n_rep
@@ -252,7 +260,7 @@ def replace_sdpa_with_coreml_sdpa(module: torch.nn.Module):
             setattr(
                 module,
                 name,
-                SDPACoreML(child.dim, child.head_dim, child.n_rep),
+                SDPACoreML(child.kv_cache, child.dim, child.head_dim, child.n_rep),
             )
         else:
             replace_sdpa_with_coreml_sdpa(child)