| |
|
|
|
|
| import torch |
| import triton |
| import triton.language as tl |
|
|
| from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices, solve_tril |
| from fla.ops.utils.op import exp |
| from fla.utils import autotune_cache_kwargs |
|
|
|
|
| @triton.heuristics({ |
| 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, |
| }) |
| @triton.autotune( |
| configs=[ |
| triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) |
| for BK in [32, 64] |
| for num_warps in [1, 2, 4, 8] |
| for num_stages in [2, 3, 4] |
| ], |
| key=["BC"], |
| **autotune_cache_kwargs, |
| ) |
| @triton.jit(do_not_specialize=['T']) |
| def chunk_kda_fwd_kernel_intra_sub_inter( |
| q, |
| k, |
| g, |
| beta, |
| Aqk, |
| Akk, |
| scale, |
| cu_seqlens, |
| chunk_indices, |
| T, |
| H: tl.constexpr, |
| K: tl.constexpr, |
| BT: tl.constexpr, |
| BC: tl.constexpr, |
| BK: tl.constexpr, |
| NC: tl.constexpr, |
| IS_VARLEN: tl.constexpr, |
| ): |
| i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| i_b, i_h = i_bh // H, i_bh % H |
| i_i, i_j = i_c // NC, i_c % NC |
| if IS_VARLEN: |
| i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) |
| bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) |
| T = eos - bos |
| else: |
| bos, eos = i_b * T, i_b * T + T |
|
|
| if i_t * BT + i_i * BC >= T: |
| return |
| if i_i <= i_j: |
| return |
|
|
| q += (bos * H + i_h) * K |
| k += (bos * H + i_h) * K |
| g += (bos * H + i_h) * K |
| Aqk += (bos * H + i_h) * BT |
| Akk += (bos * H + i_h) * BT |
|
|
| p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) |
| b_b = tl.load(p_b, boundary_check=(0,)) |
|
|
| b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) |
| b_Akk = tl.zeros([BC, BC], dtype=tl.float32) |
| for i_k in range(tl.cdiv(K, BK)): |
| p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| b_kt = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) |
| p_gk = tl.make_block_ptr(g, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) |
|
|
| o_k = i_k * BK + tl.arange(0, BK) |
| m_k = o_k < K |
| |
| b_gn = tl.load(g + (i_t * BT + i_i * BC) * H*K + o_k, mask=m_k, other=0) |
| |
| b_g = tl.load(p_g, boundary_check=(0, 1)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) |
| |
| b_gk = tl.load(p_gk, boundary_check=(0, 1)) |
| b_kt = tl.load(b_kt, boundary_check=(0, 1)) |
| |
| b_ktg = b_kt * exp(b_gn[:, None] - b_gk) |
| b_Akk += tl.dot(b_k, b_ktg) |
|
|
| b_q = tl.load(p_q, boundary_check=(0, 1)) |
| b_qg = b_q * exp(b_g - b_gn[None, :]) * scale |
| b_Aqk += tl.dot(b_qg, b_ktg) |
|
|
| b_Akk *= b_b[:, None] |
|
|
| p_Akk = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) |
| tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) |
| p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) |
| tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) |
|
|
|
|
| @triton.heuristics({ |
| 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, |
| }) |
| @triton.autotune( |
| configs=[ |
| triton.Config({}, num_warps=num_warps) |
| for num_warps in [1, 2, 4, 8] |
| ], |
| key=["BK", "BT"], |
| **autotune_cache_kwargs, |
| ) |
| @triton.jit(do_not_specialize=['T']) |
| def chunk_kda_fwd_kernel_intra_sub_intra( |
| q, |
| k, |
| g, |
| beta, |
| Aqk, |
| Akk, |
| scale, |
| cu_seqlens, |
| chunk_indices, |
| T, |
| H: tl.constexpr, |
| K: tl.constexpr, |
| BT: tl.constexpr, |
| BC: tl.constexpr, |
| BK: tl.constexpr, |
| IS_VARLEN: tl.constexpr, |
| ): |
| i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| i_b, i_h = i_bh // H, i_bh % H |
| if IS_VARLEN: |
| i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) |
| bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) |
| T = eos - bos |
| else: |
| bos, eos = i_b * T, i_b * T + T |
|
|
| if i_t * BT + i_i * BC >= T: |
| return |
|
|
| o_i = tl.arange(0, BC) |
| o_k = tl.arange(0, BK) |
| m_k = o_k < K |
| m_A = (i_t * BT + i_i * BC + o_i) < T |
| o_A = (bos + i_t * BT + i_i * BC + o_i) * H*BT + i_h * BT + i_i * BC |
|
|
| p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) |
| p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)) |
| b_g = tl.load(p_g, boundary_check=(0, 1)) |
|
|
| p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h |
| b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] |
|
|
| p_kt = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k |
| p_gk = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k |
|
|
| for j in range(0, min(BC, T - i_t * BT - i_i * BC)): |
| b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) |
| b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) |
| b_ktg = b_kt[None, :] * exp(b_g - b_gk[None, :]) |
| b_Aqk = tl.sum(b_q * b_ktg, 1) |
| b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.) |
| b_Akk = tl.sum(b_k * b_ktg, 1) |
| b_Akk = tl.where(o_i > j, b_Akk, 0.) |
| tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) |
| tl.store(Akk + o_A + j, b_Akk, mask=m_A) |
| p_kt += H*K |
| p_gk += H*K |
|
|
|
|
| @triton.heuristics({ |
| 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, |
| }) |
| @triton.autotune( |
| configs=[ |
| triton.Config({}, num_warps=num_warps, num_stages=num_stages) |
| for num_warps in [1, 2, 4, 8] |
| for num_stages in [2, 3, 4] |
| ], |
| key=['BK', 'NC', 'BT'], |
| **autotune_cache_kwargs, |
| ) |
| @triton.jit(do_not_specialize=['B', 'T']) |
| def chunk_kda_bwd_kernel_intra( |
| q, |
| k, |
| g, |
| beta, |
| dAqk, |
| dAkk, |
| dq, |
| dq2, |
| dk, |
| dk2, |
| dg, |
| db, |
| cu_seqlens, |
| chunk_indices, |
| B, |
| T, |
| H: tl.constexpr, |
| K: tl.constexpr, |
| BT: tl.constexpr, |
| BC: tl.constexpr, |
| BK: tl.constexpr, |
| NC: tl.constexpr, |
| IS_VARLEN: tl.constexpr, |
| ): |
| i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| i_b, i_h = i_bh // H, i_bh % H |
| i_k, i_i = i_kc // NC, i_kc % NC |
|
|
| all = B * T |
| if IS_VARLEN: |
| i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) |
| bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) |
| else: |
| bos, eos = i_b * T, i_b * T + T |
| T = eos - bos |
| if i_t * BT + i_i * BC >= T: |
| return |
|
|
| o_k = i_k * BK + tl.arange(0, BK) |
| m_k = o_k < K |
|
|
| q += (bos * H + i_h) * K |
| k += (bos * H + i_h) * K |
| g += (bos * H + i_h) * K |
| beta += bos * H + i_h |
|
|
| dAqk += (bos * H + i_h) * BT |
| dAkk += (bos * H + i_h) * BT |
| dq += (bos * H + i_h) * K |
| dq2 += (bos * H + i_h) * K |
| dk += (bos * H + i_h) * K |
| dk2 += (bos * H + i_h) * K |
| dg += (bos * H + i_h) * K |
| db += (i_k * all + bos) * H + i_h |
|
|
| p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| b_g = tl.load(p_g, boundary_check=(0, 1)) |
|
|
| p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) |
| b_b = tl.load(p_b, boundary_check=(0,)) |
|
|
| b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) |
| b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) |
| if i_i > 0: |
| p_gn = g + (i_t * BT + i_i * BC) * H*K + o_k |
| |
| b_gn = tl.load(p_gn, mask=m_k, other=0) |
| for i_j in range(0, i_i): |
| p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) |
| p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) |
| |
| b_k = tl.load(p_k, boundary_check=(0, 1)) |
| b_gk = tl.load(p_gk, boundary_check=(0, 1)) |
| b_kg = b_k * exp(b_gn[None, :] - b_gk) |
| |
| b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) |
| b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) |
| |
| b_dq2 += tl.dot(b_dAqk, b_kg) |
| b_dk2 += tl.dot(b_dAkk, b_kg) |
| b_dq2 *= exp(b_g - b_gn[None, :]) |
| b_dk2 *= exp(b_g - b_gn[None, :]) |
|
|
| o_i = tl.arange(0, BC) |
| m_dA = (i_t * BT + i_i * BC + o_i) < T |
| o_dA = (i_t * BT + i_i * BC + o_i) * H*BT + i_i * BC |
| p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k |
| p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k |
|
|
| p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| b_q = tl.load(p_q, boundary_check=(0, 1)) |
| b_k = tl.load(p_k, boundary_check=(0, 1)) |
|
|
| for j in range(0, min(BC, T - i_t * BT - i_i * BC)): |
| |
| b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) |
| b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) |
| |
| b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) |
| b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) |
| |
| m_i = o_i[:, None] >= j |
| |
| b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kj[None, :] * exp(b_g - b_gkj[None, :]), 0.) |
| b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * exp(b_g - b_gkj[None, :]), 0.) |
|
|
| p_kj += H*K |
| p_gkj += H*K |
| b_db = tl.sum(b_dk2 * b_k, 1) |
| b_dk2 *= b_b[:, None] |
|
|
| p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_dq2 = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) |
|
|
| b_dg = b_q * b_dq2 |
| b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) |
| tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) |
|
|
| tl.debug_barrier() |
| b_dkt = tl.zeros([BC, BK], dtype=tl.float32) |
|
|
| NC = min(NC, tl.cdiv(T - i_t * BT, BC)) |
| if i_i < NC - 1: |
| p_gn = g + (min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + o_k |
| |
| b_gn = tl.load(p_gn, mask=m_k, other=0) |
| for i_j in range(i_i + 1, NC): |
| p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) |
| p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) |
| p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) |
| p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) |
| p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) |
| |
| b_b = tl.load(p_b, boundary_check=(0,)) |
| |
| b_q = tl.load(p_q, boundary_check=(0, 1)) |
| b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] |
| b_gk = tl.load(p_gk, boundary_check=(0, 1)) |
| |
| b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) |
| b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) |
|
|
| o_j = i_t * BT + i_j * BC + o_i |
| m_j = o_j < T |
| |
| b_qg = b_q * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) |
| b_kbg = b_kb * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) |
| |
| |
| b_dkt += tl.dot(b_dAqk, b_qg) |
| b_dkt += tl.dot(b_dAkk, b_kbg) |
| b_dkt *= exp(b_gn[None, :] - b_g) |
| o_dA = (i_t * BT + i_i * BC) * H*BT + i_i * BC + o_i |
| p_qj = q + (i_t * BT + i_i * BC) * H*K + o_k |
| p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k |
| p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k |
| p_bj = beta + (i_t * BT + i_i * BC) * H |
|
|
| for j in range(0, min(BC, T - i_t * BT - i_i * BC)): |
| |
| b_dAqk = tl.load(dAqk + o_dA + j * H*BT) |
| b_dAkk = tl.load(dAkk + o_dA + j * H*BT) |
| |
| b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) |
| b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) |
| b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) |
| |
| m_i = o_i[:, None] <= j |
| b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * exp(b_gkj[None, :] - b_g), 0.) |
| b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kbj[None, :] * exp(b_gkj[None, :] - b_g), 0.) |
|
|
| p_qj += H*K |
| p_kj += H*K |
| p_gkj += H*K |
| p_bj += H |
| p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_dk2 = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
| p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) |
|
|
| b_dg += (b_dk2 - b_dkt) * b_k |
| b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) |
| b_dk2 += b_dkt |
|
|
| tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) |
| tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) |
|
|
|
|
| def chunk_kda_fwd_intra( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| gk: torch.Tensor | None = None, |
| beta: torch.Tensor | None = None, |
| scale: float | None = None, |
| cu_seqlens: torch.LongTensor | None = None, |
| chunk_size: int = 64, |
| output_dtype: torch.dtype = torch.float32, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| r""" |
| Args: |
| q (torch.Tensor): |
| The query tensor of shape `[B, T, H, K]`. |
| k (torch.Tensor): |
| The key tensor of shape `[B, T, H, K]`. |
| gk (torch.Tensor): |
| The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. |
| beta (torch.Tensor): |
| The beta tensor of shape `[B, T, H]`. Default: `None`. |
| scale (Optional[float]): |
| The scale factor. Default: `None`. |
| cu_seqlens (torch.LongTensor): |
| The cumulative sequence lengths of the input tensor. |
| Default: None |
| chunk_size (int): |
| The chunk size. Default: 64. |
| output_dtype (torch.dtype): |
| The dtype of the output tensor. Default: `torch.float32` |
| |
| Returns: |
| Aqk (torch.Tensor): |
| The intra Aqk tensor of shape `[B, T, H, BT]` where `BT` is the chunk size. |
| Akk (torch.Tensor): |
| The intra Akk tensor of shape `[B, T, H, BT]` where `BT` is the chunk size. |
| """ |
| B, T, H, K = k.shape |
| assert K <= 256 |
| BT = chunk_size |
| chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None |
| NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) |
|
|
| BC = min(16, BT) |
| NC = triton.cdiv(BT, BC) |
| BK = max(triton.next_power_of_2(K), 16) |
|
|
| Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) |
| Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) |
| grid = (NT, NC * NC, B * H) |
| chunk_kda_fwd_kernel_intra_sub_inter[grid]( |
| q=q, |
| k=k, |
| g=gk, |
| beta=beta, |
| Aqk=Aqk, |
| Akk=Akk, |
| scale=scale, |
| cu_seqlens=cu_seqlens, |
| chunk_indices=chunk_indices, |
| T=T, |
| H=H, |
| K=K, |
| BT=BT, |
| BC=BC, |
| NC=NC, |
| ) |
|
|
| grid = (NT, NC, B * H) |
| chunk_kda_fwd_kernel_intra_sub_intra[grid]( |
| q=q, |
| k=k, |
| g=gk, |
| beta=beta, |
| Aqk=Aqk, |
| Akk=Akk, |
| scale=scale, |
| cu_seqlens=cu_seqlens, |
| chunk_indices=chunk_indices, |
| T=T, |
| H=H, |
| K=K, |
| BT=BT, |
| BC=BC, |
| BK=BK, |
| ) |
| Akk = solve_tril( |
| A=Akk, |
| cu_seqlens=cu_seqlens, |
| output_dtype=k.dtype, |
| ) |
| return Aqk, Akk |
|
|
|
|
| def chunk_kda_bwd_intra( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| g: torch.Tensor, |
| beta: torch.Tensor, |
| dAqk: torch.Tensor, |
| dAkk: torch.Tensor, |
| dq: torch.Tensor, |
| dk: torch.Tensor, |
| db: torch.Tensor, |
| dg: torch.Tensor, |
| cu_seqlens: torch.LongTensor | None = None, |
| chunk_size: int = 64, |
| ): |
| B, T, H, K = k.shape |
| BT = chunk_size |
| BC = min(16, BT) |
| BK = min(64, triton.next_power_of_2(K)) |
|
|
| chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None |
| NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) |
| NC = triton.cdiv(BT, BC) |
| NK = triton.cdiv(K, BK) |
|
|
| dq2 = torch.empty_like(q) |
| dk2 = torch.empty_like(k) |
| db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) |
| dg2 = torch.empty_like(dg, dtype=torch.float) |
| grid = (NK * NC, NT, B * H) |
| chunk_kda_bwd_kernel_intra[grid]( |
| q=q, |
| k=k, |
| g=g, |
| beta=beta, |
| dAqk=dAqk, |
| dAkk=dAkk, |
| dq=dq, |
| dq2=dq2, |
| dk=dk, |
| dk2=dk2, |
| dg=dg2, |
| db=db2, |
| cu_seqlens=cu_seqlens, |
| chunk_indices=chunk_indices, |
| B=B, |
| T=T, |
| H=H, |
| K=K, |
| BT=BT, |
| BC=BC, |
| BK=BK, |
| NC=NC, |
| ) |
| dq = dq2 |
| dk = dk2 |
| db = db2.sum(0).add_(db) |
| dg = chunk_local_cumsum(dg2.add_(dg), chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens) |
|
|
| return dq, dk2, db, dg |
|
|
