| |
|
|
| 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 |
|
|
