| |
|
|
| 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', |
| ] |
|
|
