allow for ablating fine block selection for negative control

Author: lucidrains

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -277,7 +277,11 @@ def __init__(
 
         self.selection_block_size = selection_block_size
 
-        assert num_selected_blocks > 0
+        assert num_selected_blocks >= 0
+
+        if num_selected_blocks == 0:
+            print(f'`num_selected_blocks` should be set greater than 0, unless if you are ablating it for experimental purposes')
+
         self.num_selected_blocks = num_selected_blocks
 
         # they combine the three sparse branches through a learned combine with sigmoid activation

pyproject.toml

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

tests/test_sparse_attn.py

@@ -10,12 +10,14 @@
 @pytest.mark.parametrize('seq_len', (1, 4, 31, 32, 120))
 @pytest.mark.parametrize('kv_heads', (8, 4))
 @pytest.mark.parametrize('selection_block_size', (8, 4, 2))
+@pytest.mark.parametrize('num_selected_block', (0, 2))
 @pytest.mark.parametrize('query_heads_share_selected_kv', (False, True))
 def test_sparse_attn(
     use_diff_topk,
     seq_len,
     kv_heads,
     selection_block_size,
+    num_selected_block,
     query_heads_share_selected_kv
 ):
     attn = SparseAttention(
@@ -26,7 +28,7 @@ def test_sparse_attn(
         sliding_window_size = 2,
         compress_block_size = 4,
         selection_block_size = selection_block_size,
-        num_selected_blocks = 2,
+        num_selected_blocks = num_selected_block,
         use_diff_topk = use_diff_topk,
         query_heads_share_selected_kv = query_heads_share_selected_kv
     )