make inference line up with selection block size > compress block size

Author: lucidrains

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -461,16 +461,12 @@ def forward_inference(
         importance_scores = csim[..., self.num_mem_compress_kv:]
 
         num_compress_blocks = importance_scores.shape[-1]
+        num_compress_per_fine = self.selection_block_size // self.compress_block_size
 
         if self.compress_block_size != self.selection_block_size:
-            compress_seq_len = num_compress_blocks * self.compress_block_size
-
-            importance_scores = repeat(importance_scores, '... j -> ... (bsz j)', bsz = self.compress_block_size)
-
-            fine_seq_len = round_down_mult(compress_seq_len, self.selection_block_size)
-
-            importance_scores = importance_scores[..., :fine_seq_len]
-            importance_scores = reduce(importance_scores, '... (bsz j) -> ... j', 'mean', bsz = self.selection_block_size)
+            compress_seq_len = round_down_mult(num_compress_blocks, num_compress_per_fine)
+            importance_scores = importance_scores[..., :compress_seq_len]
+            importance_scores = reduce(importance_scores, '... (j num_compress_per_fine) -> ... j', 'mean', num_compress_per_fine = num_compress_per_fine)
 
         num_fine_blocks = importance_scores.shape[-1]
         num_selected = min(self.num_selected_blocks, num_fine_blocks)
@@ -490,7 +486,9 @@ def forward_inference(
             if self.query_heads_share_selected_kv:
                 importance_scores = reduce(importance_scores, 'b (h grouped_queries) ... -> b h ...', 'mean', grouped_queries = self.num_grouped_queries)
 
+            importance_scores = F.pad(importance_scores, (1, 0), value = -1e3)
             importance_scores = importance_scores.softmax(dim = -1)
+            importance_scores = importance_scores[..., 1:]
 
             sel_scores, sel_indices = importance_scores.topk(num_selected, dim = -1)
 
@@ -689,26 +687,24 @@ def forward(
 
             if self.compress_block_size != self.selection_block_size:
 
-                compress_seq_len = num_compress_blocks * self.compress_block_size
-
-                importance_scores = repeat(importance_scores, '... j -> ... (j block_size)', block_size = self.compress_block_size)
+                num_compress_per_fine = self.selection_block_size // self.compress_block_size
 
-                padding = fine_divisible_seq_len - compress_seq_len
+                round_down_score_len = round_down_mult(importance_scores.shape[-1], num_compress_per_fine)
+                importance_scores = importance_scores[..., :round_down_score_len]
 
-                fine_query_seq_len = importance_scores.shape[-2]
-                fine_query_padding = fine_divisible_seq_len - importance_scores.shape[-2]
+                if not is_empty(importance_scores):
+                    importance_scores = reduce(importance_scores, '... (j num_compress_per_fine) -> ... j', 'mean', num_compress_per_fine = num_compress_per_fine)
 
-                importance_scores = F.pad(importance_scores, (0, padding))
+                    i, j = importance_scores.shape[-2:]
 
-                # mask out the diagonal since block causal is included by default for fine attending
+                    # mask out block diagonal
 
-                block_causal_mask = torch.ones((num_fine_blocks,) * 2, device = device, dtype = torch.bool).tril(-1)
-                block_causal_mask = repeat(block_causal_mask, 'i j -> (i n1) (j n2)', n1 = self.selection_block_size, n2 = self.selection_block_size)
-                block_causal_mask = block_causal_mask[:fine_query_seq_len]
+                    q_seq = arange(i, device = device) // self.selection_block_size
+                    k_seq = arange(j, device = device)
 
-                importance_scores = importance_scores.masked_fill(~block_causal_mask, max_neg_value(csim))
+                    block_diagonal_mask = einx.equal('i, j -> i j', q_seq, k_seq)
 
-                importance_scores = reduce(importance_scores, '... (j block_size) -> ... j', 'mean', block_size = self.selection_block_size)
+                    importance_scores = importance_scores.masked_fill(block_diagonal_mask, max_neg_value(csim))
 
             importance_scores = F.pad(importance_scores, (1, 0), value = -1e3)
             importance_scores = importance_scores.softmax(dim = -1)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "native-sparse-attention-pytorch"
-version = "0.1.23"
+version = "0.1.24"
 description = "Native Sparse Attention"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

tests/test_sparse_attn.py

@@ -50,10 +50,12 @@ def test_sparse_attn(
 @pytest.mark.parametrize('seq_len', (2, 8, 16))
 @pytest.mark.parametrize('num_selected_blocks', (0, 2))
 @pytest.mark.parametrize('compress_block_overlap_len', (0, 2))
+@pytest.mark.parametrize('selection_block_size', (5, 10, 15))
 def test_inference(
     seq_len,
     num_selected_blocks,
-    compress_block_overlap_len
+    compress_block_overlap_len,
+    selection_block_size
 ):
 
     attn = SparseAttention(
@@ -63,7 +65,7 @@ def test_inference(
         causal = True,
         sliding_window_size = 2,
         compress_block_size = 5,
-        selection_block_size = 10,
+        selection_block_size = selection_block_size,
         num_selected_blocks = num_selected_blocks,
         compress_block_overlap_len = compress_block_overlap_len
     )