[Gated DeltaNet] optimize UT transform

fla/ops/common/chunk_scaled_dot_kkt.py

@@ -8,10 +8,12 @@
 import triton.language as tl
 
 from fla.ops.common.utils import prepare_chunk_indices
+from fla.ops.utils.op import safe_exp
 
 
 @triton.heuristics({
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+    'USE_G': lambda args: args['g_cumsum'] is not None
 })
 @triton.autotune(
     configs=[
@@ -26,7 +28,9 @@
 def chunk_scaled_dot_kkt_fwd_kernel(
     k,
     beta,
+    g_cumsum,
     A,
+    Ag,
     cu_seqlens,
     chunk_indices,
     T,
@@ -35,6 +39,7 @@ def chunk_scaled_dot_kkt_fwd_kernel(
     BT: tl.constexpr,
     BK: tl.constexpr,
     IS_VARLEN: tl.constexpr,
+    USE_G: tl.constexpr,
 ):
     i_t, i_bh = tl.program_id(0), tl.program_id(1)
     i_b, i_h = i_bh // H, i_bh % H
@@ -60,10 +65,19 @@ def chunk_scaled_dot_kkt_fwd_kernel(
     p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0))
     tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))
 
+    if USE_G:
+        p_g = tl.make_block_ptr(g_cumsum + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+        b_g = tl.load(p_g, boundary_check=(0,))
+        b_g_diff = b_g[:, None] - b_g[None, :]
+        b_Ag = b_A * safe_exp(b_g_diff)
+        p_Ag = tl.make_block_ptr(Ag + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0))
+        tl.store(p_Ag, b_Ag.to(p_Ag.dtype.element_ty), boundary_check=(0, 1))
+
 
 def chunk_scaled_dot_kkt_fwd(
     k: torch.Tensor,
     beta: torch.Tensor,
+    g_cumsum: Optional[torch.Tensor],
     cu_seqlens: Optional[torch.LongTensor],
     chunk_size: int = 64,
     output_dtype: torch.dtype = torch.float32
@@ -92,15 +106,18 @@ def chunk_scaled_dot_kkt_fwd(
     chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
     A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)
+    Ag = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) if g_cumsum is not None else None
     chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)](
         k=k,
         beta=beta,
+        g_cumsum=g_cumsum,
         A=A,
+        Ag=Ag,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
         T=T,
         H=H,
         K=K,
         BT=BT,
     )
-    return A
+    return A, Ag

fla/ops/delta_rule/chunk.py

@@ -5,7 +5,6 @@
 from typing import Optional
 
 import torch
-import triton
 from einops import rearrange
 
 from fla.modules.l2norm import l2norm_bwd, l2norm_fwd
@@ -24,38 +23,32 @@ def chunk_delta_rule_fwd(
     initial_state: torch.Tensor,
     output_final_state: bool,
     cu_seqlens: Optional[torch.LongTensor] = None,
-    chunk_size: int = 64
 ):
-    T = q.shape[1]
-    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
     # obtain WY representation. u is actually the new v.
     w, u, A = fwd_prepare_wy_repr(
         k=k,
         v=v,
         beta=beta,
         cu_seqlens=cu_seqlens,
-        chunk_size=BT
     )
-
     h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
         k=k,
         w=w,
         u=u,
         g=None,
         initial_state=initial_state,
         output_final_state=output_final_state,
-        cu_seqlens=cu_seqlens,
-        chunk_size=BT
+        cu_seqlens=cu_seqlens
     )
+
     o = chunk_fwd_o(
         q=q,
         k=k,
         v=v_new,
         h=h,
         g=None,
         scale=scale,
-        cu_seqlens=cu_seqlens,
-        chunk_size=BT
+        cu_seqlens=cu_seqlens
     )
     return o, A, final_state
 
@@ -71,17 +64,13 @@ def chunk_delta_rule_bwd(
     do: torch.Tensor,
     dht: torch.Tensor,
     cu_seqlens: Optional[torch.LongTensor] = None,
-    chunk_size: int = 64
 ):
