Merge branch 'main' of https://github.com/fla-org/native-sparse-attention into main

Author: Hanyuezhuohua

README.md

@@ -36,16 +36,11 @@ To test the correctness of NSA:
 pytest tests/test_nsa.py
 ```
 
-To validate the correctness of NSA with top‑k selection (ignoring the output from the compressed attention), run the command below. Please note that the initial trial may take some time as the kernel compiles, but subsequent runs will be faster.
+To validate the correctness of NSA with top‑k selection, run the command below. Please note that the initial trial may take some time as the kernel compiles, but subsequent runs will be faster.
 ```py
 pytest tests/test_nsa_with_compression.py
 ```
 
-To verify the correctness of the top‑k selection, where sampling Q and K from a uniform distribution produces similar importance scores (resulting in slight variations in the top‑k selection), we validate this component separately. To run the test, execute:
-```py
-pytest tests/test_topk.py
-```
-
 To measure the efficiency of NSA:
 ```py
 python benchmark/benchmark_nsa.py

native_sparse_attention/ops/parallel.py

@@ -1367,7 +1367,7 @@ def parallel_nsa_bwd(
     dk = dk.sum(0)
     return dq, dk, dv
 
-
+@torch.compile
 class ParallelNSAFunction(torch.autograd.Function):
 
     @staticmethod
@@ -1448,7 +1448,6 @@ def parallel_nsa_compression(
     )
 
 
-@torch.compile
 def parallel_nsa(
     q: torch.Tensor,
     k: torch.Tensor,

tests/test_nsa.py

@@ -52,9 +52,12 @@ def test_parallel(
     torch.manual_seed(42)
     os.environ['TRITON_F32_DEFAULT'] = 'ieee'
 
-    q = torch.randn((B, T, HQ, D), dtype=dtype, device='cuda').requires_grad_(True)
-    k = torch.randn((B, T, H, D), dtype=dtype, device='cuda').requires_grad_(True)
-    v = torch.randn((B, T, H, D), dtype=dtype, device='cuda').requires_grad_(True)
+    perm_q = torch.randperm(T, device='cuda')
+    perm_k = torch.randperm(T, device='cuda')
+    perm_v = torch.randperm(T, device='cuda')
+    q = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda')[perm_q].view(1, T, 1, 1).expand(B, T, HQ, D).clone().requires_grad_(True)
+    k = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda')[perm_k].view(1, T, 1, 1).expand(B, T, H, D).clone().requires_grad_(True)
+    v = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda')[perm_v].view(1, T, 1, 1).expand(B, T, H, D).clone().requires_grad_(True)
     g_slc = torch.rand((B, T, HQ), dtype=dtype, device='cuda').requires_grad_(True)
     g_swa = torch.rand((B, T, HQ), dtype=dtype, device='cuda').requires_grad_(True)
     do = torch.randn((B, T, HQ, D), dtype=dtype, device='cuda')
@@ -147,9 +150,12 @@ def test_parallel_varlen(
         torch.tensor([T], dtype=torch.long)
     ], 0).cuda().sort()[0]
     # seq-first required for inputs with variable lengths
-    q = torch.randn((1, T, HQ, D), dtype=dtype, device='cuda').requires_grad_(True)
-    k = torch.randn((1, T, H, D), dtype=dtype, device='cuda').requires_grad_(True)
-    v = torch.randn((1, T, H, D), dtype=dtype, device='cuda').requires_grad_(True)
+    perm_q = torch.randperm(T, device='cuda')
+    perm_k = torch.randperm(T, device='cuda')
+    perm_v = torch.randperm(T, device='cuda')
+    q = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda')[perm_q].view(1, T, 1, 1).expand(1, T, HQ, D).clone().requires_grad_(True)
+    k = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda')[perm_k].view(1, T, 1, 1).expand(1, T, H, D).clone().requires_grad_(True)
+    v = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda')[perm_v].view(1, T, 1, 1).expand(1, T, H, D).clone().requires_grad_(True)
     g_slc = torch.rand((1, T, HQ), dtype=dtype, device='cuda').requires_grad_(True)
     g_swa = torch.rand((1, T, HQ), dtype=dtype, device='cuda').requires_grad_(True)
     do = torch.randn((1, T, HQ, D), dtype=dtype, device='cuda')

tests/test_nsa_with_compression.py

@@ -5,7 +5,7 @@
 import pytest
 import torch
 
-from native_sparse_attention.ops.naive import naive_nsa
+from native_sparse_attention.ops.naive import naive_nsa_with_compression
 from native_sparse_attention.ops.parallel import parallel_nsa_with_compression
 
 
@@ -50,9 +50,13 @@ def test_parallel(
     torch.manual_seed(42)
     os.environ['TRITON_F32_DEFAULT'] = 'ieee'
 
-    q = torch.randn((B, T, HQ, D), dtype=dtype, device='cuda').requires_grad_(True)
-    k = torch.randn((B, T, H, D), dtype=dtype, device='cuda').requires_grad_(True)
-    v = torch.randn((B, T, H, D), dtype=dtype, device='cuda').requires_grad_(True)
+
+    perm_q = torch.randperm(T, device='cuda')
+    perm_k = torch.randperm(T, device='cuda')
+    perm_v = torch.randperm(T, device='cuda')
+    q = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda').view(1, T, 1, 1).expand(B, T, HQ, D).clone().requires_grad_(True)
+    k = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda').view(1, T, 1, 1).expand(B, T, H, D).clone().requires_grad_(True)
+    v = torch.linspace(0, 1, steps=T, dtype=dtype, device='cuda').view(1, T, 1, 1).expand(B, T, H, D).clone().requires_grad_(True)
     g_cmp = torch.rand((B, T, HQ), dtype=dtype, device='cuda').requires_grad_(True)
     g_slc = torch.rand((B, T, HQ), dtype=dtype, device='cuda').requires_grad_(True)
     g_swa = torch.rand((B, T, HQ), dtype=dtype, device='cuda').requires_grad_(True)
@@ -77,27 +81,31 @@ def test_parallel(
     tri_dk, k.grad = k.grad.clone(), None
     tri_dv, v.grad = v.grad.clone(), None
     tri_dg_slc, g_slc.grad = g_slc.grad.clone(), None
+    tri_dg_cmp, g_cmp.grad = g_cmp.grad.clone(), None
     if window_size > 0:
         tri_dg_swa, g_swa.grad = g_swa.grad.clone(), None
 
-    ref = naive_nsa(
+    ref, ref_topk = naive_nsa_with_compression(
         q=q,
         k=k,
         v=v,
+        g_cmp=g_cmp,
         g_slc=g_slc,
         g_swa=g_swa,
-        block_indices=tri_topk,
         block_counts=block_counts,
         block_size=block_size,
         window_size=window_size,
         scale=scale
     )
 
+    print((ref_topk != tri_topk[:, :, :, :ref_topk.shape[-1]]).float().mean())
+
     ref.backward(do)
     ref_dq, q.grad = q.grad.clone(), None
     ref_dk, k.grad = k.grad.clone(), None
     ref_dv, v.grad = v.grad.clone(), None
     ref_dg_slc, g_slc.grad = g_slc.grad.clone(), None
+    ref_dg_cmp, g_cmp.grad = g_cmp.grad.clone(), None
     if window_size > 0:
         ref_dg_swa, g_swa.grad = g_swa.grad.clone(), None
 
@@ -106,6 +114,7 @@ def test_parallel(
     assert_close("dk", ref_dk, tri_dk, 0.005)
     assert_close("dv", ref_dv, tri_dv, 0.005)
     assert_close("dg_slc", ref_dg_slc, tri_dg_slc, 0.005)
+    assert_close("dg_cmp", ref_dg_cmp, tri_dg_cmp, 0.005)
     if window_size > 0:
         assert_close("dg_swa", ref_dg_swa, tri_dg_swa, 0.005)