File size: 4,955 Bytes
5ebdfe3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
import torch
import triton
import triton.language as tl
from fla.ops.utils.op import exp
from fla.utils import input_guard
@triton.jit
def mesa_net_decoding_one_step_kernel(
q,
k,
v,
g,
o,
lamb,
beta,
prev_h_kk,
prev_h_kv,
curr_h_kk,
curr_h_kv,
B: tl.constexpr,
H: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
MAX_CG_STEP: tl.constexpr,
):
i_nh = tl.program_id(0)
i_h = i_nh % H
o_k = tl.arange(0, BK)
o_v = tl.arange(0, BV)
p_q = q + i_nh * K + o_k
p_k = k + i_nh * K + o_k
p_v = v + i_nh * V + o_v
p_beta = beta + i_nh
p_g = g + i_nh
p_lamb = lamb + i_h * K + o_k
b_g = exp(tl.load(p_g).to(tl.float32))
b_beta = tl.load(p_beta).to(tl.float32)
mask_k = o_k < K
mask_v = o_v < V
mask_kk = mask_k[:, None] & mask_k[None, :]
mask_kv = mask_k[:, None] & mask_v[None, :]
b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
b_lamb = tl.load(p_lamb, mask=mask_k, other=0).to(tl.float32)
p_hkk_prev = prev_h_kk + i_nh * K * K + o_k[:, None] * K + o_k[None, :]
b_h_kk = tl.load(p_hkk_prev, mask=mask_kk, other=0).to(tl.float32)
b_h_kk = b_h_kk * b_g + (b_k * b_beta)[:, None] * b_k[None, :]
p_hkk_curr = curr_h_kk + i_nh * K * K + o_k[:, None] * K + o_k[None, :]
tl.store(p_hkk_curr, b_h_kk.to(p_hkk_curr.dtype.element_ty), mask=mask_kk)
p_hkv_prev = prev_h_kv + i_nh * K * V + o_k[:, None] * V + o_v[None, :]
b_h_kv = tl.load(p_hkv_prev, mask=mask_kv, other=0).to(tl.float32)
b_h_kv = b_h_kv * b_g + (b_k * b_beta)[:, None] * b_v[None, :]
p_hkv_curr = curr_h_kv + i_nh * K * V + o_k[:, None] * V + o_v[None, :]
tl.store(p_hkv_curr, b_h_kv.to(p_hkv_curr.dtype.element_ty), mask=mask_kv)
diag_mask = tl.arange(0, BK)[:, None] == tl.arange(0, BK)[None, :]
diag_mask = diag_mask & mask_kk
b_h_kk_diag = tl.sum(tl.where(diag_mask, b_h_kk, 0.0), axis=1)
b_x = b_q / (b_h_kk_diag + b_lamb + 1e-5)
b_Hx = tl.sum(b_h_kk * b_x[:, None], axis=0)
b_r = b_q - b_Hx - b_lamb * b_x
b_p = tl.zeros([BK], dtype=tl.float32)
b_p += b_r
delta_old = tl.sum(b_r * b_r)
for i_iter in range(MAX_CG_STEP):
b_Ap = tl.sum(b_h_kk * b_p[:, None], axis=0) + b_lamb * b_p
pAp = tl.sum(b_p * b_Ap)
alpha = delta_old / (pAp + 1e-5)
b_x = b_x + alpha * b_p
b_r = b_r - alpha * b_Ap
delta_new = tl.sum(b_r * b_r)
beta_cg = delta_new / (delta_old + 1e-5)
b_p = b_r + beta_cg * b_p
delta_old = delta_new
b_o = tl.sum(b_h_kv * b_x[:, None], axis=0)
p_o = o + i_nh * V + o_v
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
@input_guard
def mesa_net_decoding_one_step(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
lamb: torch.Tensor,
beta: torch.Tensor,
prev_h_kk: torch.Tensor,
prev_h_kv: torch.Tensor,
max_CG_iteration: int = 30,
):
"""
Triton implementation of Mesa Net CG one step
Args:
q (torch.Tensor):
query tensor [B, H, K]
k (torch.Tensor):
key tensor [B, H, K]
v (torch.Tensor):
value tensor [B, H, V]
g (torch.Tensor):
gate tensor [B, H]
lamb (torch.Tensor):
lambda tensor [H, K]
beta (torch.Tensor):
beta tensor [B, H]
prev_h_kk (torch.Tensor):
previous hidden state KK [B, H, K, K]
prev_h_kv (torch.Tensor):
previous hidden state KV [B, H, K, V]
max_CG_iteration (int):
maximum CG iterations
Returns:
o (torch.Tensor):
output tensor [B, H, V]
h_kk_new (torch.Tensor):
updated hidden state KK [B, H, K, K]
h_kv_new (torch.Tensor):
updated hidden state KV [B, H, K, V]
"""
B, H, K, V = *q.shape, v.shape[-1]
o = torch.empty((B, H, V), dtype=q.dtype, device=q.device)
curr_h_kk = torch.empty_like(prev_h_kk)
curr_h_kv = torch.empty_like(prev_h_kv)
BK = max(triton.next_power_of_2(K), 16)
BV = max(triton.next_power_of_2(V), 16)
assert BK <= 128 and BV <= 128, "BK and BV must be less than or equal to 128"
grid = (B * H,)
mesa_net_decoding_one_step_kernel[grid](
q=q,
k=k,
v=v,
g=g,
o=o,
lamb=lamb,
beta=beta,
prev_h_kk=prev_h_kk,
prev_h_kv=prev_h_kv,
curr_h_kk=curr_h_kk,
curr_h_kv=curr_h_kv,
B=B,
H=H,
K=K,
V=V,
BK=BK,
BV=BV,
MAX_CG_STEP=max_CG_iteration,
num_warps=4 if BK <= 64 else 8,
)
return o, curr_h_kk, curr_h_kv
|