[Executorch][llm] Add ring buffer based kv cache and mask calculation to MHA

Pull Request resolved: #10609

Leveraging previous work now we allow MHA to have ring buffer cache. If ring buffer cache is used
then we query the mask from kv cache and use that for sdpa instead of using precalculated mask.

In this process we had to adjsut ring buffer implementation to allow keeping the context of
full sliding window. See code for comment.
ghstack-source-id: 283404675
@exported-using-ghexport

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

Author: kimishpatel

examples/models/llama/attention.py

@@ -123,14 +123,12 @@ def __init__(
         head_dim: int,
         n_rep: int,
         max_context_len: int,
-        enable_dynamic_shape: bool,
     ):
         super().__init__()
         self.dim = dim
         self.head_dim = head_dim
         self.n_rep = n_rep
         self.max_context_len = max_context_len
-        self.enable_dynamic_shape = enable_dynamic_shape
 
     def forward(
         self,
@@ -142,21 +140,12 @@ def forward(
         seqlen,
         mask: torch.Tensor,
     ) -> torch.Tensor:
-        if self.enable_dynamic_shape:
-            start_pos = input_pos[-1].item()
-            torch._check_is_size(start_pos)
-            torch._check(start_pos < self.max_context_len)
-            seq_length = q.size(2)
-            # pyre-ignore: Incompatible parameter type [6]
-            attn_mask = mask.narrow(0, start_pos, seq_length)
-        else:
-            attn_mask = mask[None, None, input_pos]
 
         # TODO(kimishpatel): This should not be necessary because scaled_dot_product_attention
         # can natively support GQA now. But needs enable_gqa=True
         k = k.repeat_interleave(self.n_rep, dim=1)
         v = v.repeat_interleave(self.n_rep, dim=1)
-        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0)
+        y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
 
         return y.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim)
 
@@ -236,21 +225,79 @@ def __init__(
         enable_dynamic_shape: bool,
         dtype=torch.float32,
     ):
+        self.window_size = max_context_length
+        """
+        Reason why we want the kv cache size to be twice the context length:
+        Sliding window attention without ringbuffer
+        pos   0  1  2  3  4  5  6  7  8  9  10
+        0     x  0  0  0  0  0  0  0  0  0  0
+        1     x  x  0  0  0  0  0  0  0  0  0
+        2     x  x  x  0  0  0  0  0  0  0  0
+        3     x  x  x  x  0  0  0  0  0  0  0
+        4     0  x  x  x  x  0  0  0  0  0  0
+        5     0  0  x  x  x  x  0  0  0  0  0
+        6     0  0  0  x  x  x  x  0  0  0  0
+        7     0  0  0  0  x  x  x  x  0  0  0
+        8     0  0  0  0  0  x  x  x  x  0  0
+        9     0  0  0  0  0  0  x  x  x  x  0
+        10    0  0  0  0  0  0  0  x  x  x  x
+
+        So when doing attention for pos = 5 and seq_len = 4 our attention
+        mask would be
+        5     0  0  x  x  x  x  0  0  0  0  0
+        6     0  0  0  x  x  x  x  0  0  0  0
+        7     0  0  0  0  x  x  x  x  0  0  0
+        8     0  0  0  0  0  x  x  x  x  0  0
+        Thus tok at pos = 5 is able to attend to tokens at pos 2, 3 and 4.
+        This is how training is done.
+
+        Now lets consider ring kv cache of size 4. When we are at pos = 5
+        before updating the kv cache, state of the kv cache would be
+        [4 1 2 3]. That is we evicted token at pos = 0 out. Now during
+        attention calculation at pos = 5 seq len = 4, we will update cache and
+        new pos in the cache would be [8 5 6 7]. So note that 5 can now only attend
+        to itself. Not 2, 3 and 4 as you would have during training.
+        So not having kept 2, 3 and 4 in cache means we will have divergent behavior.
+        Worst case of this would have been when update it equal to the length of
+        the cache. like in our case pos = 5 seq len = 4.
+        Thus we need to have a cache that is larger. How much larger, as much as
+        the sliding window size. So twice the max_context_length.
+        How would that have helped. Lets see. At pos = 5 our cache would have
+        [0, 1, 2, 3, 4, NA, NA, NA] After cache update we would have
+        [8, 1, 2, 3, 4, 5, 6, 7]. We kicked out token at pos = 0. However, the
+        current step still has access to [pos - sliding_window_size, pos] tokens.
+        
+        To make sure we dont over attend, i.e. we dont have pos = 5
+        to attend to pos = 1, mask calculaton has to account for the sliding window
+        size.
+        """
         super().__init__(
             max_batch_size,
-            max_context_length,
+            max_context_length * 2,
             n_heads,
             head_dim,
             enable_dynamic_shape,
             dtype,
         )
