From 0c9799ca3129814131ac98ed80f6302412758050 Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 01:26:59 +0000 Subject: [PATCH 01/10] [KDA] Speed up chunk_kda by introducing lowerbound gate Co-authored-by: Wang Yucheng Co-authored-by: yzhangcs Co-authored-by: Chen Guangyu --- fla/ops/kda/chunk.py | 176 +++++--- fla/ops/kda/chunk_bwd.py | 3 +- fla/ops/kda/chunk_intra.py | 516 +++++++++++++++------- fla/ops/kda/chunk_intra_token_parallel.py | 8 +- fla/ops/kda/gate.py | 2 +- fla/ops/kda/wy_fast.py | 12 +- tests/ops/test_kda.py | 107 +++-- 7 files changed, 565 insertions(+), 259 deletions(-) diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py index 13c0e85d44..336c2383f7 100644 --- a/fla/ops/kda/chunk.py +++ b/fla/ops/kda/chunk.py @@ -1,4 +1,5 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Related files are modified and supported by the Moonshot AI Team import torch @@ -7,10 +8,11 @@ from fla.ops.gla.chunk import chunk_gla_fwd_o_gk from fla.ops.kda.chunk_bwd import chunk_kda_bwd_dAv, chunk_kda_bwd_wy_dqkg_fused from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra, chunk_kda_fwd_intra -from fla.ops.kda.gate import kda_gate_bwd, kda_gate_fwd +from fla.ops.kda.gate import kda_gate_bwd, kda_gate_chunk_cumsum from fla.ops.kda.wy_fast import 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 @@ -26,8 +28,11 @@ def chunk_kda_fwd( cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, chunk_size: int = 64, + safe_gate: bool = False, + disable_recompute: bool = False ): - w, u, kg, Aqk, Akk = chunk_kda_fwd_intra( + # qg = None if disable_recompute is False + w, u, qg, kg, Aqk, Akk = chunk_kda_fwd_intra( q=q, k=k, v=v, @@ -37,6 +42,8 @@ def chunk_kda_fwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + safe_gate=safe_gate, + disable_recompute=disable_recompute ) h, v_new, final_state = chunk_gated_delta_rule_fwd_h( k=kg, @@ -62,7 +69,11 @@ def chunk_kda_fwd( chunk_indices=chunk_indices, use_exp2=True, ) - return o, Aqk, Akk, final_state + if not disable_recompute: + # Delete to save memory + w, u, qg, kg, v_new, h = None, None, None, None, None, None + + return o, Aqk, Akk, final_state, w, u, qg, kg, v_new, h def chunk_kda_bwd( @@ -80,30 +91,36 @@ def chunk_kda_bwd( cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, chunk_size: int = 64, + safe_gate: bool = False, + disable_recompute: bool = False, + **kwargs, ): - # w = Akk @ (k * beta) - # u = Akk @ (v * beta) - w, u, qg, kg = recompute_w_u_fwd( - q=q, - k=k, - v=v, - beta=beta, - A=Akk, - gk=g, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, v_new, _ = chunk_gated_delta_rule_fwd_h( - k=kg, - w=w, - u=u, - gk=g, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - use_exp2=True, - ) + if not disable_recompute: + # w = Akk @ (k * beta) + # u = Akk @ (v * beta) + w, u, qg, kg = recompute_w_u_fwd( + q=q, + k=k, + v=v, + beta=beta, + A=Akk, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + else: + w, u, qg, kg, v_new, h = kwargs["w"], kwargs["u"], kwargs["qg"], kwargs["kg"], kwargs["v_new"], kwargs["h"] # dAqk = do @ v.T # dv = A @ do dAqk, dv = chunk_kda_bwd_dAv( @@ -149,6 +166,7 @@ def chunk_kda_bwd( chunk_size=chunk_size, chunk_indices=chunk_indices, ) + dq, dk, db, dg = chunk_kda_bwd_intra( q=q, k=k, @@ -163,6 +181,7 @@ def chunk_kda_bwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + safe_gate=safe_gate ) return dq, dk, dv, db, dg, dh0 @@ -186,30 +205,43 @@ def forward( use_qk_l2norm_in_kernel: bool = False, use_gate_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + safe_gate: bool = False, + lower_bound: float | None = None, + disable_recompute: bool = False ): + chunk_size = 64 g_org = None + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None if use_gate_in_kernel: g_org = g - g = kda_gate_fwd( + if safe_gate: + assert lower_bound is not None, "lower_bound must be set when use sage_gate" + g = kda_gate_chunk_cumsum( g=g_org, A_log=A_log, dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + else: + g = chunk_local_cumsum( + g=g, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices ) q_rstd, k_rstd = None, None if use_qk_l2norm_in_kernel: q, q_rstd = l2norm_fwd(q) k, k_rstd = l2norm_fwd(k) - chunk_size = 64 - g = chunk_local_cumsum( - g=g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices - ) - o, Aqk, Akk, final_state = chunk_kda_fwd( + (o, Aqk, Akk, final_state, w, u, qg, kg, v_new, h) = chunk_kda_fwd( q=q, k=k, v=v, @@ -220,16 +252,24 @@ def forward( output_final_state=output_final_state, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + safe_gate=safe_gate, + disable_recompute=disable_recompute ) - if use_gate_in_kernel: - g = None + + if disable_recompute is False and use_gate_in_kernel: + g = None # type: ignore ctx.save_for_backward( - q, q_rstd, k, k_rstd, v, g, g_org, beta, A_log, dt_bias, Aqk, Akk, initial_state, cu_seqlens, chunk_indices + q, q_rstd, k, k_rstd, v, g, g_org, beta, A_log, dt_bias, Aqk, Akk, + w, u, qg, kg, v_new, h, + initial_state, cu_seqlens, chunk_indices ) ctx.chunk_size = chunk_size + ctx.safe_gate = safe_gate ctx.scale = scale + ctx.lower_bound = lower_bound ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.disable_recompute = disable_recompute return o.to(q.dtype), final_state @staticmethod @@ -240,21 +280,21 @@ def backward( do: torch.Tensor, dht: torch.Tensor, ): - (q, q_rstd, k, k_rstd, v, g, g_org, beta, A_log, dt_bias, Aqk, Akk, initial_state, cu_seqlens, chunk_indices) = ( + (q, q_rstd, k, k_rstd, v, g, g_org, beta, A_log, dt_bias, Aqk, Akk, + w, u, qg, kg, v_new, h, + initial_state, cu_seqlens, chunk_indices) = ( ctx.saved_tensors ) - if ctx.use_gate_in_kernel: - g = kda_gate_fwd( + if ctx.disable_recompute is False and ctx.use_gate_in_kernel: + g = kda_gate_chunk_cumsum( g=g_org, A_log=A_log, dt_bias=dt_bias, - ) - g = chunk_local_cumsum( - g=g, - chunk_size=ctx.chunk_size, scale=RCP_LN2, + chunk_size=ctx.chunk_size, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices + chunk_indices=chunk_indices, + lower_bound=ctx.lower_bound ) dq, dk, dv, db, dg, dh0 = chunk_kda_bwd( q=q, @@ -271,22 +311,43 @@ def backward( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=ctx.chunk_size, + safe_gate=ctx.safe_gate, + disable_recompute=ctx.disable_recompute, + w=w, u=u, qg=qg, kg=kg, v_new=v_new, h=h ) if ctx.use_qk_l2norm_in_kernel: dq = l2norm_bwd(q, q_rstd, dq) dk = l2norm_bwd(k, k_rstd, dk) dA, dbias = None, None + if ctx.use_gate_in_kernel: + dg = chunk_local_cumsum( + dg, + chunk_size=ctx.chunk_size, + reverse=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) dg, dA, dbias = kda_gate_bwd( g=g_org, A_log=A_log, dt_bias=dt_bias, dyg=dg, + lower_bound=ctx.lower_bound ) dA = dA.to(A_log) if dt_bias is not None: dbias = dbias.to(dt_bias) - return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), dA, dbias, None, dh0, None, None, None, None, None + else: + dg = chunk_local_cumsum( + dg, + chunk_size=ctx.chunk_size, + reverse=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return (dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), dA, dbias, None, dh0, + None, None, None, None, None, None, None, None) @torch.compiler.disable @@ -302,7 +363,10 @@ def chunk_kda( use_qk_l2norm_in_kernel: bool = False, use_gate_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + safe_gate: bool = False, + lower_bound: float | None = None, + disable_recompute: bool = False, **kwargs, ): r""" @@ -339,8 +403,9 @@ def chunk_kda( cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. - chunk_indices (torch.LongTensor): - Chunk indices used for variable-length training, + cu_seqlens_cpu (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. Returns: o (torch.Tensor): @@ -408,10 +473,8 @@ def chunk_kda( A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." - assert k.shape[-1] <= 256, "Currently we only support key headdim <=256 for KDA :-(" assert beta.shape == q.shape[:3], "beta must be of shape (batch size, seq len, num of head)." assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." - if scale is None: scale = k.shape[-1] ** -0.5 o, final_state = ChunkKDAFunction.apply( @@ -428,6 +491,9 @@ def chunk_kda( use_qk_l2norm_in_kernel, use_gate_in_kernel, cu_seqlens, - chunk_indices, + cu_seqlens_cpu, + safe_gate, + lower_bound, + disable_recompute ) return o, final_state diff --git a/fla/ops/kda/chunk_bwd.py b/fla/ops/kda/chunk_bwd.py index c5e9db367d..14d2184a58 100644 --- a/fla/ops/kda/chunk_bwd.py +++ b/fla/ops/kda/chunk_bwd.py @@ -217,8 +217,7 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) - - tl.debug_barrier() + tl.debug_barrier() # DO NOT REMOVE THIS LINE! if i_k == 0: p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) p_dv2 = tl.make_block_ptr(dv2, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) diff --git a/fla/ops/kda/chunk_intra.py b/fla/ops/kda/chunk_intra.py index 8757e8cc0a..f41fd391bf 100644 --- a/fla/ops/kda/chunk_intra.py +++ b/fla/ops/kda/chunk_intra.py @@ -6,12 +6,12 @@ from fla.ops.kda.chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel from fla.ops.kda.wy_fast import recompute_w_u_fwd -from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices -from fla.ops.utils.op import exp2 -from fla.utils import IS_TF32_SUPPORTED, autotune_cache_kwargs +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp2, gather +from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs if IS_TF32_SUPPORTED: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32x3') + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') else: SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') @@ -51,6 +51,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( BC: tl.constexpr, BK: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, ): """ Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. @@ -121,8 +122,6 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) - b_k1, b_g1 = b_k0, b_g0 - b_k2, b_g2 = b_k0, b_g0 if i_tc1 < T: p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) @@ -141,59 +140,59 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) - if i_tc2 < T: - p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_g2 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn2 = tl.load(g + i_tc2 * H*K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) - b_qg2 = b_q2 * b_gqn2 - b_kg2 = b_k2 * b_gqn2 - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) - b_Aqk20 += tl.dot(b_qg2, b_kgt) - b_Akk20 += tl.dot(b_kg2, b_kgt) - # [BC, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) - # [BC, BC] - b_Aqk21 += tl.dot(b_qg2, b_kgt) - b_Akk21 += tl.dot(b_kg2, b_kgt) - - if i_tc3 < T: - p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_g3 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn3 = tl.load(g + i_tc3 * H*K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) - b_qg3 = b_q3 * b_gqn3 - b_kg3 = b_k3 * b_gqn3 - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) - # [BC, BC] - b_Aqk30 += tl.dot(b_qg3, b_kgt) - b_Akk30 += tl.dot(b_kg3, b_kgt) - # [BK, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) - # [BC, BC] - b_Aqk31 += tl.dot(b_qg3, b_kgt) - b_Akk31 += tl.dot(b_kg3, b_kgt) - # [BK, BC] - b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) - # [BC, BC] - b_Aqk32 += tl.dot(b_qg3, b_kgt) - b_Akk32 += tl.dot(b_kg3, b_kgt) + if i_tc2 < T: + p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_g2 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_g3 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) ################################################################################ # save off-diagonal Aqk blocks and prepare Akk @@ -241,45 +240,47 @@ def chunk_kda_fwd_kernel_inter_solve_fused( ################################################################################ # forward substitution on diagonals ################################################################################ - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - b_Ai00 = -tl.where(m_A, b_Ai00, 0) - b_Ai11 = -tl.where(m_A, b_Ai11, 0) - b_Ai22 = -tl.where(m_A, b_Ai22, 0) - b_Ai33 = -tl.where(m_A, b_Ai33, 0) - - for i in range(2, min(BC, T - i_tc0)): - b_a00 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) - b_a00 = tl.where(o_i < i, b_a00, 0.) - b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) - b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) - for i in range(BC + 2, min(2*BC, T - i_tc0)): - b_a11 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) - b_a11 = tl.where(o_i < i - BC, b_a11, 0.) - b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) - b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) - for i in range(2*BC + 2, min(3*BC, T - i_tc0)): - b_a22 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) - b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) - b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) - b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) - for i in range(3*BC + 2, min(4*BC, T - i_tc0)): - b_a33 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) - b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) - b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) - b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) - - b_Ai00 += m_I - b_Ai11 += m_I - b_Ai22 += m_I - b_Ai33 += m_I + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2*BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2*BC + 2, min(3*BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) + for i in range(3*BC + 2, min(4*BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I ################################################################################ # compute merged inverse using off-diagonals ################################################################################ - # we used tf32x3 to maintain matrix inverse's precision whenever possible. + # we used tf32 to maintain matrix inverse's precision whenever possible. b_Ai10 = -tl.dot( tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), b_Ai00, @@ -381,6 +382,8 @@ def chunk_kda_bwd_kernel_intra( BK: tl.constexpr, NC: tl.constexpr, IS_VARLEN: tl.constexpr, + SAFE_GATE: tl.constexpr, + USE_GATHER: tl.constexpr, ): i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H @@ -427,7 +430,7 @@ def chunk_kda_bwd_kernel_intra( if i_i > 0: p_gn = g + i_ti * H*K + o_k # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) + b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] for i_j in range(0, i_i): p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) @@ -436,14 +439,14 @@ def chunk_kda_bwd_kernel_intra( # [BC, BK] b_k = tl.load(p_k, boundary_check=(0, 1)) b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_k * exp2(b_gn[None, :] - b_gk) + b_kg = b_k * exp2(b_gn - b_gk) # [BC, BC] b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) # [BC, BK] b_dq2 += tl.dot(b_dAqk, b_kg) b_dk2 += tl.dot(b_dAkk, b_kg) - b_gqn = exp2(b_g - b_gn[None, :]) + b_gqn = exp2(b_g - b_gn) b_dq2 *= b_gqn b_dk2 *= b_gqn @@ -458,22 +461,48 @@ def chunk_kda_bwd_kernel_intra( b_q = tl.load(p_q, boundary_check=(0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC] - b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) - b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) - # [BK] - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] >= j - # [BC, BK] - b_kgj = b_kj[None, :] * exp2(b_g - b_gkj[None, :]) - b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kgj, 0.) - b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kgj, 0.) - - p_kj += H*K - p_gkj += H*K + if SAFE_GATE: + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) + b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + + m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_qk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.) + b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.) + b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.) + exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) + exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) + + b_k_exp_diag_qk = b_k * exp_neg_b_g_diag_qk + b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + else: + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC] + b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) + b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) + # [BK] + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] >= j + # [BC, BK] + b_gqk = exp2(b_g - b_gkj[None, :]) + b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kj[None, :] * b_gqk, 0.) + b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * b_gqk, 0.) + + p_kj += H*K + p_gkj += H*K + b_db = tl.sum(b_dk2 * b_k, 1) b_dk2 *= b_b[:, None] @@ -493,7 +522,7 @@ def chunk_kda_bwd_kernel_intra( if i_i < NC - 1: p_gn = g + (min(i_ti + BC, T) - 1) * H*K + o_k # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) + b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] for i_j in range(i_i + 1, NC): p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) @@ -514,35 +543,70 @@ def chunk_kda_bwd_kernel_intra( o_j = i_t * BT + i_j * BC + o_i m_j = o_j < T # [BC, BK] - b_gkn = tl.where(m_j[:, None], exp2(b_gk - b_gn[None, :]), 0) - b_qg = b_q * b_gkn - b_kbg = b_kb * b_gkn + b_gkn = exp2(b_gk - b_gn) + b_qg = b_q * tl.where(m_j[:, None], b_gkn, 0) + b_kbg = b_kb * tl.where(m_j[:, None], b_gkn, 0) # [BC, BK] - b_dkt += tl.dot(b_dAqk, b_qg) + tl.dot(b_dAkk, b_kbg) - b_dkt *= exp2(b_gn[None, :] - b_g) + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dkt += tl.dot(b_dAqk, b_qg) + b_dkt += tl.dot(b_dAkk, b_kbg) + b_dkt *= exp2(b_gn - b_g) o_dA = i_ti * H*BT + i_i * BC + o_i p_qj = q + i_ti * H*K + o_k p_kj = k + i_ti * H*K + o_k p_gkj = g + i_ti * H*K + o_k p_bj = beta + i_ti * H - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_dAqk = tl.load(dAqk + o_dA + j * H*BT) - b_dAkk = tl.load(dAkk + o_dA + j * H*BT) - # [BK,] - b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) - b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] <= j - b_gkq = exp2(b_gkj[None, :] - b_g) - b_dkt += tl.where(m_i, (b_dAkk[:, None] * b_kbj[None, :] + b_dAqk[:, None] * b_qj[None, :]) * b_gkq, 0.) - - p_qj += H*K - p_kj += H*K - p_gkj += H*K - p_bj += H + if SAFE_GATE: + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) + b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + + m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_kk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_kk = tl.where(m_i_diag_kk, b_dAqk_diag_kk, 0.) + b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.) + # ensure numerical stability + b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn, 0.) + exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) + exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) + + b_q_exp = b_q * exp_b_g_diag_kk + b_kb_exp = b_k * b_b[:, None] * exp_b_g_diag_kk + + b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk + b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk + else: + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dAqk = tl.load(dAqk + o_dA + j * H*BT) + b_dAkk = tl.load(dAkk + o_dA + j * H*BT) + # [BK,] + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] <= j + b_gkq = exp2(b_gkj[None, :] - b_g) + b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * b_gkq, 0.) + b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kbj[None, :] * b_gkq, 0.) + + p_qj += H*K + p_kj += H*K + p_gkj += H*K + p_bj += H p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) p_dk2 = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) @@ -556,6 +620,122 @@ def chunk_kda_bwd_kernel_intra( tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT", "BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * H + i_h) * K + beta = beta + bos * H + i_h + Aqk = Aqk + (bos * H + i_h) * BT + Akk = Akk + (bos * H + i_h) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = b_g - b_gn + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (H*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * H*BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + def chunk_kda_fwd_intra( q: torch.Tensor, k: torch.Tensor, @@ -566,6 +746,8 @@ def chunk_kda_fwd_intra( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, + disable_recompute: bool = False, ): B, T, H, K = k.shape BT = chunk_size @@ -573,6 +755,7 @@ def chunk_kda_fwd_intra( if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) # Akk must be zero-initialized - kernel only writes lower triangular @@ -581,18 +764,42 @@ def chunk_kda_fwd_intra( Akkd = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32) # Step 1: Run token_parallel first to compute diagonal blocks into Akkd (fp32) - Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( - q=q, - k=k, - gk=gk, - beta=beta, - Aqk=Aqk, - Akk=Akkd, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - sub_chunk_size=BC, - ) + # Step 1: compute diagonal blocks into Akk_diag (fp32) + if safe_gate: + grid = (NT, NC, B * H) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=IS_GATHER_SUPPORTED, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) grid = (NT, B * H) chunk_kda_fwd_kernel_inter_solve_fused[grid]( @@ -611,17 +818,19 @@ def chunk_kda_fwd_intra( K=K, BT=BT, BC=BC, + USE_SAFE_GATE=safe_gate, ) - w, u, _, kg = recompute_w_u_fwd( + w, u, qg, kg = recompute_w_u_fwd( k=k, v=v, beta=beta, A=Akk, + q=q if disable_recompute else None, gk=gk, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, ) - return w, u, kg, Aqk, Akk + return w, u, qg, kg, Aqk, Akk def chunk_kda_bwd_intra( @@ -638,6 +847,7 @@ def chunk_kda_bwd_intra( cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, chunk_size: int = 64, + safe_gate: bool = False, ): B, T, H, K = k.shape BT = chunk_size @@ -679,16 +889,12 @@ def chunk_kda_bwd_intra( BC=BC, BK=BK, NC=NC, + SAFE_GATE=safe_gate, + USE_GATHER=IS_GATHER_SUPPORTED, ) dq = dq2 dk = dk2 db = db2.sum(0).add_(db) - dg = chunk_local_cumsum( - dg2, - chunk_size=chunk_size, - reverse=True, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) + dg = dg2 return dq, dk, db, dg diff --git a/fla/ops/kda/chunk_intra_token_parallel.py b/fla/ops/kda/chunk_intra_token_parallel.py index b43bf0a891..3640cce972 100644 --- a/fla/ops/kda/chunk_intra_token_parallel.py +++ b/fla/ops/kda/chunk_intra_token_parallel.py @@ -103,13 +103,15 @@ def chunk_kda_fwd_kernel_intra_token_parallel( b_kj = tl.load(p_kj, boundary_check=(0, 1)).to(tl.float32) b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) - b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) + b_kgj = b_kj * exp2(b_g - b_gj) + + b_kgj = tl.where(m_k[None, :], b_kgj, 0.0) # [BH] b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale - b_Akk = tl.where(j < i_t, tl.sum(b_k * b_kgj, axis=1), 0.0) + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) tl.store(Aqk + i_t * H*BT + (i_hg * BH + o_h) * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_h) - tl.store(Akk + i_t * H*BC + (i_hg * BH + o_h) * BC + j % BC, b_Akk.to(Akk.dtype.element_ty), mask=m_h) + tl.store(Akk + i_t * H*BC + (i_hg * BH + o_h) * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_h) def chunk_kda_fwd_intra_token_parallel( diff --git a/fla/ops/kda/gate.py b/fla/ops/kda/gate.py index 7353dadaba..70c16bb426 100644 --- a/fla/ops/kda/gate.py +++ b/fla/ops/kda/gate.py @@ -1,5 +1,5 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -# This file is modified and supported by the Moonshot Team +# This file is modified and supported by the Moonshot AI Team import torch import torch.nn.functional as F diff --git a/fla/ops/kda/wy_fast.py b/fla/ops/kda/wy_fast.py index 622ae9eace..828f3349f0 100644 --- a/fla/ops/kda/wy_fast.py +++ b/fla/ops/kda/wy_fast.py @@ -19,13 +19,13 @@ triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] - for DOT_PRECISION in (["tf32x3", "ieee"] if IS_TF32_SUPPORTED else ["ieee"]) + for DOT_PRECISION in (["tf32", "tf32x3", "ieee"] if IS_TF32_SUPPORTED else ["ieee"]) ], key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) -def recompute_w_u_fwd_kernel( +def recompute_w_u_fwd_kda_kernel( q, k, qg, @@ -113,7 +113,7 @@ def recompute_w_u_fwd_kernel( **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_bwd_kernel( +def prepare_wy_repr_bwd_kda_kernel( k, v, beta, @@ -218,7 +218,7 @@ def recompute_w_u_fwd( gk: torch.Tensor | None = None, cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: B, T, H, K, V = *k.shape, v.shape[-1] BT = A.shape[-1] BK = 64 @@ -232,7 +232,7 @@ def recompute_w_u_fwd( u = torch.empty_like(v) qg = torch.empty_like(q) if q is not None else None kg = torch.empty_like(k) if gk is not None else None - recompute_w_u_fwd_kernel[(NT, B*H)]( + recompute_w_u_fwd_kda_kernel[(NT, B*H)]( q=q, k=k, qg=qg, @@ -283,7 +283,7 @@ def prepare_wy_repr_bwd( dg2 = torch.empty_like(gk, dtype=torch.float) dA = torch.empty_like(A, dtype=torch.float) db = torch.empty_like(beta, dtype=torch.float) - prepare_wy_repr_bwd_kernel[(NT, B * H)]( + prepare_wy_repr_bwd_kda_kernel[(NT, B * H)]( k=k, v=v, beta=beta, diff --git a/tests/ops/test_kda.py b/tests/ops/test_kda.py index 5c7c31fe9c..cfe994be3c 100644 --- a/tests/ops/test_kda.py +++ b/tests/ops/test_kda.py @@ -1,6 +1,5 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -import os import pytest import torch @@ -148,22 +147,24 @@ def test_fused_recurrent( "use_qk_l2norm_in_kernel", "use_gate_in_kernel", "dtype", - "tma", + "safe_gate", + "disable_recompute", ), [ pytest.param( *test, - id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-qk_l2norm{}-gate{}-dtype{}-tma{}".format(*test), + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-qk_l2norm{}-gate{}-dtype{}-safe_gate{}-disable_recompute{}".format( + *test), ) for test in [ - (1, 63, 1, 64, 1, 1, 0, False, False, torch.float16, True), - (2, 500, 3, 60, 1, 1, 0, False, False, torch.float16, True), - (2, 1000, 3, 64, 0.1, 1, 0.5, False, False, torch.float16, False), - (3, 1024, 4, 100, 1, 0.1, 0, False, False, torch.float16, False), - (4, 1024, 4, 128, 0.1, 1, 0, False, False, torch.float16, True), - (4, 1024, 4, 128, 0.1, 1, 0, True, False, torch.float16, True), - (2, 1500, 4, 128, 0.1, 10, 0, False, True, torch.float16, False), - (4, 2048, 8, 64, 0.1, 1, 0, False, True, torch.float16, True), + (1, 63, 1, 64, 1, 1, 0, False, False, torch.float16, True, False), + (2, 500, 3, 60, 1, 1, 0, False, False, torch.float16, True, True), + (2, 1000, 3, 64, 0.1, 1, 0.5, False, False, torch.float16, False, True), + (3, 1024, 4, 100, 1, 0.1, 0, False, False, torch.float16, False, False), + (4, 1024, 4, 128, 0.1, 1, 0, False, False, torch.float16, True, True), + (4, 1024, 4, 128, 0.1, 1, 0, True, False, torch.float16, True, False), + (2, 1500, 4, 128, 0.1, 10, 0, False, True, torch.float16, False, True), + (4, 2048, 8, 64, 0.1, 1, 0, False, True, torch.float16, True, True), ] ], ) @@ -178,23 +179,29 @@ def test_chunk( use_qk_l2norm_in_kernel: bool, use_gate_in_kernel: bool, dtype: torch.dtype, - tma: bool, + safe_gate: bool, + disable_recompute: bool, ): torch.manual_seed(42) - if not tma: - os.environ["FLA_USE_TMA"] = "0" - else: - os.environ["FLA_USE_TMA"] = "1" q = torch.rand(B, T, H, D, dtype=dtype) k = torch.rand(B, T, H, D, dtype=dtype) v = torch.rand(B, T, H, D, dtype=dtype) - g = torch.randn(B, T, H, D, dtype=torch.float) + g = torch.randn(B, T, H, D, dtype=torch.float if not use_gate_in_kernel else dtype) if use_gate_in_kernel: A_log = torch.randn(H, dtype=torch.float) dt_bias = torch.randn(H * D, dtype=torch.float) else: g = F.logsigmoid(g) / gate_logit_normalizer g = g * (torch.rand_like(g) > mask_p) + if safe_gate: + lower_bound = -5.0 + if not use_gate_in_kernel: + g = g.clamp(-5, 0) + naive_kda_gate_fn = naive_kda_lowerbound_gate + else: + lower_bound = None + naive_kda_gate_fn = naive_kda_gate + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() h0 = torch.randn(B, H, D, D, dtype=torch.float32) if use_gate_in_kernel: @@ -208,7 +215,7 @@ def test_chunk( q=F.normalize(q.clone(), p=2, dim=-1), k=F.normalize(k.clone(), p=2, dim=-1), v=v.clone(), - g=(naive_kda_gate(g, A_log, dt_bias) if use_gate_in_kernel else g.clone()), + g=(naive_kda_gate_fn(g, A_log, dt_bias) if use_gate_in_kernel else g.clone()), beta=beta.clone(), scale=scale, initial_state=h0.clone(), @@ -234,6 +241,9 @@ def test_chunk( output_final_state=True, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_gate_in_kernel=use_gate_in_kernel, + safe_gate=safe_gate, + lower_bound=lower_bound, + disable_recompute=disable_recompute ) ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) if use_gate_in_kernel: @@ -250,21 +260,21 @@ def test_chunk( assert_close("dg", ref_dg, tri_dg, 0.02) assert_close("db", ref_db, tri_db, 0.02) if use_gate_in_kernel: - assert_close("dA", ref_dA, tri_dA, 0.001) - assert_close("dbias", ref_dbias, tri_dbias, 0.001) + assert_close("dA", ref_dA, tri_dA, 0.003) + assert_close("dbias", ref_dbias, tri_dbias, 0.008) assert_close("dh0", ref_dh0, tri_dh0, 0.008) @pytest.mark.parametrize( - ("H", "D", "mask_p", "cu_seqlens", "dtype", "use_tma"), + ("H", "D", "mask_p", "cu_seqlens", "dtype", "use_gate_in_kernel", "safe_gate", "disable_recompute"), [ - pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}-tma{}".format(*test)) + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}-gate{}-safe_gate{}-disable_recompute{}".format(*test)) for test in [ - (4, 60, 0.1, [0, 15], torch.float16, False), - (4, 64, 0.9, [0, 256, 500, 1000], torch.float16, False), - (4, 128, 0.5, [0, 256, 500, 1000], torch.float16, False), - (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16, False), - (4, 256, 0, [0, 100, 300, 1200, 3000, 4096], torch.float16, False), + (4, 60, 0.1, [0, 15], torch.float16, True, False, False), + (4, 64, 0.9, [0, 256, 500, 1000], torch.float16, True, False, False), + (4, 128, 0.5, [0, 256, 500, 1000], torch.float16, False, False, False), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16, True, False, False), + (4, 256, 0, [0, 100, 300, 1200, 3000, 4096], torch.float16, False, True, True), ] ], ) @@ -274,16 +284,14 @@ def test_chunk_varlen( mask_p: float, cu_seqlens: list[int], dtype: torch.dtype, - use_tma: bool, + use_gate_in_kernel: bool, + safe_gate: bool, + disable_recompute: bool, ): - if not use_tma: - os.environ["FLA_USE_TMA"] = "0" - else: - os.environ["FLA_USE_TMA"] = "1" torch.manual_seed(42) - os.environ["TRITON_F32_DEFAULT"] = "ieee" # randomly split the sequence into N segments cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + cu_seqlens_cpu = cu_seqlens.cpu() T = cu_seqlens[-1] N = len(cu_seqlens) - 1 @@ -291,14 +299,25 @@ def test_chunk_varlen( q = torch.randn((1, T, H, D), dtype=dtype) k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) v = torch.randn((1, T, H, D), dtype=dtype) - g = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)) + g = torch.randn(1, T, H, D, dtype=torch.float if not use_gate_in_kernel else dtype) + if use_gate_in_kernel: + A_log = torch.log(torch.randn(1, 1, H, 1, dtype=torch.float32, device=device).uniform_(1, 16)) + dt_bias = torch.randn(H * D, dtype=torch.float32, device=device) + else: + g = F.logsigmoid(g) + g = g * (torch.rand_like(g) > mask_p) mask = torch.rand_like(g) > mask_p g = g * mask + (~mask) * (-1000) + if safe_gate: + assert use_gate_in_kernel is False + g = g.clamp(-5, 0) beta = torch.rand(1, T, H, dtype=dtype).sigmoid() h0 = torch.randn((N, H, D, D), dtype=torch.float32) q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, beta, h0)) + if use_gate_in_kernel: + A_log, dt_bias = map(lambda x: x.to(device).requires_grad_(), (A_log, dt_bias)) do = torch.randn_like(v) dht = torch.rand_like(h0) @@ -308,13 +327,22 @@ def test_chunk_varlen( v=v.clone(), g=g.clone(), beta=beta.clone(), + A_log=(A_log.clone() if use_gate_in_kernel else None), + dt_bias=(dt_bias.clone() if use_gate_in_kernel else None), initial_state=h0.clone(), output_final_state=True, cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + use_gate_in_kernel=use_gate_in_kernel, + safe_gate=safe_gate, + disable_recompute=disable_recompute ) ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None + if use_gate_in_kernel: + tri_dA, A_log.grad = A_log.grad, None + tri_dbias, dt_bias.grad = dt_bias.grad, None ref = [] ref_ht = [] @@ -324,7 +352,8 @@ def test_chunk_varlen( k=k[:, cu_seqlens[i]: cu_seqlens[i + 1]], # k is already normalized v=v[:, cu_seqlens[i]: cu_seqlens[i + 1]], beta=beta[:, cu_seqlens[i]: cu_seqlens[i + 1]], - g=g[:, cu_seqlens[i]: cu_seqlens[i + 1]], + g=(naive_kda_gate(g[:, cu_seqlens[i]: cu_seqlens[i + 1]].to(torch.float), A_log.to(torch.float), + dt_bias.to(torch.float)) if use_gate_in_kernel else g[:, cu_seqlens[i]: cu_seqlens[i + 1]]), initial_state=h0[i], output_final_state=True, ) @@ -335,7 +364,9 @@ def test_chunk_varlen( ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad - + if use_gate_in_kernel: + ref_dA, A_log.grad = A_log.grad, None + ref_dbias, dt_bias.grad = dt_bias.grad, None assert_close("o", ref, tri, 0.005) assert_close("ht", ref_ht, tri_ht, 0.005) assert_close("dq", ref_dq, tri_dq, 0.007) @@ -344,7 +375,9 @@ def test_chunk_varlen( assert_close("dg", ref_dg, tri_dg, 0.015) assert_close("db", ref_db, tri_db, 0.015) assert_close("dh0", ref_dh0, tri_dh0, 0.007) - os.environ["FLA_USE_TMA"] = "0" + if use_gate_in_kernel: + assert_close("dA", ref_dA, tri_dA, 0.008) + assert_close("dbias", ref_dbias, tri_dbias, 0.005) @pytest.mark.parametrize( From 97c085f58eeaff850f5f687179c2adbaee8576ed Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 01:40:30 +0000 Subject: [PATCH 02/10] add doc strings --- fla/ops/kda/chunk.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py index 336c2383f7..e90adadc7c 100644 --- a/fla/ops/kda/chunk.py +++ b/fla/ops/kda/chunk.py @@ -406,6 +406,18 @@ def chunk_kda( cu_seqlens_cpu (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. + safe_gate (bool): + Whether the kernel can assume the input gate values `g` are in a safe range. + When `True`, the kernel can use M=16 TensorCore acceleration. + The safe range is approximately [-5, 0). Default: `False`. + lower_bound (Optional[float]): + Lower bound for the forget gate activation function when `use_gate_in_kernel=True`. + This parameter modifies the internal forget gate activation and is recommended + to be set to `-5` when `safe_gate` is enabled. Default: `None`. + disable_recompute (bool): + Whether to disable gradient recomputation in the kernel. When `True`, the kernel + will save all intermediate activations for backward pass, which is beneficial + for training small models at the cost of increased memory usage. Default: `False`. Returns: o (torch.Tensor): @@ -473,6 +485,7 @@ def chunk_kda( A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." + assert k.shape[-1] <= 256, "Currently we only support key headdim <=256 for KDA :-(" assert beta.shape == q.shape[:3], "beta must be of shape (batch size, seq len, num of head)." assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." if scale is None: From e9a06058f2a3ba099bcff028adf08f0e06e487b5 Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 01:41:43 +0000 Subject: [PATCH 03/10] fix typos --- fla/ops/kda/chunk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py index e90adadc7c..ed7f769c46 100644 --- a/fla/ops/kda/chunk.py +++ b/fla/ops/kda/chunk.py @@ -217,7 +217,7 @@ def forward( if use_gate_in_kernel: g_org = g if safe_gate: - assert lower_bound is not None, "lower_bound must be set when use sage_gate" + assert lower_bound is not None, "lower_bound must be set when use safe_gate" g = kda_gate_chunk_cumsum( g=g_org, A_log=A_log, From 21a9a0b20d3f3c5683ae0683f59eeddaf0db79ef Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 01:45:49 +0000 Subject: [PATCH 04/10] remove lines Signed-off-by: lizhiyuan --- fla/ops/kda/chunk.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py index ed7f769c46..a2de0f1194 100644 --- a/fla/ops/kda/chunk.py +++ b/fla/ops/kda/chunk.py @@ -335,9 +335,6 @@ def backward( dyg=dg, lower_bound=ctx.lower_bound ) - dA = dA.to(A_log) - if dt_bias is not None: - dbias = dbias.to(dt_bias) else: dg = chunk_local_cumsum( dg, From e71d48fcc3cb8f7de3deb40fcf0144617a3c24c3 Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 01:56:45 +0000 Subject: [PATCH 05/10] clean --- fla/ops/kda/chunk.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py index a2de0f1194..664402e85b 100644 --- a/fla/ops/kda/chunk.py +++ b/fla/ops/kda/chunk.py @@ -354,8 +354,8 @@ def chunk_kda( v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, use_gate_in_kernel: bool = False, From 970f78cfce72cf7a8b0351b4d17d45d1dffaf1bb Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 02:23:26 +0000 Subject: [PATCH 06/10] Add check --- fla/ops/kda/chunk.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py index 664402e85b..523a9944c4 100644 --- a/fla/ops/kda/chunk.py +++ b/fla/ops/kda/chunk.py @@ -480,6 +480,11 @@ def chunk_kda( if use_gate_in_kernel: assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") + if safe_gate: + if lower_bound is None: + raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") + if not (-5 <= lower_bound < 0): + raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." assert k.shape[-1] <= 256, "Currently we only support key headdim <=256 for KDA :-(" From 3896ef426b8f67e00137d5496f1097fa1a653f15 Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 02:35:45 +0000 Subject: [PATCH 07/10] update bmk --- benchmarks/ops/benchmark_kda.py | 24 ++++++------------------ 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/benchmarks/ops/benchmark_kda.py b/benchmarks/ops/benchmark_kda.py index a5801c8a1c..7558cdb079 100644 --- a/benchmarks/ops/benchmark_kda.py +++ b/benchmarks/ops/benchmark_kda.py @@ -1,5 +1,4 @@ -import os import torch import triton @@ -38,21 +37,11 @@ def benchmark(T, provider): dtype = torch.bfloat16 B, H, D = 1, 16, 128 - # Set TMA environment variable based on provider - original_tma_env = os.environ.get('FLA_USE_TMA', '0') - - if provider.endswith('_no_tma'): - os.environ['FLA_USE_TMA'] = '0' - provider_base = provider.replace('_no_tma', '') - else: - os.environ['FLA_USE_TMA'] = '1' - provider_base = provider - quantiles = [0.5, 0.2, 0.8] results = 0, 0, 0 do = torch.randn(B, T, H, D, dtype=dtype, device=device) - if provider_base == 'gdn': + if provider == 'gdn': q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) @@ -69,7 +58,7 @@ def benchmark(T, provider): )[0].backward(do), quantiles=quantiles, ) - elif provider_base == 'attn': + elif provider == 'attn': q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) @@ -81,7 +70,7 @@ def benchmark(T, provider): ).backward(do), quantiles=quantiles, ) - elif provider_base == 'comba': + elif provider == 'comba': q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) p = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) @@ -100,7 +89,7 @@ def benchmark(T, provider): )[0].backward(do), quantiles=quantiles, ) - elif provider_base == 'kda': + elif provider == 'kda': q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) @@ -114,10 +103,11 @@ def benchmark(T, provider): g=g, beta=beta, use_qk_l2norm_in_kernel=True, + safe_gate=True, # Plese Carefully read doc strings )[0].backward(do), quantiles=quantiles, ) - elif provider_base == 'dplr': + elif provider == 'dplr': q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) a = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) @@ -137,8 +127,6 @@ def benchmark(T, provider): quantiles=quantiles, ) - # Restore original TMA environment variable - os.environ['FLA_USE_TMA'] = original_tma_env return results From 2e003e35597b3d9b1b579a5d41ced6a16f01f0b9 Mon Sep 17 00:00:00 2001 From: lizhiyuan Date: Tue, 30 Dec 2025 03:33:28 +0000 Subject: [PATCH 08/10] add dA warning --- tests/ops/test_kda.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/ops/test_kda.py b/tests/ops/test_kda.py index cfe994be3c..6719c266e5 100644 --- a/tests/ops/test_kda.py +++ b/tests/ops/test_kda.py @@ -260,7 +260,7 @@ def test_chunk( assert_close("dg", ref_dg, tri_dg, 0.02) assert_close("db", ref_db, tri_db, 0.02) if use_gate_in_kernel: - assert_close("dA", ref_dA, tri_dA, 0.003) + assert_close("dA", ref_dA, tri_dA, 0.003, warning=True) assert_close("dbias", ref_dbias, tri_dbias, 0.008) assert_close("dh0", ref_dh0, tri_dh0, 0.008) @@ -376,7 +376,7 @@ def test_chunk_varlen( assert_close("db", ref_db, tri_db, 0.015) assert_close("dh0", ref_dh0, tri_dh0, 0.007) if use_gate_in_kernel: - assert_close("dA", ref_dA, tri_dA, 0.008) + assert_close("dA", ref_dA, tri_dA, 0.008, warning=True) assert_close("dbias", ref_dbias, tri_dbias, 0.005) From 86867caabd0ab2c7c1222bb4e14df72388c57b67 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E8=87=B4=E8=BF=9C?= Date: Tue, 30 Dec 2025 15:07:09 +0800 Subject: [PATCH 09/10] [KDA]: Support FP16 Gate (train/flash-linear-attention!43) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit load g 的时候手动 to fp32，解耦对 fp32 g 的需求 --- fla/ops/common/chunk_delta_h.py | 20 ++++++++++---------- fla/ops/gla/chunk.py | 2 +- fla/ops/kda/chunk_bwd.py | 4 ++-- fla/ops/kda/chunk_intra.py | 12 ++++++------ fla/ops/kda/wy_fast.py | 4 ++-- 5 files changed, 21 insertions(+), 21 deletions(-) diff --git a/fla/ops/common/chunk_delta_h.py b/fla/ops/common/chunk_delta_h.py index 3b2ca04cca..0cdda2724c 100644 --- a/fla/ops/common/chunk_delta_h.py +++ b/fla/ops/common/chunk_delta_h.py @@ -146,7 +146,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( m_t = (i_t * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + bos * H + last_idx * H + i_h) 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 = tl.load(p_g, boundary_check=(0,)).to(tl.float32) if USE_EXP2: b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) @@ -163,28 +163,28 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k1, mask=(o_k1 < K), other=0.) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) if USE_EXP2: b_h1 *= exp2(b_gk_last1)[:, None] else: b_h1 *= exp(b_gk_last1)[:, None] if K > 64: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k2, mask=(o_k2 < K), other=0.) + b_gk_last2 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) if USE_EXP2: b_h2 *= exp2(b_gk_last2)[:, None] else: b_h2 *= exp(b_gk_last2)[:, None] if K > 128: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k3, mask=(o_k3 < K), other=0.) + b_gk_last3 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) if USE_EXP2: b_h3 *= exp2(b_gk_last3)[:, None] else: b_h3 *= exp(b_gk_last3)[:, None] if K > 192: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k4, mask=(o_k4 < K), other=0.) + b_gk_last4 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) if USE_EXP2: b_h4 *= exp2(b_gk_last4)[:, None] else: @@ -335,7 +335,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( if USE_G: bg_last = tl.load(g + (bos + last_idx) * H + i_h) 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 = tl.load(p_g, boundary_check=(0,)).to(tl.float32) if USE_EXP2: bg_last_exp = exp2(bg_last) b_g_exp = exp2(b_g) @@ -354,7 +354,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + last_idx * H*K + o_k1, mask=(o_k1 < K), other=0.) + b_gk_last1 = tl.load(gk + last_idx * H*K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) if K > 64: @@ -362,7 +362,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + last_idx * H*K + o_k2, mask=(o_k2 < K), other=0.) + b_gk_last2 = tl.load(gk + last_idx * H*K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) if K > 128: @@ -370,7 +370,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + last_idx * H*K + o_k3, mask=(o_k3 < K), other=0.) + b_gk_last3 = tl.load(gk + last_idx * H*K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) if K > 192: @@ -378,7 +378,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_k = tl.load(p_k, boundary_check=(0, 1)) if USE_GK: o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + last_idx * H*K + o_k4, mask=(o_k4 < K), other=0.) + b_gk_last4 = tl.load(gk + last_idx * H*K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) if USE_G: diff --git a/fla/ops/gla/chunk.py b/fla/ops/gla/chunk.py index 1f0a6c1c09..2e8bdd45a3 100644 --- a/fla/ops/gla/chunk.py +++ b/fla/ops/gla/chunk.py @@ -349,7 +349,7 @@ def chunk_gla_fwd_kernel_o( # [BT, BK] b_q = tl.load(p_q, boundary_check=(0, 1)) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) # [BT, BK] if USE_EXP2: b_qg = (b_q * exp2(b_g)).to(b_q.dtype) diff --git a/fla/ops/kda/chunk_bwd.py b/fla/ops/kda/chunk_bwd.py index 14d2184a58..34c61439a4 100644 --- a/fla/ops/kda/chunk_bwd.py +++ b/fla/ops/kda/chunk_bwd.py @@ -188,10 +188,10 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) p_gn = g + (min(T, i_t * BT + BT) - 1) * H*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0) + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) b_dq = tl.zeros([BT, BK], dtype=tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) diff --git a/fla/ops/kda/chunk_intra.py b/fla/ops/kda/chunk_intra.py index f41fd391bf..8cd30016aa 100644 --- a/fla/ops/kda/chunk_intra.py +++ b/fla/ops/kda/chunk_intra.py @@ -420,7 +420,7 @@ def chunk_kda_bwd_kernel_intra( db += (i_k * all + bos) * H + i_h p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) @@ -430,7 +430,7 @@ def chunk_kda_bwd_kernel_intra( if i_i > 0: p_gn = g + i_ti * H*K + o_k # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(0, i_i): p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) @@ -522,7 +522,7 @@ def chunk_kda_bwd_kernel_intra( if i_i < NC - 1: p_gn = g + (min(i_ti + BC, T) - 1) * H*K + o_k # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(i_i + 1, NC): p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) @@ -535,7 +535,7 @@ def chunk_kda_bwd_kernel_intra( # [BC, BK] b_q = tl.load(p_q, boundary_check=(0, 1)) b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) # [BC, BC] b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) @@ -562,7 +562,7 @@ def chunk_kda_bwd_kernel_intra( b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) else: p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) b_q = tl.load(p_q, boundary_check=(0, 1)) p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) @@ -697,7 +697,7 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( # current block, keep numerical stability by subtracting the left boundary # less than 85 to avoid overflow in exp2 - b_gm = b_g - b_gn + b_gm = (b_g - b_gn).to(tl.float32) b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) diff --git a/fla/ops/kda/wy_fast.py b/fla/ops/kda/wy_fast.py index 828f3349f0..94aa1441ff 100644 --- a/fla/ops/kda/wy_fast.py +++ b/fla/ops/kda/wy_fast.py @@ -79,7 +79,7 @@ def recompute_w_u_fwd_kda_kernel( b_kb = b_k * b_b[:, None] p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) b_kb *= exp2(b_gk) if STORE_QG: p_q = tl.make_block_ptr(q + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) @@ -91,7 +91,7 @@ def recompute_w_u_fwd_kda_kernel( last_idx = min(i_t * BT + BT, T) - 1 o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.) + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.).to(tl.float32) b_kg = b_k * tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], exp2(b_gn[None, :] - b_gk), 0) p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) From 990535b7a614bd7e249cacce3918c833429c5950 Mon Sep 17 00:00:00 2001 From: Zhiyuan Li Date: Tue, 30 Dec 2025 21:47:22 +0800 Subject: [PATCH 10/10] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20chunk=5Fdelta=5Fh.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- fla/ops/common/chunk_delta_h.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fla/ops/common/chunk_delta_h.py b/fla/ops/common/chunk_delta_h.py index 0cdda2724c..c57536af74 100644 --- a/fla/ops/common/chunk_delta_h.py +++ b/fla/ops/common/chunk_delta_h.py @@ -144,7 +144,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( last_idx = min((i_t + 1) * BT, T) - 1 if USE_G: m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) 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,)).to(tl.float32) if USE_EXP2: @@ -333,7 +333,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( last_idx = min((i_t + 1) * BT, T) - 1 if USE_G: - bg_last = tl.load(g + (bos + last_idx) * H + i_h) + bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) 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,)).to(tl.float32) if USE_EXP2: