first handle when sequence length has number of compressed windows less than the number of selected blocks, but it still breaks for no blocks (seq length less than compress block size)

Author: lucidrains

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -243,7 +243,9 @@ def forward(
 
         importance_scores = cattn[..., num_mem_compress_kv:]
 
-        selected_importance_values, selected_block_indices = importance_scores.topk(self.num_selected_blocks, dim = -1)
+        topk = min(self.num_selected_blocks, importance_scores.shape[-1])
+
+        selected_importance_values, selected_block_indices = importance_scores.topk(topk, dim = -1)
 
         if self.use_diff_topk:
             gates = selected_importance_values + (1. - selected_importance_values).detach()

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "native-sparse-attention-pytorch"
-version = "0.0.15"
+version = "0.0.16"
 description = "Native Sparse Attention"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_sparse_attn.py

@@ -7,8 +7,10 @@
 from native_sparse_attention_pytorch import SparseAttention
 
 @pytest.mark.parametrize('use_diff_topk', (False, True))
+@pytest.mark.parametrize('seq_len', (4, 31, 32, 120))
 def test_sparse_attn(
-    use_diff_topk
+    use_diff_topk,
+    seq_len
 ):
     attn = SparseAttention(
         dim = 512,
@@ -21,7 +23,7 @@ def test_sparse_attn(
         use_diff_topk = use_diff_topk
     )
 
-    tokens = torch.randn(2, 31, 512)
+    tokens = torch.randn(2, seq_len, 512)
 
     attended = attn(tokens)