Update DSV4 sparse attention

Author: yinli-systems

tests/unit_tests/cpu/test_deepseek_v4_dsa.py

@@ -12,12 +12,16 @@
 the previous index-based formulation for every compression ratio.
 """
 
+import inspect
 import unittest
 
 import torch
 from torch.nn.attention.flex_attention import BlockMask
 
-from torchtitan.models.deepseek_v4.attention import DSV4FlexInnerAttention
+from torchtitan.models.deepseek_v4.attention import (
+    CompressedSparseAttention,
+    DSV4FlexInnerAttention,
+)
 from torchtitan.models.deepseek_v4.compressor import Indexer
 
 
@@ -92,6 +96,107 @@ def build_dsa(ratio, window_size, block_size=128):
 
 
 class TestDSABlockMask(unittest.TestCase):
+    def test_csa_gather_attention_reference_forward_and_backward(self):
+        torch.manual_seed(0)
+        seqlen, n_heads, head_dim = 16, 2, 8
+        ratio, index_topk = 4, 3
+        cfg = CompressedSparseAttention.Config(
+            block_size=8,
+            window_size=5,
+            compress_ratio=ratio,
+            softmax_scale=head_dim**-0.5,
+            index_topk=index_topk,
+        )
+        attention = CompressedSparseAttention(cfg)
+        inputs = [
+            torch.randn(seqlen, n_heads, head_dim, requires_grad=True),
+            torch.randn(seqlen, head_dim, requires_grad=True),
+            torch.randn(seqlen // ratio, head_dim, requires_grad=True),
+            torch.randn(seqlen, 3, 6, requires_grad=True),
+            torch.randn(seqlen // ratio, 6, requires_grad=True),
+            torch.randn(seqlen, 3, requires_grad=True),
+            torch.randn(n_heads, requires_grad=True),
+        ]
+
+        output = attention(*inputs)
+        self.assertEqual(output.shape, (seqlen, n_heads, head_dim))
+        self.assertTrue(torch.isfinite(output).all())
+        output.square().mean().backward()
+        for tensor in (*inputs[:3], inputs[-1]):
+            self.assertIsNotNone(tensor.grad)
+            self.assertTrue(torch.isfinite(tensor.grad).all())
+
+    def test_static_mask_does_not_capture_token_pair_tensor(self):
+        seqlen = 32768
+        dsa = build_dsa(ratio=128, window_size=2048, block_size=128)
+        block_mask = dsa.build_block_mask(seqlen=seqlen, device="cpu")
+
+        closure = inspect.getclosurevars(block_mask.mask_mod)
+        captured_tensors = [
+            value
+            for value in closure.nonlocals.values()
+            if isinstance(value, torch.Tensor)
+        ]
+        self.assertEqual(captured_tensors, [])
+
+        metadata = [
+            block_mask.kv_num_blocks,
+            block_mask.kv_indices,
+            block_mask.q_num_blocks,
+            block_mask.q_indices,
+        ]
+        metadata_bytes = sum(x.numel() * x.element_size() for x in metadata)
+        self.assertLess(metadata_bytes, 1024 * 1024)
+
+    def test_static_masks_match_index_formulation(self):
+        device = torch.device("cpu")
+        cases = [
+            (128, 17, 0, 32),
+            (128, 63, 1, (32, 16)),
+            (256, 64, 128, 32),
+            (512, 128, 128, (64, 32)),
+        ]
+        for seqlen, window_size, ratio, block_size in cases:
+            with self.subTest(
+                seqlen=seqlen,
+                window_size=window_size,
+                ratio=ratio,
+                block_size=block_size,
+            ):
+                n_cmp = seqlen // ratio if ratio > 1 else 0
+                sink_idx = seqlen + n_cmp
+                win = window_idxs(window_size, 1, seqlen, device)
+                compress = (
+                    compress_idxs(ratio, 1, seqlen, device, seqlen)
+                    if ratio > 1
+                    else torch.empty((1, seqlen, 0), dtype=torch.int64)
+                )
+                selected = (
+                    torch.cat([win, compress], dim=-1) if compress.size(-1) else win
+                )
+
+                dsa = build_dsa(ratio, window_size, block_size)
+                block_mask = dsa.build_block_mask(seqlen=seqlen, device=device)
+                expected = old_attended(selected, sink_idx)
+                actual = new_attended(block_mask, seqlen, n_cmp)
+                self.assertEqual(expected, actual)
+
+                bq, bk = (
+                    block_size
+                    if isinstance(block_size, tuple)
+                    else (block_size, block_size)
+                )
+                for q_block in range(seqlen // bq):
+                    q_slice = expected[q_block * bq : (q_block + 1) * bq]
+                    expected_blocks = {
+                        kv_idx // bk for attended in q_slice for kv_idx in attended
+                    }
+                    num_blocks = int(block_mask.kv_num_blocks[0, 0, q_block])
+                    actual_blocks = set(
+                        block_mask.kv_indices[0, 0, q_block, :num_blocks].tolist()
+                    )
+                    self.assertEqual(expected_blocks, actual_blocks)
+
     def test_attended_sets_match_old_formulation(self):
         torch.manual_seed(0)
         device = torch.device("cpu")

tests/unit_tests/gpu/test_deepseek_v4_attention_cache.py

@@ -0,0 +1,133 @@
+# 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.
+
+"""Compare cached and freshly built DeepSeek V4 SWA/HCA attention masks."""
+
+import unittest
+from unittest import mock
+
+import torch
+
+from torchtitan.models.deepseek_v4 import model_registry
+from torchtitan.models.deepseek_v4.attention import (
+    dsv4_mask_key,
+    DSV4FlexInnerAttention,
+)
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable")
+class TestDeepSeekV4AttentionCache(unittest.TestCase):
+    def _make_model(self):
+        # Only mask construction and parameter-free inner attention are used.
+        with torch.device("meta"):
+            model = model_registry("debugmodel", seq_len=512).model.build()
+        return model
+
+    def _assert_outputs_and_grads_match(self, inner, cached_mask, seqlen):
+        device = torch.device("cuda")
+        num_heads, head_dim = 2, 64
+        inputs = [
+            torch.randn(seqlen, num_heads, head_dim, device=device, requires_grad=True),
+            torch.randn(seqlen, head_dim, device=device, requires_grad=True),
+        ]
+        if inner.compress_ratio > 1:
+            inputs.append(
+                torch.randn(
+                    seqlen // inner.compress_ratio,
+                    head_dim,
+                    device=device,
+                    requires_grad=True,
+                )
+            )
+        inputs.append(torch.randn(num_heads, device=device, requires_grad=True))
+        reference_inputs = [x.detach().clone().requires_grad_(True) for x in inputs]
+
+        fresh_mask = DSV4FlexInnerAttention.build_block_mask(
+            inner, seqlen=seqlen, device=device
+        )
+        self.assertIsNot(cached_mask, fresh_mask)
+        cached_output = inner(*inputs, attention_masks=cached_mask)
+        reference_output = inner(*reference_inputs, attention_masks=fresh_mask)
+        grad_output = torch.randn_like(cached_output)
+        cached_grads = torch.autograd.grad(cached_output, inputs, grad_output)
+        reference_grads = torch.autograd.grad(
+            reference_output, reference_inputs, grad_output
+        )
+
+        torch.testing.assert_close(
+            cached_output, reference_output, atol=1e-5, rtol=1e-5
+        )
+        names = ["q", "swa_k"]
+        if inner.compress_ratio > 1:
+            names.append("cmp_k")
+        names.append("attn_sink")
+        for name, actual, expected in zip(
+            names, cached_grads, reference_grads, strict=True
+        ):
+            torch.testing.assert_close(
+                actual, expected, atol=1e-5, rtol=1e-5, msg=f"grad[{name}] mismatch"
+            )
+
+    def test_cached_masks_reused_across_sequence_lengths(self):
+        model = self._make_model()
+        masks_by_length = {}
+        build_mask = DSV4FlexInnerAttention.build_block_mask
+
+        for seqlen, expected_builds in [(256, 2), (256, 0), (512, 2), (256, 0)]:
+            with self.subTest(seqlen=seqlen, expected_builds=expected_builds):
+                positions = torch.arange(seqlen, device="cuda")
+                with mock.patch.object(
+                    DSV4FlexInnerAttention,
+                    "build_block_mask",
+                    autospec=True,
+                    side_effect=build_mask,
+                ) as builder:
+                    masks = model.get_attention_masks(positions)
+                self.assertEqual(builder.call_count, expected_builds)
+                self.assertEqual(set(masks), {"swa", "hca_128"})
+
+                if seqlen in masks_by_length:
+                    for key, mask in masks.items():
+                        self.assertIs(mask, masks_by_length[seqlen][key])
+                else:
+                    for other_masks in masks_by_length.values():
+                        for key, mask in masks.items():
+                            self.assertIsNot(mask, other_masks[key])
+                    masks_by_length[seqlen] = masks
+                self.assertEqual(len(model.mask_cache), 2 * len(masks_by_length))
+
+    def test_cached_masks_match_fresh_outputs_and_grads(self):
+        torch.manual_seed(0)
+        model = self._make_model()
+        seqlen = 128
+        masks = model.get_attention_masks(torch.arange(seqlen, device="cuda"))
+        checked = set()
+        for layer in model.layers.values():
+            inner = layer.attention.inner_attention
+            key = dsv4_mask_key(inner.compress_ratio)
+            if key is not None and key not in checked:
+                with self.subTest(mask=key):
+                    self._assert_outputs_and_grads_match(inner, masks[key], seqlen)
+                checked.add(key)
+        self.assertEqual(checked, {"swa", "hca_128"})
+
+    def test_rejects_incompatible_layers_before_deduplication(self):
+        for field, value in [("window_size", 32), ("block_size", 64)]:
+            for populate_cache in (False, True):
+                with self.subTest(field=field, populate_cache=populate_cache):
+                    model = self._make_model()
+                    positions = torch.arange(256, device="cuda")
+                    if populate_cache:
+                        model.get_attention_masks(positions)
+                    inner = model.layers["1"].attention.inner_attention
+                    self.assertNotEqual(getattr(inner, field), value)
+                    setattr(inner, field, value)
+                    with self.assertRaises(AssertionError):
+                        model.get_attention_masks(positions)
+
+
+if __name__ == "__main__":
+    unittest.main()