# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang import math import warnings import torch import triton import triton.language as tl from . import invcum try: from flash_attn import flash_attn_func, flash_attn_varlen_func except ImportError: warnings.warn( "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", category=ImportWarning, ) flash_attn_func = None from fla.layers.utils import pad_input, unpad_input BLOCK_SIZE_C = 512 def parallel_deltaformer_chunk_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, u: torch.Tensor, qk_scale: float, beta: torch.Tensor, ): C, H, D = q.size() T, _H, _D = k.size() __C, __H = beta.size() assert H == _H and D == _D and H == __H and __C == C w = torch.empty(C, H, C, device=q.device, dtype=q.dtype) lse = torch.empty(C, H, device=q.device, dtype=torch.float) parallel_deltaformer_kernel(q, k, v, u, w, lse, qk_scale, beta) return w, lse def parallel_deltaformer_bwd_u_chunk( q: torch.Tensor, k: torch.Tensor, lse: torch.Tensor, grad_v: torch.Tensor, fa_scale: float, beta: torch.Tensor, ): C, H, D = q.size() T, _H, _D = k.size() grad_u = torch.empty_like(q) def grid(META): return (triton.cdiv(C, META['BLOCK_C']), H) parallel_deltaformer_bwd_kernel_u[grid]( grad_u, q, k, grad_v, lse, beta, H, T, C, D, fa_scale, ) return grad_u def parallel_deltaformer_bwd_qk( q: torch.Tensor, k: torch.Tensor, u: torch.Tensor, lse: torch.Tensor, grad_v: torch.Tensor, qk_scale: float, fa_scale: float, beta: torch.Tensor, ): T, H, D = k.size() row_dot_sum = torch.empty_like(lse) def grid_bp(META): return (triton.cdiv(T, META['BLOCK_C']), H) parallel_deltaformer_bwd_kernel_row_sum[grid_bp]( row_dot_sum, q, k, grad_v, u, lse, H, T, D, fa_scale, ) grad_k = torch.empty_like(k) grad_q = torch.empty_like(q) parallel_deltaformer_bwd_kernel_qk[grid_bp]( grad_q, grad_k, q, k, grad_v, u, lse, beta, row_dot_sum, H, T, D, fa_scale, qk_scale, ) return grad_q, grad_k, row_dot_sum def parallel_deltaformer_kernel( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, u: torch.Tensor, w: torch.Tensor, lse: torch.Tensor, qk_scale: float, beta: torch.Tensor, ) -> None: C, H, D = q.size() T, _H, _D = k.size() def grid(META): return (triton.cdiv(C, META['BLOCK_C']), H) parallel_deltaformer_fwd_kernel[grid]( q, k, v, u, w, lse, beta, H, T, C, D, qk_scale, ) def _config_deltaformer(): return [ triton.Config({'BLOCK_C': BC, 'BLOCK_T': BT}, num_stages=ns, num_warps=nw) for BC in [128, 64] for BT in [64, 32] for ns in [3, 2] for nw in [8, 4] ] @triton.autotune(configs=_config_deltaformer(), key=['C', 'D']) @triton.jit def parallel_deltaformer_fwd_kernel( q_ptr, k_ptr, v_ptr, u_ptr, w_ptr, lse_ptr, beta_ptr, H, T, C, D: tl.constexpr, qk_scale: float, BLOCK_C: tl.constexpr, BLOCK_T: tl.constexpr, ): pid_c = tl.program_id(axis=0) pid_h = tl.program_id(axis=1) rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C colid_block = tl.arange(0, BLOCK_T) rowmax = tl.zeros([BLOCK_C], dtype=tl.float32) - float('inf') rowsum = tl.zeros([BLOCK_C], dtype=tl.float32) + 1 acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) q_blk_ptr = tl.make_block_ptr( base=q_ptr + pid_h * D, shape=(C, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) q = tl.load(q_blk_ptr, boundary_check=(0,)) for kv_i in range(0, T, BLOCK_T): k_blk_ptr = tl.make_block_ptr( base=k_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_T), order=(0, 1), ) k = tl.load(k_blk_ptr, boundary_check=(1,)) qk = tl.dot(q, k) * qk_scale if kv_i >= T - C: mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) qk = tl.where(mask, -1e6, qk) rowmax_i = tl.maximum(rowmax, tl.max(qk, axis=1)) qk -= rowmax_i[:, None] p = tl.math.exp2(qk) rowsum_i = tl.sum(p, axis=1) alpha = tl.math.exp2(rowmax - rowmax_i) rowsum = rowsum * alpha + rowsum_i acc = acc * alpha[:, None] rowmax = rowmax_i if kv_i < T - C: u_blk_ptr = tl.make_block_ptr( base=u_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(kv_i, 0), block_shape=(BLOCK_T, D), order=(1, 0), ) u = tl.load(u_blk_ptr, boundary_check=(0,)) acc = tl.dot(p.to(u_ptr.dtype.element_ty), u, acc) lse = rowmax + tl.math.log2(rowsum) lse_block_ptr = lse_ptr + pid_h + rowid_block * H lse_mask = rowid_block < C tl.store(lse_block_ptr, lse, mask=lse_mask) v_ptr = tl.make_block_ptr( base=v_ptr + pid_h * D, shape=(C, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) acc = acc / rowsum[:, None] beta_ptr = tl.make_block_ptr( base=beta_ptr + pid_h, shape=(C,), strides=(H,), offsets=(pid_c * BLOCK_C,), block_shape=(BLOCK_C,), order=(0,), ) beta = tl.load(beta_ptr, boundary_check=(0,)) acc = acc * beta[:, None] v = tl.load(v_ptr, boundary_check=(0,)) u = v - acc.to(v_ptr.dtype.element_ty) u_block_ptr = tl.make_block_ptr( base=u_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(T - C + pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) tl.store(u_block_ptr, u, boundary_check=(0, 1)) for kv_i in range(T - C, T, BLOCK_T): k_blk_ptr = tl.make_block_ptr( base=k_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_T), order=(0, 1), ) k = tl.load(k_blk_ptr, boundary_check=(1,)) qk = tl.dot(q, k) * qk_scale mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) qk -= rowmax[:, None] p = tl.math.exp2(qk) / rowsum[:, None] p = tl.where(mask, 0, p) w_blk_ptr = tl.make_block_ptr( base=w_ptr + pid_h * C, shape=(C, C), strides=(H * C, 1), offsets=(pid_c * BLOCK_C, kv_i - (T - C)), block_shape=(BLOCK_C, BLOCK_T), order=(1, 0), ) tl.store(w_blk_ptr, p.to(w_ptr.dtype.element_ty), boundary_check=(0, 1)) @triton.autotune(configs=_config_deltaformer(), key=['C', 'D']) @triton.jit def parallel_deltaformer_bwd_kernel_u( o_ptr, q_ptr, k_ptr, v_ptr, lse_ptr, beta_ptr, H, T, C, D: tl.constexpr, fa_scale, BLOCK_C: tl.constexpr, BLOCK_T: tl.constexpr, ): pid_c = tl.program_id(axis=0) pid_h = tl.program_id(axis=1) acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) q_blk_ptr = tl.make_block_ptr( base=q_ptr + pid_h * D, shape=(C, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) q = tl.load(q_blk_ptr, boundary_check=(0,)) for kv_i in range(0, T, BLOCK_T): k_blk_ptr = tl.make_block_ptr( base=k_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_T), order=(0, 1), ) k = tl.load(k_blk_ptr, boundary_check=(1,)) qk = tl.dot(q, k) * fa_scale lse_blk_ptr = tl.make_block_ptr( base=lse_ptr + pid_h, shape=(T,), strides=(H,), offsets=(kv_i,), block_shape=(BLOCK_T,), order=(0,), ) lse = tl.load(lse_blk_ptr, boundary_check=(0,)) beta_blk_ptr = tl.make_block_ptr( base=beta_ptr + pid_h, shape=(T,), strides=(H,), offsets=(kv_i,), block_shape=(BLOCK_T,), order=(0,), ) beta = tl.load(beta_blk_ptr, boundary_check=(0,)) p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] v_blk_ptr = tl.make_block_ptr( base=v_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(kv_i, 0), block_shape=(BLOCK_T, D), order=(1, 0), ) v = tl.load(v_blk_ptr, boundary_check=(0,)) acc = tl.dot(p.to(v_ptr.dtype.element_ty), v, acc) o_blk_ptr = tl.make_block_ptr( base=o_ptr + pid_h * D, shape=(C, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) tl.store(o_blk_ptr, acc.to(o_ptr.dtype.element_ty), boundary_check=(0,)) @triton.autotune(configs=_config_deltaformer(), key=['T', 'D']) @triton.jit def parallel_deltaformer_bwd_kernel_row_sum( row_dot_ptr, q_ptr, k_ptr, grad_v_ptr, u_ptr, lse_ptr, H, T, D: tl.constexpr, fa_scale, BLOCK_C: tl.constexpr, BLOCK_T: tl.constexpr, ): pid_c = tl.program_id(axis=0) pid_h = tl.program_id(axis=1) rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C colid_block = tl.arange(0, BLOCK_T) acc = tl.zeros([BLOCK_C], dtype=tl.float32) k_row_blk_ptr = tl.make_block_ptr( base=q_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) lse_blk_ptr = tl.make_block_ptr( base=lse_ptr + pid_h, shape=(T,), strides=(H,), offsets=(pid_c * BLOCK_C,), block_shape=(BLOCK_C,), order=(0,), ) lse = tl.load(lse_blk_ptr, boundary_check=(0,)) grad_v_blk_ptr = tl.make_block_ptr( base=grad_v_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) for kv_i in range(0, (pid_c + 1) * BLOCK_C, BLOCK_T): k_blk_ptr = tl.make_block_ptr( base=k_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_T), order=(0, 1), ) k = tl.load(k_blk_ptr, boundary_check=(1,)) qk = tl.dot(k_row, k) * fa_scale p = tl.math.exp2(qk - lse[:, None]) u_blk_ptr = tl.make_block_ptr( base=u_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_T), order=(0, 1), ) ut = tl.load(u_blk_ptr, boundary_check=(1,)) dp = tl.dot(grad_v_row, ut) if kv_i + BLOCK_T >= pid_c * BLOCK_C: mask = (rowid_block[:, None] <= colid_block[None, :] + kv_i) p = tl.where(mask, 0., p) dp = tl.where(mask, 0., dp) acc += tl.sum(p * dp, axis=1) row_dot_block_ptr = tl.make_block_ptr( base=row_dot_ptr + pid_h, shape=(T,), strides=(H,), offsets=(pid_c * BLOCK_C,), block_shape=(BLOCK_C,), order=(0,), ) tl.store(row_dot_block_ptr, acc, boundary_check=(0,)) @triton.autotune(configs=[triton.Config({'BLOCK_C': BC}, num_stages=ns, num_warps=nw) for BC in [64, 32] for ns in [4, 3] for nw in [4]], key=['T', 'D']) @triton.jit def parallel_deltaformer_bwd_kernel_qk( grad_q_ptr, grad_k_ptr, q_ptr, k_ptr, grad_v_ptr, u_ptr, lse_ptr, beta_ptr, row_dot_ptr, H, T, D: tl.constexpr, fa_scale: tl.constexpr, qk_scale: tl.constexpr, BLOCK_C: tl.constexpr, ): pid_c = tl.program_id(axis=0) pid_h = tl.program_id(axis=1) block_i = tl.arange(0, BLOCK_C) acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) k_row_blk_ptr = tl.make_block_ptr( base=q_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) lse_blk_ptr = tl.make_block_ptr( base=lse_ptr + pid_h, shape=(T,), strides=(H,), offsets=(pid_c * BLOCK_C,), block_shape=(BLOCK_C,), order=(0,), ) lse = tl.load(lse_blk_ptr, boundary_check=(0,)) beta_blk_ptr = tl.make_block_ptr( base=beta_ptr + pid_h, shape=(T,), strides=(H,), offsets=(pid_c * BLOCK_C,), block_shape=(BLOCK_C,), order=(0,), ) beta = tl.load(beta_blk_ptr, boundary_check=(0,)) grad_v_blk_ptr = tl.make_block_ptr( base=grad_v_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) row_dot_blk_ptr = tl.make_block_ptr( base=row_dot_ptr + pid_h, shape=(T,), strides=(H,), offsets=(pid_c * BLOCK_C,), block_shape=(BLOCK_C,), order=(0,), ) row_dot_row = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) for kv_i in range(0, pid_c * BLOCK_C, BLOCK_C): k_blk_ptr = tl.make_block_ptr( base=k_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_C), order=(0, 1), ) kt = tl.load(k_blk_ptr, boundary_check=(1,)) qk = tl.dot(k_row, kt) * fa_scale p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] u_blk_ptr = tl.make_block_ptr( base=u_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_C), order=(0, 1), ) ut = tl.load(u_blk_ptr) dp = tl.dot(grad_v_row, ut) da = p * (dp - row_dot_row[:, None]) k = tl.trans(kt, 1, 0) acc = tl.dot(da.to(k.dtype), k, acc) k_row_blk_ptr = tl.make_block_ptr( base=k_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(pid_c * BLOCK_C, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) k_row_true = tl.load(k_row_blk_ptr, boundary_check=(0,)) qk = tl.dot(k_row, tl.trans(k_row_true, 1, 0)) * fa_scale p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] u_blk_ptr = tl.make_block_ptr( base=u_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, pid_c * BLOCK_C), block_shape=(D, BLOCK_C), order=(0, 1), ) ut = tl.load(u_blk_ptr) dp = tl.dot(grad_v_row, ut) dpm = dp - row_dot_row[:, None] mask = block_i[None, :] < block_i[:, None] p = tl.where(mask, p, 0.) dpm = tl.where(mask, dpm, 0.) da = p * dpm daat = da acc = tl.dot(daat.to(k_row.dtype), k_row_true, acc) grad_q_blk_ptr = tl.make_block_ptr( base=grad_q_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(BLOCK_C * pid_c, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) acc = acc * qk_scale tl.store(grad_q_blk_ptr, acc.to(grad_q_ptr.dtype.element_ty), boundary_check=(0,)) daat = tl.trans(da, 1, 0) acc = tl.dot(daat.to(k_row.dtype), k_row) k_row = k_row_true nu = -tl.trans(ut, 1, 0) for kv_i in range((pid_c + 1) * BLOCK_C, T, BLOCK_C): k_blk_ptr = tl.make_block_ptr( base=q_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_C), order=(0, 1), ) kt = tl.load(k_blk_ptr, boundary_check=(1,)) lse_blk_ptr = tl.make_block_ptr( base=lse_ptr + pid_h, shape=(T,), strides=(H,), offsets=(kv_i,), block_shape=(BLOCK_C,), order=(0,), ) lse = tl.load(lse_blk_ptr, boundary_check=(0,)) beta_blk_ptr = tl.make_block_ptr( base=beta_ptr + pid_h, shape=(T,), strides=(H,), offsets=(kv_i,), block_shape=(BLOCK_C,), order=(0,), ) beta = tl.load(beta_blk_ptr, boundary_check=(0,)) qk = tl.dot(k_row, kt) * fa_scale p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] grad_vt_blk_ptr = tl.make_block_ptr( base=grad_v_ptr + pid_h * D, shape=(D, T), strides=(1, H * D), offsets=(0, kv_i), block_shape=(D, BLOCK_C), order=(0, 1), ) grad_vt = tl.load(grad_vt_blk_ptr, boundary_check=(1,)) row_dot_blk_ptr = tl.make_block_ptr( base=row_dot_ptr + pid_h, shape=(T,), strides=(H,), offsets=(kv_i,), block_shape=(BLOCK_C,), order=(0,), ) row_dot = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) dp = tl.dot(nu, grad_vt) da = p * (dp - row_dot[None, :]) k = tl.trans(kt, 1, 0) acc = tl.dot(da.to(k.dtype), k, acc) grad_k_blk_ptr = tl.make_block_ptr( base=grad_k_ptr + pid_h * D, shape=(T, D), strides=(H * D, 1), offsets=(BLOCK_C * pid_c, 0), block_shape=(BLOCK_C, D), order=(1, 0), ) acc = acc * qk_scale tl.store(grad_k_blk_ptr, acc.to(grad_k_ptr.dtype.element_ty), boundary_check=(0,)) class ParallelDeltaformerFunction(torch.autograd.Function): @staticmethod def forward( ctx, qo: torch.Tensor, ko: torch.Tensor, vo: torch.Tensor, betao: torch.Tensor | None = None, C: int = BLOCK_SIZE_C, cu_seqlens: torch.LongTensor | None = None, ): B, T, H, D = ko.size() C = min(C, T) ctx.C = C ctx.cu_seqlens = cu_seqlens if cu_seqlens is not None: need_aux = qo.requires_grad or ko.requires_grad or vo.requires_grad or (betao is not None and betao.requires_grad) u, ws, lses = ParallelDeltaformerFunction._forward_impl( qo, ko, vo, betao, C, need_aux=need_aux, cu_seqlens=cu_seqlens) saved_beta = betao if betao is not None else torch.ones(B, T, H, device=ko.device, dtype=ko.dtype) ctx.beta_is_none = betao is None if need_aux: ctx.save_for_backward(qo, ko, vo, u, ws, lses, saved_beta) else: ctx.save_for_backward() return u u, ws, lses = ParallelDeltaformerFunction._forward_impl(qo, ko, vo, betao, C, need_aux=True) saved_beta = betao if betao is not None else torch.ones(B, T, H, device=ko.device, dtype=ko.dtype) ctx.save_for_backward(qo, ko, vo, u, ws, lses, saved_beta) ctx.beta_is_none = betao is None return u @staticmethod def backward( ctx, grad_u: torch.Tensor, ): if getattr(ctx, 'cu_seqlens', None) is not None: cu = ctx.cu_seqlens qo, ko, vo, u_full, ws, lses, betao = ctx.saved_tensors B, T_max, H, D = ko.size() qk_scale = 1.0 / math.sqrt(D) fa_scale = qk_scale / math.log(2) dq = torch.zeros_like(qo) dk = torch.zeros_like(ko) dv = torch.zeros_like(vo) dbeta = None if ctx.beta_is_none else torch.zeros_like(betao) C = ctx.C N = len(cu) - 1 chunk_bases = [] total = 0 lengths = [] for b in range(N): L = int(cu[b + 1].item() - cu[b].item()) lengths.append(L) chunk_bases.append(total) if L > 0: total += (L + C - 1) // C for b in range(N): L = lengths[b] if L == 0: continue base = chunk_bases[b] seq_start = int(cu[b].item()) seq_end = seq_start + L q_seq = qo[0, seq_start:seq_end, :, :] k_seq = ko[0, seq_start:seq_end, :, :] u_seq = u_full[0, seq_start:seq_end, :, :] beta_seq = betao[0, seq_start:seq_end, :] lse_seq = lses[0, seq_start:seq_end, :] go_seq = grad_u[0, seq_start:seq_end, :, :] gv_seq = torch.zeros_like(u_seq) start = ((L - 1) // C) * C for i_local in range(start, -1, -C): Ci = min(C, L - i_local) i0 = i_local i1 = i_local + Ci do = go_seq[i0:i1, :, :] if i_local < L - C: qi = k_seq[i0:i1, :, :] ki = q_seq[i1:L, :, :] lse_tail = lse_seq[i1:L, :] beta_tail = beta_seq[i1:L, :] du_tail = parallel_deltaformer_bwd_u_chunk(qi, ki, lse_tail, gv_seq[i1:L, :, :], fa_scale, beta_tail) do = do - du_tail Wpad = ws[base + (i_local // C)] W = Wpad[:Ci, :, :Ci] W_t = W.transpose(0, 1).contiguous() du_chunk = invcum.backward_x(do.transpose(0, 1).contiguous(), W_t).transpose(0, 1).contiguous() gv_seq[i0:i1, :, :].copy_(du_chunk) gq, gk, gbeta = parallel_deltaformer_bwd_qk(q_seq, k_seq, u_seq, lse_seq, gv_seq, qk_scale, fa_scale, beta_seq) dq[0, seq_start:seq_end, :, :].copy_(gq) dk[0, seq_start:seq_end, :, :].copy_(gk) dv[0, seq_start:seq_end, :, :].copy_(gv_seq) if dbeta is not None: dbeta[0, seq_start:seq_end, :].copy_(gbeta) return dq, dk, dv, dbeta, None, None qo, ko, vo, u, ws, lses, betao = ctx.saved_tensors C = ctx.C B, T, H, D = ko.size() grad_q = torch.zeros_like(qo) grad_k = torch.zeros_like(ko) grad_v = torch.zeros_like(vo) grad_beta_out = None if ctx.beta_is_none else torch.zeros_like(betao) qk_scale = 1.0 / math.sqrt(D) fa_scale = qk_scale / math.log(2) chunk_base = 0 for b in range(B): grad_v_seq = torch.empty(T, H, D, device=ko.device, dtype=ko.dtype) for i in range(T - C, -1, -C): Ci = min(C, T - i) do = grad_u[b, i:i + Ci, :, :] if i < T - C: qi = ko[b, i:i + Ci, :, :] ki = qo[b, i + Ci:, :, :] lse = lses[b, i + Ci:, :] if not ctx.beta_is_none: beta_single = betao[b, i + Ci:, :] else: beta_single = torch.ones(T - i - Ci, H, device=ko.device, dtype=ko.dtype) du = parallel_deltaformer_bwd_u_chunk(qi, ki, lse, grad_v_seq[i + Ci:, :, :], fa_scale, beta_single) do = grad_u[b, i:i + Ci, :, :] - du W = ws[chunk_base + (i // C)][:Ci, :, :Ci] W_t = W.transpose(0, 1).contiguous() du = invcum.backward_x(do.transpose(0, 1).contiguous(), W_t).transpose(0, 1).contiguous() grad_v_seq[i:i + Ci, :, :].copy_(du) q_seq = qo[b] k_seq = ko[b] u_seq = u[b] lse_seq = lses[b] beta_seq = betao[b] if not ctx.beta_is_none else torch.ones(T, H, device=ko.device, dtype=ko.dtype) gq, gk, gbeta = parallel_deltaformer_bwd_qk(q_seq, k_seq, u_seq, lse_seq, grad_v_seq, qk_scale, fa_scale, beta_seq) grad_q[b].copy_(gq) grad_k[b].copy_(gk) grad_v[b].copy_(grad_v_seq) if not ctx.beta_is_none: grad_beta_out[b].copy_(gbeta) chunk_base += (T + C - 1) // C return grad_q, grad_k, grad_v, grad_beta_out, None, None @staticmethod def _forward_impl( qo: torch.Tensor, ko: torch.Tensor, vo: torch.Tensor, betao: torch.Tensor | None, C: int, need_aux: bool, cu_seqlens: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: B, T_max, H, D = ko.size() C = min(C, T_max) qk_scale = 1.0 / math.sqrt(D) fa_scale = qk_scale / math.log(2) if cu_seqlens is None: if betao is None: beta_full = torch.ones(B, T_max, H, device=ko.device, dtype=ko.dtype) else: beta_full = betao u_full = torch.empty_like(vo) if need_aux: total_chunks = B * ((T_max + C - 1) // C) ws = torch.empty(total_chunks, C, H, C, device=ko.device, dtype=ko.dtype) lses = torch.empty(B, T_max, H, device=ko.device, dtype=torch.float) chunk_base = 0 else: ws = None lses = None chunk_base = 0 for b in range(B): for i in range(0, T_max, C): Ci = min(C, T_max - i) qi = qo[b, i:i + Ci, :, :] ki = ko[b, :i + Ci, :, :] vi = vo[b, i:i + Ci, :, :] ui_prev = u_full[b, :i + Ci, :, :] betai = beta_full[b, i:i + Ci, :] w, lse_chunk = parallel_deltaformer_chunk_fwd(qi, ki, vi, ui_prev, fa_scale, betai) w = w * betai.unsqueeze(-1) if need_aux: wpad = torch.zeros(C, H, C, device=ko.device, dtype=ko.dtype) wpad[:Ci, :, :Ci].copy_(w) ws[chunk_base + (i // C)].copy_(wpad) lses[b, i:i + Ci, :].copy_(lse_chunk) u_chunk_view = u_full[b, i:i + Ci, :, :] w_t = w.transpose(0, 1).contiguous() u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous() invcum.forward_inplace(u_chunk_view_t, w_t) u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1)) chunk_base += (T_max + C - 1) // C return u_full, ws, lses N = len(cu_seqlens) - 1 assert cu_seqlens.dim() == 1 and cu_seqlens.size(0) == N + 1, "cu_seqlens must be [N+1]" device = ko.device dtype_k = ko.dtype if betao is None: beta_full = torch.ones(B, T_max, H, device=device, dtype=dtype_k) else: beta_full = betao u_full = torch.empty_like(vo) if need_aux: total_chunks = sum((max(0, int(cu_seqlens[b + 1].item() - cu_seqlens[b].item())) + C - 1) // C for b in range(N)) ws = torch.empty(total_chunks, C, H, C, device=device, dtype=dtype_k) lses = torch.empty(B, T_max, H, device=device, dtype=torch.float) chunk_base = 0 else: ws = None lses = None chunk_base = 0 for b in range(N): seq_start = int(cu_seqlens[b].item()) seq_end = int(cu_seqlens[b + 1].item()) L = max(0, seq_end - seq_start) if L == 0: continue for i_local in range(0, L, C): Ci = min(C, L - i_local) li0 = i_local li1 = i_local + Ci abs_start = seq_start + li0 abs_end = seq_start + li1 abs_context_end = seq_start + li1 qi = qo[0, abs_start:abs_end, :, :] ki = ko[0, seq_start:abs_context_end, :, :] vi = vo[0, abs_start:abs_end, :, :] ui_prev = u_full[0, seq_start:abs_context_end, :, :] betai = beta_full[0, abs_start:abs_end, :] w, lse_chunk = parallel_deltaformer_chunk_fwd(qi, ki, vi, ui_prev, fa_scale, betai) w = w * betai.unsqueeze(-1) if need_aux: wpad = torch.zeros(C, H, C, device=device, dtype=dtype_k) wpad[:Ci, :, :Ci].copy_(w) ws[chunk_base + (i_local // C)].copy_(wpad) lses[0, abs_start:abs_end, :].copy_(lse_chunk) u_chunk_view = u_full[0, abs_start:abs_end, :, :] w_t = w.transpose(0, 1).contiguous() u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous() invcum.forward_inplace(u_chunk_view_t, w_t) u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1)) chunk_base += (L + C - 1) // C return u_full, ws, lses def deltaformer_attn( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, beta: torch.Tensor | None = None, attention_mask: torch.LongTensor | None = None, cu_seqlens: torch.LongTensor | None = None, C: int = BLOCK_SIZE_C, ) -> torch.Tensor: if flash_attn_func is None: raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") B, T, H, D = k.shape C = min(C, T) u = ParallelDeltaformerFunction.apply(q, k, v, beta, C, cu_seqlens) if attention_mask is not None: q_padded, (k_padded, u_padded), indices_q, cu_seqlens_lens, max_seq_lens = unpad_input(q, (k, u), attention_mask, T) cu_seqlens_q, cu_seqlens_k = cu_seqlens_lens max_seqlen_q, max_seqlen_k = max_seq_lens o = flash_attn_varlen_func( q_padded, k_padded, u_padded, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, causal=True, window_size=(-1, -1), ) o = pad_input(o, indices_q, B, T) elif cu_seqlens is not None: max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item()) o = flash_attn_varlen_func( q.squeeze(0), k.squeeze(0), u.squeeze(0), cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens, max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen, causal=True, window_size=(-1, -1), ).unsqueeze(0) else: o = flash_attn_func(q, k, u, causal=True, window_size=(-1, -1)) return o __all__ = [ 'deltaformer_attn', ]