-        self.cache_positions_manager = CachePositionsManager(max_context_length)
+        self.cache_positions_manager = CachePositionsManager(self.max_context_length)
+        self.is_ring_buffer = True
+
+    def create_causal_mask_for_ring_buffer(self, start_pos, seq_len):
+        pos_q = start_pos + torch.arange(seq_len, dtype=torch.long).view(-1, 1)
+        cache_positions = self.cache_positions_manager.cache_positions
+        delta = pos_q - cache_positions
+        attn_mask = (cache_positions >= 0) & (delta >= 0) & (delta < self.window_size)
+        attn_mask = torch.where(attn_mask == True, 0, float("-inf"))  # noqa E712
+        return attn_mask
 
     def update(
         self, input_pos: torch.Tensor, k_val: torch.Tensor, v_val: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         # input_pos: [S], k_val: [B, H, S, D]
         seq_len = k_val.size(2)
+        assert seq_len <= self.k_cache.size(
+            2
+        ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(2)})"
         indices = self.cache_positions_manager.calculate_positions_and_update_indices(
             input_pos, seq_len
         )
@@ -286,6 +333,7 @@ def __init__(self, args: ModelArgs, layer_id: int, rope: Rope):
         self.attention_qkv_bias = args.attention_qkv_bias
         self.use_qk_norm = args.use_qk_norm
         self.qk_norm_before_rope = args.qk_norm_before_rope
+        self.enable_dynamic_shape = args.enable_dynamic_shape
 
         if self.use_qk_norm:
             q_norm_dim = self.head_dim
@@ -331,7 +379,6 @@ def __init__(self, args: ModelArgs, layer_id: int, rope: Rope):
                 head_dim=self.head_dim,
                 n_rep=self.n_rep,
                 max_context_len=self.max_context_len,
