diff --git a/examples/models/llama/attention.py b/examples/models/llama/attention.py index 10861a4189d..98a838bfa30 100644 --- a/examples/models/llama/attention.py +++ b/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 ) diff --git a/examples/models/llama/llama_transformer.py b/examples/models/llama/llama_transformer.py index b893516aa5f..73117826708 100644 --- a/examples/models/llama/llama_transformer.py +++ b/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 diff --git a/examples/models/llama/tests/BUCK b/examples/models/llama/tests/BUCK index 563c2f95435..74430d9e306 100644 --- a/examples/models/llama/tests/BUCK +++ b/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 = [ diff --git a/examples/models/llama/tests/test_ring_kv_cache.py b/examples/models/llama/tests/test_ring_kv_cache.py index 0a10e39cf16..f923a1bf2bf 100644 --- a/examples/models/llama/tests/test_ring_kv_cache.py +++ b/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( diff --git a/examples/models/llama/tests/test_transformer_block.py b/examples/models/llama/tests/test_transformer_block.py new file mode 100644 index 00000000000..0dbcd1b60fe --- /dev/null +++ b/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))