| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang |
| 4 | # |
| 5 | # This file contains code copied from the flash-linear-attention project. |
| 6 | # The original source code was licensed under the MIT license and included |
| 7 | # the following copyright notice: |
| 8 | # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang |
| 9 | |
| 10 | # ruff: noqa: E501 |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm.triton_utils import tl, triton |
| 15 | |
| 16 | from .index import prepare_chunk_indices |
| 17 | |
| 18 | |
| 19 | @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) |
| 20 | @triton.autotune( |
| 21 | configs=[ |
| 22 | triton.Config({}, num_warps=num_warps, num_stages=num_stages) |
| 23 | for num_warps in [2, 4, 8] |
| 24 | for num_stages in [2, 3, 4] |
| 25 | ], |
| 26 | key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], |
| 27 | ) |
| 28 | @triton.jit(do_not_specialize=["T"]) |
| 29 | def recompute_w_u_fwd_kernel( |
| 30 | k, |
| 31 | v, |
| 32 | beta, |
| 33 | w, |
| 34 | u, |
| 35 | A, |
| 36 | g, |
| 37 | cu_seqlens, |
| 38 | chunk_indices, |
| 39 | T, |
| 40 | H: tl.constexpr, |
| 41 | Hg: tl.constexpr, |
| 42 | K: tl.constexpr, |
| 43 | V: tl.constexpr, |
| 44 | BT: tl.constexpr, |
| 45 | BK: tl.constexpr, |
| 46 | BV: tl.constexpr, |
| 47 | IS_VARLEN: tl.constexpr, |
| 48 | ): |
| 49 | i_t, i_bh = tl.program_id(0), tl.program_id(1) |
| 50 | i_b, i_h = i_bh // H, i_bh % H |
| 51 | if IS_VARLEN: |
| 52 | i_n, i_t = ( |
| 53 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 54 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 55 | ) |
| 56 | bos, eos = ( |
| 57 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 58 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 59 | ) |
| 60 | T = eos - bos |
| 61 | else: |
| 62 | bos, eos = i_b * T, i_b * T + T |
| 63 | p_beta = tl.make_block_ptr( |
| 64 | beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) |
| 65 | ) |
| 66 | p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) |
| 67 | p_A = tl.make_block_ptr( |
| 68 | A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) |
| 69 | ) |
| 70 | b_beta = tl.load(p_beta, boundary_check=(0,)) |
| 71 | b_A = tl.load(p_A, boundary_check=(0, 1)) |
| 72 | b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) |
| 73 | |
| 74 | for i_v in range(tl.cdiv(V, BV)): |
| 75 | p_v = tl.make_block_ptr( |
| 76 | v + (bos * H + i_h) * V, |
| 77 | (T, V), |
| 78 | (H * V, 1), |
| 79 | (i_t * BT, i_v * BV), |
| 80 | (BT, BV), |
| 81 | (1, 0), |
| 82 | ) |
| 83 | p_u = tl.make_block_ptr( |
| 84 | u + (bos * H + i_h) * V, |
| 85 | (T, V), |
| 86 | (H * V, 1), |
| 87 | (i_t * BT, i_v * BV), |
| 88 | (BT, BV), |
| 89 | (1, 0), |
| 90 | ) |
| 91 | b_v = tl.load(p_v, boundary_check=(0, 1)) |
| 92 | b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) |
| 93 | b_u = tl.dot(b_A, b_vb, allow_tf32=False) |
| 94 | tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) |
| 95 | |
| 96 | for i_k in range(tl.cdiv(K, BK)): |
| 97 | p_k = tl.make_block_ptr( |
| 98 | k + (bos * Hg + i_h // (H // Hg)) * K, |
| 99 | (T, K), |
| 100 | (Hg * K, 1), |
| 101 | (i_t * BT, i_k * BK), |
| 102 | (BT, BK), |
| 103 | (1, 0), |
| 104 | ) |
| 105 | p_w = tl.make_block_ptr( |
| 106 | w + (bos * H + i_h) * K, |
| 107 | (T, K), |
| 108 | (H * K, 1), |
| 109 | (i_t * BT, i_k * BK), |
| 110 | (BT, BK), |
| 111 | (1, 0), |
| 112 | ) |
| 113 | b_k = tl.load(p_k, boundary_check=(0, 1)) |
| 114 | b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) |
| 115 | b_w = tl.dot(b_A, b_kb) |
| 116 | tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) |
| 117 | |
| 118 | |
| 119 | def recompute_w_u_fwd( |
| 120 | k: torch.Tensor, |
| 121 | v: torch.Tensor, |
| 122 | beta: torch.Tensor, |
| 123 | g_cumsum: torch.Tensor, |
| 124 | A: torch.Tensor, |
| 125 | cu_seqlens: torch.Tensor | None, |
| 126 | chunk_indices: torch.Tensor | None = None, |
| 127 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 128 | B, T, Hg, K, V = *k.shape, v.shape[-1] |
| 129 | H = v.shape[-2] |
| 130 | BT = A.shape[-1] |
| 131 | |
| 132 | if chunk_indices is None and cu_seqlens is not None: |
| 133 | chunk_indices = prepare_chunk_indices(cu_seqlens, BT) |
| 134 | NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) |
| 135 | BK = 64 |
| 136 | BV = 64 |
| 137 | u = torch.empty_like(v) |
| 138 | w = k.new_empty(B, T, H, K) |
| 139 | recompute_w_u_fwd_kernel[(NT, B * H)]( |
| 140 | k=k, |
| 141 | v=v, |
| 142 | beta=beta, |
| 143 | w=w, |
| 144 | u=u, |
| 145 | A=A, |
| 146 | g=g_cumsum, |
| 147 | cu_seqlens=cu_seqlens, |
| 148 | chunk_indices=chunk_indices, |
| 149 | T=T, |
| 150 | H=H, |
| 151 | Hg=Hg, |
| 152 | K=K, |
| 153 | V=V, |
| 154 | BT=BT, |
| 155 | BK=BK, |
| 156 | BV=BV, |
| 157 | ) |
| 158 | return w, u |