Fix LGA cache devices and skipped-attention residuals (pytorch#22087)

Differential Revision: D117006258

Pull Request resolved: pytorch#22087

Author: YIWENX14

examples/models/llama/attention.py

@@ -131,7 +131,7 @@ def update(
             torch._check(start_pos < self.max_context_length)
             dim_to_slice = 2
             seq_length = k_val.size(dim_to_slice)
-            indices = torch.arange(seq_length) + start_pos
+            indices = torch.arange(seq_length, device=self.k_cache.device) + start_pos
             self.k_cache.index_copy_(dim_to_slice, indices, k_val)
             self.v_cache.index_copy_(dim_to_slice, indices, v_val)
             return self.k_cache, self.v_cache
@@ -181,7 +181,9 @@ def forward(
 def _create_causal_mask_for_ring_buffer(
     cache_positions, window_size, start_pos, seq_len
 ):
-    pos_q = start_pos + torch.arange(seq_len, dtype=torch.long).view(-1, 1)
+    pos_q = start_pos + torch.arange(
+        seq_len, dtype=torch.long, device=cache_positions.device
+    ).view(-1, 1)
     delta = pos_q - cache_positions
     attn_mask = (cache_positions >= 0) & (delta >= 0) & (delta < window_size)
     attn_mask = torch.where(attn_mask == True, 0, float("-inf"))  # noqa E712
@@ -239,11 +241,18 @@ def calculate_positions_and_update_indices(self, input_pos: torch.Tensor, seq_le
         """
         start_pos = input_pos[0].item()
         torch._check_is_size(start_pos)
-        orig_indices = torch.arange(seq_len, dtype=torch.long) + start_pos
+        device = self.cache_positions.device
+        orig_indices = (
+            torch.arange(seq_len, dtype=torch.long, device=device) + start_pos
+        )
         indices = orig_indices % self.max_context_length
 
-        full_t = torch.full((self.max_context_length,), -1, dtype=torch.long)
-        arange_tensor = torch.arange(self.max_context_length, dtype=torch.long)
+        full_t = torch.full(
+            (self.max_context_length,), -1, dtype=torch.long, device=device
+        )
+        arange_tensor = torch.arange(
+            self.max_context_length, dtype=torch.long, device=device
+        )
         cache_positions = torch.where(
             arange_tensor < start_pos, self.cache_positions, full_t
         )

examples/models/llama/llama_transformer.py

@@ -267,17 +267,22 @@ def from_type(cls, layer_id, args, rope) -> "TransformerBlock":
         attention = cls(args, layer_id, rope, **args.attention_kwargs)
         return TransformerBlock(args, attention, mlp_type=mlp_type, layer_id=layer_id)
 
+    def _apply_attention_residual(self, x, attention_output):
+        if isinstance(self.attention, AttentionSkip):
+            if not self.use_residual_gate:
+                return x
+            attention_output = torch.zeros_like(x)
+        if self.use_residual_gate:
+            if hasattr(self, "post_attn_norm"):
+                attention_output = self.post_attn_norm(attention_output)
+            return self.add_attn(stream=x, branch=attention_output)
+        return x + attention_output
+
     def forward(self, x, freqs_cos, freqs_sin, attn_options: ForwardOptions):  # x: 1xN
-        h, attn_options_update = self.attention(
+        attention_output, attn_options_update = self.attention(
             self.attention_norm(x), freqs_cos, freqs_sin, **attn_options
         )
-        if not isinstance(self.attention, AttentionSkip):
-            if self.use_residual_gate:
-                if hasattr(self, "post_attn_norm"):
-                    h = self.post_attn_norm(h)
-                h = self.add_attn(stream=x, branch=h)
-            else:
-                h = x + h
+        h = self._apply_attention_residual(x, attention_output)
 
         if self.mlp_type == "skip":
             out = h

examples/models/llama/tests/BUCK

@@ -67,6 +67,17 @@ fbcode_target(_kind = python_unittest,
     ],
 )
 
+fbcode_target(_kind = python_unittest,
+    name = "test_transformer_block",
+    srcs = [
+        "test_transformer_block.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/examples/models/llama:llama_transformer",
+    ],
+)
+
 fbcode_target(_kind = python_unittest,
     name = "test_ring_attention",
     srcs = [

examples/models/llama/tests/test_ring_kv_cache.py

@@ -7,7 +7,7 @@
 import unittest
 
 import torch
-from executorch.examples.models.llama.attention import RingKVCache
+from executorch.examples.models.llama.attention import KVCache, RingKVCache
 
 
 class TestRingKVCache(unittest.TestCase):
@@ -20,6 +20,80 @@ def setUp(self):
         self.enable_dynamic_shape = True
         self.dtype = torch.float32
 
+    def _require_usable_cuda(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA is not available")
+        try:
+            torch.zeros(1, device="cuda").cpu()
+        except (RuntimeError, torch.AcceleratorError) as error:
+            self.skipTest(f"CUDA kernels are not usable: {error}")
+
+    def test_dynamic_kv_cache_update_on_cuda(self):
+        self._require_usable_cuda()
+        cache = KVCache(
+            max_batch_size=1,
+            max_context_length=self.max_context_length,
+            n_heads=self.n_heads,
+            head_dim=self.head_dim,
+            enable_dynamic_shape=True,
+            dtype=self.dtype,
+        ).cuda()
+        input_pos = torch.tensor([2], dtype=torch.long, device="cuda")
+        k_val = torch.randn(
+            1, self.n_heads, 3, self.head_dim, device="cuda", dtype=self.dtype
+        )
+        v_val = torch.randn_like(k_val)
+
+        k_out, v_out = cache.update(input_pos, k_val, v_val)
+
+        self.assertEqual(k_out.device.type, "cuda")
+        self.assertEqual(v_out.device.type, "cuda")
+        torch.testing.assert_close(k_out[:, :, 2:5], k_val)
+        torch.testing.assert_close(v_out[:, :, 2:5], v_val)
+
+    def test_ring_cache_positions_and_mask_on_cuda(self):
+        self._require_usable_cuda()
+        cache = RingKVCache(
+            max_batch_size=1,
+            max_context_length=self.max_context_length,
+            n_heads=self.n_heads,
+            head_dim=self.head_dim,
+            enable_dynamic_shape=True,
+            dtype=self.dtype,
+        ).cuda()
+        input_pos = torch.tensor([0], dtype=torch.long, device="cuda")
+        seq_len = 3
+        k_val = torch.randn(
+            1,
+            self.n_heads,
+            seq_len,
+            self.head_dim,
+            device="cuda",
+            dtype=self.dtype,
+        )
+        v_val = torch.randn_like(k_val)
+
+        cache.update(input_pos, k_val, v_val)
+        mask = cache.create_causal_mask_for_ring_buffer(start_pos=0, seq_len=seq_len)
+
+        expected_positions = torch.tensor(
+            [0, 1, 2] + [-1] * 13, dtype=torch.long, device="cuda"
+        )
+        expected_mask = torch.full(
+            (seq_len, 16), float("-inf"), device="cuda", dtype=self.dtype
+        )
+        expected_mask[0, 0] = 0
+        expected_mask[1, :2] = 0
+        expected_mask[2, :3] = 0
+        self.assertEqual(
+            cache.cache_positions_manager.cache_positions.device.type, "cuda"
+        )
+        torch.testing.assert_close(
+            cache.cache_positions_manager.cache_positions, expected_positions
+        )
+        self.assertEqual(mask.device.type, "cuda")
+        torch.testing.assert_close(mask, expected_mask)
+
     def test_basic_update(self):
         """Test basic update functionality of RingKVCache."""
         cache = RingKVCache(

examples/models/llama/tests/test_transformer_block.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.examples.models.llama.attention import AttentionSkip
+from executorch.examples.models.llama.llama_transformer import TransformerBlock
+from executorch.examples.models.llama.model_args import ModelArgs
+
+
+class TestTransformerBlock(unittest.TestCase):
+    def test_attention_skip_applies_gated_zero_branch_residual(self) -> None:
+        args = ModelArgs(
+            dim=4,
+            hidden_dim=8,
+            n_heads=1,
+            n_kv_heads=1,
+            head_dim=4,
+            use_residual_gate=True,
+        )
+        block = TransformerBlock(
+            args,
+            AttentionSkip(),
+            mlp_type="skip",
+            layer_id=2,
+        ).eval()
+        x = torch.randn(1, 3, args.dim)
+        freqs = torch.empty(0)
+
+        output, update = block(x, freqs, freqs, {})
+        assert block.add_attn is not None
+        expected = block.add_attn(stream=x, branch=torch.zeros_like(x))
+
+        self.assertIsNone(update)
+        torch.testing.assert_close(output, expected)
+        self.assertFalse(torch.equal(output, x))