diff --git a/fla/ops/common/backends/tilelang/__init__.py b/fla/ops/common/backends/tilelang/__init__.py index d0d411a3b2..7d710dda0a 100644 --- a/fla/ops/common/backends/tilelang/__init__.py +++ b/fla/ops/common/backends/tilelang/__init__.py @@ -65,15 +65,25 @@ def chunk_bwd_dqkwg_verifier( return False, "TileLang backend only supports gated case (g != None)" if g_gamma is not None: return False, "TileLang backend does not support g_gamma" + if q.dtype not in (torch.float16, torch.bfloat16): + return False, f"TileLang backend only supports fp16/bf16 model tensors, got {q.dtype}" + if chunk_size != 64: + return False, f"TileLang backend requires chunk_size == 64, got {chunk_size}" + if state_v_first: + return False, "TileLang backend does not support state_v_first" if v.shape[2] % k.shape[2] != 0: return False, ( f"TileLang backend requires num_v_heads (HV={v.shape[2]}) to be divisible by " f"num_qk_heads (H={k.shape[2]}); HV % H must be 0 for GVA" ) - if h.dtype != q.dtype: + if h.dtype != dh.dtype: return False, ( - f"TileLang backend requires h.dtype == q.dtype (got h={h.dtype}, q={q.dtype}); " - "e.g. simple_gla's bwd keeps h/dh in fp32 for h·dh reduction precision" + f"TileLang backend requires h.dtype == dh.dtype (got h={h.dtype}, dh={dh.dtype})" + ) + if h.dtype not in (q.dtype, torch.float32): + return False, ( + f"TileLang backend requires h/dh dtype to be q.dtype or fp32 " + f"(got h={h.dtype}, q={q.dtype})" ) return True, None diff --git a/fla/ops/common/backends/tilelang/chunk_bwd.py b/fla/ops/common/backends/tilelang/chunk_bwd.py index bf7576c90a..53f6e84961 100644 --- a/fla/ops/common/backends/tilelang/chunk_bwd.py +++ b/fla/ops/common/backends/tilelang/chunk_bwd.py @@ -5,6 +5,8 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors +from pathlib import Path + import tilelang import tilelang.language as T import torch @@ -13,11 +15,14 @@ from fla.ops.utils import prepare_chunk_indices from fla.utils import check_shared_mem +_CUDA126_FP8_E8M0_STUB = Path(__file__).with_name("cuda126_fp8_e8m0_stub.cuh") +_TILELANG_COMPILE_FLAGS = ["-include", str(_CUDA126_FP8_E8M0_STUB)] + @tilelang.jit(pass_configs={ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) +}, compile_flags=_TILELANG_COMPILE_FLAGS) def _build_kernel( B, H, @@ -31,6 +36,7 @@ def _build_kernel( hD1, hD2, dtype_str, + state_dtype_str, USE_G, USE_DW, STATE_V_FIRST, @@ -39,6 +45,7 @@ def _build_kernel( ): dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} _dtype = dtype_map[dtype_str] + _state_dtype = dtype_map[state_dtype_str] NV = tilelang.cdiv(V, BV) threads = num_warps * 32 tile_hD1, tile_hD2 = (BV, BK) if STATE_V_FIRST else (BK, BV) @@ -53,6 +60,7 @@ def _build_kernel( _hD1, _hD2, _thD1, _thD2 = hD1, hD2, tile_hD1, tile_hD2 _threads = threads _USE_G, _USE_DW = USE_G, USE_DW + _CAST_STATE_FOR_MMA = state_dtype_str != dtype_str _TS, _VAR = STATE_V_FIRST, IS_VARLEN # T, NT, total_h are dynamic (vary with sequence length, no recompilation). @@ -91,8 +99,11 @@ def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, # -- shared tiles -- s_v = T.alloc_shared((_BT, _BV), _dtype) s_do = T.alloc_shared((_BT, _BV), _dtype) - s_h = T.alloc_shared((_thD1, _thD2), _dtype) - s_dh = T.alloc_shared((_thD1, _thD2), _dtype) + s_h = T.alloc_shared((_thD1, _thD2), _state_dtype) + s_dh = T.alloc_shared((_thD1, _thD2), _state_dtype) + if _CAST_STATE_FOR_MMA: + s_h_mma = T.alloc_shared((_thD1, _thD2), _dtype) + s_dh_mma = T.alloc_shared((_thD1, _thD2), _dtype) # dg_last accumulator (shared scalar, accumulated across V-loop) if _USE_G: @@ -105,18 +116,23 @@ def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, for i_v_py in T.Pipelined(_NV, num_stages=2): v_off_c = i_v_py * _BV - T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v) - T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do) + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) if _TS: - T.copy(h[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_h) - T.copy(dh[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_dh) + T.copy(h[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_h, disable_tma=True) + T.copy(dh[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_dh, disable_tma=True) else: - T.copy(h[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h) - T.copy(dh[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh) + T.copy(h[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h, disable_tma=True) + T.copy(dh[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh, disable_tma=True) T.gemm(s_do, s_v, b_ds, transpose_B=True) + if _CAST_STATE_FOR_MMA: + for _i, _j in T.Parallel(_thD1, _thD2): + s_h_mma[_i, _j] = T.cast(s_h[_i, _j], _dtype) + s_dh_mma[_i, _j] = T.cast(s_dh[_i, _j], _dtype) + # h·dh reduction must precede the gemms that consume s_h/s_dh: # the downstream gemms are what the pipeline recognizes as consumers, # so they act as the barrier against the next iter's prefetch. @@ -131,19 +147,33 @@ def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, s_dg_last_acc[0] = s_dg_last_acc[0] + f_hdh_scalar[0] if _TS: - T.gemm(s_do, s_h, b_dq) - T.gemm(s_v, s_dh, b_dk) + if _CAST_STATE_FOR_MMA: + T.gemm(s_do, s_h_mma, b_dq) + T.gemm(s_v, s_dh_mma, b_dk) + else: + T.gemm(s_do, s_h, b_dq) + T.gemm(s_v, s_dh, b_dk) else: - T.gemm(s_do, s_h, b_dq, transpose_B=True) - T.gemm(s_v, s_dh, b_dk, transpose_B=True) + if _CAST_STATE_FOR_MMA: + T.gemm(s_do, s_h_mma, b_dq, transpose_B=True) + T.gemm(s_v, s_dh_mma, b_dk, transpose_B=True) + else: + T.gemm(s_do, s_h, b_dq, transpose_B=True) + T.gemm(s_v, s_dh, b_dk, transpose_B=True) if _USE_DW: s_dv = T.alloc_shared((_BT, _BV), _dtype) - T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_dv) + T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_dv, disable_tma=True) if _TS: - T.gemm(s_dv, s_h, b_dw) + if _CAST_STATE_FOR_MMA: + T.gemm(s_dv, s_h_mma, b_dw) + else: + T.gemm(s_dv, s_h, b_dw) else: - T.gemm(s_dv, s_h, b_dw, transpose_B=True) + if _CAST_STATE_FOR_MMA: + T.gemm(s_dv, s_h_mma, b_dw, transpose_B=True) + else: + T.gemm(s_dv, s_h, b_dw, transpose_B=True) # ========== store dw (negated, with varlen boundary mask) ========== if _USE_DW: @@ -161,8 +191,8 @@ def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, # q/k are indexed by the shared qk-head, not the value-head. s_q = T.alloc_shared((_BT, _BK), _dtype) s_k = T.alloc_shared((_BT, _BK), _dtype) - T.copy(q[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_q) - T.copy(k[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_k) + T.copy(q[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_k, disable_tma=True) # ========== USE_G path ========== if _USE_G: @@ -262,8 +292,8 @@ def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, def kernel( q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), dq: T.Tensor(dqk_s, _dtype), + h: T.Tensor(h_s, _state_dtype), do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _state_dtype), dq: T.Tensor(dqk_s, _dtype), dk: T.Tensor(dqk_s, _dtype), dw: T.Tensor(dqk_s, _dtype), dv: T.Tensor(v_s, _dtype), dg: T.Tensor(dg_s, T.float32), cu_seqlens: T.Tensor((Ncu_d,), T.int32), @@ -285,8 +315,8 @@ def kernel( def kernel( q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), dq: T.Tensor(dqk_s, _dtype), + h: T.Tensor(h_s, _state_dtype), do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _state_dtype), dq: T.Tensor(dqk_s, _dtype), dk: T.Tensor(dqk_s, _dtype), dw: T.Tensor(dqk_s, _dtype), dv: T.Tensor(v_s, _dtype), dg: T.Tensor(dg_s, T.float32), scale: T.float32, @@ -304,39 +334,1354 @@ def kernel( return kernel -def chunk_bwd_dqkwg_tilelang( - q, - k, - v, - do, - h, - dh, - w=None, - g=None, - g_gamma=None, - dv=None, - scale=None, - state_v_first=False, - cu_seqlens=None, - chunk_size=64, - chunk_indices=None, +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}, compile_flags=_TILELANG_COMPILE_FLAGS) +def _build_kernel_k_inner( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + STATE_V_FIRST, + num_warps=4, ): - B, T, H, K = k.shape - HV, V = v.shape[2], v.shape[-1] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - IS_VARLEN = cu_seqlens is not None + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + _dtype = dtype_map[dtype_str] + _state_dtype = dtype_map[state_dtype_str] + NV = tilelang.cdiv(V, BV) + threads = num_warps * 32 - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NK = triton.cdiv(K, BK) - if scale is None: - scale = K ** -0.5 + _B, _H, _HV, _K, _V = B, H, HV, K, V + _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV + tile_hD1, tile_hD2 = (BV, BK) if STATE_V_FIRST else (BK, BV) + _hD1, _hD2, _thD1, _thD2 = hD1, hD2, tile_hD1, tile_hD2 + _threads = threads + _CAST_STATE_FOR_MMA = state_dtype_str != dtype_str + _TS = STATE_V_FIRST - USE_G = g is not None - USE_DW = w is not None + T_d, total_h_d = T.dynamic("T, total_h") + + qk_s = (_B, T_d, _H, _K) + dqk_s = (_B, T_d, _HV, _K) + v_s = (_B, T_d, _HV, _V) + h_s = (total_h_d, _hD1, _hD2) + g_s = (_B, T_d, _HV) + dg_s = (_B, T_d, _HV) + + @T.prim_func + def kernel( + q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), + h: T.Tensor(h_s, _state_dtype), do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _state_dtype), dq: T.Tensor(dqk_s, _dtype), + dk: T.Tensor(dqk_s, _dtype), dg: T.Tensor(dg_s, T.float32), + scale: T.float32, + ): + with T.Kernel(T.ceildiv(T_d, _BT), _B * _HV, threads=_threads) as (i_t, i_bh): + i_b = i_bh // _HV + i_h = i_bh % _HV + NT_local = T.ceildiv(T_d, _BT) + h_idx = (i_b * NT_local + i_t) * _HV + i_h + t_s = i_t * _BT + + s_g = T.alloc_shared((_BT,), T.float32) + T.copy(g[i_b, t_s:t_s + _BT, i_h], s_g, disable_tma=True) + last_pos = T.max(0, T.min(_BT, T_d - i_t * _BT) - 1) + g_last = s_g[last_pos] + + s_v = T.alloc_shared((_BT, _BV), _dtype) + s_do = T.alloc_shared((_BT, _BV), _dtype) + b_ds = T.alloc_fragment((_BT, _BT), T.float32) + T.clear(b_ds) + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + T.gemm(s_do, s_v, b_ds, transpose_B=True) + + s_ds = T.alloc_shared((_BT, _BT), _dtype) + f_ds = T.alloc_fragment((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + causal = (_i >= _j) & ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + f_ds[_i, _j] = T.if_then_else( + causal, + T.cast(b_ds[_i, _j] * T.exp2(s_g[_i] - s_g[_j]) * scale, _dtype), + T.cast(0, _dtype)) + T.copy(f_ds, s_ds) + + s_dg_acc = T.alloc_shared((_BT,), T.float32) + for _i in T.Parallel(_BT): + s_dg_acc[_i] = 0.0 + T.sync_threads() + + for i_k_py in T.serial(_NK): + k_off = i_k_py * _BK + + b_dq = T.alloc_fragment((_BT, _BK), T.float32) + b_dk = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dq) + T.clear(b_dk) + + s_h = T.alloc_shared((_thD1, _thD2), _state_dtype) + s_dh = T.alloc_shared((_thD1, _thD2), _state_dtype) + if _CAST_STATE_FOR_MMA: + s_h_mma = T.alloc_shared((_thD1, _thD2), _dtype) + s_dh_mma = T.alloc_shared((_thD1, _thD2), _dtype) + + s_hdh_acc = T.alloc_shared((1,), T.float32) + for _i in T.Parallel(1): + s_hdh_acc[0] = 0.0 + T.sync_threads() + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + if _TS: + T.copy(h[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_h, disable_tma=True) + T.copy(dh[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_dh, disable_tma=True) + else: + T.copy(h[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h, disable_tma=True) + T.copy(dh[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh, disable_tma=True) + + if _CAST_STATE_FOR_MMA: + for _i, _j in T.Parallel(_thD1, _thD2): + s_h_mma[_i, _j] = T.cast(s_h[_i, _j], _dtype) + s_dh_mma[_i, _j] = T.cast(s_dh[_i, _j], _dtype) + + f_hdh = T.alloc_fragment((_thD1, _thD2), T.float32) + for _i, _j in T.Parallel(_thD1, _thD2): + f_hdh[_i, _j] = T.cast(s_h[_i, _j], T.float32) * T.cast(s_dh[_i, _j], T.float32) + f_hdh_row = T.alloc_fragment((_thD1,), T.float32) + T.reduce_sum(f_hdh, f_hdh_row, dim=1) + f_hdh_scalar = T.alloc_fragment((1,), T.float32) + T.reduce_sum(f_hdh_row, f_hdh_scalar, dim=0) + s_hdh_acc[0] = s_hdh_acc[0] + f_hdh_scalar[0] + + if _TS: + if _CAST_STATE_FOR_MMA: + T.gemm(s_do, s_h_mma, b_dq) + T.gemm(s_v, s_dh_mma, b_dk) + else: + T.gemm(s_do, s_h, b_dq) + T.gemm(s_v, s_dh, b_dk) + else: + if _CAST_STATE_FOR_MMA: + T.gemm(s_do, s_h_mma, b_dq, transpose_B=True) + T.gemm(s_v, s_dh_mma, b_dk, transpose_B=True) + else: + T.gemm(s_do, s_h, b_dq, transpose_B=True) + T.gemm(s_v, s_dh, b_dk, transpose_B=True) + + s_q = T.alloc_shared((_BT, _BK), _dtype) + s_k = T.alloc_shared((_BT, _BK), _dtype) + T.copy(q[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k, disable_tma=True) + + for _i, _j in T.Parallel(_BT, _BK): + b_dq[_i, _j] = b_dq[_i, _j] * T.exp2(s_g[_i]) * scale + + for _i, _j in T.Parallel(_BT, _BK): + m_t = (i_t * _BT + _i) < T_d + b_dk[_i, _j] = T.if_then_else( + m_t, + b_dk[_i, _j] * T.exp2(-s_g[_i] + g_last), + 0.0) + + b_dg_last = T.alloc_var(T.float32) + b_dg_last = s_hdh_acc[0] * T.exp2(g_last) + + f_prod2 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg2 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod2, f_dg2, dim=1) + f_dkk_scalar = T.alloc_fragment((1,), T.float32) + T.reduce_sum(f_dg2, f_dkk_scalar, dim=0) + b_dg_last = b_dg_last + f_dkk_scalar[0] + + T.gemm(s_ds, s_k, b_dq) + T.gemm(s_ds, s_q, b_dk, transpose_A=True) + + f_prod1 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod1[_i, _j] = b_dq[_i, _j] * T.cast(s_q[_i, _j], T.float32) + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg1 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod1, f_dg1, dim=1) + T.reduce_sum(f_prod2, f_dg2, dim=1) + + for _i in T.Parallel(_BT): + val = T.alloc_var(T.float32) + val = f_dg1[_i] - f_dg2[_i] + val = T.if_then_else(_i == last_pos, val + b_dg_last, val) + s_dg_acc[_i] = s_dg_acc[_i] + val + + f_out = T.alloc_fragment((_BT, _BK), _dtype) + s_out = T.alloc_shared((_BT, _BK), _dtype) + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dq[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dq[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dk[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + + for _i in T.Parallel(_BT): + if (i_t * _BT + _i) < T_d: + dg[i_b, t_s + _i, i_h] = s_dg_acc[_i] + + return kernel + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}, compile_flags=_TILELANG_COMPILE_FLAGS) +def _build_kernel_dv_v_first( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=4, +): + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + _dtype = dtype_map[dtype_str] + _state_dtype = dtype_map[state_dtype_str] + NV = tilelang.cdiv(V, BV) + threads = num_warps * 32 + + _B, _H, _HV, _K, _V = B, H, HV, K, V + _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV + _hD1, _hD2 = hD1, hD2 + _threads = threads + _G = HV // H + _CAST_STATE_FOR_MMA = state_dtype_str != dtype_str + + T_d, total_h_d = T.dynamic("T, total_h") + + qk_s = (_B, T_d, _H, _K) + v_s = (_B, T_d, _HV, _V) + h_s = (total_h_d, _hD1, _hD2) + g_s = (_B, T_d, _HV) + + @T.prim_func + def kernel( + q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), + g: T.Tensor(g_s, T.float32), do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _state_dtype), dv: T.Tensor(v_s, _dtype), + scale: T.float32, + ): + with T.Kernel(_NV, T.ceildiv(T_d, _BT), _B * _HV, threads=_threads) as (i_v, i_t, i_bh): + i_b = i_bh // _HV + i_h = i_bh % _HV + i_hqk = i_h // _G + NT_local = T.ceildiv(T_d, _BT) + h_idx = (i_b * NT_local + i_t) * _HV + i_h + t_s = i_t * _BT + v_off = i_v * _BV + + s_g = T.alloc_shared((_BT,), T.float32) + T.copy(g[i_b, t_s:t_s + _BT, i_h], s_g, disable_tma=True) + last_pos = T.max(0, T.min(_BT, T_d - i_t * _BT) - 1) + g_last = s_g[last_pos] + + b_A = T.alloc_fragment((_BT, _BT), T.float32) + b_dv = T.alloc_fragment((_BT, _BV), T.float32) + T.clear(b_A) + T.clear(b_dv) + + s_q = T.alloc_shared((_BT, _BK), _dtype) + s_k = T.alloc_shared((_BT, _BK), _dtype) + s_dh = T.alloc_shared((_BV, _BK), _state_dtype) + if _CAST_STATE_FOR_MMA: + s_dh_mma = T.alloc_shared((_BV, _BK), _dtype) + + for i_k_py in T.Pipelined(_NK, num_stages=2): + k_off = i_k_py * _BK + + T.copy(q[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_k, disable_tma=True) + T.copy(dh[h_idx, v_off:v_off + _BV, k_off:k_off + _BK], s_dh, disable_tma=True) + + T.gemm(s_k, s_q, b_A, transpose_B=True) + + if _CAST_STATE_FOR_MMA: + for _i, _j in T.Parallel(_BV, _BK): + s_dh_mma[_i, _j] = T.cast(s_dh[_i, _j], _dtype) + T.gemm(s_k, s_dh_mma, b_dv, transpose_B=True) + else: + T.gemm(s_k, s_dh, b_dv, transpose_B=True) + + for _i, _j in T.Parallel(_BT, _BV): + m_t = (i_t * _BT + _i) < T_d + b_dv[_i, _j] = T.if_then_else( + m_t, + b_dv[_i, _j] * T.exp2(-s_g[_i] + g_last), + 0.0) + + f_A = T.alloc_fragment((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + m_i = (i_t * _BT + _i) < T_d + m_j = (i_t * _BT + _j) < T_d + causal = (_i <= _j) & m_i & m_j + f_A[_i, _j] = T.if_then_else( + causal, + T.cast(b_A[_i, _j] * T.exp2(s_g[_j] - s_g[_i]) * scale, _dtype), + T.cast(0, _dtype)) + + s_A = T.alloc_shared((_BT, _BT), _dtype) + s_do = T.alloc_shared((_BT, _BV), _dtype) + T.copy(f_A, s_A) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_do, disable_tma=True) + T.gemm(s_A, s_do, b_dv) + + f_out = T.alloc_fragment((_BT, _BV), _dtype) + s_out = T.alloc_shared((_BT, _BV), _dtype) + for _i, _j in T.Parallel(_BT, _BV): + f_out[_i, _j] = T.cast(b_dv[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BV): + if (i_t * _BT + _i) < T_d: + dv[i_b, t_s + _i, i_h, v_off + _j] = s_out[_i, _j] + + return kernel + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}, compile_flags=_TILELANG_COMPILE_FLAGS) +def _build_kernel_k_inner_shadow_state( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + num_warps=4, +): + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + _dtype = dtype_map[dtype_str] + NV = tilelang.cdiv(V, BV) + threads = num_warps * 32 + + _B, _H, _HV, _K, _V = B, H, HV, K, V + _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV + _hD1, _hD2 = hD1, hD2 + _threads = threads + + T_d, total_h_d = T.dynamic("T, total_h") + + qk_s = (_B, T_d, _H, _K) + dqk_s = (_B, T_d, _HV, _K) + v_s = (_B, T_d, _HV, _V) + h_s = (total_h_d, _hD1, _hD2) + hdh_s = (total_h_d,) + g_s = (_B, T_d, _HV) + dg_s = (_B, T_d, _HV) + + @T.prim_func + def kernel( + q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), + h_mma: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), + dh_mma: T.Tensor(h_s, _dtype), hdh_last: T.Tensor(hdh_s, T.float32), + dq: T.Tensor(dqk_s, _dtype), dk: T.Tensor(dqk_s, _dtype), + dg: T.Tensor(dg_s, T.float32), scale: T.float32, + ): + with T.Kernel(T.ceildiv(T_d, _BT), _B * _HV, threads=_threads) as (i_t, i_bh): + i_b = i_bh // _HV + i_h = i_bh % _HV + NT_local = T.ceildiv(T_d, _BT) + h_idx = (i_b * NT_local + i_t) * _HV + i_h + t_s = i_t * _BT + + s_g = T.alloc_shared((_BT,), T.float32) + T.copy(g[i_b, t_s:t_s + _BT, i_h], s_g, disable_tma=True) + last_pos = T.max(0, T.min(_BT, T_d - i_t * _BT) - 1) + g_last = s_g[last_pos] + hdh_scaled = hdh_last[h_idx] * T.exp2(g_last) + + s_v = T.alloc_shared((_BT, _BV), _dtype) + s_do = T.alloc_shared((_BT, _BV), _dtype) + b_ds = T.alloc_fragment((_BT, _BT), T.float32) + T.clear(b_ds) + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + T.gemm(s_do, s_v, b_ds, transpose_B=True) + + s_ds = T.alloc_shared((_BT, _BT), _dtype) + f_ds = T.alloc_fragment((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + causal = (_i >= _j) & ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + f_ds[_i, _j] = T.if_then_else( + causal, + T.cast(b_ds[_i, _j] * T.exp2(s_g[_i] - s_g[_j]) * scale, _dtype), + T.cast(0, _dtype)) + T.copy(f_ds, s_ds) + + s_dg_acc = T.alloc_shared((_BT,), T.float32) + for _i in T.Parallel(_BT): + s_dg_acc[_i] = 0.0 + T.sync_threads() + + for i_k_py in T.serial(_NK): + k_off = i_k_py * _BK + + b_dq = T.alloc_fragment((_BT, _BK), T.float32) + b_dk = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dq) + T.clear(b_dk) + + s_h_mma = T.alloc_shared((_BK, _BV), _dtype) + s_dh_mma = T.alloc_shared((_BK, _BV), _dtype) + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + T.copy(h_mma[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h_mma, disable_tma=True) + T.copy(dh_mma[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh_mma, disable_tma=True) + + T.gemm(s_do, s_h_mma, b_dq, transpose_B=True) + T.gemm(s_v, s_dh_mma, b_dk, transpose_B=True) + + s_q = T.alloc_shared((_BT, _BK), _dtype) + s_k = T.alloc_shared((_BT, _BK), _dtype) + T.copy(q[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k, disable_tma=True) + + for _i, _j in T.Parallel(_BT, _BK): + b_dq[_i, _j] = b_dq[_i, _j] * T.exp2(s_g[_i]) * scale + + for _i, _j in T.Parallel(_BT, _BK): + m_t = (i_t * _BT + _i) < T_d + b_dk[_i, _j] = T.if_then_else( + m_t, + b_dk[_i, _j] * T.exp2(-s_g[_i] + g_last), + 0.0) + + b_dg_last = T.alloc_var(T.float32) + b_dg_last = T.if_then_else(i_k_py == 0, hdh_scaled, 0.0) + + f_prod2 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg2 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod2, f_dg2, dim=1) + f_dkk_scalar = T.alloc_fragment((1,), T.float32) + T.reduce_sum(f_dg2, f_dkk_scalar, dim=0) + b_dg_last = b_dg_last + f_dkk_scalar[0] + + T.gemm(s_ds, s_k, b_dq) + T.gemm(s_ds, s_q, b_dk, transpose_A=True) + + f_prod1 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod1[_i, _j] = b_dq[_i, _j] * T.cast(s_q[_i, _j], T.float32) + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg1 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod1, f_dg1, dim=1) + T.reduce_sum(f_prod2, f_dg2, dim=1) + + for _i in T.Parallel(_BT): + val = T.alloc_var(T.float32) + val = f_dg1[_i] - f_dg2[_i] + val = T.if_then_else(_i == last_pos, val + b_dg_last, val) + s_dg_acc[_i] = s_dg_acc[_i] + val + + f_out = T.alloc_fragment((_BT, _BK), _dtype) + s_out = T.alloc_shared((_BT, _BK), _dtype) + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dq[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dq[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dk[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + + for _i in T.Parallel(_BT): + if (i_t * _BT + _i) < T_d: + dg[i_b, t_s + _i, i_h] = s_dg_acc[_i] + + return kernel + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}, compile_flags=_TILELANG_COMPILE_FLAGS) +def _build_kernel_k_inner_dh_shadow( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=4, +): + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + _dtype = dtype_map[dtype_str] + _state_dtype = dtype_map[state_dtype_str] + NV = tilelang.cdiv(V, BV) + threads = num_warps * 32 + + _B, _H, _HV, _K, _V = B, H, HV, K, V + _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV + _hD1, _hD2 = hD1, hD2 + _threads = threads + _CAST_H_FOR_MMA = state_dtype_str != dtype_str + + T_d, total_h_d = T.dynamic("T, total_h") + + qk_s = (_B, T_d, _H, _K) + dqk_s = (_B, T_d, _HV, _K) + v_s = (_B, T_d, _HV, _V) + h_s = (total_h_d, _hD1, _hD2) + hdh_s = (total_h_d,) + g_s = (_B, T_d, _HV) + dg_s = (_B, T_d, _HV) + + @T.prim_func + def kernel( + q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), + h: T.Tensor(h_s, _state_dtype), do: T.Tensor(v_s, _dtype), + dh_mma: T.Tensor(h_s, _dtype), hdh_last: T.Tensor(hdh_s, T.float32), + dq: T.Tensor(dqk_s, _dtype), dk: T.Tensor(dqk_s, _dtype), + dg: T.Tensor(dg_s, T.float32), scale: T.float32, + ): + with T.Kernel(T.ceildiv(T_d, _BT), _B * _HV, threads=_threads) as (i_t, i_bh): + i_b = i_bh // _HV + i_h = i_bh % _HV + NT_local = T.ceildiv(T_d, _BT) + h_idx = (i_b * NT_local + i_t) * _HV + i_h + t_s = i_t * _BT + + s_g = T.alloc_shared((_BT,), T.float32) + T.copy(g[i_b, t_s:t_s + _BT, i_h], s_g, disable_tma=True) + last_pos = T.max(0, T.min(_BT, T_d - i_t * _BT) - 1) + g_last = s_g[last_pos] + hdh_scaled = hdh_last[h_idx] * T.exp2(g_last) + + s_v = T.alloc_shared((_BT, _BV), _dtype) + s_do = T.alloc_shared((_BT, _BV), _dtype) + b_ds = T.alloc_fragment((_BT, _BT), T.float32) + T.clear(b_ds) + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + T.gemm(s_do, s_v, b_ds, transpose_B=True) + + s_ds = T.alloc_shared((_BT, _BT), _dtype) + f_ds = T.alloc_fragment((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + causal = (_i >= _j) & ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + f_ds[_i, _j] = T.if_then_else( + causal, + T.cast(b_ds[_i, _j] * T.exp2(s_g[_i] - s_g[_j]) * scale, _dtype), + T.cast(0, _dtype)) + T.copy(f_ds, s_ds) + + s_dg_acc = T.alloc_shared((_BT,), T.float32) + for _i in T.Parallel(_BT): + s_dg_acc[_i] = 0.0 + T.sync_threads() + + for i_k_py in T.serial(_NK): + k_off = i_k_py * _BK + + b_dq = T.alloc_fragment((_BT, _BK), T.float32) + b_dk = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dq) + T.clear(b_dk) + + s_h = T.alloc_shared((_BK, _BV), _state_dtype) + s_dh_mma = T.alloc_shared((_BK, _BV), _dtype) + if _CAST_H_FOR_MMA: + s_h_mma = T.alloc_shared((_BK, _BV), _dtype) + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + T.copy(h[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h, disable_tma=True) + T.copy(dh_mma[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh_mma, disable_tma=True) + + if _CAST_H_FOR_MMA: + for _i, _j in T.Parallel(_BK, _BV): + s_h_mma[_i, _j] = T.cast(s_h[_i, _j], _dtype) + + if _CAST_H_FOR_MMA: + T.gemm(s_do, s_h_mma, b_dq, transpose_B=True) + else: + T.gemm(s_do, s_h, b_dq, transpose_B=True) + T.gemm(s_v, s_dh_mma, b_dk, transpose_B=True) + + s_q = T.alloc_shared((_BT, _BK), _dtype) + s_k = T.alloc_shared((_BT, _BK), _dtype) + T.copy(q[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k, disable_tma=True) + + for _i, _j in T.Parallel(_BT, _BK): + b_dq[_i, _j] = b_dq[_i, _j] * T.exp2(s_g[_i]) * scale + + for _i, _j in T.Parallel(_BT, _BK): + m_t = (i_t * _BT + _i) < T_d + b_dk[_i, _j] = T.if_then_else( + m_t, + b_dk[_i, _j] * T.exp2(-s_g[_i] + g_last), + 0.0) + + b_dg_last = T.alloc_var(T.float32) + b_dg_last = T.if_then_else(i_k_py == 0, hdh_scaled, 0.0) + + f_prod2 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg2 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod2, f_dg2, dim=1) + f_dkk_scalar = T.alloc_fragment((1,), T.float32) + T.reduce_sum(f_dg2, f_dkk_scalar, dim=0) + b_dg_last = b_dg_last + f_dkk_scalar[0] + + T.gemm(s_ds, s_k, b_dq) + T.gemm(s_ds, s_q, b_dk, transpose_A=True) + + f_prod1 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod1[_i, _j] = b_dq[_i, _j] * T.cast(s_q[_i, _j], T.float32) + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg1 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod1, f_dg1, dim=1) + T.reduce_sum(f_prod2, f_dg2, dim=1) + + for _i in T.Parallel(_BT): + val = T.alloc_var(T.float32) + val = f_dg1[_i] - f_dg2[_i] + val = T.if_then_else(_i == last_pos, val + b_dg_last, val) + s_dg_acc[_i] = s_dg_acc[_i] + val + + f_out = T.alloc_fragment((_BT, _BK), _dtype) + s_out = T.alloc_shared((_BT, _BK), _dtype) + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dq[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dq[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dk[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + + for _i in T.Parallel(_BT): + if (i_t * _BT + _i) < T_d: + dg[i_b, t_s + _i, i_h] = s_dg_acc[_i] + + return kernel + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}, compile_flags=_TILELANG_COMPILE_FLAGS) +def _build_kernel_k_inner_dh_shadow_terminal_dot( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=4, +): + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + _dtype = dtype_map[dtype_str] + _state_dtype = dtype_map[state_dtype_str] + NV = tilelang.cdiv(V, BV) + threads = num_warps * 32 + + _B, _H, _HV, _K, _V = B, H, HV, K, V + _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV + _hD1, _hD2 = hD1, hD2 + _threads = threads + _CAST_H_FOR_MMA = state_dtype_str != dtype_str + + T_d, total_h_d = T.dynamic("T, total_h") + + qk_s = (_B, T_d, _H, _K) + dqk_s = (_B, T_d, _HV, _K) + v_s = (_B, T_d, _HV, _V) + h_s = (total_h_d, _hD1, _hD2) + g_s = (_B, T_d, _HV) + dg_s = (_B, T_d, _HV) + + @T.prim_func + def kernel( + q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), + h: T.Tensor(h_s, _state_dtype), do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _state_dtype), dh_mma: T.Tensor(h_s, _dtype), + dq: T.Tensor(dqk_s, _dtype), dk: T.Tensor(dqk_s, _dtype), + dg: T.Tensor(dg_s, T.float32), scale: T.float32, + ): + with T.Kernel(T.ceildiv(T_d, _BT), _B * _HV, threads=_threads) as (i_t, i_bh): + i_b = i_bh // _HV + i_h = i_bh % _HV + NT_local = T.ceildiv(T_d, _BT) + h_idx = (i_b * NT_local + i_t) * _HV + i_h + t_s = i_t * _BT + + s_g = T.alloc_shared((_BT,), T.float32) + T.copy(g[i_b, t_s:t_s + _BT, i_h], s_g, disable_tma=True) + last_pos = T.max(0, T.min(_BT, T_d - i_t * _BT) - 1) + g_last = s_g[last_pos] + + s_v = T.alloc_shared((_BT, _BV), _dtype) + s_do = T.alloc_shared((_BT, _BV), _dtype) + b_ds = T.alloc_fragment((_BT, _BT), T.float32) + T.clear(b_ds) + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + T.gemm(s_do, s_v, b_ds, transpose_B=True) + + s_ds = T.alloc_shared((_BT, _BT), _dtype) + f_ds = T.alloc_fragment((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + causal = (_i >= _j) & ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + f_ds[_i, _j] = T.if_then_else( + causal, + T.cast(b_ds[_i, _j] * T.exp2(s_g[_i] - s_g[_j]) * scale, _dtype), + T.cast(0, _dtype)) + T.copy(f_ds, s_ds) + + s_dg_acc = T.alloc_shared((_BT,), T.float32) + for _i in T.Parallel(_BT): + s_dg_acc[_i] = 0.0 + T.sync_threads() + + for i_k_py in T.serial(_NK): + k_off = i_k_py * _BK + + b_dq = T.alloc_fragment((_BT, _BK), T.float32) + b_dk = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dq) + T.clear(b_dk) + + s_h = T.alloc_shared((_BK, _BV), _state_dtype) + s_dh = T.alloc_shared((_BK, _BV), _state_dtype) + s_dh_mma = T.alloc_shared((_BK, _BV), _dtype) + if _CAST_H_FOR_MMA: + s_h_mma = T.alloc_shared((_BK, _BV), _dtype) + + s_hdh_acc = T.alloc_shared((1,), T.float32) + for _i in T.Parallel(1): + s_hdh_acc[0] = 0.0 + T.sync_threads() + + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off_c = i_v_py * _BV + + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do, disable_tma=True) + T.copy(h[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h, disable_tma=True) + T.copy(dh[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh, disable_tma=True) + T.copy(dh_mma[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh_mma, disable_tma=True) + + if _CAST_H_FOR_MMA: + for _i, _j in T.Parallel(_BK, _BV): + s_h_mma[_i, _j] = T.cast(s_h[_i, _j], _dtype) + + f_hdh = T.alloc_fragment((_BK, _BV), T.float32) + for _i, _j in T.Parallel(_BK, _BV): + f_hdh[_i, _j] = T.cast(s_h[_i, _j], T.float32) * T.cast(s_dh[_i, _j], T.float32) + f_hdh_row = T.alloc_fragment((_BK,), T.float32) + T.reduce_sum(f_hdh, f_hdh_row, dim=1) + f_hdh_scalar = T.alloc_fragment((1,), T.float32) + T.reduce_sum(f_hdh_row, f_hdh_scalar, dim=0) + s_hdh_acc[0] = s_hdh_acc[0] + f_hdh_scalar[0] + + if _CAST_H_FOR_MMA: + T.gemm(s_do, s_h_mma, b_dq, transpose_B=True) + else: + T.gemm(s_do, s_h, b_dq, transpose_B=True) + T.gemm(s_v, s_dh_mma, b_dk, transpose_B=True) + + s_q = T.alloc_shared((_BT, _BK), _dtype) + s_k = T.alloc_shared((_BT, _BK), _dtype) + T.copy(q[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k, disable_tma=True) + + for _i, _j in T.Parallel(_BT, _BK): + b_dq[_i, _j] = b_dq[_i, _j] * T.exp2(s_g[_i]) * scale + + for _i, _j in T.Parallel(_BT, _BK): + m_t = (i_t * _BT + _i) < T_d + b_dk[_i, _j] = T.if_then_else( + m_t, + b_dk[_i, _j] * T.exp2(-s_g[_i] + g_last), + 0.0) + + b_dg_last = T.alloc_var(T.float32) + b_dg_last = s_hdh_acc[0] * T.exp2(g_last) + + f_prod2 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg2 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod2, f_dg2, dim=1) + f_dkk_scalar = T.alloc_fragment((1,), T.float32) + T.reduce_sum(f_dg2, f_dkk_scalar, dim=0) + b_dg_last = b_dg_last + f_dkk_scalar[0] + + T.gemm(s_ds, s_k, b_dq) + T.gemm(s_ds, s_q, b_dk, transpose_A=True) + + f_prod1 = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_prod1[_i, _j] = b_dq[_i, _j] * T.cast(s_q[_i, _j], T.float32) + f_prod2[_i, _j] = b_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_dg1 = T.alloc_fragment((_BT,), T.float32) + T.reduce_sum(f_prod1, f_dg1, dim=1) + T.reduce_sum(f_prod2, f_dg2, dim=1) + + for _i in T.Parallel(_BT): + val = T.alloc_var(T.float32) + val = f_dg1[_i] - f_dg2[_i] + val = T.if_then_else(_i == last_pos, val + b_dg_last, val) + s_dg_acc[_i] = s_dg_acc[_i] + val + + f_out = T.alloc_fragment((_BT, _BK), _dtype) + s_out = T.alloc_shared((_BT, _BK), _dtype) + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dq[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dq[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dk[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + + for _i in T.Parallel(_BT): + if (i_t * _BT + _i) < T_d: + dg[i_b, t_s + _i, i_h] = s_dg_acc[_i] + + return kernel + + +def _can_use_k_inner_dqkwg( + q, k, v, h, dh, w, g, g_gamma, dv, state_v_first, cu_seqlens, chunk_size, chunk_indices, K, V, NK +): + return ( + g is not None + and g_gamma is None + and w is None + and dv is None + and cu_seqlens is None + and chunk_indices is None + and not state_v_first + and chunk_size == 64 + and q.dtype in (torch.float16, torch.bfloat16) + and h.dtype == torch.float32 + and dh.dtype == torch.float32 + and h.dtype == dh.dtype + and v.shape[2] == k.shape[2] + and K >= 128 + and V >= 128 + and NK >= 2 + ) + + +def chunk_bwd_dqkwg_tilelang_k_inner( + q, + k, + v, + do, + h, + dh, + g, + scale=None, + chunk_size=64, +): + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + BT = chunk_size + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) + dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) + dg = torch.empty(B, T, HV, dtype=torch.float32, device=q.device) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + state_dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[h.dtype] + + num_warps = 4 if min(K, V) >= 64 else 2 + kernel = _build_kernel_k_inner( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + False, + num_warps=num_warps, + ) + kernel(q, k, v, g, h_flat, do, dh_flat, dq, dk, dg, scale) + return dq, dk, None, dg + + +def chunk_bwd_dqkwg_tilelang_k_inner_v_first( + q, + k, + v, + do, + h, + dh, + g, + scale=None, + chunk_size=64, +): + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + BT = chunk_size + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) + dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) + dg = torch.empty(B, T, HV, dtype=torch.float32, device=q.device) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + state_dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[h.dtype] + + num_warps = 4 if min(K, V) >= 64 else 2 + kernel = _build_kernel_k_inner( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + True, + num_warps=num_warps, + ) + kernel(q, k, v, g, h_flat, do, dh_flat, dq, dk, dg, scale) + return dq, dk, None, dg + + +def chunk_bwd_dqkwg_tilelang_k_inner_v_first_d256( + q, + k, + v, + do, + h, + dh, + g, + scale=None, + chunk_size=64, +): + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + if K != 256 or V != 256: + raise ValueError(f"D256 V-first TileLang K-inner requires K == V == 256, got K={K}, V={V}") + BT = chunk_size + BK = 32 + BV = 128 + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) + dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) + dg = torch.empty(B, T, HV, dtype=torch.float32, device=q.device) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + state_dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[h.dtype] + + kernel = _build_kernel_k_inner( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + True, + num_warps=4, + ) + kernel(q, k, v, g, h_flat, do, dh_flat, dq, dk, dg, scale) + return dq, dk, None, dg + + +def chunk_bwd_dv_tilelang_v_first_d256( + q, + k, + g, + do, + dh, + scale=None, + chunk_size=64, +): + B, T, H, K = k.shape + HV, V = do.shape[2], do.shape[-1] + if K != 256 or V != 256: + raise ValueError(f"D256 V-first TileLang dV requires K == V == 256, got K={K}, V={V}") + BT = chunk_size + BK = 64 + BV = 128 + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dv = torch.empty_like(do) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + hD1, hD2 = dh_flat.shape[-2], dh_flat.shape[-1] + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + state_dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[dh.dtype] + + kernel = _build_kernel_dv_v_first( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=4, + ) + kernel(q, k, g, do, dh_flat, dv, scale) + return dv + + +def chunk_bwd_dqkwg_tilelang_k_inner_shadow_state( + q, + k, + v, + do, + h_mma, + dh_mma, + hdh_last, + g, + scale=None, + chunk_size=64, +): + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + BT = chunk_size + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) + dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) + dg = torch.empty(B, T, HV, dtype=torch.float32, device=q.device) + + h_mma_flat = h_mma.reshape(-1, h_mma.shape[-2], h_mma.shape[-1]) + dh_mma_flat = dh_mma.reshape(-1, dh_mma.shape[-2], dh_mma.shape[-1]) + hdh_last_flat = hdh_last.reshape(-1) + hD1, hD2 = h_mma_flat.shape[-2], h_mma_flat.shape[-1] + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + + num_warps = 4 if min(K, V) >= 64 else 2 + kernel = _build_kernel_k_inner_shadow_state( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + num_warps=num_warps, + ) + kernel(q, k, v, g, h_mma_flat, do, dh_mma_flat, hdh_last_flat, dq, dk, dg, scale) + return dq, dk, None, dg + + +def chunk_bwd_dqkwg_tilelang_k_inner_dh_shadow( + q, + k, + v, + do, + h, + dh_mma, + hdh_last, + g, + scale=None, + chunk_size=64, +): + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + BT = chunk_size + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) + dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) + dg = torch.empty(B, T, HV, dtype=torch.float32, device=q.device) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_mma_flat = dh_mma.reshape(-1, dh_mma.shape[-2], dh_mma.shape[-1]) + hdh_last_flat = hdh_last.reshape(-1) + hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + state_dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[h.dtype] + + num_warps = 4 if min(K, V) >= 64 else 2 + kernel = _build_kernel_k_inner_dh_shadow( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=num_warps, + ) + kernel(q, k, v, g, h_flat, do, dh_mma_flat, hdh_last_flat, dq, dk, dg, scale) + return dq, dk, None, dg + + +def chunk_bwd_dqkwg_tilelang_k_inner_dh_shadow_terminal_dot( + q, + k, + v, + do, + h, + dh, + dh_mma, + g, + scale=None, + chunk_size=64, +): + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + BT = chunk_size + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) + dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) + dg = torch.empty(B, T, HV, dtype=torch.float32, device=q.device) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + dh_mma_flat = dh_mma.reshape(-1, dh_mma.shape[-2], dh_mma.shape[-1]) + hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + state_dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[h.dtype] + + num_warps = 4 if min(K, V) >= 64 else 2 + kernel = _build_kernel_k_inner_dh_shadow_terminal_dot( + B, + H, + HV, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=num_warps, + ) + kernel(q, k, v, g, h_flat, do, dh_flat, dh_mma_flat, dq, dk, dg, scale) + return dq, dk, None, dg + + +def chunk_bwd_dqkwg_tilelang( + q, + k, + v, + do, + h, + dh, + w=None, + g=None, + g_gamma=None, + dv=None, + scale=None, + state_v_first=False, + cu_seqlens=None, + chunk_size=64, + chunk_indices=None, +): + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + IS_VARLEN = cu_seqlens is not None + + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + USE_G = g is not None + USE_DW = w is not None + + if _can_use_k_inner_dqkwg( + q, k, v, h, dh, w, g, g_gamma, dv, state_v_first, cu_seqlens, chunk_size, chunk_indices, K, V, NK + ): + return chunk_bwd_dqkwg_tilelang_k_inner( + q=q, + k=k, + v=v, + do=do, + h=h, + dh=dh, + g=g, + scale=scale, + chunk_size=chunk_size, + ) # Outputs — kernel writes dq/dk at value-head (HV) granularity; reduce to # qk-head (H) below when GVA is active. dg stays at HV (per value-head). @@ -349,6 +1694,9 @@ def chunk_bwd_dqkwg_tilelang( dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] + if h.dtype != dh.dtype: + raise ValueError(f"TileLang chunk_bwd_dqkwg requires h.dtype == dh.dtype, got {h.dtype} and {dh.dtype}") + state_dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[h.dtype] # Cache key: B, H, HV, tile sizes, flags. T is dynamic (no recompilation for different seq lengths). # Small head dims (< 64) cannot be warp-partitioned across 4 warps by TileLang's @@ -367,6 +1715,7 @@ def chunk_bwd_dqkwg_tilelang( hD1, hD2, dtype_str, + state_dtype_str, USE_G, USE_DW, state_v_first, diff --git a/fla/ops/common/backends/tilelang/cuda126_fp8_e8m0_stub.cuh b/fla/ops/common/backends/tilelang/cuda126_fp8_e8m0_stub.cuh new file mode 100644 index 0000000000..b33b2c2c67 --- /dev/null +++ b/fla/ops/common/backends/tilelang/cuda126_fp8_e8m0_stub.cuh @@ -0,0 +1,73 @@ +#pragma once + +// TileLang 0.1.9 assumes CUDA 12.6 exposes the FP8 e8m0 CUDA C++ symbols. +// The CUDA 12.6 toolkit paired with this verifier does not define them, yet +// TileLang includes its FP8 helper header even for bf16-only kernels. These +// declarations let non-e8m0 kernels compile; the stubs must never be used for +// real e8m0 computation. + +#include +#include + +#if defined(__CUDACC_VER_MAJOR__) && defined(__CUDACC_VER_MINOR__) && \ + (__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ < 8) + +struct __CUDA_ALIGN__(1) __nv_fp8_e8m0 { + __nv_fp8_storage_t __x; + + __host__ __device__ __nv_fp8_e8m0() = default; + __host__ __device__ explicit __nv_fp8_e8m0(__nv_fp8_storage_t x) : __x(x) {} +}; + +__host__ __device__ __forceinline__ __nv_bfloat16_raw +__nv_cvt_e8m0_to_bf16raw(const __nv_fp8_storage_t) { + __nv_bfloat16_raw out; + out.x = 0U; + return out; +} + +__host__ __device__ __forceinline__ __nv_bfloat162_raw +__nv_cvt_e8m0x2_to_bf162raw(const __nv_fp8x2_storage_t) { + __nv_bfloat162_raw out; + out.x = 0U; + out.y = 0U; + return out; +} + +__host__ __device__ __forceinline__ __nv_fp8_storage_t +__nv_cvt_bfloat16raw_to_e8m0( + const __nv_bfloat16_raw, + const __nv_saturation_t, + const cudaRoundMode) { + return 0U; +} + +__host__ __device__ __forceinline__ __nv_fp8x2_storage_t +__nv_cvt_bfloat162raw_to_e8m0x2( + const __nv_bfloat162_raw, + const __nv_saturation_t, + const cudaRoundMode) { + return 0U; +} + +__host__ __device__ __forceinline__ __nv_fp8_storage_t +__nv_cvt_float_to_e8m0(const float, const __nv_saturation_t, const cudaRoundMode) { + return 0U; +} + +__host__ __device__ __forceinline__ __nv_fp8x2_storage_t +__nv_cvt_float2_to_e8m0x2(const float2, const __nv_saturation_t, const cudaRoundMode) { + return 0U; +} + +__host__ __device__ __forceinline__ __nv_fp8_storage_t +__nv_cvt_double_to_e8m0(const double, const __nv_saturation_t, const cudaRoundMode) { + return 0U; +} + +__host__ __device__ __forceinline__ __nv_fp8x2_storage_t +__nv_cvt_double2_to_e8m0x2(const double2, const __nv_saturation_t, const cudaRoundMode) { + return 0U; +} + +#endif diff --git a/fla/ops/common/chunk_h.py b/fla/ops/common/chunk_h.py index 2ade63a76c..dba26c2c7d 100644 --- a/fla/ops/common/chunk_h.py +++ b/fla/ops/common/chunk_h.py @@ -19,6 +19,7 @@ @triton.heuristics({ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'STORE_MMA_STATE': lambda args: args['h_mma'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( @@ -37,6 +38,7 @@ def chunk_fwd_kernel_h( k, v, h, + h_mma, g, g_gamma, gk, @@ -59,6 +61,7 @@ def chunk_fwd_kernel_h( USE_GV: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, + STORE_MMA_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr, ): @@ -103,12 +106,22 @@ def chunk_fwd_kernel_h( if STATE_V_FIRST: p_h = h + o_h + o_v[:, None] * K + o_k[None, :] m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + if STORE_MMA_STATE: + p_h_mma = h_mma + o_h + o_v[:, None] * K + o_k[None, :] else: p_h = h + o_h + o_k[:, None] * V + o_v[None, :] m_h = (o_k[:, None] < K) & (o_v[None, :] < V) + if STORE_MMA_STATE: + p_h_mma = h_mma + o_h + o_k[:, None] * V + o_v[None, :] if i_t % NTS == 0: tl.store(p_h, (tl.trans(b_h) if STATE_V_FIRST else b_h).to(p_h.dtype.element_ty), mask=m_h) + if STORE_MMA_STATE: + tl.store( + p_h_mma, + (tl.trans(b_h) if STATE_V_FIRST else b_h).to(p_h_mma.dtype.element_ty), + mask=m_h, + ) # [BK, BT] b_k = tl.load(p_k, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) # [BT, BV] @@ -162,6 +175,8 @@ def chunk_fwd_kernel_h( @triton.heuristics({ 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'STORE_MMA_STATE': lambda args: args['dh_mma'] is not None, + 'FUSE_HDH_LAST': lambda args: args['h_for_hdh'] is not None and args['hdh_last'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @triton.autotune( @@ -184,6 +199,9 @@ def chunk_bwd_kernel_dh( gv, do, dh, + dh_mma, + h_for_hdh, + hdh_last, dht, dh0, cu_seqlens, @@ -205,6 +223,8 @@ def chunk_bwd_kernel_dh( USE_GV: tl.constexpr, STORE_INITIAL_STATE_GRADIENT: tl.constexpr, USE_FINAL_STATE_GRADIENT: tl.constexpr, + STORE_MMA_STATE: tl.constexpr, + FUSE_HDH_LAST: tl.constexpr, IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr, ): @@ -245,12 +265,35 @@ def chunk_bwd_kernel_dh( if STATE_V_FIRST: p_dh = dh + o_dh + o_v[:, None] * K + o_k[None, :] m_dh = (o_v[:, None] < V) & (o_k[None, :] < K) + if STORE_MMA_STATE: + p_dh_mma = dh_mma + o_dh + o_v[:, None] * K + o_k[None, :] + if FUSE_HDH_LAST: + p_h_for_hdh = h_for_hdh + o_dh + o_v[:, None] * K + o_k[None, :] else: p_dh = dh + o_dh + o_k[:, None] * V + o_v[None, :] m_dh = (o_k[:, None] < K) & (o_v[None, :] < V) + if STORE_MMA_STATE: + p_dh_mma = dh_mma + o_dh + o_k[:, None] * V + o_v[None, :] + if FUSE_HDH_LAST: + p_h_for_hdh = h_for_hdh + o_dh + o_k[:, None] * V + o_v[None, :] if i_t % (BS // BT) == 0: tl.store(p_dh, (tl.trans(b_dh) if STATE_V_FIRST else b_dh).to(p_dh.dtype.element_ty), mask=m_dh) + if STORE_MMA_STATE: + tl.store( + p_dh_mma, + (tl.trans(b_dh) if STATE_V_FIRST else b_dh).to(p_dh_mma.dtype.element_ty), + mask=m_dh, + ) + if FUSE_HDH_LAST: + b_h_for_hdh = tl.load(p_h_for_hdh, mask=m_dh, other=0.0).to(tl.float32) + if STATE_V_FIRST: + b_h_for_hdh = tl.trans(b_h_for_hdh) + tl.atomic_add( + hdh_last + (boh + i_s) * H + i_h, + tl.sum(b_h_for_hdh * b_dh), + sem='relaxed', + ) last_idx = min(i_t * BT + BT, T) - 1 o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -303,6 +346,45 @@ def chunk_bwd_kernel_dh( tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=(o_k[:, None] < K) & (o_v[None, :] < V)) +@triton.jit +def _chunk_hdh_last_kernel( + h, + dh, + hdh_last, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i_h = tl.program_id(0) + offs = tl.arange(0, BLOCK) + acc = tl.zeros((BLOCK,), dtype=tl.float32) + + for base in range(0, D, BLOCK): + idx = base + offs + mask = idx < D + b_h = tl.load(h + i_h * D + idx, mask=mask, other=0.).to(tl.float32) + b_dh = tl.load(dh + i_h * D + idx, mask=mask, other=0.).to(tl.float32) + acc += b_h * b_dh + + tl.store(hdh_last + i_h, tl.sum(acc, axis=0)) + + +def chunk_hdh_last(h: torch.Tensor, dh: torch.Tensor) -> torch.Tensor: + h_flat = h.reshape(-1, h.shape[-2] * h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2] * dh.shape[-1]) + D = h_flat.shape[1] + block = min(triton.next_power_of_2(D), 1024) + hdh_last = torch.empty(h_flat.shape[0], dtype=torch.float32, device=h.device) + _chunk_hdh_last_kernel[(h_flat.shape[0],)]( + h_flat, + dh_flat, + hdh_last, + D, + BLOCK=block, + num_warps=8 if block >= 1024 else 4, + ) + return hdh_last + + def chunk_fwd_h( k: torch.Tensor, v: torch.Tensor, @@ -317,7 +399,8 @@ def chunk_fwd_h( chunk_size: int = 64, split_size: int | None = None, states_in_fp32: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: + output_mma_state: bool = False, +) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = chunk_size BS = BT if split_size is None else split_size @@ -332,12 +415,14 @@ def chunk_fwd_h( # `state_v_first` stores the states in V-first `[V, K]` layout instead of `[K, V]` state_shape = (V, K) if state_v_first else (K, V) h = k.new_empty(B, NS, H, *state_shape, dtype=k.dtype if not states_in_fp32 else torch.float) + h_mma = k.new_empty(B, NS, H, *state_shape, dtype=k.dtype) if output_mma_state else None ht = k.new_empty(N, H, *state_shape, dtype=torch.float) if output_final_state else None def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) chunk_fwd_kernel_h[grid]( k=k, v=v, h=h, + h_mma=h_mma, g=g, g_gamma=g_gamma, gk=gk, @@ -358,6 +443,8 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), USE_GV=gv is not None, STATE_V_FIRST=state_v_first, ) + if output_mma_state: + return h, ht, h_mma return h, ht @@ -378,7 +465,9 @@ def chunk_bwd_dh( chunk_size: int = 64, split_size: int | None = None, states_in_fp32: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: + output_mma_state: bool = False, + h_for_hdh: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HQ = q.shape[2] BT = chunk_size @@ -396,6 +485,8 @@ def chunk_bwd_dh( # `state_v_first` stores the states in V-first `[V, K]` layout instead of `[K, V]` state_shape = (V, K) if state_v_first else (K, V) dh = k.new_empty(B, NS, HQ, *state_shape, dtype=k.dtype if not states_in_fp32 else torch.float) + dh_mma = k.new_empty(B, NS, HQ, *state_shape, dtype=q.dtype) if output_mma_state else None + hdh_last = q.new_zeros(B, NS, H, dtype=torch.float32) if h_for_hdh is not None else None dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) @@ -407,6 +498,9 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), gv=gv, do=do, dh=dh, + dh_mma=dh_mma, + h_for_hdh=h_for_hdh, + hdh_last=hdh_last, dht=dht, dh0=dh0, cu_seqlens=cu_seqlens, @@ -426,4 +520,8 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), USE_GV=gv is not None, STATE_V_FIRST=state_v_first, ) + if output_mma_state: + if hdh_last is not None: + return dh, dh0, dh_mma, hdh_last + return dh, dh0, dh_mma return dh, dh0 diff --git a/fla/ops/delta_rule/backends/__init__.py b/fla/ops/delta_rule/backends/__init__.py new file mode 100644 index 0000000000..f2fbd01989 --- /dev/null +++ b/fla/ops/delta_rule/backends/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Backend implementations for delta-rule ops.""" diff --git a/fla/ops/delta_rule/backends/tilelang/__init__.py b/fla/ops/delta_rule/backends/tilelang/__init__.py new file mode 100644 index 0000000000..38ca0913fb --- /dev/null +++ b/fla/ops/delta_rule/backends/tilelang/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""TileLang backend helpers for delta-rule ops.""" + +from fla.ops.delta_rule.backends.tilelang.chunk_bwd import ( + chunk_delta_rule_wy_dqkw_fused_tilelang, +) + +__all__ = ["chunk_delta_rule_wy_dqkw_fused_tilelang"] diff --git a/fla/ops/delta_rule/backends/tilelang/chunk_bwd.py b/fla/ops/delta_rule/backends/tilelang/chunk_bwd.py new file mode 100644 index 0000000000..046bce5e5d --- /dev/null +++ b/fla/ops/delta_rule/backends/tilelang/chunk_bwd.py @@ -0,0 +1,348 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""TileLang fused delta-rule WY + dq/dk backward experiment. + +This ports the op-local Triton fused-WY algebra into a one-program-per-chunk +TileLang kernel. The D128 route removes the separate ``chunk_bwd_dqkwg`` and +``prepare_wy_repr_bwd`` launches for dense ungated delta rule while keeping the +current D256 Triton route as fallback unless explicitly opted in for comparison. +""" + +from pathlib import Path + +import tilelang +import tilelang.language as T +import torch + +_COMMON_TILELANG_DIR = Path(__file__).parents[3] / "common" / "backends" / "tilelang" +_CUDA126_FP8_E8M0_STUB = _COMMON_TILELANG_DIR / "cuda126_fp8_e8m0_stub.cuh" +_TILELANG_COMPILE_FLAGS = ["-include", str(_CUDA126_FP8_E8M0_STUB)] + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}, compile_flags=_TILELANG_COMPILE_FLAGS) +def _build_delta_rule_wy_dqkw_fused_kernel( + B, + H, + K, + V, + BT, + BK, + BV, + hD1, + hD2, + dtype_str, + num_warps=4, +): + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + _dtype = dtype_map[dtype_str] + _B, _H, _K, _V = B, H, K, V + _BT, _BK, _BV = BT, BK, BV + _NK = tilelang.cdiv(K, BK) + _NV = tilelang.cdiv(V, BV) + _hD1, _hD2 = hD1, hD2 + _threads = num_warps * 32 + + T_d, total_h_d = T.dynamic("T, total_h") + + qk_s = (_B, T_d, _H, _K) + v_s = (_B, T_d, _H, _V) + beta_s = (_B, T_d, _H) + A_s = (_B, T_d, _H, _BT) + h_s = (total_h_d, _hD1, _hD2) + + @T.prim_func + def kernel( + q: T.Tensor(qk_s, _dtype), + k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), + v_new: T.Tensor(v_s, _dtype), + beta: T.Tensor(beta_s, _dtype), + A: T.Tensor(A_s, _dtype), + h: T.Tensor(h_s, _dtype), + do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _dtype), + dv: T.Tensor(v_s, _dtype), + dq: T.Tensor(qk_s, _dtype), + dk: T.Tensor(qk_s, _dtype), + dv2: T.Tensor(v_s, _dtype), + dbeta: T.Tensor(beta_s, _dtype), + scale: T.float32, + ): + with T.Kernel(T.ceildiv(T_d, _BT), _B * _H, threads=_threads) as (i_t, i_bh): + i_b = i_bh // _H + i_h = i_bh % _H + NT_local = T.ceildiv(T_d, _BT) + h_idx = (i_b * NT_local + i_t) * _H + i_h + t_s = i_t * _BT + + s_beta = T.alloc_shared((_BT,), T.float32) + for _i in T.Parallel(_BT): + valid = (i_t * _BT + _i) < T_d + s_beta[_i] = T.if_then_else( + valid, + T.cast(beta[i_b, t_s + _i, i_h], T.float32), + 0.0, + ) + + s_A_src = T.alloc_shared((_BT, _BT), _dtype) + T.copy(A[i_b, t_s:t_s + _BT, i_h, 0:_BT], s_A_src, disable_tma=True) + s_A = T.alloc_shared((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + valid = ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + s_A[_i, _j] = T.if_then_else(valid, s_A_src[_j, _i], T.cast(0, _dtype)) + + b_ds = T.alloc_fragment((_BT, _BT), T.float32) + b_dA = T.alloc_fragment((_BT, _BT), T.float32) + b_dbeta = T.alloc_fragment((_BT,), T.float32) + T.clear(b_ds) + T.clear(b_dA) + T.clear(b_dbeta) + + s_do = T.alloc_shared((_BT, _BV), _dtype) + s_v_new = T.alloc_shared((_BT, _BV), _dtype) + s_v = T.alloc_shared((_BT, _BV), _dtype) + s_v_beta = T.alloc_shared((_BT, _BV), _dtype) + s_dv = T.alloc_shared((_BT, _BV), _dtype) + b_dvb = T.alloc_fragment((_BT, _BV), T.float32) + f_dvv = T.alloc_fragment((_BT, _BV), T.float32) + f_row_v = T.alloc_fragment((_BT,), T.float32) + f_dv2 = T.alloc_fragment((_BT, _BV), _dtype) + s_dv2 = T.alloc_shared((_BT, _BV), _dtype) + + for i_v in T.Pipelined(_NV, num_stages=2): + v_off = i_v * _BV + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_do, disable_tma=True) + T.copy(v_new[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v_new, disable_tma=True) + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v, disable_tma=True) + T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_dv, disable_tma=True) + for _i, _j in T.Parallel(_BT, _BV): + s_v_beta[_i, _j] = T.cast(T.cast(s_v[_i, _j], T.float32) * s_beta[_i], _dtype) + + T.gemm(s_do, s_v_new, b_ds, transpose_B=True) + T.gemm(s_dv, s_v_beta, b_dA, transpose_B=True) + + T.clear(b_dvb) + T.gemm(s_A, s_dv, b_dvb) + for _i, _j in T.Parallel(_BT, _BV): + f_dv2[_i, _j] = T.cast(b_dvb[_i, _j] * s_beta[_i], _dtype) + f_dvv[_i, _j] = b_dvb[_i, _j] * T.cast(s_v[_i, _j], T.float32) + T.copy(f_dv2, s_dv2) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BV): + if (i_t * _BT + _i) < T_d: + dv2[i_b, t_s + _i, i_h, v_off + _j] = s_dv2[_i, _j] + T.reduce_sum(f_dvv, f_row_v, dim=1) + for _i in T.Parallel(_BT): + b_dbeta[_i] = b_dbeta[_i] + f_row_v[_i] + + s_ds = T.alloc_shared((_BT, _BT), _dtype) + f_ds = T.alloc_fragment((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + valid = ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + causal = (_i >= _j) & valid + f_ds[_i, _j] = T.if_then_else(causal, T.cast(b_ds[_i, _j], _dtype), T.cast(0, _dtype)) + T.copy(f_ds, s_ds) + + s_q = T.alloc_shared((_BT, _BK), _dtype) + s_k = T.alloc_shared((_BT, _BK), _dtype) + s_h = T.alloc_shared((_BK, _BV), _dtype) + s_dh = T.alloc_shared((_BK, _BV), _dtype) + s_dw = T.alloc_shared((_BT, _BK), _dtype) + s_k_beta = T.alloc_shared((_BT, _BK), _dtype) + s_out = T.alloc_shared((_BT, _BK), _dtype) + f_out = T.alloc_fragment((_BT, _BK), _dtype) + f_dkk = T.alloc_fragment((_BT, _BK), T.float32) + f_row_k = T.alloc_fragment((_BT,), T.float32) + + for i_k in T.serial(_NK): + k_off = i_k * _BK + T.copy(q[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k, disable_tma=True) + for _i, _j in T.Parallel(_BT, _BK): + s_k_beta[_i, _j] = T.cast(T.cast(s_k[_i, _j], T.float32) * s_beta[_i], _dtype) + + b_dq = T.alloc_fragment((_BT, _BK), T.float32) + b_dk = T.alloc_fragment((_BT, _BK), T.float32) + b_dk_ds = T.alloc_fragment((_BT, _BK), T.float32) + b_dw = T.alloc_fragment((_BT, _BK), T.float32) + b_dk_beta = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dq) + T.clear(b_dk) + T.clear(b_dk_ds) + T.clear(b_dw) + + for i_v in T.Pipelined(_NV, num_stages=2): + v_off = i_v * _BV + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_do, disable_tma=True) + T.copy(v_new[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v_new, disable_tma=True) + T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_dv, disable_tma=True) + T.copy(h[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_h, disable_tma=True) + T.copy(dh[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_dh, disable_tma=True) + + T.gemm(s_do, s_h, b_dq, transpose_B=True) + T.gemm(s_v_new, s_dh, b_dk, transpose_B=True) + T.gemm(s_dv, s_h, b_dw, transpose_B=True) + + T.gemm(s_ds, s_k, b_dq) + T.gemm(s_ds, s_q, b_dk_ds, transpose_A=True) + + for _i, _j in T.Parallel(_BT, _BK): + b_dq[_i, _j] = b_dq[_i, _j] * scale + b_dk[_i, _j] = b_dk[_i, _j] + b_dk_ds[_i, _j] * scale + s_dw[_i, _j] = T.cast(-b_dw[_i, _j], _dtype) + + T.gemm(s_dw, s_k_beta, b_dA, transpose_B=True) + + T.clear(b_dk_beta) + T.gemm(s_A, s_dw, b_dk_beta) + for _i, _j in T.Parallel(_BT, _BK): + b_dk[_i, _j] = b_dk[_i, _j] + b_dk_beta[_i, _j] * s_beta[_i] + f_dkk[_i, _j] = b_dk_beta[_i, _j] * T.cast(s_k[_i, _j], T.float32) + T.reduce_sum(f_dkk, f_row_k, dim=1) + for _i in T.Parallel(_BT): + b_dbeta[_i] = b_dbeta[_i] + f_row_k[_i] + + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dq[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dq[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + for _i, _j in T.Parallel(_BT, _BK): + f_out[_i, _j] = T.cast(b_dk[_i, _j], _dtype) + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + + s_dA = T.alloc_shared((_BT, _BT), T.float32) + T.copy(b_dA, s_dA) + for _i, _j in T.Parallel(_BT, _BT): + valid = ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + strict_lower = (_i > _j) & valid + s_dA[_i, _j] = T.if_then_else(strict_lower, s_dA[_i, _j], 0.0) + + s_dA_dtype = T.alloc_shared((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + s_dA_dtype[_i, _j] = T.cast(s_dA[_i, _j], _dtype) + + b_dA2 = T.alloc_fragment((_BT, _BT), T.float32) + T.clear(b_dA2) + T.gemm(s_dA_dtype, s_A, b_dA2) + for _i, _j in T.Parallel(_BT, _BT): + s_dA_dtype[_i, _j] = T.cast(b_dA2[_i, _j], _dtype) + + b_dA3 = T.alloc_fragment((_BT, _BT), T.float32) + T.clear(b_dA3) + T.gemm(s_A, s_dA_dtype, b_dA3) + + s_dA_final = T.alloc_shared((_BT, _BT), _dtype) + for _i, _j in T.Parallel(_BT, _BT): + valid = ((i_t * _BT + _i) < T_d) & ((i_t * _BT + _j) < T_d) + strict_lower = (_i > _j) & valid + s_dA_final[_i, _j] = T.if_then_else( + strict_lower, + T.cast(-b_dA3[_i, _j], _dtype), + T.cast(0, _dtype), + ) + + s_dk_prev = T.alloc_shared((_BT, _BK), _dtype) + for i_k in T.serial(_NK): + k_off = i_k * _BK + T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k, disable_tma=True) + T.copy(dk[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_dk_prev, disable_tma=True) + for _i, _j in T.Parallel(_BT, _BK): + s_k_beta[_i, _j] = T.cast(T.cast(s_k[_i, _j], T.float32) * s_beta[_i], _dtype) + + b_dk_beta = T.alloc_fragment((_BT, _BK), T.float32) + b_dk_extra = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dk_beta) + T.clear(b_dk_extra) + T.gemm(s_dA_final, s_k, b_dk_beta) + T.gemm(s_dA_final, s_k_beta, b_dk_extra, transpose_A=True) + + for _i, _j in T.Parallel(_BT, _BK): + f_dkk[_i, _j] = b_dk_beta[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_out[_i, _j] = T.cast( + T.cast(s_dk_prev[_i, _j], T.float32) + + b_dk_extra[_i, _j] + + b_dk_beta[_i, _j] * s_beta[_i], + _dtype, + ) + T.reduce_sum(f_dkk, f_row_k, dim=1) + for _i in T.Parallel(_BT): + b_dbeta[_i] = b_dbeta[_i] + f_row_k[_i] + + T.copy(f_out, s_out) + T.sync_threads() + for _i, _j in T.Parallel(_BT, _BK): + if (i_t * _BT + _i) < T_d: + dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] + + for _i in T.Parallel(_BT): + if (i_t * _BT + _i) < T_d: + dbeta[i_b, t_s + _i, i_h] = T.cast(b_dbeta[_i], _dtype) + + return kernel + + +def chunk_delta_rule_wy_dqkw_fused_tilelang( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dv: torch.Tensor, + scale: float | None = None, + chunk_size: int = 64, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K = k.shape + V = v.shape[-1] + if K != V or K not in (128, 256): + raise ValueError(f"TileLang fused delta WY/DQKW requires K == V in {{128, 256}}, got K={K}, V={V}") + if scale is None: + scale = K ** -0.5 + + BT = chunk_size + BK = 64 + BV = 64 + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16'}[q.dtype] + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv2 = torch.empty_like(v) + dbeta = torch.empty_like(beta) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] + + kernel = _build_delta_rule_wy_dqkw_fused_kernel( + B, + H, + K, + V, + BT, + BK, + BV, + hD1, + hD2, + dtype_str, + num_warps=4, + ) + kernel(q, k, v, v_new, beta, A, h_flat, do, dh_flat, dv, dq, dk, dv2, dbeta, scale) + return dq, dk, dv2, dbeta diff --git a/fla/ops/delta_rule/backends/triton/__init__.py b/fla/ops/delta_rule/backends/triton/__init__.py new file mode 100644 index 0000000000..600f5d01ce --- /dev/null +++ b/fla/ops/delta_rule/backends/triton/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Triton backend helpers for delta-rule ops.""" + +from fla.ops.delta_rule.backends.triton.chunk_bwd import ( + chunk_delta_rule_wy_dqkw_fused_triton, +) + +__all__ = ["chunk_delta_rule_wy_dqkw_fused_triton"] diff --git a/fla/ops/delta_rule/backends/triton/chunk_bwd.py b/fla/ops/delta_rule/backends/triton/chunk_bwd.py new file mode 100644 index 0000000000..9278e2c056 --- /dev/null +++ b/fla/ops/delta_rule/backends/triton/chunk_bwd.py @@ -0,0 +1,236 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Fused delta-rule WY + dq/dk backward experiment. + +This is an op-local producer-consumer fusion for the dense ungated delta-rule +backward. It keeps the WY consumer of ``dw = -(dv @ h)`` in the same program +that produces ``dw`` so the wide ``dw`` tensor is never materialized to HBM. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import safe_dot + + +@triton.jit(do_not_specialize=["T"]) +def _chunk_delta_rule_wy_dqkw_fused_kernel( + q, + k, + v, + v_new, + beta, + A, + h, + do, + dh, + dv, + dq, + dk, + dv2, + dbeta, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + i_t = tl.program_id(0) + i_bh = tl.program_id(1).to(tl.int64) + i_b = i_bh // H + i_h = i_bh % H + + NT = tl.cdiv(T, BT) + i_tg = (i_b * NT + i_t).to(tl.int64) + bos = (i_b * T).to(tl.int64) + t0 = i_t * BT + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + v_new += (bos * H + i_h) * V + beta += bos * H + i_h + A += (bos * H + i_h) * BT + h += (i_tg * H + i_h) * K * V + do += (bos * H + i_h) * V + dh += (i_tg * H + i_h) * K * V + dv += (bos * H + i_h) * V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dv2 += (bos * H + i_h) * V + dbeta += bos * H + i_h + + o_t = t0 + tl.arange(0, BT) + m_t = o_t < T + + p_beta = tl.make_block_ptr(beta, (T,), (H,), (t0,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + # Existing prepare_wy_repr_bwd consumes A transposed. + p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, t0), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + b_ds = tl.zeros([BT, BT], dtype=tl.float32) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + b_dbeta = tl.zeros([BT], dtype=tl.float32) + + # K-independent work: local dS plus the V-side WY backward contribution. + for i_v in range(tl.cdiv(V, BV)): + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + + b_ds += tl.dot(b_do, tl.trans(b_v_new)) + b_dA += tl.dot(b_dv, tl.trans(b_v)) + + b_dvb = tl.dot(b_A, b_dv) + b_dv2 = b_dvb * b_beta[:, None] + b_dbeta += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) + + m_lower = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]) + b_ds = tl.where(m_lower, b_ds, 0).to(q.dtype.element_ty) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (t0, i_k * BK), (BT, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (t0, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + + 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)) + + b_dq += tl.dot(b_ds, b_k) + b_dk += tl.dot(tl.trans(b_ds), b_q) * scale + b_dq *= scale + + b_dw = -b_dw.to(b_A.dtype) + b_dA += tl.dot(b_dw, tl.trans(b_k.to(b_A.dtype))) + + b_dk_beta = tl.dot(b_A, b_dw) + b_dbeta += tl.sum(b_dk_beta * b_k, 1) + b_dk += b_dk_beta * b_beta[:, None] + + p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (t0, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (t0, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + m_strict = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t[None, :]) + b_dA = tl.where(m_strict, b_dA * b_beta[None, :], 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_strict, -b_dA, 0) + + # Final transformed-dA terms for dk/dbeta. This stays in the fused launch, + # avoiding the separate prepare_wy_repr_bwd + dk.add_ HBM round trip. + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (t0, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (t0, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.load(p_dk, boundary_check=(0, 1)) + b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) + + b_dk_beta = tl.dot(b_dA.to(b_k.dtype), b_k) + b_dbeta += tl.sum(b_dk_beta * b_k, 1) + b_dk += safe_dot(tl.trans(b_dA.to(b_k.dtype)), b_k_beta, allow_tf32=False) + b_dk += b_dk_beta * b_beta[:, None] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + p_dbeta = tl.make_block_ptr(dbeta, (T,), (H,), (t0,), (BT,), (0,)) + tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) + + +def chunk_delta_rule_wy_dqkw_fused_triton( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dv: torch.Tensor, + scale: float | None = None, + chunk_size: int = 64, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K = k.shape + V = v.shape[-1] + BT = chunk_size + BK = 64 + BV = 64 + if scale is None: + scale = K ** -0.5 + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv2 = torch.empty_like(v) + dbeta = torch.empty_like(beta) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + grid = (triton.cdiv(T, BT), B * H) + _chunk_delta_rule_wy_dqkw_fused_kernel[grid]( + q, + k, + v, + v_new, + beta, + A, + h_flat, + do, + dh_flat, + dv, + dq, + dk, + dv2, + dbeta, + scale, + T, + H, + K, + V, + BT, + BK, + BV, + num_warps=4, + num_stages=3, + ) + return dq, dk, dv2, dbeta diff --git a/fla/ops/delta_rule/chunk.py b/fla/ops/delta_rule/chunk.py index d69b598ae9..5b22f315dc 100644 --- a/fla/ops/delta_rule/chunk.py +++ b/fla/ops/delta_rule/chunk.py @@ -5,16 +5,103 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors +import os + import torch from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.backends.tilelang import TileLangBackend 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.delta_rule.backends.tilelang import chunk_delta_rule_wy_dqkw_fused_tilelang +from fla.ops.delta_rule.backends.triton import chunk_delta_rule_wy_dqkw_fused_triton from fla.ops.delta_rule.wy_fast import prepare_wy_repr_bwd, prepare_wy_repr_fwd, recompute_w_u_fwd from fla.ops.utils.index import prepare_chunk_indices from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +def _can_use_fused_wy_dqkw( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + dh: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + initial_state: torch.Tensor | None, + dht: torch.Tensor | None, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None, + chunk_size: int, +) -> bool: + return ( + os.environ.get("FLA_DELTA_RULE_FUSED_WY", "1") != "0" + and os.environ.get("FLA_TILELANG") == "1" + and q.is_cuda + and cu_seqlens is None + and chunk_indices is None + and initial_state is None + and dht is None + and chunk_size == 64 + and q.dtype in (torch.float16, torch.bfloat16) + and q.dtype == k.dtype == v.dtype == beta.dtype == do.dtype == dv.dtype + and h.dtype == dh.dtype == q.dtype + and q.shape == k.shape + and v.shape == do.shape == dv.shape + and q.shape[2] == v.shape[2] + and q.shape[-1] == v.shape[-1] + and q.shape[-1] == 256 + and q.shape[1] % 64 == 0 + and A.shape[-1] == 64 + ) + + +def _can_use_tilelang_wy_dqkw_fused( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + dh: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + initial_state: torch.Tensor | None, + dht: torch.Tensor | None, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None, + chunk_size: int, +) -> bool: + if not ( + os.environ.get("FLA_DELTA_RULE_FUSED_WY", "1") != "0" + and os.environ.get("FLA_DELTA_RULE_TILELANG_WY", "1") != "0" + and TileLangBackend.is_available() + and TileLangBackend.is_enabled() + and q.is_cuda + and cu_seqlens is None + and chunk_indices is None + and initial_state is None + and dht is None + and chunk_size == 64 + and q.dtype in (torch.float16, torch.bfloat16) + and q.dtype == k.dtype == v.dtype == beta.dtype == do.dtype == dv.dtype + and h.dtype == dh.dtype == q.dtype + and q.shape == k.shape + and v.shape == do.shape == dv.shape + and q.shape[2] == v.shape[2] + and q.shape[-1] == v.shape[-1] + and q.shape[1] % 64 == 0 + and A.shape[-1] == 64 + ): + return False + head_dim = q.shape[-1] + if head_dim == 128: + return True + return head_dim == 256 and os.environ.get("FLA_DELTA_RULE_TILELANG_WY_D256", "0") == "1" + + def chunk_delta_rule_fwd( q: torch.Tensor, k: torch.Tensor, @@ -119,6 +206,68 @@ def chunk_delta_rule_bwd( chunk_indices=chunk_indices, chunk_size=chunk_size, ) + if _can_use_tilelang_wy_dqkw_fused( + q=q, + k=k, + v=v, + beta=beta, + A=A, + h=h, + dh=dh, + do=do, + dv=dv, + initial_state=initial_state, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ): + dq, dk, dv, db = chunk_delta_rule_wy_dqkw_fused_tilelang( + q=q, + k=k, + v=v, + v_new=v_new, + beta=beta, + A=A, + h=h, + do=do, + dh=dh, + dv=dv, + scale=scale, + chunk_size=chunk_size, + ) + return dq, dk, dv, db, dh0 + if _can_use_fused_wy_dqkw( + q=q, + k=k, + v=v, + beta=beta, + A=A, + h=h, + dh=dh, + do=do, + dv=dv, + initial_state=initial_state, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ): + dq, dk, dv, db = chunk_delta_rule_wy_dqkw_fused_triton( + q=q, + k=k, + v=v, + v_new=v_new, + beta=beta, + A=A, + h=h, + do=do, + dh=dh, + dv=dv, + scale=scale, + chunk_size=chunk_size, + ) + return dq, dk, dv, db, dh0 dq, dk, dw, _ = chunk_bwd_dqkwg( q=q, k=k, diff --git a/fla/ops/gated_delta_rule/wy_fast.py b/fla/ops/gated_delta_rule/wy_fast.py index 63821cd7ee..15263ccc09 100644 --- a/fla/ops/gated_delta_rule/wy_fast.py +++ b/fla/ops/gated_delta_rule/wy_fast.py @@ -5,6 +5,8 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors +import os + import torch import triton import triton.language as tl @@ -220,7 +222,9 @@ def prepare_wy_repr_bwd_kernel( b_A = tl.zeros([BT, BT], dtype=tl.float32) b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty) + # keep this barrier: triton-ascend can misorder b_dA before the dk load-modify-store without it tl.debug_barrier() + for i_k in range(tl.cdiv(K, BK)): o_k = i_k * BK + tl.arange(0, BK) m_k = m_t[:, None] & (o_k[None, :] < K) @@ -307,13 +311,33 @@ def prepare_wy_repr_bwd( 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) + + if k.dtype is torch.float32 and os.environ.get('TRITON_F32_DEFAULT') == 'ieee': + return _prepare_wy_repr_bwd_split( + k=k, + v=v, + beta=beta, + A=A, + dw=dw, + du=du, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + NT=NT, + ) + CONST_TILING = 64 if check_shared_mem() else 32 BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + return_dg = g is not None + if g is None: + # exp2(0) is the no-g path; reuse the g-specialized compiler artifact. + g = torch.zeros_like(beta) + dk = k.new_empty(B, T, HV, K) dv = torch.empty_like(v) - dg = torch.empty_like(g) if g is not None else None + dg = torch.empty_like(g) db = torch.empty_like(beta) prepare_wy_repr_bwd_kernel[(NT, B * HV)]( k=k, @@ -338,6 +362,112 @@ def prepare_wy_repr_bwd( BK=BK, BV=BV, ) + if H != HV: + dk = dk.view(B, T, H, HV // H, K).sum(3) + if not return_dg: + dg = None + return dk, dv, db, dg + + +def _prepare_wy_repr_bwd_split( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + g: torch.Tensor | None, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None, + NT: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + from fla.ops.gated_delta_rule.backends.triton_ascend import wy_fast as wy_split + + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] + BT = A.shape[-1] + BK = min(max(triton.next_power_of_2(K), 16), 32) + BV = min(max(triton.next_power_of_2(V), 16), 32) + use_g = g is not None + is_varlen = cu_seqlens is not None + + dk = k.new_empty(B, T, HV, K) + dv = torch.empty_like(v) + dg = torch.empty_like(g) if use_g else None + db = torch.empty_like(beta) + dA_scr = torch.empty_like(A, dtype=torch.float32) + dA_mid = torch.empty_like(A, dtype=torch.float32) + dA_out = torch.empty_like(A, dtype=torch.float32) + a2_scr = torch.empty_like(A, dtype=torch.float32) + col_acc_scr = torch.empty( + B, + triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices), + HV, + BT, + dtype=torch.float32, + device=k.device, + ) + g_arg = g if use_g else beta + dg_arg = dg if use_g else beta + grid = (NT, B * HV) + base = dict( + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + BT=BT, + IS_VARLEN=is_varlen, + NT_OFFSET=0, + BH_OFFSET=0, + num_warps=2, + ) + + wy_split.prepare_wy_repr_bwd_k_npu[grid]( + k=k, beta=beta, g=g_arg, A=A, dw=dw, + dk=dk, dA_scr=dA_scr, db=db, dg=dg_arg, + H=H, HV=HV, K=K, BK=BK, USE_G=use_g, + **base, + ) + wy_split.prepare_wy_repr_bwd_v_npu[grid]( + v=v, beta=beta, A=A, du=du, dv=dv, dA_scr=dA_scr, db=db, + HV=HV, V=V, BV=BV, + **base, + ) + wy_split.prepare_wy_repr_bwd_da_mask_npu[grid]( + dA_scr=dA_scr, + HV=HV, + **base, + ) + wy_split.prepare_wy_repr_bwd_da_dot1_npu[grid]( + A=A, dA_scr=dA_scr, dA_mid=dA_mid, + HV=HV, + **base, + ) + wy_split.prepare_wy_repr_bwd_da_dot2_npu[grid]( + A=A, dA_mid=dA_mid, dA_out=dA_out, + HV=HV, + **base, + ) + if use_g: + wy_split.prepare_wy_repr_bwd_da_gate_npu[grid]( + g=g_arg, dA_out=dA_out, + HV=HV, BC=16, + **base, + ) + wy_split.prepare_wy_repr_bwd_finalize_k_npu[grid]( + k=k, beta=beta, dA_out=dA_out, dk=dk, db=db, + H=H, HV=HV, K=K, BK=BK, + **base, + ) + if use_g: + wy_split.prepare_wy_repr_bwd_finalize_a2_npu[grid]( + k=k, beta=beta, a2_scr=a2_scr, + H=H, HV=HV, K=K, BK=BK, + **base, + ) + wy_split.prepare_wy_repr_bwd_finalize_dg_npu[grid]( + dA_out=dA_out, a2_scr=a2_scr, dg=dg_arg, col_acc_scr=col_acc_scr, + HV=HV, BC=16, + **base, + ) if H != HV: dk = dk.view(B, T, H, HV // H, K).sum(3) return dk, dv, db, dg diff --git a/fla/ops/gla/backends/__init__.py b/fla/ops/gla/backends/__init__.py new file mode 100644 index 0000000000..5ea4566637 --- /dev/null +++ b/fla/ops/gla/backends/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Backend implementations for GLA ops.""" diff --git a/fla/ops/gla/backends/tilelang/__init__.py b/fla/ops/gla/backends/tilelang/__init__.py new file mode 100644 index 0000000000..65542e4f5f --- /dev/null +++ b/fla/ops/gla/backends/tilelang/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""TileLang backends for GLA ops.""" diff --git a/fla/ops/gla/backends/tilelang/chunk_bwd.py b/fla/ops/gla/backends/tilelang/chunk_bwd.py new file mode 100644 index 0000000000..a1a44937d7 --- /dev/null +++ b/fla/ops/gla/backends/tilelang/chunk_bwd.py @@ -0,0 +1,261 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from pathlib import Path + +import tilelang +import tilelang.language as T +import torch +import triton + +_CUDA126_FP8_E8M0_STUB = ( + Path(__file__).parents[3] + / "common" + / "backends" + / "tilelang" + / "cuda126_fp8_e8m0_stub.cuh" +) +_TILELANG_COMPILE_FLAGS = ["-include", str(_CUDA126_FP8_E8M0_STUB)] + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}, compile_flags=_TILELANG_COMPILE_FLAGS) +def _build_chunk_gla_fused_bwd_k_tile( + B, + H, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=4, +): + dtype_map = {"float16": T.float16, "bfloat16": T.bfloat16, "float32": T.float32} + _dtype = dtype_map[dtype_str] + _state_dtype = dtype_map[state_dtype_str] + NV = tilelang.cdiv(V, BV) + threads = num_warps * 32 + + _B, _H, _K, _V = B, H, K, V + _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV + _hD1, _hD2 = hD1, hD2 + _threads = threads + _CAST_STATE_FOR_MMA = state_dtype_str != dtype_str + + T_d, total_h_d = T.dynamic("T, total_h") + + qk_s = (_B, T_d, _H, _K) + v_s = (_B, T_d, _H, _V) + g_s = (_B, T_d, _H, _K) + a_s = (_B, T_d, _H, _BT) + h_s = (total_h_d, _hD1, _hD2) + + @T.prim_func + def kernel( + q: T.Tensor(qk_s, _dtype), + k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), + g: T.Tensor(g_s, T.float32), + h: T.Tensor(h_s, _state_dtype), + do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _state_dtype), + dA: T.Tensor(a_s, T.float32), + dq: T.Tensor(qk_s, T.float32), + dk: T.Tensor(qk_s, T.float32), + dg: T.Tensor(g_s, T.float32), + scale: T.float32, + ): + with T.Kernel(_NK, T.ceildiv(T_d, _BT), _B * _H, threads=_threads) as (i_k, i_t, i_bh): + i_b = i_bh // _H + i_h = i_bh % _H + NT_local = T.ceildiv(T_d, _BT) + h_idx = (i_b * NT_local + i_t) * _H + i_h + t_s = i_t * _BT + k_off = i_k * _BK + last_pos = T.max(0, T.min(_BT, T_d - i_t * _BT) - 1) + + b_dq_inter = T.alloc_fragment((_BT, _BK), T.float32) + b_dk_inter = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dq_inter) + T.clear(b_dk_inter) + + s_dgk = T.alloc_shared((_BK,), T.float32) + for _j in T.Parallel(_BK): + s_dgk[_j] = 0.0 + + s_v = T.alloc_shared((_BT, _BV), _dtype) + s_do = T.alloc_shared((_BT, _BV), _dtype) + s_h = T.alloc_shared((_BK, _BV), _state_dtype) + s_dh = T.alloc_shared((_BK, _BV), _state_dtype) + if _CAST_STATE_FOR_MMA: + s_h_mma = T.alloc_shared((_BK, _BV), _dtype) + s_dh_mma = T.alloc_shared((_BK, _BV), _dtype) + + # Inter-chunk contribution: do @ h, v @ dh, and per-K h*dh. + for i_v_py in T.Pipelined(_NV, num_stages=2): + v_off = i_v_py * _BV + T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v, disable_tma=True) + T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_do, disable_tma=True) + T.copy(h[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_h, disable_tma=True) + T.copy(dh[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_dh, disable_tma=True) + + if _CAST_STATE_FOR_MMA: + for _i, _j in T.Parallel(_BK, _BV): + s_h_mma[_i, _j] = T.cast(s_h[_i, _j], _dtype) + s_dh_mma[_i, _j] = T.cast(s_dh[_i, _j], _dtype) + T.gemm(s_do, s_h_mma, b_dq_inter, transpose_B=True) + T.gemm(s_v, s_dh_mma, b_dk_inter, transpose_B=True) + else: + T.gemm(s_do, s_h, b_dq_inter, transpose_B=True) + T.gemm(s_v, s_dh, b_dk_inter, transpose_B=True) + + f_hdh = T.alloc_fragment((_BK, _BV), T.float32) + for _i, _j in T.Parallel(_BK, _BV): + f_hdh[_i, _j] = T.cast(s_h[_i, _j], T.float32) * T.cast(s_dh[_i, _j], T.float32) + f_hdh_k = T.alloc_fragment((_BK,), T.float32) + T.reduce_sum(f_hdh, f_hdh_k, dim=1) + for _j in T.Parallel(_BK): + s_dgk[_j] = s_dgk[_j] + f_hdh_k[_j] + + s_q = T.alloc_shared((_BT, _BK), _dtype) + s_k = T.alloc_shared((_BT, _BK), _dtype) + s_g = T.alloc_shared((_BT, _BK), T.float32) + T.copy(q[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_q, disable_tma=True) + T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k, disable_tma=True) + T.copy(g[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_g, disable_tma=True) + + s_inter_dq = T.alloc_shared((_BT, _BK), T.float32) + s_inter_dk = T.alloc_shared((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + b_dq_inter[_i, _j] = b_dq_inter[_i, _j] * T.exp2(s_g[_i, _j]) * scale + b_dk_inter[_i, _j] = b_dk_inter[_i, _j] * T.exp2(s_g[last_pos, _j] - s_g[_i, _j]) + T.copy(b_dq_inter, s_inter_dq) + T.copy(b_dk_inter, s_inter_dk) + + for _j in T.Parallel(_BK): + s_dgk[_j] = s_dgk[_j] * T.exp2(s_g[last_pos, _j]) + + f_kdk = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_kdk[_i, _j] = s_inter_dk[_i, _j] * T.cast(s_k[_i, _j], T.float32) + f_kdk_t = T.alloc_fragment((_BK, _BT), T.float32) + for _i, _j in T.Parallel(_BK, _BT): + f_kdk_t[_i, _j] = f_kdk[_j, _i] + f_kdk_col = T.alloc_fragment((_BK,), T.float32) + T.reduce_sum(f_kdk_t, f_kdk_col, dim=1) + for _j in T.Parallel(_BK): + s_dgk[_j] = s_dgk[_j] + f_kdk_col[_j] + + # Intra-chunk contribution, fused into the same consumer: + # dq += exp2(g) * (dA @ (k * exp2(-g))) + # dk += exp2(-g) * (dA.T @ (q * exp2(g))) + b_dq = T.alloc_fragment((_BT, _BK), T.float32) + b_dk = T.alloc_fragment((_BT, _BK), T.float32) + T.clear(b_dq) + T.clear(b_dk) + s_dA = T.alloc_shared((_BT, _BT), T.float32) + s_kg = T.alloc_shared((_BT, _BK), T.float32) + s_qg = T.alloc_shared((_BT, _BK), T.float32) + T.copy(dA[i_b, t_s:t_s + _BT, i_h, 0:_BT], s_dA, disable_tma=True) + for _i, _j in T.Parallel(_BT, _BK): + s_kg[_i, _j] = T.cast(s_k[_i, _j], T.float32) * T.exp2(-s_g[_i, _j]) + s_qg[_i, _j] = T.cast(s_q[_i, _j], T.float32) * T.exp2(s_g[_i, _j]) + T.gemm(s_dA, s_kg, b_dq) + T.gemm(s_dA, s_qg, b_dk, transpose_A=True) + + for _i, _j in T.Parallel(_BT, _BK): + b_dq[_i, _j] = b_dq[_i, _j] * T.exp2(s_g[_i, _j]) + s_inter_dq[_i, _j] + b_dk[_i, _j] = b_dk[_i, _j] * T.exp2(-s_g[_i, _j]) + s_inter_dk[_i, _j] + + f_dg = T.alloc_fragment((_BT, _BK), T.float32) + for _i, _j in T.Parallel(_BT, _BK): + f_dg[_i, _j] = ( + T.cast(s_q[_i, _j], T.float32) * b_dq[_i, _j] + - T.cast(s_k[_i, _j], T.float32) * b_dk[_i, _j] + ) + f_dg_t = T.alloc_fragment((_BK, _BT), T.float32) + for _i, _j in T.Parallel(_BK, _BT): + f_dg_t[_i, _j] = f_dg[_j, _i] + f_dg_col = T.alloc_fragment((_BK,), T.float32) + T.reduce_sum(f_dg_t, f_dg_col, dim=1) + + s_dg_raw = T.alloc_shared((_BT, _BK), T.float32) + T.copy(f_dg, s_dg_raw) + T.cumsum(src=f_dg, dim=0) + for _i, _j in T.Parallel(_BT, _BK): + f_dg[_i, _j] = s_dg_raw[_i, _j] - f_dg[_i, _j] + f_dg_col[_j] + s_dgk[_j] + + T.copy(b_dq, dq[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK]) + T.copy(b_dk, dk[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK]) + T.copy(f_dg, dg[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK]) + + return kernel + + +def chunk_gla_bwd_dqkg_fused_tilelang( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dA: torch.Tensor, + scale: float | None = None, + chunk_size: int = 64, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K = k.shape + V = v.shape[-1] + BT = chunk_size + BK = 64 + BV = 64 + NK = triton.cdiv(K, BK) + if scale is None: + scale = K ** -0.5 + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + dg = torch.empty_like(g, dtype=torch.float) + + h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) + dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) + hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] + dtype_str = { + torch.float16: "float16", + torch.bfloat16: "bfloat16", + torch.float32: "float32", + }[q.dtype] + state_dtype_str = { + torch.float16: "float16", + torch.bfloat16: "bfloat16", + torch.float32: "float32", + }[h.dtype] + + kernel = _build_chunk_gla_fused_bwd_k_tile( + B, + H, + K, + V, + BT, + BK, + BV, + NK, + hD1, + hD2, + dtype_str, + state_dtype_str, + num_warps=4, + ) + kernel(q, k, v, g, h_flat, do, dh_flat, dA, dq, dk, dg, scale) + return dq, dk, dg diff --git a/fla/ops/gla/chunk.py b/fla/ops/gla/chunk.py index c55bc8b001..8f84b313ed 100644 --- a/fla/ops/gla/chunk.py +++ b/fla/ops/gla/chunk.py @@ -9,6 +9,7 @@ import triton import triton.language as tl +from fla.ops.common.backends.tilelang import TileLangBackend from fla.ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h from fla.ops.utils import prepare_chunk_indices from fla.ops.utils.cache import fla_cache_autotune @@ -21,6 +22,52 @@ BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] +def _can_use_tilelang_fused_bwd_k_tile( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor | None, + h: torch.Tensor, + dh: torch.Tensor, + do: torch.Tensor, + initial_state: torch.Tensor | None, + dht: torch.Tensor | None, + state_v_first: bool, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None, +) -> bool: + B, T, H, K = k.shape + HV, V = v.shape[2], v.shape[-1] + return ( + TileLangBackend.is_available() + and TileLangBackend.is_enabled() + and q.is_cuda + and cu_seqlens is None + and chunk_indices is None + and initial_state is None + and dht is None + and not state_v_first + and chunk_size == 64 + and T % chunk_size == 0 + and q.dtype in (torch.float16, torch.bfloat16) + and k.dtype == q.dtype + and v.dtype == q.dtype + and do.dtype == q.dtype + and g_cumsum is not None + and g_cumsum.dtype == torch.float32 + and h.dtype == torch.float32 + and dh.dtype == torch.float32 + and H == HV + and K == V + and K >= 128 + and V >= 128 + and K % 64 == 0 + and V % 64 == 0 + and B > 0 + ) + + def _prune_gla_bwd_configs(configs, nargs, **kwargs): # Keep a tile only if it leaves headroom below its dim, or is the smallest # option (so small dims still autotune); this avoids a software-pipelined @@ -1310,31 +1357,62 @@ def chunk_gla_bwd( chunk_size=chunk_size, chunk_indices=chunk_indices, ) - dq, dk = chunk_gla_bwd_dqk_intra( - q=q, - k=k, - g=g_cumsum, - dA=dA, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - dq, dk, dg = chunk_gla_bwd_dqkg( + if _can_use_tilelang_fused_bwd_k_tile( q=q, k=k, v=v, + g_cumsum=g_cumsum, h=h, - g=g_cumsum, - do=do, dh=dh, - dq=dq, - dk=dk, - scale=scale, + do=do, + initial_state=initial_state, + dht=dht, + state_v_first=state_v_first, cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) + ): + from fla.ops.gla.backends.tilelang.chunk_bwd import ( + chunk_gla_bwd_dqkg_fused_tilelang, + ) + dq, dk, dg = chunk_gla_bwd_dqkg_fused_tilelang( + q=q, + k=k, + v=v, + h=h, + g=g_cumsum, + do=do, + dh=dh, + dA=dA, + scale=scale, + chunk_size=chunk_size, + ) + else: + dq, dk = chunk_gla_bwd_dqk_intra( + q=q, + k=k, + g=g_cumsum, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dq, dk, dg = chunk_gla_bwd_dqkg( + q=q, + k=k, + v=v, + h=h, + g=g_cumsum, + do=do, + dh=dh, + dq=dq, + dk=dk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + ) return dq, dk, dv, dg, dh0 diff --git a/fla/ops/simple_gla/chunk.py b/fla/ops/simple_gla/chunk.py index 90302fdffb..50b96983ae 100644 --- a/fla/ops/simple_gla/chunk.py +++ b/fla/ops/simple_gla/chunk.py @@ -7,11 +7,220 @@ import torch +from fla.ops.common.backends.tilelang import TileLangBackend from fla.ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv, chunk_fwd_o from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices from fla.ops.utils.constant import RCP_LN2 -from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, input_guard + + +def _can_use_shadow_state_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None, + g_gamma: torch.Tensor | None, + initial_state: torch.Tensor | None, + do: torch.Tensor, + dht: torch.Tensor | None, + state_v_first: bool, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None, +) -> bool: + _, _, H, K = k.shape + HQ = q.shape[2] + HV, V = v.shape[2], v.shape[-1] + return ( + TileLangBackend.is_available() + and TileLangBackend.is_enabled() + and g is not None + and g_gamma is None + and initial_state is None + and dht is None + and cu_seqlens is None + and chunk_indices is None + and not state_v_first + and chunk_size == 64 + and q.dtype in (torch.float16, torch.bfloat16) + and k.dtype == q.dtype + and v.dtype == q.dtype + and do.dtype == q.dtype + and HQ == H + and H == HV + and K == V + and K >= 128 + and V >= 128 + and K % 64 == 0 + and V % 64 == 0 + ) + + +def _can_use_direct_mixed_state_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None, + g_gamma: torch.Tensor | None, + initial_state: torch.Tensor | None, + do: torch.Tensor, + dht: torch.Tensor | None, + state_v_first: bool, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None, +) -> bool: + return _can_use_shadow_state_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + +def _can_use_v_first_direct_state_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None, + g_gamma: torch.Tensor | None, + initial_state: torch.Tensor | None, + do: torch.Tensor, + dht: torch.Tensor | None, + state_v_first: bool, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None, +) -> bool: + K = k.shape[-1] + V = v.shape[-1] + return ( + _can_use_shadow_state_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + and K in (128, 256) + and V in (128, 256) + ) + + +def _can_use_v_first_d256_matured_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None, + g_gamma: torch.Tensor | None, + initial_state: torch.Tensor | None, + do: torch.Tensor, + dht: torch.Tensor | None, + state_v_first: bool, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None, +) -> bool: + K = k.shape[-1] + V = v.shape[-1] + if not q.is_cuda: + return False + device_index = q.device.index if q.device.index is not None else torch.cuda.current_device() + return ( + check_shared_mem('hopper', device_index) + and _can_use_shadow_state_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + and K == 256 + and V == 256 + ) + + +def _can_use_dh_shadow_state_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None, + g_gamma: torch.Tensor | None, + initial_state: torch.Tensor | None, + do: torch.Tensor, + dht: torch.Tensor | None, + state_v_first: bool, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None, +) -> bool: + return _can_use_shadow_state_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + +def _can_use_dh_shadow_terminal_dot_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None, + g_gamma: torch.Tensor | None, + initial_state: torch.Tensor | None, + do: torch.Tensor, + dht: torch.Tensor | None, + state_v_first: bool, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None, +) -> bool: + return _can_use_shadow_state_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) def chunk_simple_gla_fwd( @@ -74,65 +283,284 @@ def chunk_simple_gla_bwd( chunk_indices: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: # (SY 09/22) states_in_fp32 seems not affecting the error of dg but for safety, set to True - h, _ = chunk_fwd_h( + use_v_first_d256_matured = _can_use_v_first_d256_matured_dqkwg( + q=q, k=k, v=v, g=g, g_gamma=g_gamma, - gk=None, - gv=None, - h0=initial_state, - output_final_state=False, + initial_state=initial_state, + do=do, + dht=dht, + state_v_first=state_v_first, cu_seqlens=cu_seqlens, chunk_size=chunk_size, - states_in_fp32=True, - state_v_first=state_v_first, + chunk_indices=chunk_indices, ) - dh, dh0 = chunk_bwd_dh( + use_v_first_direct_state = _can_use_v_first_direct_state_dqkwg( q=q, k=k, v=v, g=g, g_gamma=g_gamma, - gk=None, - gv=None, + initial_state=initial_state, do=do, - h0=initial_state, dht=dht, - scale=scale, + state_v_first=state_v_first, cu_seqlens=cu_seqlens, chunk_size=chunk_size, - states_in_fp32=True, + chunk_indices=chunk_indices, + ) and not use_v_first_d256_matured + use_dh_shadow_terminal_dot = _can_use_dh_shadow_terminal_dot_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, state_v_first=state_v_first, - ) - dq, dk, _, dg = chunk_bwd_dqkwg( + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) and not use_v_first_d256_matured and not use_v_first_direct_state + use_dh_shadow_state = _can_use_dh_shadow_state_dqkwg( q=q, k=k, v=v, g=g, g_gamma=g_gamma, - h=h, + initial_state=initial_state, do=do, - dh=dh, - scale=scale, + dht=dht, + state_v_first=state_v_first, cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + ) and not use_v_first_d256_matured and not use_v_first_direct_state and not use_dh_shadow_terminal_dot + use_direct_mixed_state = _can_use_direct_mixed_state_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) and ( + not use_v_first_d256_matured + and not use_v_first_direct_state + and not use_dh_shadow_terminal_dot + and not use_dh_shadow_state ) - dv = chunk_bwd_dv( + use_shadow_state = _can_use_shadow_state_dqkwg( q=q, k=k, + v=v, g=g, g_gamma=g_gamma, + initial_state=initial_state, do=do, - dh=dh, - scale=scale, + dht=dht, + state_v_first=state_v_first, cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - state_v_first=state_v_first, + ) and ( + not use_v_first_d256_matured + and not use_v_first_direct_state + and not use_dh_shadow_terminal_dot + and not use_dh_shadow_state + and not use_direct_mixed_state + ) + internal_state_v_first = state_v_first or use_v_first_d256_matured or use_v_first_direct_state + h_result = chunk_fwd_h( + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=None, + gv=None, + h0=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=True, + state_v_first=internal_state_v_first, + output_mma_state=use_shadow_state, + ) + if use_shadow_state: + h, _, h_mma = h_result + else: + h, _ = h_result + + dh_result = chunk_bwd_dh( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=None, + gv=None, + do=do, + h0=initial_state, + dht=dht, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=True, + state_v_first=internal_state_v_first, + output_mma_state=use_shadow_state or use_dh_shadow_state or use_dh_shadow_terminal_dot, + h_for_hdh=h if (use_shadow_state or use_dh_shadow_state) else None, ) + if use_shadow_state or use_dh_shadow_state: + dh, dh0, dh_mma, hdh_last = dh_result + elif use_dh_shadow_terminal_dot: + dh, dh0, dh_mma = dh_result + else: + dh, dh0 = dh_result + + if use_v_first_d256_matured: + from fla.ops.common.backends.tilelang.chunk_bwd import ( + chunk_bwd_dqkwg_tilelang_k_inner_v_first_d256, + ) + dq, dk, _, dg = chunk_bwd_dqkwg_tilelang_k_inner_v_first_d256( + q=q, + k=k, + v=v, + do=do, + h=h, + dh=dh, + g=g, + scale=scale, + chunk_size=chunk_size, + ) + elif use_v_first_direct_state: + from fla.ops.common.backends.tilelang.chunk_bwd import ( + chunk_bwd_dqkwg_tilelang_k_inner_v_first, + ) + dq, dk, _, dg = chunk_bwd_dqkwg_tilelang_k_inner_v_first( + q=q, + k=k, + v=v, + do=do, + h=h, + dh=dh, + g=g, + scale=scale, + chunk_size=chunk_size, + ) + elif use_dh_shadow_terminal_dot: + from fla.ops.common.backends.tilelang.chunk_bwd import ( + chunk_bwd_dqkwg_tilelang_k_inner_dh_shadow_terminal_dot, + ) + dq, dk, _, dg = chunk_bwd_dqkwg_tilelang_k_inner_dh_shadow_terminal_dot( + q=q, + k=k, + v=v, + do=do, + h=h, + dh=dh, + dh_mma=dh_mma, + g=g, + scale=scale, + chunk_size=chunk_size, + ) + elif use_dh_shadow_state: + from fla.ops.common.backends.tilelang.chunk_bwd import ( + chunk_bwd_dqkwg_tilelang_k_inner_dh_shadow, + ) + dq, dk, _, dg = chunk_bwd_dqkwg_tilelang_k_inner_dh_shadow( + q=q, + k=k, + v=v, + do=do, + h=h, + dh_mma=dh_mma, + hdh_last=hdh_last, + g=g, + scale=scale, + chunk_size=chunk_size, + ) + elif use_shadow_state: + from fla.ops.common.backends.tilelang.chunk_bwd import ( + chunk_bwd_dqkwg_tilelang_k_inner_shadow_state, + ) + dq, dk, _, dg = chunk_bwd_dqkwg_tilelang_k_inner_shadow_state( + q=q, + k=k, + v=v, + do=do, + h_mma=h_mma, + dh_mma=dh_mma, + hdh_last=hdh_last, + g=g, + scale=scale, + chunk_size=chunk_size, + ) + elif use_direct_mixed_state: + from fla.ops.common.backends.tilelang.chunk_bwd import ( + chunk_bwd_dqkwg_tilelang_k_inner, + ) + dq, dk, _, dg = chunk_bwd_dqkwg_tilelang_k_inner( + q=q, + k=k, + v=v, + do=do, + h=h, + dh=dh, + g=g, + scale=scale, + chunk_size=chunk_size, + ) + else: + dq, dk, _, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + h=h, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + ) + if use_v_first_d256_matured: + from fla.ops.common.backends.tilelang.chunk_bwd import ( + chunk_bwd_dv_tilelang_v_first_d256, + ) + dv = chunk_bwd_dv_tilelang_v_first_d256( + q=q, + k=k, + g=g, + do=do, + dh=dh, + scale=scale, + chunk_size=chunk_size, + ) + else: + dv = chunk_bwd_dv( + q=q, + k=k, + g=g, + g_gamma=g_gamma, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + state_v_first=internal_state_v_first, + ) return dq, dk, dv, dg, dh0