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/source_transformation/quantized_kv_cache.py

@@ -118,12 +118,12 @@ def update(self, input_pos, k_val, v_val):
                 v_zero_points, self.v_cache_zero_points, start_pos
             )
         else:
-            self.k_cache[:, :, input_pos] = quantized_k_val
-            self.k_cache_scales[:, :, input_pos] = k_scales
-            self.k_cache_zero_points[:, :, input_pos] = k_zero_points
-            self.v_cache[:, :, input_pos] = quantized_v_val
-            self.v_cache_scales[:, :, input_pos] = v_scales
-            self.v_cache_zero_points[:, :, input_pos] = v_zero_points
+            self.k_cache[:, input_pos] = quantized_k_val
+            self.k_cache_scales[:, input_pos] = k_scales
+            self.k_cache_zero_points[:, input_pos] = k_zero_points
+            self.v_cache[:, input_pos] = quantized_v_val
+            self.v_cache_scales[:, input_pos] = v_scales
+            self.v_cache_zero_points[:, input_pos] = v_zero_points
 
         k_out = torch.ops.quantized_decomposed.dequantize_per_token(
             self.k_cache,
@@ -149,8 +149,8 @@ def update(self, input_pos, k_val, v_val):
             _ = torch.ops.llama.update_cache(k_val, k_out, start_pos)
             _ = torch.ops.llama.update_cache(v_val, v_out, start_pos)
         else:
-            k_out[:, :, input_pos] = k_val
-            v_out[:, :, input_pos] = v_val
+            k_out[:, input_pos] = k_val
+            v_out[:, input_pos] = v_val
 
         return k_out.transpose(1, 2), v_out.transpose(1, 2)
 

examples/models/llama/source_transformation/test_quantized_kv_cache.py

@@ -24,16 +24,12 @@ def _init_cache(self):
             self.max_seq_len,
             self.n_kv_heads,
             self.head_dim,
-            self.transpose_kv_cache,
             self.enable_dynamic_shape,
             dtype=self.dtype,
         )
 
     def _init_kv(self):
-        if self.transpose_kv_cache:
-            shape = (1, self.n_kv_heads, self.seq_len, self.head_dim)
-        else:
-            shape = (1, self.seq_len, self.n_kv_heads, self.head_dim)
+        shape = (1, self.n_kv_heads, self.seq_len, self.head_dim)
         k = torch.rand(shape, dtype=self.dtype)
         v = torch.rand(shape, dtype=self.dtype)
         return k, v
@@ -45,11 +41,11 @@ def setUp(self):
         self.n_kv_heads = 8
         self.head_dim = 17
         self.enable_dynamic_shape = False
-        self.transpose_kv_cache = False
         self.dtype = torch.float32
 
-    def _test_simple_update_fetch(self, is_transposed=False, is_dynamic_shape=False):
-        self.transpose_kv_cache = is_transposed
+    def _test_simple_update_fetch(
+        self, is_dynamic_shape=False, use_custom_update_cache_op=False
+    ):
         self.enable_dynamic_shape = is_dynamic_shape
         input_pos = torch.tensor([0, 1, 2])
         self.seq_len = input_pos.size(0)
@@ -64,10 +60,7 @@ def _test_simple_update_fetch(self, is_transposed=False, is_dynamic_shape=False)
         )
 
         def index(t, input_pos):
-            if self.transpose_kv_cache:
-                return t[:, :, input_pos, :]
-            else:
-                return t[:, input_pos, :, :]
+            return t[:, :, input_pos, :]
 
         sliced_k_cache = index(updated_k_cache, input_pos)
         sliced_v_cache = index(updated_v_cache, input_pos)
@@ -115,14 +108,16 @@ def index(t, input_pos):
             atol=1e-02,
         )
 
-    def test_simple_update_fetch_not_transposed(self):
+    def test_simple_update_fetch(self):
         self._test_simple_update_fetch()
 
-    def test_simple_update_fetch_not_transposed_dynamic_shape(self):
-        self._test_simple_update_fetch(is_dynamic_shape=True)
+    def test_simple_update_fetch_use_custom_op(self):
+        self._test_simple_update_fetch(use_custom_update_cache_op=True)
 
-    def test_simple_update_fetch_transposed(self):
-        self._test_simple_update_fetch(is_transposed=True)
+    def test_simple_update_fetch_dynamic_shape(self):
+        self._test_simple_update_fetch(is_dynamic_shape=True)
 
-    def test_simple_update_fetch_transposed_dynamic_shape(self):
-        self._test_simple_update_fetch(is_transposed=True, is_dynamic_shape=True)
+    def test_simple_update_fetch_dynamic_shape_use_custom_op(self):
+        self._test_simple_update_fetch(
+            is_dynamic_shape=True, use_custom_update_cache_op=True
+        )