some intermittent issue with flex attention on sample, just disable and plan on writing own triton kernel

Author: lucidrains

native_sparse_attention_pytorch/transformer.py

@@ -221,7 +221,7 @@ def sample(
                 out,
                 cache = cache,
                 return_cache = True,
-                disable_flex = not is_first,
+                disable_flex = True,
                 disable_triton_kernel = not is_first
             )
 

pyproject.toml

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