fixed, resuming with backwards gqa..

lucidrains · lucidrains · commit 6e52c4bddfe3 · 2025-02-26T16:15:44.000Z
diff --git a/native_sparse_attention_pytorch/triton_native_sparse_attention.py b/native_sparse_attention_pytorch/triton_native_sparse_attention.py
@@ -144,7 +144,7 @@ def forward_kernel(
 
     lse_ptrs = (
         Lse +
-        (off_hb + offs_lse_qh[:, None]) * seqlen_q_rounded +
+        offs_qh[:, None] * seqlen_q_rounded +
         offs_m[None, :]
     )
 
diff --git a/test_triton_nsa.py b/test_triton_nsa.py
@@ -93,7 +93,7 @@ def regular_attend(
 
 fine_block_size = 16
 
-q = torch.randn(1, 2, 512, 64).cuda()
+q = torch.randn(1, 4, 512, 64).cuda()
 k = torch.randn(1, 2, 512, 64).cuda()
 v = torch.randn(1, 2, 512, 64).cuda()
 

Original file line number	Diff line number	Diff line change
`@@ -144,7 +144,7 @@ def forward_kernel(`
`144`	`144`
`145`	`145`	`lse_ptrs = (`
`146`	`146`	`Lse +`
`147`		`- (off_hb + offs_lse_qh[:, None]) * seqlen_q_rounded +`
	`147`	`+ offs_qh[:, None] * seqlen_q_rounded +`
`148`	`148`	`offs_m[None, :]`
`149`	`149`	`)`
`150`	`150`