fix flex sliding window attn

Author: lucidrains

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -41,18 +41,18 @@
 
 # flex attn sliding attention mask
 
+
 def create_sliding_mask(seq_len, window_size, causal = True):
+
     def sliding_mask(_, __, q_idx, kv_idx):
 
         distance = q_idx - kv_idx
-        mask = distance <= window_size
+        backward_sliding_mask = distance <= window_size
 
-        if causal:
-            mask = mask & q_idx >= kv_idx
-        else:
-            mask = mask & (distance >= -window_size)
+        forward_distance = 0 if causal else -window_size
+        forward_sliding_mask = distance >= forward_distance
 
-        return mask
+        return backward_sliding_mask & forward_sliding_mask
 
     block_mask = create_block_mask(sliding_mask, B = None, H = None, Q_LEN = seq_len, KV_LEN = seq_len, _compile = True)
     return block_mask

pyproject.toml

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