| |
|
|
|
|
| import torch |
| import triton |
| import triton.language as tl |
|
|
| from fla.ops.utils.cumsum import chunk_global_cumsum |
| from fla.ops.utils.op import exp |
| from fla.utils import autotune_cache_kwargs, check_shared_mem |
|
|
|
|
| @triton.heuristics({ |
| 'USE_G': lambda args: args['g_cumsum'] is not None, |
| }) |
| @triton.autotune( |
| configs=[ |
| triton.Config({}, num_warps=num_warps, num_stages=num_stages) |
| for num_warps in [1, 2, 4] + ([] if check_shared_mem('hopper') else [8]) |
| for num_stages in [2, 3, 4, 5] |
| ], |
| key=['H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'], |
| **autotune_cache_kwargs, |
| ) |
| @triton.jit |
| def naive_attn_decoding_kernel( |
| q, |
| k, |
| v, |
| o, |
| g_cumsum, |
| scale, |
| gate_scale, |
| cu_seqlens, |
| T, |
| B: tl.constexpr, |
| H: tl.constexpr, |
| HQ: tl.constexpr, |
| G: tl.constexpr, |
| K: tl.constexpr, |
| V: tl.constexpr, |
| BS: tl.constexpr, |
| BK: tl.constexpr, |
| BV: tl.constexpr, |
| USE_G: tl.constexpr, |
| ): |
| i_v, i_bh = tl.program_id(0), tl.program_id(1) |
| i_b, i_hq = i_bh // HQ, i_bh % HQ |
| i_h = i_hq // G |
|
|
| bos, eos = tl.load(cu_seqlens + i_b).to(tl.int32), tl.load(cu_seqlens + i_b + 1).to(tl.int32) |
| T = eos - bos |
|
|
| p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1, ), (0, ), (BK,), (0,)) |
| p_o = tl.make_block_ptr(o + i_bh * V, (V,), (1, ), (0, ), (BV,), (0,)) |
|
|
| b_q = tl.load(p_q, boundary_check=(0,)) |
| b_q = (b_q * scale).to(b_q.dtype) |
|
|
| b_o = tl.zeros([BV ], dtype=tl.float32) |
|
|
| b_m = tl.full([1], float('-inf'), dtype=tl.float32) |
| b_acc = tl.zeros([1], dtype=tl.float32) |
|
|
| if USE_G: |
| p_g = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (T-1,), (1,), (0,)) |
| b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) |
| else: |
| b_gq = None |
|
|
| for i_s in range(0, T, BS): |
| p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_s, 0), (BS, BK), (1, 0)) |
| p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) |
| |
| b_k = tl.load(p_k, boundary_check=(0, 1)) |
| |
| b_v = tl.load(p_v, boundary_check=(0, 1)) |
| |
| b_s = tl.sum(b_q[None, :] * b_k, 1) |
|
|
| mask = i_s + tl.arange(0, BS) < T |
| b_s = tl.where(mask, b_s, float('-inf')) |
|
|
| if USE_G: |
| p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) |
| b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) |
| b_s += (b_gq - b_gk) * gate_scale |
| |
| b_m, b_mp = tl.maximum(b_m, tl.max(b_s)), b_m |
| b_r = exp(b_mp - b_m) |
| |
| b_p = exp(b_s - b_m) |
|
|
| |
| b_acc = b_acc * b_r + tl.sum(b_p, 0) |
| |
| b_o = b_o * b_r + tl.sum(b_p[:, None] * b_v, 0) |
| b_mp = b_m |
| b_o = b_o / b_acc |
| tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, )) |
|
|
|
|
| def attn_decoding_one_step( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| g: torch.Tensor | None = None, |
| scale: float | None = None, |
| cu_seqlens: torch.LongTensor = None, |
| do_gate_scale: bool = False, |
| ): |
| r""" |
| Args: |
| q (torch.Tensor): |
| query of shape `[1, B, HQ, K]`. |
| k (torch.Tensor): |
| keys of shape `[1, T, H, K]`. |
| GQA will be applied if HQ is divisible by H. T is the cumulative length for all batch. |
| v (torch.Tensor): |
| values of shape `[1, T, H, V]`. |
| g (Optional[torch.Tensor]): |
| log decay factors of shape `[1, T, H]`. Default: `None`. |
| scale (Optional[float]): |
| Scale factor for attention scores. |
| If not provided, it will default to `1 / sqrt(K)`. Default: `None`. |
| cu_seqlens (torch.LongTensor): |
| Cumulative sequence lengths of shape `[N+1]` used for variable-length training, |
| consistent with the FlashAttention API. |
| do_gate_scale (bool): |
| Whether to apply gate scale. Default: `False`. If `True`, the attention scale will also be applied |
| to the gating bias term in Forgetting Transformer or PaTH-FoX. |
| |
| Returns: |
| o (torch.Tensor): |
| Outputs of shape `[B, 1, HQ, V]`. |
| """ |
| assert cu_seqlens is not None, "The cu_seqlens must be provided for varlen decoding" |
| B, T, H, K, V = *k.shape, v.shape[-1] |
| N = len(cu_seqlens) - 1 |
| HQ = q.shape[2] |
| G = HQ // H |
| if scale is None: |
| scale = K ** -0.5 |
|
|
| BK = max(triton.next_power_of_2(K), 16) |
| if check_shared_mem('hopper', q.device.index): |
| BS = min(64, max(16, triton.next_power_of_2(T))) |
| BV = min(256, max(16, triton.next_power_of_2(V))) |
| elif check_shared_mem('ampere', q.device.index): |
| BS = min(32, max(16, triton.next_power_of_2(T))) |
| BV = min(128, max(16, triton.next_power_of_2(V))) |
| else: |
| BS = min(32, max(16, triton.next_power_of_2(T))) |
| BV = min(64, max(16, triton.next_power_of_2(V))) |
| g_cumsum = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, output_dtype=torch.float32) if g is not None else None |
| NV = triton.cdiv(V, BV) |
| o = torch.empty(*q.shape[:-1], V, dtype=v.dtype, device=q.device) |
| gate_scale = 1.0 if not do_gate_scale else scale |
|
|
| grid = (NV, N * HQ) |
| naive_attn_decoding_kernel[grid]( |
| q=q, |
| k=k, |
| v=v, |
| o=o, |
| g_cumsum=g_cumsum, |
| scale=scale, |
| gate_scale=gate_scale, |
| cu_seqlens=cu_seqlens, |
| B=B, |
| T=T, |
| H=H, |
| HQ=HQ, |
| G=G, |
| K=K, |
| V=V, |
| BS=BS, |
| BK=BK, |
| BV=BV, |
| ) |
| return o |
|
|