-                enable_dynamic_shape=args.enable_dynamic_shape,
             )
 
     def forward(
@@ -368,8 +415,22 @@ def forward(
 
         if self.use_kv_cache:
             assert input_pos is not None
+            if self.enable_dynamic_shape:
+                start_pos = input_pos[-1].item()
+                torch._check_is_size(start_pos)
+                torch._check(start_pos < self.max_context_len)
+                seq_length = q.size(2)
+                # pyre-ignore: Incompatible parameter type [6]
+                attn_mask = self.mask.narrow(0, start_pos, seq_length)
+            else:
+                # mask is always 2D
+                attn_mask = self.mask[input_pos]
             k, v = self.kv_cache.update(input_pos, k, v)
-            output = self.SDPA(input_pos, q, k, v, bsz, seqlen, self.mask)
+            if getattr(self.kv_cache, "is_ring_buffer", False):
+                attn_mask = self.kv_cache.create_causal_mask_for_ring_buffer(
+                    input_pos[0].item(), seqlen
+                )
+            output = self.SDPA(input_pos, q, k, v, bsz, seqlen, attn_mask)
             return self.wo(output), None
 
         # grouped multiquery attention: expand out keys and values

examples/models/llama/source_transformation/sdpa.py

@@ -22,15 +22,11 @@ class SDPACustom(torch.nn.Module):
     def __init__(
         self,
         dim: int,
-        max_context_len,
-        enable_dynamic_shape,
         use_attention_mask: bool = False,
     ):
         super().__init__()
         self.dim = dim
-        self.max_context_len = max_context_len
         self.use_attention_mask = use_attention_mask
-        self.enable_dynamic_shape = enable_dynamic_shape
 
     def forward(
         self,
@@ -42,16 +38,6 @@ def forward(
         seqlen,
         mask,
     ):
-        if self.use_attention_mask:
-            if self.enable_dynamic_shape:
-                start_pos = input_pos[-1].item()
-                torch._check_is_size(start_pos)
-                torch._check(start_pos < self.max_context_len)
-                seq_length = q.size(2)
-                mask = mask.narrow(0, start_pos, seq_length)
-            else:
-                mask = mask[input_pos]
-
         q = q.transpose(1, 2)  # (bs, seqlen, n_local_heads, head_dim)
         k = k.transpose(1, 2)
         v = v.transpose(1, 2)
@@ -96,8 +82,6 @@ def _replace_sdpa_with_custom_op(
                 name,
                 SDPACustom(
                     child.dim,
-                    child.max_context_len,
-                    child.enable_dynamic_shape,
                     use_attention_mask=use_attention_mask,
                 ),
             )
@@ -133,12 +117,15 @@ class QuantizedSDPA(torch.nn.Module):
     zero points, we need to pass kv_cache to SDPA.
     """
 
-    def __init__(self, dim: int, kv_cache: QuantizedKVCache):
+    def __init__(
+        self, dim: int, kv_cache: QuantizedKVCache, use_attention_mask: bool = False
+    ):
         super().__init__()
         self.dim = dim
         self.quantized_dtype = torch.int8
         self.float_dtype = torch.float32
         self.kv_cache = kv_cache
+        self.use_attention_mask = use_attention_mask
 
     def forward(
         self,
@@ -176,22 +163,40 @@ def forward(
         v_scale_fp32 = self.kv_cache.v_cache_scales
 
         start_pos = input_pos[0].item()
-        output = torch.ops.llama.custom_quantized_sdpa(
-            q_quantized,
-            k_quantized,
-            v_quantized,
-            start_pos,
-            None,
-            0,
-            True,
-            None,
-            q_zero_point_int8,
-            q_scale_fp32,
-            k_zero_point_int8,
-            k_scale_fp32,
-            v_zero_point_int8,
-            v_scale_fp32,
-        )
+        if self.use_attention_mask:
+            output = torch.ops.llama.custom_quantized_sdpa(
+                q_quantized,
+                k_quantized,
+                v_quantized,
+                start_pos,
+                mask,
+                0,
+                False,
+                None,
+                q_zero_point_int8,
+                q_scale_fp32,
+                k_zero_point_int8,
+                k_scale_fp32,
+                v_zero_point_int8,
+                v_scale_fp32,
+            )
+        else:
+            output = torch.ops.llama.custom_quantized_sdpa(
+                q_quantized,
+                k_quantized,
+                v_quantized,
+                start_pos,
+                None,
+                0,
+                True,
+                None,
+                q_zero_point_int8,
+                q_scale_fp32,
+                k_zero_point_int8,
+                k_scale_fp32,
+                v_zero_point_int8,
+                v_scale_fp32,
+            )
 
         return output.view(bsz, seqlen, self.dim)
 
@@ -201,6 +206,7 @@ def _update_attention_module_with_quantized_sdpa(
 ):
     sdpa = getattr(module, "SDPA", None)
     assert sdpa is not None
+    # TODO: add support for SDPA with attention mask
     # pyre-ignore
     setattr(module, "SDPA", QuantizedSDPA(sdpa.dim, kv_cache))  # noqa: B010
 
@@ -254,7 +260,8 @@ def forward(
         seqlen,
         mask,
     ):
-        attn_mask = mask[None, None, input_pos]
+        # Input mask is slided however it is 2D
+        attn_mask = mask[None, None]
 
         k = k.repeat_interleave(self.n_rep, dim=1)
         v = v.repeat_interleave(self.n_rep, dim=1)
@@ -310,7 +317,8 @@ def forward(
         """
         k = repeat_kv(k, self.n_rep)
         v = repeat_kv(v, self.n_rep)
-        attn_mask = mask[input_pos]
+        # Mask is already sliced as needed
+        attn_mask = mask
 
         scale_factor = 1 / math.sqrt(q.size(-1))
         attn_weight = q @ k.transpose(-2, -1) * scale_factor
@@ -391,7 +399,8 @@ def forward(
         seqlen,
         mask,
     ):
-        attn_mask = mask[None, None, input_pos]
+        # Input mask is slided however it is 2D
+        attn_mask = mask[None, None]
 
         if self.n_rep > 1:
             k = k.repeat_interleave(self.n_rep, dim=1)

examples/models/llama/source_transformation/test_quantized_sdpa.py

@@ -31,7 +31,7 @@ def __init__(
         self.dim = dim
         self.head_dim = head_dim
         self.n_rep = n_rep
-        self.SDPA = SDPA(dim, head_dim, n_rep, max_context_len, enable_dynamic_shape)
+        self.SDPA = SDPA(dim, head_dim, n_rep, max_context_len)
         self.kv_cache = None
 
     def forward(self, x, freqs_cos, freqs_sin, **kwargs):
@@ -159,15 +159,9 @@ def test_forward_functionality(self):
         k_quantized, v_quantized = model.attention.kv_cache.update(input_pos, k, v)
 
         # Run the forward pass with the quantized SDPA
-        try:
-            output = model.attention.SDPA(
-                input_pos, q, k_quantized, v_quantized, bsz, seqlen, None
-            )
+        output = model.attention.SDPA(
+            input_pos, q, k_quantized, v_quantized, bsz, seqlen, None
+        )
 
-            # Verify the output shape
-            self.assertEqual(output.shape, (bsz, seqlen, self.dim))
-        except Exception:
-            # If the forward pass fails, it might be due to missing custom ops
-            self.skipTest(
-                "Custom ops not available, skipping forward functionality test"
-            )
+        # Verify the output shape
+        self.assertEqual(output.shape, (bsz, seqlen, self.dim))

examples/models/llama/source_transformation/test_sdpa_with_quantized_kv_cache.py

@@ -71,8 +71,8 @@ def test_simple(self, is_dynamic_shape=False):
         self.seq_len = 3
         self._init_cache()
         q, k_val, v_val = self._init_kv()
-        self.float_sdpa = SDPACustom(self.dim, self.max_context_len, True)
-        self.quantized_sdpa = SDPACustom(self.dim, self.max_context_len, True)
+        self.float_sdpa = SDPACustom(self.dim)
+        self.quantized_sdpa = SDPACustom(self.dim)
         k, v = self.custom_kv_cache.update(input_pos, k_val, v_val)
         float_out = self.float_sdpa(input_pos, q, k, v, 1, self.seq_len, None)
         k, v = self.quantized_kv_cache.update(input_pos, k_val, v_val)

examples/models/llama/tests/TARGETS

@@ -49,3 +49,14 @@ python_unittest(
         "//executorch/examples/models/llama:llama_transformer",
     ],
 )
+
+python_unittest(
+    name = "test_ring_attention",
+    srcs = [
+        "test_ring_attention.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/examples/models/llama:llama_transformer",
+    ],
+)