[GDN] Add exp2 support across chunk kernels for improved performance

fla/ops/comba/chunk.py

@@ -1,6 +1,5 @@
 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
 
-
 import torch
 
 from fla.modules.l2norm import l2norm_bwd, l2norm_fwd
@@ -9,6 +8,7 @@
 from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
 from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o
 from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices, solve_tril
+from fla.ops.utils.constant import RCP_LN2
 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
 
 
@@ -25,7 +25,13 @@ def chunk_comba_fwd(
     cu_seqlens: torch.LongTensor | None = None,
     chunk_indices: torch.LongTensor | None = None,
 ):
-    g0, g = chunk_comba_cumsum_scalar_fwd(g, chunk_size=64, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices)
+    g0, g = chunk_comba_cumsum_scalar_fwd(
+        g,
+        chunk_size=64,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        scale=RCP_LN2,
+    )
     # obtain WY representation. u is actually the new v.
     A = chunk_scaled_dot_comba_pkt_fwd(
         k=k,
@@ -36,6 +42,7 @@ def chunk_comba_fwd(
         cu_seqlens=cu_seqlens,
         output_dtype=torch.float32,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     A = solve_tril(
         A=A,
@@ -51,6 +58,7 @@ def chunk_comba_fwd(
         g_cumsum=g0,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
         k=k,
@@ -61,6 +69,7 @@ def chunk_comba_fwd(
         output_final_state=output_final_state,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     o = chunk_fwd_o(
         q=q,
@@ -71,6 +80,7 @@ def chunk_comba_fwd(
         scale=scale,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     return g0, g, o, A, final_state
 
@@ -99,6 +109,7 @@ def chunk_comba_bwd(
         g_cumsum=g0,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     h, v_new, _ = chunk_gated_delta_rule_fwd_h(
         k=k,
@@ -109,6 +120,7 @@ def chunk_comba_bwd(
         output_final_state=False,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     dv = chunk_bwd_dv_local(
         q=q,
@@ -118,6 +130,7 @@ def chunk_comba_bwd(
         scale=scale,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(
         q=q,
@@ -131,6 +144,7 @@ def chunk_comba_bwd(
         scale=scale,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     dq, dk, dw, dg = chunk_bwd_dqkwg(
         q=q,
@@ -145,6 +159,7 @@ def chunk_comba_bwd(
         scale=scale,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     dk2, dv, dp, db, dg0, dg2 = prepare_wy_repr_bwd(
         k=k,
@@ -158,6 +173,7 @@ def chunk_comba_bwd(
         du=dv,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        use_exp2=True,
     )
     dk.add_(dk2)
     dg.add_(dg2)
@@ -362,3 +378,11 @@ def chunk_comba(
         cu_seqlens_cpu,
     )
     return o, final_state
+        cu_seqlens,
+        cu_seqlens_cpu,
+    )
+    return o, final_state
+        cu_seqlens,
+        cu_seqlens_cpu,
+    )
+    return o, final_state

fla/ops/comba/fused_recurrent.py

@@ -1,11 +1,10 @@
 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
 
-
 import torch
 import triton
 import triton.language as tl
 
-from fla.ops.utils.op import exp
+from fla.ops.utils.op import exp, exp2
 from fla.utils import input_guard
 
 
@@ -40,6 +39,7 @@ def fused_recurrent_comba_fwd_kernel(
     IS_BETA_HEADWISE: tl.constexpr,  # whether beta is headwise vector or scalar,
     USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
     IS_VARLEN: tl.constexpr,
+    USE_EXP2: tl.constexpr,
 ):
     i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
     i_n, i_hv = i_nh // HV, i_nh % HV
@@ -89,7 +89,10 @@ def fused_recurrent_comba_fwd_kernel(
         # [BV]
         b_v -= tl.sum(b_h * b_p[:, None], 0)
         # [BK, BV]
-        b_h *= exp(b_g)
+        if USE_EXP2:
+            b_h *= exp2(b_g)
+        else:
+            b_h *= exp(b_g)
         if IS_BETA_HEADWISE:
             b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
         else:
@@ -126,6 +129,7 @@ def fused_recurrent_comba_fwd(
     output_final_state: bool,
     use_qk_l2norm_in_kernel: bool = False,
     cu_seqlens: torch.LongTensor | None = None,
+    use_exp2: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     B, T, H, K, V = *k.shape, v.shape[-1]
     HV = v.shape[2]
@@ -165,6 +169,7 @@ def fused_recurrent_comba_fwd(
         BV=BV,
         IS_BETA_HEADWISE=beta.ndim == v.ndim,
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        USE_EXP2=use_exp2,
         num_warps=num_warps,
         num_stages=num_stages,
     )

fla/ops/comba/utils.py

@@ -8,6 +8,7 @@
 
 
 @triton.heuristics({
+    'HAS_SCALE': lambda args: args['scale'] is not None,
     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
 })
 @triton.autotune(
@@ -23,14 +24,15 @@ def chunk_comba_cumsum_scalar_fwd_kernel(
     g,
     g0,
     g1,
+    scale,
     cu_seqlens,
     chunk_indices,
     T,
     B: tl.constexpr,
     H: tl.constexpr,
     BT: tl.constexpr,
+    HAS_SCALE: tl.constexpr,
     IS_VARLEN: tl.constexpr,
-    HEAD_FIRST: tl.constexpr,
 ):
     i_t, i_bh = tl.program_id(0), tl.program_id(1)
     i_b, i_h = i_bh // H, i_bh % H
@@ -41,16 +43,13 @@ def chunk_comba_cumsum_scalar_fwd_kernel(
     else:
         bos, eos = i_b * T, i_b * T + T
 
-    if HEAD_FIRST:
-        p_g = tl.make_block_ptr(g + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
-        p_g0 = tl.make_block_ptr(g0 + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
-        p_g1 = tl.make_block_ptr(g1 + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
-    else:
-        p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
-        p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
-        p_g1 = tl.make_block_ptr(g1 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    p_g1 = tl.make_block_ptr(g1 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
     # [BT]
     b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32)
+    if HAS_SCALE:
+        b_g = b_g * scale
     b_g1 = tl.cumsum(b_g, axis=0)
     b_g0 = b_g1 - b_g
     tl.store(p_g0, b_g0.to(p_g0.dtype.element_ty), boundary_check=(0,))
@@ -61,14 +60,11 @@ def chunk_comba_cumsum_scalar_fwd(
     g: torch.Tensor,
     chunk_size: int,
     cu_seqlens: torch.Tensor | None = None,
-    head_first: bool = False,
     output_dtype: torch.dtype | None = torch.float,
     chunk_indices: torch.LongTensor | None = None,
+    scale: float | None = None,
 ) -> torch.Tensor:
-    if head_first:
-        B, H, T = g.shape
-    else:
-        B, T, H = g.shape
+    B, T, H = g.shape
     assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2"
     BT = chunk_size
     if chunk_indices is None and cu_seqlens is not None:
@@ -80,13 +76,13 @@ def chunk_comba_cumsum_scalar_fwd(
         g,
         g0,
         g1,
+        scale,
         cu_seqlens,
         chunk_indices,
         T=T,
         B=B,
         H=H,
         BT=BT,
-        HEAD_FIRST=head_first,
     )
     return g0, g1
 
@@ -113,7 +109,6 @@ def chunk_comba_cumsum_scalar_bwd_kernel(
     H: tl.constexpr,
     BT: tl.constexpr,
     IS_VARLEN: tl.constexpr,
-    HEAD_FIRST: tl.constexpr,
 ):
     i_t, i_bh = tl.program_id(0), tl.program_id(1)
     i_b, i_h = i_bh // H, i_bh % H
@@ -124,12 +119,8 @@ def chunk_comba_cumsum_scalar_bwd_kernel(
     else:
         bos, eos = i_b * T, i_b * T + T
 
-    if HEAD_FIRST:
-        p_dg0 = tl.make_block_ptr(dg0 + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
-        p_dgr = tl.make_block_ptr(dgr + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
-    else:
-        p_dg0 = tl.make_block_ptr(dg0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
-        p_dgr = tl.make_block_ptr(dgr + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    p_dg0 = tl.make_block_ptr(dg0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    p_dgr = tl.make_block_ptr(dgr + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
     # [BT]
     """
     b_dg:   1,2,3,4
@@ -149,14 +140,10 @@ def chunk_comba_cumsum_scalar_bwd(
     dg0: torch.Tensor,
     chunk_size: int,
     cu_seqlens: torch.Tensor | None = None,
-    head_first: bool = False,
     output_dtype: torch.dtype | None = torch.float,
     chunk_indices: torch.LongTensor | None = None,
 ) -> torch.Tensor:
-    if head_first:
-        B, H, T = dg0.shape
-    else:
-        B, T, H = dg0.shape
+    B, T, H = dg0.shape
     assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2"
     BT = chunk_size
     if chunk_indices is None and cu_seqlens is not None:
@@ -173,6 +160,5 @@ def chunk_comba_cumsum_scalar_bwd(
         B=B,
         H=H,
         BT=BT,
-        HEAD_FIRST=head_first,
     )
     return dg