[DeltaFormer] Fixed testing ops error (#602)

* [Stick-Breaking Attention] Add Model

* Revert "[Stick-Breaking Attention] Add Model"

This reverts commit db7411a.

* [deltaformer] fixed ops test error

* Test under fp16

* added varlen ops test & passed

* Update naive.py

---------

Co-authored-by: Yu Zhang <[email protected]>

Author: Nathancgy

fla/ops/deltaformer/naive.py

@@ -36,16 +36,34 @@ def tril_softmax(scores: torch.Tensor, strict: bool = True) -> torch.Tensor:
     return probs
 
 
-def naive_deltaformer_attn(
+def naive_causal_attention_bhtd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+) -> torch.Tensor:
+    B, H, T, D = q.shape
+    qk_scale = 1.0 / math.sqrt(D)
+    scores = torch.matmul(q, k.transpose(-1, -2)) * qk_scale  # [B, H, T, T]
+    causal_mask = torch.triu(torch.ones(T, T, device=q.device), diagonal=1).bool()
+    scores = scores.masked_fill(causal_mask, float('-inf'))
+    attn_weights = torch.softmax(scores, dim=-1)  # [B, H, T, T]
+    o = torch.matmul(attn_weights, v)  # [B, H, T, D]
+
+    return o
+
+
+def naive_deltaformer_attn_head_first(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
     beta: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """
-    Naive reference implementation of DeltaFormer pre-attention.
+    Naive reference implementation of DeltaFormer attention for head-first format.
 
-    Computes u[i] = v[i] - beta[i] * sum_{j<i} softmax(q[i] @ k[:i]^T) @ u[:i]
+    Two-stage process:
+    1. Computes u[i] = v[i] - beta[i] * sum_{j<i} softmax(q[i] @ k[:i]^T) @ u[:i]
+    2. Applies causal attention: o = causal_attn(q, k, u)
 
     Args:
         q: [B, H, T, D]
@@ -54,7 +72,7 @@ def naive_deltaformer_attn(
         beta: [B, H, T] or None (defaults to ones)
 
     Returns:
-        u: [B, H, T, D]
+        o: [B, H, T, D]
     """
     assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "q,k,v must be [B,H,T,D]"
     B, H, T, D = q.shape
@@ -83,8 +101,49 @@ def naive_deltaformer_attn(
             weighted_sum = (w.unsqueeze(-1) * u_prev).sum(dim=-2)  # [B,H,D]
             u_t = vf[:, :, t, :] - betaf[:, :, t].unsqueeze(-1) * weighted_sum
         u_list.append(u_t)
-    u = torch.stack(u_list, dim=2)
-    return u.to(orig_dtype)
+    u = torch.stack(u_list, dim=2)  # [B,H,T,D]
+
+    o = naive_causal_attention_bhtd(q, k, u.to(orig_dtype))
+    return o.to(orig_dtype)
+
+
+def naive_deltaformer_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Naive reference implementation of DeltaFormer attention for sequence-first format.
+
+    Args:
+        q: [B, T, H, D]
+        k: [B, T, H, D]
+        v: [B, T, H, D]
+        beta: [B, T, H] or None (defaults to ones)
+
+    Returns:
+        o: [B, T, H, D]
+    """
+    assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "q,k,v must be [B,T,H,D]"
+    B, T, H, D = q.shape
+    assert k.shape == (B, T, H, D) and v.shape == (B, T, H, D)
+
+    q_bhtd = q.transpose(1, 2)  # [B, T, H, D] -> [B, H, T, D]
+    k_bhtd = k.transpose(1, 2)  # [B, T, H, D] -> [B, H, T, D]
+    v_bhtd = v.transpose(1, 2)  # [B, T, H, D] -> [B, H, T, D]
+
+    if beta is not None:
+        assert beta.shape == (B, T, H)
+        beta_bhtd = beta.transpose(1, 2)  # [B, T, H] -> [B, H, T]
+    else:
+        beta_bhtd = None
+
+    o_bhtd = naive_deltaformer_attn_head_first(q_bhtd, k_bhtd, v_bhtd, beta_bhtd)
+
+    o_bthd = o_bhtd.transpose(1, 2)  # [B, H, T, D] -> [B, T, H, D]
+
+    return o_bthd
 
 
 __all__ = [

fla/ops/deltaformer/parallel.py

@@ -300,7 +300,7 @@ def parallel_deltaformer_bwd_kernel_u(
         block_shape=(BLOCK_C, D),
         order=(1, 0),
     )
-    q = tl.load(q_blk_ptr)
+    q = tl.load(q_blk_ptr, boundary_check=(0,))
 
     for kv_i in range(0, T, BLOCK_T):
         k_blk_ptr = tl.make_block_ptr(
@@ -311,7 +311,7 @@ def parallel_deltaformer_bwd_kernel_u(
             block_shape=(D, BLOCK_T),
             order=(0, 1),
         )
-        k = tl.load(k_blk_ptr)
+        k = tl.load(k_blk_ptr, boundary_check=(1,))
         qk = tl.dot(q, k) * fa_scale
 
         lse_blk_ptr = tl.make_block_ptr(
@@ -322,7 +322,7 @@ def parallel_deltaformer_bwd_kernel_u(
             block_shape=(BLOCK_T,),
             order=(0,),
         )
-        lse = tl.load(lse_blk_ptr)
+        lse = tl.load(lse_blk_ptr, boundary_check=(0,))
         beta_blk_ptr = tl.make_block_ptr(
             base=beta_ptr + pid_h,
             shape=(T,),
@@ -331,7 +331,7 @@ def parallel_deltaformer_bwd_kernel_u(
             block_shape=(BLOCK_T,),
             order=(0,),
         )
-        beta = tl.load(beta_blk_ptr)
+        beta = tl.load(beta_blk_ptr, boundary_check=(0,))
 
         p = tl.math.exp2(qk - lse[None, :]) * beta[None, :]
 
@@ -343,7 +343,7 @@ def parallel_deltaformer_bwd_kernel_u(
             block_shape=(BLOCK_T, D),
             order=(1, 0),
         )
-        v = tl.load(v_blk_ptr)
+        v = tl.load(v_blk_ptr, boundary_check=(0,))
         acc = tl.dot(p.to(v_ptr.dtype.element_ty), v, acc)
 
     o_blk_ptr = tl.make_block_ptr(
@@ -354,7 +354,7 @@ def parallel_deltaformer_bwd_kernel_u(
         block_shape=(BLOCK_C, D),
         order=(1, 0),
     )
-    tl.store(o_blk_ptr, acc.to(o_ptr.dtype.element_ty))
+    tl.store(o_blk_ptr, acc.to(o_ptr.dtype.element_ty), boundary_check=(0,))
 
 
 @triton.autotune(configs=_config_deltaformer(), key=['T', 'D'])
@@ -389,7 +389,7 @@ def parallel_deltaformer_bwd_kernel_row_sum(
         block_shape=(BLOCK_C, D),
         order=(1, 0),
     )
-    k_row = tl.load(k_row_blk_ptr)
+    k_row = tl.load(k_row_blk_ptr, boundary_check=(0,))
     lse_blk_ptr = tl.make_block_ptr(
         base=lse_ptr + pid_h,
         shape=(T,),
@@ -398,7 +398,7 @@ def parallel_deltaformer_bwd_kernel_row_sum(
         block_shape=(BLOCK_C,),
         order=(0,),
     )
-    lse = tl.load(lse_blk_ptr)
+    lse = tl.load(lse_blk_ptr, boundary_check=(0,))
     grad_v_blk_ptr = tl.make_block_ptr(
         base=grad_v_ptr + pid_h * D,
         shape=(T, D),
@@ -407,7 +407,7 @@ def parallel_deltaformer_bwd_kernel_row_sum(
         block_shape=(BLOCK_C, D),
         order=(1, 0),
     )
-    grad_v_row = -tl.load(grad_v_blk_ptr)
+    grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,))
 
     for kv_i in range(0, (pid_c + 1) * BLOCK_C, BLOCK_T):
         k_blk_ptr = tl.make_block_ptr(
@@ -418,7 +418,7 @@ def parallel_deltaformer_bwd_kernel_row_sum(
             block_shape=(D, BLOCK_T),
             order=(0, 1),
         )
-        k = tl.load(k_blk_ptr)
+        k = tl.load(k_blk_ptr, boundary_check=(1,))
         qk = tl.dot(k_row, k) * fa_scale
         p = tl.math.exp2(qk - lse[:, None])
 
@@ -430,7 +430,7 @@ def parallel_deltaformer_bwd_kernel_row_sum(
             block_shape=(D, BLOCK_T),
             order=(0, 1),
         )
-        ut = tl.load(u_blk_ptr)
+        ut = tl.load(u_blk_ptr, boundary_check=(1,))
         dp = tl.dot(grad_v_row, ut)
         if kv_i + BLOCK_T >= pid_c * BLOCK_C:
             mask = (rowid_block[:, None] <= colid_block[None, :] + kv_i)
@@ -445,7 +445,7 @@ def parallel_deltaformer_bwd_kernel_row_sum(
         block_shape=(BLOCK_C,),
         order=(0,),
     )
-    tl.store(row_dot_block_ptr, acc)
+    tl.store(row_dot_block_ptr, acc, boundary_check=(0,))
 
 
 @triton.autotune(configs=[triton.Config({'BLOCK_C': BC}, num_stages=ns, num_warps=nw)
@@ -484,7 +484,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         block_shape=(BLOCK_C, D),
         order=(1, 0),
     )
-    k_row = tl.load(k_row_blk_ptr)
+    k_row = tl.load(k_row_blk_ptr, boundary_check=(0,))
     lse_blk_ptr = tl.make_block_ptr(
         base=lse_ptr + pid_h,
         shape=(T,),
@@ -493,7 +493,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         block_shape=(BLOCK_C,),
         order=(0,),
     )
-    lse = tl.load(lse_blk_ptr)
+    lse = tl.load(lse_blk_ptr, boundary_check=(0,))
     beta_blk_ptr = tl.make_block_ptr(
         base=beta_ptr + pid_h,
         shape=(T,),
@@ -502,7 +502,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         block_shape=(BLOCK_C,),
         order=(0,),
     )
-    beta = tl.load(beta_blk_ptr)
+    beta = tl.load(beta_blk_ptr, boundary_check=(0,))
     grad_v_blk_ptr = tl.make_block_ptr(
         base=grad_v_ptr + pid_h * D,
         shape=(T, D),
@@ -511,7 +511,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         block_shape=(BLOCK_C, D),
         order=(1, 0),
     )
-    grad_v_row = -tl.load(grad_v_blk_ptr)
+    grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,))
     row_dot_blk_ptr = tl.make_block_ptr(
         base=row_dot_ptr + pid_h,
         shape=(T,),
@@ -520,7 +520,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         block_shape=(BLOCK_C,),
         order=(0,),
     )
-    row_dot_row = tl.load(row_dot_blk_ptr).to(k_ptr.dtype.element_ty)
+    row_dot_row = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty)
 
     for kv_i in range(0, pid_c * BLOCK_C, BLOCK_C):
         k_blk_ptr = tl.make_block_ptr(
@@ -531,7 +531,7 @@ def parallel_deltaformer_bwd_kernel_qk(
             block_shape=(D, BLOCK_C),
             order=(0, 1),
         )
-        kt = tl.load(k_blk_ptr)
+        kt = tl.load(k_blk_ptr, boundary_check=(1,))
         qk = tl.dot(k_row, kt) * fa_scale
         p = tl.math.exp2(qk - lse[:, None]) * beta[:, None]
 
@@ -557,7 +557,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         block_shape=(BLOCK_C, D),
         order=(1, 0),
     )
-    k_row_true = tl.load(k_row_blk_ptr)
+    k_row_true = tl.load(k_row_blk_ptr, boundary_check=(0,))
     qk = tl.dot(k_row, tl.trans(k_row_true, 1, 0)) * fa_scale
     p = tl.math.exp2(qk - lse[:, None]) * beta[:, None]
     u_blk_ptr = tl.make_block_ptr(
@@ -587,7 +587,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         order=(1, 0),
     )
     acc = acc * qk_scale
-    tl.store(grad_q_blk_ptr, acc.to(grad_q_ptr.dtype.element_ty))
+    tl.store(grad_q_blk_ptr, acc.to(grad_q_ptr.dtype.element_ty), boundary_check=(0,))
 
     daat = tl.trans(da, 1, 0)
     acc = tl.dot(daat.to(k_row.dtype), k_row)
@@ -602,7 +602,7 @@ def parallel_deltaformer_bwd_kernel_qk(
             block_shape=(D, BLOCK_C),
             order=(0, 1),
         )
-        kt = tl.load(k_blk_ptr)
+        kt = tl.load(k_blk_ptr, boundary_check=(1,))
         lse_blk_ptr = tl.make_block_ptr(
             base=lse_ptr + pid_h,
             shape=(T,),
@@ -611,7 +611,7 @@ def parallel_deltaformer_bwd_kernel_qk(
             block_shape=(BLOCK_C,),
             order=(0,),
         )
-        lse = tl.load(lse_blk_ptr)
+        lse = tl.load(lse_blk_ptr, boundary_check=(0,))
         beta_blk_ptr = tl.make_block_ptr(
             base=beta_ptr + pid_h,
             shape=(T,),
@@ -620,7 +620,7 @@ def parallel_deltaformer_bwd_kernel_qk(
             block_shape=(BLOCK_C,),
             order=(0,),
         )
-        beta = tl.load(beta_blk_ptr)
+        beta = tl.load(beta_blk_ptr, boundary_check=(0,))
         qk = tl.dot(k_row, kt) * fa_scale
         p = tl.math.exp2(qk - lse[None, :]) * beta[None, :]
 
@@ -632,7 +632,7 @@ def parallel_deltaformer_bwd_kernel_qk(
             block_shape=(D, BLOCK_C),
             order=(0, 1),
         )
-        grad_vt = tl.load(grad_vt_blk_ptr)
+        grad_vt = tl.load(grad_vt_blk_ptr, boundary_check=(1,))
         row_dot_blk_ptr = tl.make_block_ptr(
             base=row_dot_ptr + pid_h,
             shape=(T,),
@@ -641,7 +641,7 @@ def parallel_deltaformer_bwd_kernel_qk(
             block_shape=(BLOCK_C,),
             order=(0,),
         )
-        row_dot = tl.load(row_dot_blk_ptr).to(k_ptr.dtype.element_ty)
+        row_dot = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty)
         dp = tl.dot(nu, grad_vt)
         da = p * (dp - row_dot[None, :])
         k = tl.trans(kt, 1, 0)
@@ -656,7 +656,7 @@ def parallel_deltaformer_bwd_kernel_qk(
         order=(1, 0),
     )
     acc = acc * qk_scale
-    tl.store(grad_k_blk_ptr, acc.to(grad_k_ptr.dtype.element_ty))
+    tl.store(grad_k_blk_ptr, acc.to(grad_k_ptr.dtype.element_ty), boundary_check=(0,))
 
 
 class ParallelDeltaformerFunction(torch.autograd.Function):
@@ -872,6 +872,7 @@ def _forward_impl(
                     w_t = w.transpose(0, 1).contiguous()
                     u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous()
                     invcum.forward_inplace(u_chunk_view_t, w_t)
+                    u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1))
 
                 chunk_base += (T_max + C - 1) // C
 
@@ -932,6 +933,7 @@ def _forward_impl(
                 w_t = w.transpose(0, 1).contiguous()
                 u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous()
                 invcum.forward_inplace(u_chunk_view_t, w_t)
+                u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1))
 
             chunk_base += (L + C - 1) // C
 
@@ -953,7 +955,7 @@ def deltaformer_attn(
     B, T, H, D = k.shape
     C = min(C, T)
 
-    u = ParallelDeltaformerFunction.apply(k, k, v, beta, C, cu_seqlens)
+    u = ParallelDeltaformerFunction.apply(q, k, v, beta, C, cu_seqlens)
 
     if attention_mask is not None:
         q_padded, (k_padded, u_padded), indices_q, cu_seqlens_lens, max_seq_lens = unpad_input(q, (k, u), attention_mask, T)