diff --git a/fla/ops/comba/chunk.py b/fla/ops/comba/chunk.py index e43ddb0bea..fb936db925 100644 --- a/fla/ops/comba/chunk.py +++ b/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) diff --git a/fla/ops/comba/fused_recurrent.py b/fla/ops/comba/fused_recurrent.py index b7c4265fa1..13c1500909 100644 --- a/fla/ops/comba/fused_recurrent.py +++ b/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, ) diff --git a/fla/ops/comba/utils.py b/fla/ops/comba/utils.py index ee9e6de5f7..30aa090c5e 100644 --- a/fla/ops/comba/utils.py +++ b/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 diff --git a/fla/ops/comba/wy_fast.py b/fla/ops/comba/wy_fast.py index 5742704f58..b6cd1fade6 100644 --- a/fla/ops/comba/wy_fast.py +++ b/fla/ops/comba/wy_fast.py @@ -1,12 +1,11 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang - import torch import triton import triton.language as tl from fla.ops.utils import prepare_chunk_indices -from fla.ops.utils.op import exp +from fla.ops.utils.op import exp, exp2 from fla.utils import autotune_cache_kwargs, check_shared_mem @@ -41,6 +40,7 @@ def chunk_scaled_dot_comba_pkt_fwd_kernel( BK: tl.constexpr, IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H @@ -70,7 +70,10 @@ def chunk_scaled_dot_comba_pkt_fwd_kernel( p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_g0 = tl.load(p_g0, boundary_check=(0,)) b_g = tl.load(p_g, boundary_check=(0,)) - b_A = b_A * exp(b_g0[:, None] - b_g[None, :]) + if USE_EXP2: + b_A = b_A * exp2(b_g0[:, None] - b_g[None, :]) + else: + b_A = b_A * exp(b_g0[:, None] - b_g[None, :]) m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) @@ -88,6 +91,7 @@ def chunk_scaled_dot_comba_pkt_fwd( chunk_size: int = 64, output_dtype: torch.dtype = torch.float32, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> torch.Tensor: r""" Compute beta \mathcal{A}(i-1/j) * P * K^T. @@ -135,6 +139,7 @@ def chunk_scaled_dot_comba_pkt_fwd( H=H, K=K, BT=BT, + USE_EXP2=use_exp2, ) return A @@ -178,6 +183,7 @@ def prepare_wy_repr_bwd_kernel( BK: tl.constexpr, BV: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_EXP2: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H @@ -196,7 +202,10 @@ def prepare_wy_repr_bwd_kernel( b_A = tl.load(p_A, boundary_check=(0, 1)) b_beta = tl.load(p_beta, boundary_check=(0,)) b_g0 = tl.load(p_g0, boundary_check=(0,)) - b_g0_exp = tl.exp(b_g0) + if USE_EXP2: + b_g0_exp = exp2(b_g0) + else: + b_g0_exp = tl.exp(b_g0) b_g = tl.load(p_g, boundary_check=(0,)) b_dbeta = tl.zeros([BT], dtype=tl.float32) @@ -236,7 +245,10 @@ def prepare_wy_repr_bwd_kernel( b_dA = tl.where(m_A, b_dA, 0) b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - b_dA = tl.where(m_A, -b_dA * exp(b_g0[:, None] - b_g[None, :]), 0).to(k.dtype.element_ty) + if USE_EXP2: + b_dA = tl.where(m_A, -b_dA * exp2(b_g0[:, None] - b_g[None, :]), 0).to(k.dtype.element_ty) + else: + b_dA = tl.where(m_A, -b_dA * exp(b_g0[:, None] - b_g[None, :]), 0).to(k.dtype.element_ty) b_dA = b_dA.to(k.dtype.element_ty) b_A = tl.zeros([BT, BT], dtype=tl.float32) @@ -299,6 +311,7 @@ def recompute_w_u_fwd_kernel( BK: tl.constexpr, BV: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_EXP2: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H @@ -313,7 +326,10 @@ def recompute_w_u_fwd_kernel( p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) b_beta = tl.load(p_beta, boundary_check=(0,)) b_A = tl.load(p_A, boundary_check=(0, 1)) - b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + if USE_EXP2: + b_g = exp2(tl.load(p_g, boundary_check=(0,))) + else: + b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) for i_v in range(tl.cdiv(V, BV)): p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) @@ -340,6 +356,7 @@ def recompute_w_u_fwd( A: torch.Tensor, cu_seqlens: torch.LongTensor | None, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = A.shape[-1] @@ -369,6 +386,7 @@ def recompute_w_u_fwd( BT=BT, BK=BK, BV=BV, + USE_EXP2=use_exp2, ) return w, u @@ -385,6 +403,7 @@ def prepare_wy_repr_bwd( du: torch.Tensor, cu_seqlens: torch.LongTensor | None, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = 64 @@ -426,5 +445,6 @@ def prepare_wy_repr_bwd( BT=BT, BK=BK, BV=BV, + USE_EXP2=use_exp2, ) return dk, dv, dp, dbeta, dg0, dg diff --git a/fla/ops/common/chunk_o.py b/fla/ops/common/chunk_o.py index f1dea2aa62..3e1fdae2e0 100644 --- a/fla/ops/common/chunk_o.py +++ b/fla/ops/common/chunk_o.py @@ -5,7 +5,7 @@ import triton.language as tl from fla.ops.utils import prepare_chunk_indices -from fla.ops.utils.op import exp +from fla.ops.utils.op import exp, exp2 from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem('ada') else [32]) @@ -47,6 +47,7 @@ def chunk_fwd_kernel_o( BV: tl.constexpr, USE_G: tl.constexpr, USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, TRANSPOSE_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): @@ -99,14 +100,22 @@ def chunk_fwd_kernel_o( g += bos * H + i_h p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_g = tl.load(p_g, boundary_check=(0,)) - b_o = b_o * exp(b_g)[:, None] - b_A = b_A * exp(b_g[:, None] - b_g[None, :]) + if USE_EXP2: + b_o = b_o * exp2(b_g)[:, None] + b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) + else: + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * exp(b_g[:, None] - b_g[None, :]) if USE_G_GAMMA: b_gamma = tl.load(g_gamma + i_h) b_g = b_gamma * (tl.arange(0, BT) + 1) - b_o = b_o * exp(b_g)[:, None] - b_A = b_A * exp(b_g[:, None] - b_g[None, :]) + if USE_EXP2: + b_o = b_o * exp2(b_g)[:, None] + b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) + else: + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * exp(b_g[:, None] - b_g[None, :]) o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -166,6 +175,7 @@ def chunk_bwd_kernel_dqkwg( BV: tl.constexpr, USE_G: tl.constexpr, USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, USE_DW: tl.constexpr, TRANSPOSE_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, @@ -263,16 +273,25 @@ def chunk_bwd_kernel_dqkwg( p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_g = tl.load(p_g, boundary_check=(0,)) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) - b_dg_last *= exp(b_g_last) - - b_dq = b_dq * exp(b_g)[:, None] * scale + if USE_EXP2: + b_dg_last *= exp2(b_g_last) + b_dq = b_dq * exp2(b_g)[:, None] * scale + else: + b_dg_last *= exp(b_g_last) + b_dq = b_dq * exp(b_g)[:, None] * scale b_dg += tl.sum(b_dq * b_q, axis=1) - b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] + if USE_EXP2: + b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] + else: + b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] b_dg -= tl.sum(b_k * b_dk, axis=1) b_dg_last += tl.sum(b_dk * b_k) - b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale + if USE_EXP2: + b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale + else: + b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k)) b_dg += tl.sum(b_ds2, axis=1) b_dg -= tl.sum(b_ds2, axis=0) @@ -290,9 +309,14 @@ def chunk_bwd_kernel_dqkwg( tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) elif USE_G_GAMMA: - b_dq = b_dq * exp(b_g)[:, None] * scale - b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] - b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale + if USE_EXP2: + b_dq = b_dq * exp2(b_g)[:, None] * scale + b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] + b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale + else: + b_dq = b_dq * exp(b_g)[:, None] * scale + b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] + b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale b_ds = b_ds.to(b_k.dtype) # [BT, BK] b_dq += tl.dot(b_ds, b_k) @@ -345,6 +369,7 @@ def chunk_bwd_kernel_dv( BV: tl.constexpr, USE_G: tl.constexpr, USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, IS_VARLEN: tl.constexpr, ): i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -394,8 +419,12 @@ def chunk_bwd_kernel_dv( m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) if USE_G or USE_G_GAMMA: - b_A = tl.where(m_A, b_A * exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) - b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] + if USE_EXP2: + b_A = tl.where(m_A, b_A * exp2(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) + b_dv *= tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] + else: + b_A = tl.where(m_A, b_A * exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) + b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] else: b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) @@ -441,6 +470,7 @@ def chunk_bwd_kernel_dv_local( BV: tl.constexpr, USE_G: tl.constexpr, USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, USE_A: tl.constexpr, IS_VARLEN: tl.constexpr, ): @@ -480,7 +510,10 @@ def chunk_bwd_kernel_dv_local( b_q = tl.load(p_q, boundary_check=(0, 1)) b_A += tl.dot(b_k, b_q) * scale if USE_G or USE_G_GAMMA: - b_A *= exp(b_g[None, :] - b_g[:, None]) + if USE_EXP2: + b_A *= exp2(b_g[None, :] - b_g[:, None]) + else: + b_A *= exp(b_g[None, :] - b_g[:, None]) o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -506,6 +539,7 @@ def chunk_fwd_o( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, transpose_state_layout: bool = False, ) -> torch.Tensor: B, T, H, K, V = *q.shape, v.shape[-1] @@ -534,6 +568,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) K=K, V=V, BT=BT, + USE_EXP2=use_exp2, TRANSPOSE_STATE=transpose_state_layout, ) return o @@ -550,6 +585,7 @@ def chunk_bwd_dv( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> torch.Tensor: B, T, H, K, V = *k.shape, do.shape[-1] BT = chunk_size @@ -589,6 +625,7 @@ def chunk_bwd_dv( BT=BT, BK=BK, BV=BV, + USE_EXP2=use_exp2, ) return dv @@ -604,6 +641,7 @@ def chunk_bwd_dv_local( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> torch.Tensor: B, T, H, K, V = *k.shape, do.shape[-1] BT = chunk_size @@ -640,6 +678,7 @@ def chunk_bwd_dv_local( BT=BT, BK=BK, BV=BV, + USE_EXP2=use_exp2, ) return dv @@ -659,6 +698,7 @@ def chunk_bwd_dqkwg( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, transpose_state_layout: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -708,6 +748,7 @@ def chunk_bwd_dqkwg( BT=BT, BK=BK, BV=BV, + USE_EXP2=use_exp2, TRANSPOSE_STATE=transpose_state_layout, ) diff --git a/fla/ops/gated_delta_product/chunk.py b/fla/ops/gated_delta_product/chunk.py index e738022bb2..6ca1a9cdf7 100644 --- a/fla/ops/gated_delta_product/chunk.py +++ b/fla/ops/gated_delta_product/chunk.py @@ -1,6 +1,5 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang - import torch from einops import rearrange @@ -188,6 +187,7 @@ def backward( dht=dht, cu_seqlens=cu_seqlens * ctx.num_householder if cu_seqlens is not None else None, chunk_indices=chunk_indices_dp, + use_exp2=False, ) dg = rearrange(dg, 'b (l n) h -> b l n h ', n=ctx.num_householder)[:, :, 0].contiguous().to(g) else: diff --git a/fla/ops/gated_delta_rule/chunk.py b/fla/ops/gated_delta_rule/chunk.py index 8fea87503e..35b210b04f 100644 --- a/fla/ops/gated_delta_rule/chunk.py +++ b/fla/ops/gated_delta_rule/chunk.py @@ -17,6 +17,7 @@ from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 from fla.ops.utils.index import prepare_chunk_indices from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard @@ -33,9 +34,16 @@ def chunk_gated_delta_rule_fwd( cu_seqlens: torch.LongTensor | None = None, cp_context: FLACPContext | None = None, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = True, transpose_state_layout: bool = False, ): - g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + g = chunk_local_cumsum( + g, + chunk_size=64, + scale=RCP_LN2 if use_exp2 else None, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) # obtain WY representation. u is actually the new v. # fused kkt + solve_tril + recompute_w_u w, u, A = chunk_gated_delta_rule_fwd_intra( @@ -45,6 +53,7 @@ def chunk_gated_delta_rule_fwd( beta=beta, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, ) if cp_context is not None: @@ -56,6 +65,7 @@ def chunk_gated_delta_rule_fwd( cu_seqlens=cu_seqlens, initial_state=initial_state, context=cp_context, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) @@ -68,6 +78,7 @@ def chunk_gated_delta_rule_fwd( output_final_state=output_final_state, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) @@ -83,6 +94,7 @@ def chunk_gated_delta_rule_fwd( scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) return g, o, A, final_state, initial_state @@ -102,6 +114,7 @@ def chunk_gated_delta_rule_bwd( cu_seqlens: torch.LongTensor | None = None, cp_context: FLACPContext | None = None, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = True, transpose_state_layout: bool = False, ): w, u = recompute_w_u_fwd( @@ -112,6 +125,7 @@ def chunk_gated_delta_rule_bwd( g=g, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, ) if cp_context is not None: @@ -126,6 +140,7 @@ def chunk_gated_delta_rule_bwd( output_final_state=False, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) dv = chunk_bwd_dv_local( @@ -136,6 +151,7 @@ def chunk_gated_delta_rule_bwd( scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, ) if cp_context is not None: @@ -153,6 +169,7 @@ def chunk_gated_delta_rule_bwd( dht=dht, initial_state=initial_state, context=cp_context, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) @@ -168,6 +185,7 @@ def chunk_gated_delta_rule_bwd( scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) dq, dk, dw, dg = chunk_bwd_dqkwg( @@ -183,6 +201,7 @@ def chunk_gated_delta_rule_bwd( scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) dk2, dv, db, dg2 = prepare_wy_repr_bwd( @@ -195,6 +214,7 @@ def chunk_gated_delta_rule_bwd( du=dv, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, ) dk.add_(dk2) dg.add_(dg2) diff --git a/fla/ops/gated_delta_rule/chunk_fwd.py b/fla/ops/gated_delta_rule/chunk_fwd.py index 98b4edb93a..597d55da2a 100644 --- a/fla/ops/gated_delta_rule/chunk_fwd.py +++ b/fla/ops/gated_delta_rule/chunk_fwd.py @@ -6,7 +6,7 @@ from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd from fla.ops.utils import prepare_chunk_indices -from fla.ops.utils.op import exp +from fla.ops.utils.op import exp, exp2 from fla.utils import IS_TF32_SUPPORTED, autotune_cache_kwargs if IS_TF32_SUPPORTED: @@ -43,6 +43,7 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( BC: tl.constexpr, BK: tl.constexpr, USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, IS_VARLEN: tl.constexpr, ): """ @@ -163,19 +164,34 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( ############################################################################ if USE_G: - # diagonal blocks: g_diff = g_i - g_j within sub-chunk - b_A00 *= exp(b_g0[:, None] - b_g0[None, :]) - b_A11 *= exp(b_g1[:, None] - b_g1[None, :]) - b_A22 *= exp(b_g2[:, None] - b_g2[None, :]) - b_A33 *= exp(b_g3[:, None] - b_g3[None, :]) - - # off-diagonal blocks: g_diff = g_row - g_col (cross sub-chunk) - b_A10 *= exp(b_g1[:, None] - b_g0[None, :]) - b_A20 *= exp(b_g2[:, None] - b_g0[None, :]) - b_A21 *= exp(b_g2[:, None] - b_g1[None, :]) - b_A30 *= exp(b_g3[:, None] - b_g0[None, :]) - b_A31 *= exp(b_g3[:, None] - b_g1[None, :]) - b_A32 *= exp(b_g3[:, None] - b_g2[None, :]) + if USE_EXP2: + # diagonal blocks: g_diff = g_i - g_j within sub-chunk + b_A00 *= exp2(b_g0[:, None] - b_g0[None, :]) + b_A11 *= exp2(b_g1[:, None] - b_g1[None, :]) + b_A22 *= exp2(b_g2[:, None] - b_g2[None, :]) + b_A33 *= exp2(b_g3[:, None] - b_g3[None, :]) + + # off-diagonal blocks: g_diff = g_row - g_col (cross sub-chunk) + b_A10 *= exp2(b_g1[:, None] - b_g0[None, :]) + b_A20 *= exp2(b_g2[:, None] - b_g0[None, :]) + b_A21 *= exp2(b_g2[:, None] - b_g1[None, :]) + b_A30 *= exp2(b_g3[:, None] - b_g0[None, :]) + b_A31 *= exp2(b_g3[:, None] - b_g1[None, :]) + b_A32 *= exp2(b_g3[:, None] - b_g2[None, :]) + else: + # diagonal blocks: g_diff = g_i - g_j within sub-chunk + b_A00 *= exp(b_g0[:, None] - b_g0[None, :]) + b_A11 *= exp(b_g1[:, None] - b_g1[None, :]) + b_A22 *= exp(b_g2[:, None] - b_g2[None, :]) + b_A33 *= exp(b_g3[:, None] - b_g3[None, :]) + + # off-diagonal blocks: g_diff = g_row - g_col (cross sub-chunk) + b_A10 *= exp(b_g1[:, None] - b_g0[None, :]) + b_A20 *= exp(b_g2[:, None] - b_g0[None, :]) + b_A21 *= exp(b_g2[:, None] - b_g1[None, :]) + b_A30 *= exp(b_g3[:, None] - b_g0[None, :]) + b_A31 *= exp(b_g3[:, None] - b_g1[None, :]) + b_A32 *= exp(b_g3[:, None] - b_g2[None, :]) # apply beta to row dimension and mask m_d = o_i[:, None] > o_i[None, :] @@ -309,6 +325,7 @@ def chunk_gated_delta_rule_fwd_intra( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: r""" GDN intra-chunk forward: fused kkt + solve_tril + recompute_w_u. @@ -364,6 +381,7 @@ def chunk_gated_delta_rule_fwd_intra( K=K, BT=BT, BC=BC, + USE_EXP2=use_exp2, ) # Step 2: recompute_w_u @@ -375,5 +393,6 @@ def chunk_gated_delta_rule_fwd_intra( g=g, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_exp2=use_exp2, ) return w, u, A diff --git a/fla/ops/gated_delta_rule/fused_recurrent.py b/fla/ops/gated_delta_rule/fused_recurrent.py index 812c357e75..3f470a3ab0 100644 --- a/fla/ops/gated_delta_rule/fused_recurrent.py +++ b/fla/ops/gated_delta_rule/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 @@ -45,6 +44,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( IS_BETA_HEADWISE: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, + USE_EXP2: tl.constexpr, TRANSPOSE_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): @@ -109,21 +109,36 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if USE_G: b_g = tl.load(p_g).to(tl.float32) - b_h *= exp(b_g) + if USE_EXP2: + b_h *= exp2(b_g) + else: + b_h *= exp(b_g) if USE_GK: b_gk = tl.load(p_gk).to(tl.float32) - if TRANSPOSE_STATE: - b_h *= exp(b_gk[None, :]) + if USE_EXP2: + if TRANSPOSE_STATE: + b_h *= exp2(b_gk[None, :]) + else: + b_h *= exp2(b_gk[:, None]) else: - b_h *= exp(b_gk[:, None]) + if TRANSPOSE_STATE: + b_h *= exp(b_gk[None, :]) + else: + b_h *= exp(b_gk[:, None]) if USE_GV: b_gv = tl.load(p_gv).to(tl.float32) - if TRANSPOSE_STATE: - b_h *= exp(b_gv[:, None]) + if USE_EXP2: + if TRANSPOSE_STATE: + b_h *= exp2(b_gv[:, None]) + else: + b_h *= exp2(b_gv[None, :]) else: - b_h *= exp(b_gv[None, :]) + if TRANSPOSE_STATE: + b_h *= exp(b_gv[:, None]) + else: + b_h *= exp(b_gv[None, :]) if TRANSPOSE_STATE: b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1)) @@ -168,6 +183,7 @@ def fused_recurrent_gated_delta_rule_fwd( output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, transpose_state_layout: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] @@ -209,6 +225,7 @@ def fused_recurrent_gated_delta_rule_fwd( BV=BV, IS_BETA_HEADWISE=beta.ndim != v.ndim, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_EXP2=use_exp2, TRANSPOSE_STATE=transpose_state_layout, num_warps=1, num_stages=3, @@ -234,6 +251,7 @@ def forward( output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, transpose_state_layout: bool = False, ): o, final_state = fused_recurrent_gated_delta_rule_fwd( @@ -249,6 +267,7 @@ def forward( output_final_state=output_final_state, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, cu_seqlens=cu_seqlens, + use_exp2=use_exp2, transpose_state_layout=transpose_state_layout, ) @@ -277,6 +296,7 @@ def fused_recurrent_gated_delta_rule( output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, transpose_state_layout: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: r""" @@ -377,6 +397,7 @@ def fused_recurrent_gated_delta_rule( output_final_state, use_qk_l2norm_in_kernel, cu_seqlens, + use_exp2, transpose_state_layout, ) return o, final_state diff --git a/fla/ops/gated_delta_rule/wy_fast.py b/fla/ops/gated_delta_rule/wy_fast.py index ffccdda82e..52b2c92491 100644 --- a/fla/ops/gated_delta_rule/wy_fast.py +++ b/fla/ops/gated_delta_rule/wy_fast.py @@ -5,7 +5,7 @@ import triton.language as tl from fla.ops.utils import prepare_chunk_indices -from fla.ops.utils.op import exp +from fla.ops.utils.op import exp, exp2 from fla.utils import IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem if IS_NVIDIA_BLACKWELL: @@ -67,6 +67,7 @@ def recompute_w_u_fwd_kernel( BK: tl.constexpr, BV: tl.constexpr, USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, IS_VARLEN: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) @@ -93,7 +94,10 @@ def recompute_w_u_fwd_kernel( if USE_G: p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = exp(tl.load(p_g, boundary_check=(0,))) + if USE_EXP2: + b_g = exp2(tl.load(p_g, boundary_check=(0,))) + else: + b_g = exp(tl.load(p_g, boundary_check=(0,))) for i_k in range(tl.cdiv(K, BK)): p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) @@ -142,6 +146,7 @@ def prepare_wy_repr_bwd_kernel( BK: tl.constexpr, BV: tl.constexpr, USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, IS_VARLEN: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) @@ -165,7 +170,10 @@ def prepare_wy_repr_bwd_kernel( if USE_G: p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) b_g = tl.load(p_g, boundary_check=(0,)) - b_g_exp = tl.exp(b_g) + if USE_EXP2: + b_g_exp = exp2(b_g) + else: + b_g_exp = tl.exp(b_g) b_dg = tl.zeros([BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): @@ -209,7 +217,10 @@ def prepare_wy_repr_bwd_kernel( b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) if USE_G: - b_dA *= exp(b_g[:, None] - b_g[None, :]) + if USE_EXP2: + b_dA *= exp2(b_g[:, None] - b_g[None, :]) + else: + b_dA *= exp(b_g[:, None] - b_g[None, :]) b_A = tl.zeros([BT, BT], dtype=tl.float32) b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty) @@ -247,6 +258,7 @@ def recompute_w_u_fwd( g: torch.Tensor | None = None, cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = A.shape[-1] @@ -276,6 +288,7 @@ def recompute_w_u_fwd( BT=BT, BK=BK, BV=BV, + USE_EXP2=use_exp2, ) return w, u @@ -290,6 +303,7 @@ def prepare_wy_repr_bwd( g: torch.Tensor = None, cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = 64 @@ -325,6 +339,7 @@ def prepare_wy_repr_bwd( BT=BT, BK=BK, BV=BV, + USE_EXP2=use_exp2, ) return dk, dv, db, dg