gqa works for main block causal, but broken for fine selection pathway, hack on it another day

Author: lucidrains

native_sparse_attention_pytorch/transformer.py

@@ -183,7 +183,7 @@ def forward(
         ids,
         return_loss = False,
         disable_flex = False,
-        disable_triton_kernel = True
+        disable_triton_kernel = False
     ):
         if return_loss:
             ids, labels = ids[:, :-1], ids[:, 1:]

native_sparse_attention_pytorch/triton_native_sparse_attention.py

@@ -802,32 +802,50 @@ def backward_kernel_one_col_block(
         block_k = tl.load(block_k_ptrs)
         block_v = tl.load(block_v_ptrs)
 
-        q_expanded = tl.expand_dims(q, 1)
-        q_expanded = tl.broadcast_to(q_expanded, (BLOCK, 16, BLOCK_HEADDIM))
+        q_expanded = q.reshape(QUERY_HEAD_GROUPS, BLOCK, BLOCK_HEADDIM)
+        q_expanded = q_expanded.permute(1, 0, 2)
+        q_expanded = tl.expand_dims(q_expanded, 2)
+        q_expanded = tl.broadcast_to(q_expanded, (BLOCK, QUERY_HEAD_GROUPS, QUERY_EXPAND_DIM, BLOCK_HEADDIM))
+        q_expanded = q_expanded.reshape(BLOCK, 16, BLOCK_HEADDIM)
 
         block_k_permuted = tl.permute(block_k, (0, 2, 1))
         block_qk = tl.dot(q_expanded, block_k_permuted)
 
-        qk = tl.sum(block_qk, 1) / 16.
-        qk += tl.where(block_masks[:, None], 0, float("-inf"))
+        block_qk = block_qk.reshape(BLOCK, QUERY_HEAD_GROUPS, QUERY_EXPAND_DIM, BLOCK)
+        qk = tl.sum(block_qk, 2) / QUERY_EXPAND_DIM
+        qk = qk.permute(1, 0, 2)
+
+        qk += tl.where(block_masks[None, :, None], 0, float("-inf"))
+
+        qk = qk.reshape(QUERY_HEAD_GROUPS * BLOCK, BLOCK)
 
         p = tl.exp(qk * softmax_scale - lse_i[:, None])
 
         # take care of block dv
 
         block_dv = p.to(do.dtype)[:, :, None] * do[:, None, :]
-        block_dv = tl.where(block_masks[:, None, None], block_dv, 0.)
+
+        block_dv = block_dv.reshape(QUERY_HEAD_GROUPS, BLOCK, BLOCK, BLOCK_HEADDIM)
+        block_dv = tl.sum(block_dv, 0)
 
         tl.atomic_add(block_dv_ptrs, block_dv, sem = 'relaxed')
 
         # get dp
 
-        do_expanded = tl.expand_dims(do, 1)
-        do_expanded = tl.broadcast_to(do_expanded, (BLOCK, 16, BLOCK_HEADDIM))
+        do_expanded = do.reshape(QUERY_HEAD_GROUPS, BLOCK, BLOCK_HEADDIM)
+        do_expanded = do_expanded.permute(1, 0, 2)
+        do_expanded = tl.expand_dims(do_expanded, 2)
+        do_expanded = tl.broadcast_to(do_expanded, (BLOCK, QUERY_HEAD_GROUPS, QUERY_EXPAND_DIM, BLOCK_HEADDIM))
+        do_expanded = do_expanded.reshape(BLOCK, 16, BLOCK_HEADDIM)
+
         block_v = tl.permute(block_v, (0, 2, 1))
 
         dp = tl.dot(do_expanded, block_v)
-        dp = tl.sum(dp, 1) / 16.
+
+        dp = dp.reshape(BLOCK, QUERY_HEAD_GROUPS, QUERY_EXPAND_DIM, BLOCK)
+        dp = tl.sum(dp, 2) / QUERY_EXPAND_DIM
+        dp = dp.permute(1, 0, 2)
+        dp = dp.reshape(QUERY_HEAD_GROUPS * BLOCK, BLOCK)
 
         # ds
 
@@ -837,15 +855,25 @@ def backward_kernel_one_col_block(
         # block dk
 
         block_dk = ds[:, :, None] * q[:, None, :]
+        block_dk = block_dk.reshape(QUERY_HEAD_GROUPS, BLOCK, BLOCK, BLOCK_HEADDIM)
+        block_dk = tl.sum(block_dk, 0)
 
         tl.atomic_add(block_dk_ptrs, block_dk, sem = 'relaxed')
 
         # block dq
 
-        ds_expanded = tl.expand_dims(ds, 1)
-        ds_expanded = tl.broadcast_to(ds_expanded, (BLOCK, 16, BLOCK))
+        ds_expanded = ds.reshape(QUERY_HEAD_GROUPS, BLOCK, BLOCK)
+        ds_expanded = ds_expanded.permute(1, 0, 2)
+        ds_expanded = tl.expand_dims(ds_expanded, 2)
+        ds_expanded = tl.broadcast_to(ds_expanded, (BLOCK, QUERY_HEAD_GROUPS, QUERY_EXPAND_DIM, BLOCK))
+        ds_expanded = ds_expanded.reshape(BLOCK, 16, BLOCK)
+
         block_dq = tl.dot(ds_expanded, block_k)
-        block_dq = tl.sum(block_dq, 1) / 16
+
+        block_dq = block_dq.reshape(BLOCK, QUERY_HEAD_GROUPS, QUERY_EXPAND_DIM, BLOCK_HEADDIM)
+        block_dq = tl.sum(block_dq, 2) / QUERY_EXPAND_DIM
+        block_dq = block_dq.permute(1, 0, 2)
+        block_dq = block_dq.reshape(QUERY_HEAD_GROUPS * BLOCK, BLOCK_HEADDIM)
 
         dq += block_dq
 
@@ -1194,9 +1222,7 @@ def backward(self, ctx, do, _):
             out, lse, dq, dk, dv,
             block_size = block_size
         )
-
-        dk, dv = tuple(reduce(t, 'b (h g) ... -> b h ...', 'sum', g = head_groups) for t in (dk, dv))
-
+    
         return dq, dk, dv, None, None, None, None
 
 _native_sparse_attend = NSA.apply
@@ -1208,6 +1234,8 @@ def native_sparse_attend(
     fmask,
     return_lse = False
 ):
+    assert divisible_by(fq.shape[-2], block_size)
+
     out, lse = _native_sparse_attend(
         fq, fk, fv,
         block_size,

test_triton_nsa.py

@@ -93,12 +93,12 @@ def regular_attend(
 
 fine_block_size = 16
 
-q = torch.randn(1, 2, 512, 64).cuda()
-k = torch.randn(1, 2, 512, 64).cuda()
-v = torch.randn(1, 2, 512, 64).cuda()
+q = torch.randn(2, 4, 512, 64).cuda()
+k = torch.randn(2, 2, 512, 64).cuda()
+v = torch.randn(2, 2, 512, 64).cuda()
 
-indices = torch.zeros(1, 2, 512, 1).long().cuda()
-mask = torch.ones(1, 2, 512, 1).bool().cuda()
+indices = torch.zeros(2, 2, 512, 0).long().cuda()
+mask = torch.randint(0, 2, (2, 2, 512, 0)).bool().cuda()
 
 # both regular and nsa pathways `r` and `n`
 

train_triton_nsa.py

@@ -25,24 +25,25 @@
 LEARNING_RATE = 1e-4
 VALIDATE_EVERY = 100
 PRIME_LENGTH = 64
+SHOULD_GENERATE = False
 GENERATE_EVERY = 500
 GENERATE_LENGTH = 512
 SEQ_LEN = 512
 HEADS = 8
-KV_HEADS = 8
+KV_HEADS = 4
 
 USE_SPARSE_ATTN = True
 USE_TRITON_NSA = True
-USE_FLEX_FOR_FINE_SELECTION = False   # will push flex a bit, won't be efficient as each layer needs sparsity dynmically generated, but may be enough just to compare to full attention before going all-in on triton kernels
-QUERY_HEADS_SHARE_SELECTION = False  # if set to False, each query head can look at a different segment of their corresponding key / value head in GQA
+USE_FLEX_FOR_FINE_SELECTION = False  # will push flex a bit, won't be efficient as each layer needs sparsity dynmically generated, but may be enough just to compare to full attention before going all-in on triton kernels
+QUERY_HEADS_SHARE_SELECTION = True   # if set to False, each query head can look at a different segment of their corresponding key / value head in GQA
 
 # sparse attention related
 
 SLIDING_WINDOW_SIZE = 32
 COMPRESS_BLOCK_SIZE = 16
 
 FINE_BLOCK_SIZE = 16
-NUM_FINE_SELECTED = 1
+NUM_FINE_SELECTED = 0
 
 INTERPOLATED_IMPORTANCE_SCORE = False
 USE_DIFF_TOPK = True
@@ -211,7 +212,7 @@ def __getitem__(self, index):
             wandb.log(dict(valid_loss = loss.item()), step = i)
             print(f"validation loss: {loss.item():.3f}")
 
-    if i % GENERATE_EVERY == 0:
+    if SHOULD_GENERATE and i % GENERATE_EVERY == 0:
         model.eval()
 
         inp = random.choice(val_dataset)[:PRIME_LENGTH]