-    T = q.shape[1]
-    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
     w, u = fwd_recompute_w_u(
         k=k,
         v=v,
         beta=beta,
         A=A,
-        cu_seqlens=cu_seqlens,
-        chunk_size=BT
+        cu_seqlens=cu_seqlens
     )
     h, v_new, _ = chunk_gated_delta_rule_fwd_h(
         k=k,
@@ -91,16 +80,14 @@ def chunk_delta_rule_bwd(
         initial_state=initial_state,
         output_final_state=False,
         cu_seqlens=cu_seqlens,
-        chunk_size=BT
     )
     dv = chunk_bwd_dv_local(
         q=q,
         k=k,
         do=do,
         g=None,
         scale=scale,
-        cu_seqlens=cu_seqlens,
-        chunk_size=BT
+        cu_seqlens=cu_seqlens
     )
     dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(
         q=q,
@@ -112,8 +99,7 @@ def chunk_delta_rule_bwd(
         do=do,
         dv=dv,
         scale=scale,
-        cu_seqlens=cu_seqlens,
-        chunk_size=BT
+        cu_seqlens=cu_seqlens
     )
     dq, dk, dw, _ = chunk_bwd_dqkwg(
         q=q,
@@ -126,8 +112,7 @@ def chunk_delta_rule_bwd(
         dh=dh,
         g=None,
         scale=scale,
-        cu_seqlens=cu_seqlens,
-        chunk_size=BT
+        cu_seqlens=cu_seqlens
     )
     dk2, dv, db = bwd_prepare_wy_repr(
         k=k,
@@ -136,8 +121,7 @@ def chunk_delta_rule_bwd(
         A=A,
         dw=dw,
         du=dv,
-        cu_seqlens=cu_seqlens,
-        chunk_size=BT
+        cu_seqlens=cu_seqlens
     )
     dk.add_(dk2)
     return dq, dk, dv, db, dh0
@@ -160,9 +144,6 @@ def forward(
         cu_seqlens: Optional[torch.LongTensor] = None,
         use_qk_l2norm_in_kernel: bool = True
     ):
-        T = q.shape[1]
-        chunk_size = min(64, max(triton.next_power_of_2(T), 16))
-
         q_orig = q
         k_orig = k
 
@@ -179,10 +160,8 @@ def forward(
             initial_state=initial_state,
             output_final_state=output_final_state,
             cu_seqlens=cu_seqlens,
-            chunk_size=chunk_size
         )
         ctx.save_for_backward(q_orig, k_orig, v, beta, A, initial_state)
-        ctx.chunk_size = chunk_size
         ctx.scale = scale
         ctx.cu_seqlens = cu_seqlens
         ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
@@ -212,8 +191,7 @@ def backward(
             initial_state=initial_state,
             do=do,
             dht=dht,
-            cu_seqlens=ctx.cu_seqlens,
-            chunk_size=ctx.chunk_size
+            cu_seqlens=ctx.cu_seqlens
         )
         if use_qk_l2norm_in_kernel:
             dq = l2norm_bwd(q_orig, dq)

fla/ops/delta_rule/wy_fast.py

@@ -9,7 +9,7 @@
 
 from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
 from fla.ops.common.utils import prepare_chunk_indices
-from fla.ops.utils.solve_tril import solve_tril
+from fla.ops.utils import solve_tril
 from fla.utils import check_shared_mem, is_nvidia_hopper
 
 NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8]
@@ -182,28 +182,26 @@ def fwd_prepare_wy_repr(
     v: torch.Tensor,
     beta: torch.Tensor,
     cu_seqlens: Optional[torch.LongTensor],
-    chunk_size: int = 64
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    A = chunk_scaled_dot_kkt_fwd(
+    A, _ = chunk_scaled_dot_kkt_fwd(
         k=k,
         beta=beta,
+        g_cumsum=None,
         cu_seqlens=cu_seqlens,
-        chunk_size=chunk_size,
-        output_dtype=torch.float32
+        output_dtype=torch.float32,
+        chunk_size=64,
     )
     A = solve_tril(
         A=A,
         cu_seqlens=cu_seqlens,
         output_dtype=k.dtype
     )
-
     w, u = fwd_recompute_w_u(
         k=k,
         v=v,
         beta=beta,
         A=A,
         cu_seqlens=cu_seqlens,
-        chunk_size=chunk_size
     )
     return w, u, A
 
@@ -214,10 +212,9 @@ def fwd_recompute_w_u(
     beta: torch.Tensor,
     A: torch.Tensor,
     cu_seqlens: Optional[torch.LongTensor],
-    chunk_size: int
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     B, T, H, K, V = *k.shape, v.shape[-1]
-    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+    BT = 64
     CONST_TILING = 64 if check_shared_mem() else 32
     BK = min(triton.next_power_of_2(K), CONST_TILING)
     BV = min(triton.next_power_of_2(V), CONST_TILING)
@@ -254,11 +251,10 @@ def bwd_prepare_wy_repr(
     A: torch.Tensor,
     dw: torch.Tensor,
     du: torch.Tensor,
-    cu_seqlens: Optional[torch.LongTensor],
-    chunk_size: int
+    cu_seqlens: Optional[torch.LongTensor]
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     B, T, H, K, V = *k.shape, v.shape[-1]
-    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+    BT = A.shape[-1]
     CONST_TILING = 64 if check_shared_mem() else 32
     BK = min(triton.next_power_of_2(K), CONST_TILING)
     BV = min(triton.next_power_of_2(V), CONST_TILING)