more guards

Author: lucidrains

native_sparse_attention_pytorch/triton_native_sparse_attention.py

@@ -935,8 +935,17 @@ def backward_kernel_one_col_block_sparse(
         offs_m * stride_kvbl_m
     )
 
-    block_indices = tl.load(kv_block_indices_ptrs + OFF_SEL_KV_BLOCKS)
-    block_masks = tl.load(kv_block_mask_ptrs + OFF_SEL_KV_BLOCKS)
+    block_indices = tl.load(
+        kv_block_indices_ptrs + OFF_SEL_KV_BLOCKS,
+        mask = offs_m < seqlen_q,
+        other = 0.
+    )
+
+    block_masks = tl.load(
+        kv_block_mask_ptrs + OFF_SEL_KV_BLOCKS,
+        mask = offs_m < seqlen_q,
+        other = 0.
+    )
 
     blocks_offs_n = (
         block_indices[:, None] * BLOCK +

pyproject.toml

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