@triton.heuristics({
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@triton.autotune(
configs=[
triton.Config({'BK': BK}, num_warps=num_warps)
for BK in [32, 64]
for num_warps in [1, 2, 4]
],
key=["H", "HV", "K", "BC"],
)
@triton.jit(do_not_specialize=['T'])
def chunk_kda_fwd_kernel_inter_solve_fused(
q,
k,
g,
beta,
Aqk,
Akkd,
Akk,
scale,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
BT: tl.constexpr,
BC: tl.constexpr,
BK: tl.constexpr,
IS_VARLEN: tl.constexpr,
USE_SAFE_GATE: tl.constexpr,
SOLVE_TRIL_DOT_PRECISION: tl.constexpr,
):
"""
Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd.
This kernel:
1. Computes off-diagonal Aqk blocks -> writes to global
2. Computes off-diagonal Akk blocks -> keeps in registers
3. Loads diagonal Akk blocks from Akkd (fp32)
4. Does forward substitution on diagonals
5. Computes merged Akk_inv
6. Writes Akk_inv to Akk
"""
i_t, i_bh = tl.program_id(0), tl.program_id(1)
i_b, i_hv = i_bh // HV, i_bh % HV
i_h = i_hv // (HV // 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 >= T:
return
i_tc0 = i_t * BT
i_tc1 = i_t * BT + BC
i_tc2 = i_t * BT + 2 * BC
i_tc3 = i_t * BT + 3 * BC
q += (bos * H + i_h) * K
k += (bos * H + i_h) * K
g += (bos * HV + i_hv) * K
Aqk += (bos * HV + i_hv) * BT
Akk += (bos * HV + i_hv) * BT
Akkd += (bos * HV + i_hv) * BC
o_i = tl.arange(0, BC)
m_tc1 = (i_tc1 + o_i) < T
m_tc2 = (i_tc2 + o_i) < T
m_tc3 = (i_tc3 + o_i) < T
b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32)
b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32)
b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32)
b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32)
b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32)
b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32)
b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32)
b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32)
b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32)
b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32)
b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32)
b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32)
################################################################################
# off-diagonal blocks
################################################################################
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = o_k < K
p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32)
b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32)
if i_tc1 < T:
p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
# [BC, BK]
b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32)
b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32)
b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32)
# [BK]
b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
# [BC, BK]
b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0)
# [BK, BC]
b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0))
# [BC, BC]
b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt)
b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)
if i_tc2 < T:
p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
# [BC, BK]
b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32)
b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32)
b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32)
# [BK]
b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
# [BC, BK]
b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0)
b_qg2 = b_q2 * b_gqn2
b_kg2 = b_k2 * b_gqn2
# [BK, BC]
b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0))
b_Aqk20 += tl.dot(b_qg2, b_kgt)
b_Akk20 += tl.dot(b_kg2, b_kgt)
# [BC, BC]
b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1))
# [BC, BC]
b_Aqk21 += tl.dot(b_qg2, b_kgt)
b_Akk21 += tl.dot(b_kg2, b_kgt)
if i_tc3 < T:
p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
# [BC, BK]
b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32)
b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32)
b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32)
# [BK]
b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
# [BC, BK]
b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0)
b_qg3 = b_q3 * b_gqn3
b_kg3 = b_k3 * b_gqn3
# [BK, BC]
b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0))
# [BC, BC]
b_Aqk30 += tl.dot(b_qg3, b_kgt)
b_Akk30 += tl.dot(b_kg3, b_kgt)
# [BK, BC]
b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1))
# [BC, BC]
b_Aqk31 += tl.dot(b_qg3, b_kgt)
b_Akk31 += tl.dot(b_kg3, b_kgt)
# [BK, BC]
b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2))
# [BC, BC]
b_Aqk32 += tl.dot(b_qg3, b_kgt)
b_Akk32 += tl.dot(b_kg3, b_kgt)
################################################################################
# save off-diagonal Aqk blocks and prepare Akk
################################################################################
if i_tc1 < T:
p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,))
b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32)
b_Akk10 = b_Akk10 * b_b1[:, None]
if i_tc2 < T:
p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,))
b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32)
b_Akk20 = b_Akk20 * b_b2[:, None]
b_Akk21 = b_Akk21 * b_b2[:, None]
if i_tc3 < T:
p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0))
tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,))
b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32)
b_Akk30 = b_Akk30 * b_b3[:, None]
b_Akk31 = b_Akk31 * b_b3[:, None]
b_Akk32 = b_Akk32 * b_b3[:, None]
p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0))
p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0))
p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0))
p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0))
b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32)
b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32)
b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32)
b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32)
################################################################################
# forward substitution on diagonals
################################################################################
if not USE_SAFE_GATE:
m_A = o_i[:, None] > o_i[None, :]
m_I = o_i[:, None] == o_i[None, :]
b_Ai00 = -tl.where(m_A, b_Ai00, 0)
b_Ai11 = -tl.where(m_A, b_Ai11, 0)
b_Ai22 = -tl.where(m_A, b_Ai22, 0)
b_Ai33 = -tl.where(m_A, b_Ai33, 0)
for i in range(2, min(BC, T - i_tc0)):
b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
b_a00 = tl.where(o_i < i, b_a00, 0.)
b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0)
b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
for i in range(BC + 2, min(2*BC, T - i_tc0)):
b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
b_a11 = tl.where(o_i < i - BC, b_a11, 0.)
b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0)
b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11)
for i in range(2*BC + 2, min(3*BC, T - i_tc0)):
b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.)
b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0)
b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22)
for i in range(3*BC + 2, min(4*BC, T - i_tc0)):
b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.)
b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0)
b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33)
b_Ai00 += m_I
b_Ai11 += m_I
b_Ai22 += m_I
b_Ai33 += m_I
################################################################################
# compute merged inverse using off-diagonals
################################################################################
# we used tf32 to maintain matrix inverse's precision whenever possible.
b_Ai10 = -tl.dot(
tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai00,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai21 = -tl.dot(
tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai11,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai32 = -tl.dot(
tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai22,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai20 = -tl.dot(
b_Ai22,
tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai31 = -tl.dot(
b_Ai33,
tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai30 = -tl.dot(
b_Ai33,
tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION
)
################################################################################
# store full Akk_inv to Akk
################################################################################
p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0))
p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0))
p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0))
p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0))
p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0))
tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1))