move the gating back onto the selected keys for improved differentiable topk, do it for the blocked values later as well

Author: lucidrains

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -616,10 +616,10 @@ def forward(
 
             selected_importance_values, selected_block_indices = importance_scores.topk(num_selected, dim = -1)
 
+            gates = None
+
             if self.use_diff_topk:
                 gates = straight_through(selected_importance_values, 1.)
-                gates = gates.cumprod(dim = -1)[..., -1]
-                gates = repeat(gates, 'b h ... -> b (h qh) ...', qh = fine_num_grouped_queries)
 
             if self.use_triton_kernel and not disable_triton_kernel:
 
@@ -631,10 +631,13 @@ def forward(
                     fq, fk, fv,
                     self.selection_block_size,
                     selected_block_indices,
-                    fmask
+                    fmask,
+                    sel_scale = gates
                 )
 
             elif exists(fine_selection_flex_mask):
+                assert not self.use_diff_topk, 'differential topk is not available for flex attention'
+
                 # flex attention for the selection for fine attention
 
                 fine_block_mask = fine_selection_flex_mask(selected_block_indices, num_grouped_queries = fine_num_grouped_queries)
@@ -654,9 +657,6 @@ def forward(
 
                     selected_block_indices = pad_at_dim(selected_block_indices, (0, remainder), value = 0, dim = -2)
 
-                    if self.use_diff_topk:
-                        gates = pad_at_dim(gates, (0, remainder), value = 1.)
-
                 # handle block causal diagonal in the diagram, but run experiments without to see
 
                 fine_window_seq = arange(fine_divisible_seq_len, device = device) // self.selection_block_size
@@ -690,6 +690,13 @@ def forward(
                 fk = fk.gather(3, selected_block_indices)
                 fv = fv.gather(3, selected_block_indices)
 
+                # differential topk gating
+
+                if self.use_diff_topk:
+                    fk = einx.multiply('b h i sel, b h i sel j d -> b h i sel j d', gates, fk)
+
+                # merge selected key values
+
                 fk, fv = tuple(rearrange(t, 'b h i w j d -> b h i (w j) d') for t in (fk, fv))
 
                 # fine attention
@@ -712,12 +719,6 @@ def forward(
 
                 fine_attn_out = fine_attn_out[..., :seq_len, :]
 
-            # handle maybe gating
-
-            if self.use_diff_topk:
-                gates = gates[..., :seq_len]
-                fine_attn_out = einx.multiply('b h n, b h n d -> b h n d', gates, fine_attn_out)
-
         else:
             # if only first block, just do a simple block causal
 

native_sparse_attention_pytorch/triton_native_sparse_attention.py

@@ -590,6 +590,7 @@ def backward_kernel_one_col_block_sparse(
     V,
     kv_block_indices,
     kv_block_mask,
+    kv_block_grads,
     DO,
     DQ,
     DK,
@@ -619,6 +620,7 @@ def backward_kernel_one_col_block_sparse(
     BLOCK: tl.constexpr,
     QUERY_HEAD_GROUPS: tl.constexpr,
     QUERY_EXPAND_DIM: tl.constexpr,
+    RETURN_SEL_GRADS: tl.constexpr,
     OFF_SEL_KV_BLOCKS: tl.constexpr
 ):
     # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N)
@@ -638,9 +640,6 @@ def backward_kernel_one_col_block_sparse(
 
     # initialize pointers to value-like data
 
-    k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])
-    v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])
-
     q_ptrs = (
         Q +
         offs_g[None, :, None] * stride_qh +
@@ -794,9 +793,9 @@ def backward_kernel_one_col_block_sparse(
     block_qk = block_qk.reshape(BLOCK, QUERY_HEAD_GROUPS, QUERY_EXPAND_DIM, BLOCK)
     qk = tl.sum(block_qk, 2) / QUERY_EXPAND_DIM
 
-    qk += tl.where(block_masks[:, None, None], 0, float("-inf"))
+    masked_qk = qk + tl.where(block_masks[:, None, None], 0, float("-inf"))
 
-    p = tl.exp(qk * softmax_scale - lse_i[:, :, None])
+    p = tl.exp(masked_qk * softmax_scale - lse_i[:, :, None])
 
     # take care of block dv
 
@@ -823,6 +822,26 @@ def backward_kernel_one_col_block_sparse(
 
     ds = (p * (dp - Di[:, :, None]) * softmax_scale)
 
+    # maybe return gradients for better differentiable topk
+
+    if RETURN_SEL_GRADS:
+
+        kv_block_grads_ptrs = (
+            kv_block_grads +
+            offs_m * stride_kvbl_m
+        )
+
+        sel_grads = ds * qk
+        sel_grads = tl.where(block_masks[:, None, None], sel_grads, 0.)
+        sel_grads = sel_grads.reshape(BLOCK, QUERY_HEAD_GROUPS * BLOCK)
+        sel_grads = tl.sum(sel_grads, 1)
+
+        tl.atomic_add(
+            kv_block_grads_ptrs + OFF_SEL_KV_BLOCKS,
+            sel_grads,
+            sem = 'relaxed'
+        )
+
     # block dk
 
     block_dk = ds[:, :, :, None] * q[:, :, None, :].to(ds.dtype)
@@ -1145,6 +1164,7 @@ def backward_kernel(
     V,
     kv_block_indices,
     kv_block_mask,
+    kv_block_grads,
     DO,
     DQ,
     DK,
@@ -1192,19 +1212,16 @@ def backward_kernel(
     BLOCK: tl.constexpr,
     QUERY_HEAD_GROUPS: tl.constexpr,
     QUERY_EXPAND_DIM: tl.constexpr,
+    RETURN_SEL_GRADS: tl.constexpr,
     INCLUDE_BLOCK_CAUSAL: tl.constexpr
 ):
     off_hb = tl.program_id(1)
     off_b = off_hb // kv_heads
     off_h = off_hb % kv_heads
     off_qh = off_h * QUERY_HEAD_GROUPS
 
-    if INCLUDE_BLOCK_CAUSAL:
-        IS_CAUSAL = tl.program_id(0) == 0
-        OFF_SEL_KV_BLOCKS = tl.program_id(0) - 1
-    else:
-        IS_CAUSAL = False
-        OFF_SEL_KV_BLOCKS = tl.program_id(0)
+    OFF_SEL_KV_BLOCKS = tl.program_id(0) - int(INCLUDE_BLOCK_CAUSAL)
+    IS_CAUSAL = INCLUDE_BLOCK_CAUSAL and tl.program_id(0) == 0
 
     # offset pointers for batch/head
 
@@ -1220,6 +1237,7 @@ def backward_kernel(
 
     kv_block_indices += off_b * stride_kvbl_b + off_h * stride_kvbl_h
     kv_block_mask += off_b * stride_kvbl_b + off_h * stride_kvbl_h
+    kv_block_grads += off_b * stride_kvbl_b + off_h * stride_kvbl_h
 
     # pointer to row-wise quantities in value-like data
 
@@ -1283,6 +1301,7 @@ def backward_kernel(
                 V,
                 kv_block_indices,
                 kv_block_mask,
+                kv_block_grads,
                 DO,
                 DQ,
                 DK,
@@ -1312,6 +1331,7 @@ def backward_kernel(
                 BLOCK = BLOCK,
                 QUERY_HEAD_GROUPS = QUERY_HEAD_GROUPS,
                 QUERY_EXPAND_DIM = QUERY_EXPAND_DIM,
+                RETURN_SEL_GRADS = RETURN_SEL_GRADS,
                 OFF_SEL_KV_BLOCKS = OFF_SEL_KV_BLOCKS
             )
 
@@ -1320,11 +1340,13 @@ def native_sparse_attn_backward(
     q, k, v,
     kv_block_indices,
     kv_block_mask,
+    kv_block_grads,
     o,
     lse,
     dq, dk, dv,
     block_size = 128,
-    include_block_causal = True
+    include_block_causal = True,
+    return_sel_grads = False
 ):
     device = do.device
 
@@ -1387,6 +1409,7 @@ def native_sparse_attn_backward(
         v,
         kv_block_indices,
         kv_block_mask,
+        kv_block_grads,
         do,
         dq,
         dk,
@@ -1436,6 +1459,7 @@ def native_sparse_attn_backward(
         EVEN_M = divisible_by(seqlen_q, block_size),
         EVEN_N = divisible_by(seqlen_k, block_size),
         EVEN_HEADDIM = BLOCK_HEADDIM == dim,
+        RETURN_SEL_GRADS = return_sel_grads,
         INCLUDE_BLOCK_CAUSAL = include_block_causal,
         # BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
         # num_warps=num_warps,
@@ -1458,6 +1482,7 @@ def forward(
         block_size,
         selected_block_indices,
         fmask,
+        sel_scale,
         include_block_causal
     ):
         dtype = fq.dtype
@@ -1478,10 +1503,16 @@ def forward(
 
         ctx.save_for_backward(fq, fk, fv, selected_block_indices, fmask, out, lse)
 
+        return_sel_grads = exists(sel_scale)
+
+        if return_sel_grads:
+            assert (sel_scale == 1.).all(), 'for now, must be straight through as multiplier of 1.'
+
         ctx._saved_variables = (
             block_size,
             head_groups,
-            include_block_causal
+            return_sel_grads,
+            include_block_causal,
         )
 
         return out.type(dtype), lse
@@ -1495,6 +1526,7 @@ def backward(self, ctx, do, _):
         (
             block_size,
             head_groups,
+            return_sel_grads,
             include_block_causal
         ) = ctx._saved_variables
 
@@ -1503,15 +1535,23 @@ def backward(self, ctx, do, _):
         dk = torch.zeros(k.shape, dtype = torch.float32, device = device)
         dv = torch.zeros(v.shape, dtype = torch.float32, device = device)
 
+        sel_grads = torch.zeros_like(sel_block_indices).float()
+
         native_sparse_attn_backward(
             do, q, k, v,
-            sel_block_indices, mask,
+            sel_block_indices, mask, sel_grads,
             out, lse, dq, dk, dv,
             block_size = block_size,
-            include_block_causal = include_block_causal
+            include_block_causal = include_block_causal,
+            return_sel_grads = return_sel_grads
         )
 
-        return dq, dk, dv, None, None, None, None
+        ret_sel_grads = None
+
+        if return_sel_grads:
+            ret_sel_grads = sel_grads
+
+        return dq, dk, dv, None, None, None, ret_sel_grads, None
 
 _native_sparse_attend = NSA.apply
 
@@ -1531,6 +1571,7 @@ def native_sparse_attend(
     block_size: int,
     selected_block_indices: Int['b qh n sel'] | Int['b kh n sel'],
     fmask: Bool['b qh n sel'] | Bool['b kh n sel'],
+    sel_scale: Float['b kh n sel'] | Float['b qh n sel'] | None = None,
     include_block_causal = True,
     return_lse = False
 ):
@@ -1550,6 +1591,7 @@ def native_sparse_attend(
         block_size,
         selected_block_indices,
         fmask,
+        sel_scale,
         include_block_causal
     )
 

pyproject.toml

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

test_triton_nsa.py

@@ -21,21 +21,24 @@ def regular_attend(
     indices,
     mask,
     block_size,
+    sel_scale = None, 
     return_lse = False
 ):
     q_heads, seq_len, kv_heads, device = q.shape[1], q.shape[-2], k.shape[1], q.device
     assert divisible_by(q_heads, kv_heads)
 
     q, k, v = tuple(pad_to_multiple(t, block_size, dim = -2) for t in (q, k, v))
 
+    if exists(sel_scale):
+        sel_scale = pad_to_multiple(sel_scale, block_size, dim = -2)
+
     g = q_heads // kv_heads # `g` stands for `g`roups of query heads per kv head
 
     w = ceil(seq_len / block_size)
 
     q, k, v = tuple(rearrange(t, 'b h (w n) d -> b h w n d', n = block_size) for t in (q, k, v))
 
     scale = q.shape[-1] ** -0.5
-    q = q * scale
 
     q = rearrange(q, 'b (h g) ... -> b h g ...', g = g)
 
@@ -62,6 +65,10 @@ def regular_attend(
 
         bsim = rearrange(bsim, 'b h g (w i) (sel j) -> b h g w i sel j', sel = num_sel_kv_blocks, i = fine_block_size)
 
+        if exists(sel_scale):
+            sel_scale = rearrange(sel_scale, 'b h (w i) sel -> b h w i sel', i = fine_block_size)
+            bsim = einx.multiply('b h g w i sel j, b h w i sel -> b h g w i sel j', bsim, sel_scale)
+
         mask = rearrange(mask, 'b h (w i) sel -> b h 1 w i sel', i = fine_block_size)
         bsim = torch.where(mask[..., None], bsim, -torch.finfo(bsim.dtype).max)
 
@@ -78,6 +85,7 @@ def regular_attend(
 
     # attend
 
+    sim = sim * scale
     attn = sim.softmax(dim = -1)
 
     if has_sel_kv_blocks:
@@ -113,27 +121,29 @@ def regular_attend(
 
 indices = torch.randint(0, 2, (batch, kv_heads, seq_len, num_sel)).cuda()
 mask = torch.randint(0, 2, (batch, kv_heads, seq_len, num_sel)).bool().cuda()
+sel_scale = torch.ones((batch, kv_heads, seq_len, num_sel)).cuda()
 
 # both regular and nsa pathways `r` and `n`
 
-rq, rk, rv = tuple(t.clone().requires_grad_() for t in (q, k, v))
-nq, nk, nv = tuple(t.clone().requires_grad_() for t in (q, k, v))
+rq, rk, rv, rsel_scale = tuple(t.clone().requires_grad_() for t in (q, k, v, sel_scale))
+nq, nk, nv, nsel_scale = tuple(t.clone().requires_grad_() for t in (q, k, v, sel_scale))
 
 # regular forwards and backwards
 
-out, rlse = regular_attend(rq, rk, rv, indices, mask, block_size = fine_block_size, return_lse = True)
+out, rlse = regular_attend(rq, rk, rv, indices, mask, block_size = fine_block_size, sel_scale = rsel_scale, return_lse = True)
 out.sum().backward()
 
 # triton nsa forwards and backwards
 
-nsa_out, nlse = native_sparse_attend(nq, nk, nv, fine_block_size, indices, mask, return_lse = True)
+nsa_out, nlse = native_sparse_attend(nq, nk, nv, fine_block_size, indices, mask, sel_scale = nsel_scale, return_lse = True)
 nsa_out.sum().backward()
 
 # asserts
 
 assert torch.allclose(out, nsa_out, atol = 1e-2)
 assert torch.allclose(rlse, nlse, atol = 1e-2)
 
+assert torch.allclose(rsel_scale.grad, nsel_scale.grad, atol = 1e-2)
 assert torch.allclose(nv.grad, rv.grad, atol = 1e-2)
 assert torch.allclose(nq.grad, rq.grad, atol = 1e-2)
 assert torch.allclose(nk.grad, rk.grad, atol = 1e-2)