| 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 | # ruff: noqa: E501 |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm.triton_utils import tl, triton |
| 14 | |
| 15 | from .index import prepare_chunk_indices |
| 16 | from .utils import check_shared_mem, input_guard |
| 17 | |
| 18 | BS_LIST = [32, 64] if check_shared_mem() else [16, 32] |
| 19 | |
| 20 | |
| 21 | @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) |
| 22 | @triton.autotune( |
| 23 | configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], |
| 24 | key=["B", "H", "BT", "IS_VARLEN", "REVERSE"], |
| 25 | ) |
| 26 | @triton.jit(do_not_specialize=["T"]) |
| 27 | def chunk_local_cumsum_scalar_kernel( |
| 28 | s, |
| 29 | o, |
| 30 | cu_seqlens, |
| 31 | chunk_indices, |
| 32 | T, |
| 33 | B: tl.constexpr, |
| 34 | H: tl.constexpr, |
| 35 | BT: tl.constexpr, |
| 36 | REVERSE: tl.constexpr, |
| 37 | IS_VARLEN: tl.constexpr, |
| 38 | HEAD_FIRST: tl.constexpr, |
| 39 | ): |
| 40 | i_t, i_bh = tl.program_id(0), tl.program_id(1) |
| 41 | i_b, i_h = i_bh // H, i_bh % H |
| 42 | if IS_VARLEN: |
| 43 | i_n, i_t = ( |
| 44 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 45 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 46 | ) |
| 47 | bos, eos = ( |
| 48 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 49 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 50 | ) |
| 51 | T = eos - bos |
| 52 | else: |
| 53 | bos, eos = i_b * T, i_b * T + T |
| 54 | |
| 55 | if HEAD_FIRST: |
| 56 | p_s = tl.make_block_ptr( |
| 57 | s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) |
| 58 | ) |
| 59 | p_o = tl.make_block_ptr( |
| 60 | o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) |
| 61 | ) |
| 62 | else: |
| 63 | p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) |
| 64 | p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) |
| 65 | # [BT] |
| 66 | b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) |
| 67 | b_o = tl.cumsum(b_s, axis=0) |
| 68 | if REVERSE: |
| 69 | b_z = tl.sum(b_s, axis=0) |
| 70 | b_o = -b_o + b_z[None] + b_s |
| 71 | tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) |
| 72 | |
| 73 | |
| 74 | @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) |
| 75 | @triton.autotune( |
| 76 | configs=[ |
| 77 | triton.Config({"BS": BS}, num_warps=num_warps) |
| 78 | for BS in BS_LIST |
| 79 | for num_warps in [2, 4, 8] |
| 80 | ], |
| 81 | key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE"], |
| 82 | ) |
| 83 | @triton.jit(do_not_specialize=["T"]) |
| 84 | def chunk_local_cumsum_vector_kernel( |
| 85 | s, |
| 86 | o, |
| 87 | cu_seqlens, |
| 88 | chunk_indices, |
| 89 | T, |
| 90 | B: tl.constexpr, |
| 91 | H: tl.constexpr, |
| 92 | S: tl.constexpr, |
| 93 | BT: tl.constexpr, |
| 94 | BS: tl.constexpr, |
| 95 | REVERSE: tl.constexpr, |
| 96 | IS_VARLEN: tl.constexpr, |
| 97 | HEAD_FIRST: tl.constexpr, |
| 98 | ): |
| 99 | i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| 100 | i_b, i_h = i_bh // H, i_bh % H |
| 101 | if IS_VARLEN: |
| 102 | i_n, i_t = ( |
| 103 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 104 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 105 | ) |
| 106 | bos, eos = ( |
| 107 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 108 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 109 | ) |
| 110 | T = eos - bos |
| 111 | else: |
| 112 | bos, eos = i_b * T, i_b * T + T |
| 113 | |
| 114 | o_i = tl.arange(0, BT) |
| 115 | if REVERSE: |
| 116 | m_s = tl.where(o_i[:, None] <= o_i[None, :], 1.0, 0.0) |
| 117 | else: |
| 118 | m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) |
| 119 | |
| 120 | if HEAD_FIRST: |
| 121 | p_s = tl.make_block_ptr( |
| 122 | s + (bos * H + i_h * T) * S, |
| 123 | (T, S), |
| 124 | (S, 1), |
| 125 | (i_t * BT, i_s * BS), |
| 126 | (BT, BS), |
| 127 | (1, 0), |
| 128 | ) |
| 129 | p_o = tl.make_block_ptr( |
| 130 | o + (bos * H + i_h * T) * S, |
| 131 | (T, S), |
| 132 | (S, 1), |
| 133 | (i_t * BT, i_s * BS), |
| 134 | (BT, BS), |
| 135 | (1, 0), |
| 136 | ) |
| 137 | else: |
| 138 | p_s = tl.make_block_ptr( |
| 139 | s + (bos * H + i_h) * S, |
| 140 | (T, S), |
| 141 | (H * S, 1), |
| 142 | (i_t * BT, i_s * BS), |
| 143 | (BT, BS), |
| 144 | (1, 0), |
| 145 | ) |
| 146 | p_o = tl.make_block_ptr( |
| 147 | o + (bos * H + i_h) * S, |
| 148 | (T, S), |
| 149 | (H * S, 1), |
| 150 | (i_t * BT, i_s * BS), |
| 151 | (BT, BS), |
| 152 | (1, 0), |
| 153 | ) |
| 154 | # [BT, BS] |
| 155 | b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) |
| 156 | b_o = tl.dot(m_s, b_s, allow_tf32=False) |
| 157 | tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) |
| 158 | |
| 159 | |
| 160 | def chunk_local_cumsum_scalar( |
| 161 | g: torch.Tensor, |
| 162 | chunk_size: int, |
| 163 | reverse: bool = False, |
| 164 | cu_seqlens: torch.Tensor | None = None, |
| 165 | chunk_indices: torch.Tensor | None = None, |
| 166 | head_first: bool = False, |
| 167 | output_dtype: torch.dtype | None = torch.float, |
| 168 | ) -> torch.Tensor: |
| 169 | if head_first: |
| 170 | B, H, T = g.shape |
| 171 | else: |
| 172 | B, T, H = g.shape |
| 173 | assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( |
| 174 | "chunk_size must be a power of 2" |
| 175 | ) |
| 176 | if chunk_indices is None and cu_seqlens is not None: |
| 177 | chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) |
| 178 | BT = chunk_size |
| 179 | NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) |
| 180 | g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) |
| 181 | grid = (NT, B * H) |
| 182 | chunk_local_cumsum_scalar_kernel[grid]( |
| 183 | g_org, |
| 184 | g, |
| 185 | cu_seqlens, |
| 186 | chunk_indices, |
| 187 | T=T, |
| 188 | B=B, |
| 189 | H=H, |
| 190 | BT=BT, |
| 191 | HEAD_FIRST=head_first, |
| 192 | REVERSE=reverse, |
| 193 | ) |
| 194 | return g |
| 195 | |
| 196 | |
| 197 | def chunk_local_cumsum_vector( |
| 198 | g: torch.Tensor, |
| 199 | chunk_size: int, |
| 200 | reverse: bool = False, |
| 201 | cu_seqlens: torch.Tensor | None = None, |
| 202 | chunk_indices: torch.Tensor | None = None, |
| 203 | head_first: bool = False, |
| 204 | output_dtype: torch.dtype | None = torch.float, |
| 205 | ) -> torch.Tensor: |
| 206 | if head_first: |
| 207 | B, H, T, S = g.shape |
| 208 | else: |
| 209 | B, T, H, S = g.shape |
| 210 | assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( |
| 211 | "chunk_size must be a power of 2" |
| 212 | ) |
| 213 | if chunk_indices is None and cu_seqlens is not None: |
| 214 | chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) |
| 215 | BT = chunk_size |
| 216 | NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) |
| 217 | |
| 218 | g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) |
| 219 | |
| 220 | def grid(meta): |
| 221 | return (triton.cdiv(meta["S"], meta["BS"]), NT, B * H) |
| 222 | |
| 223 | # keep cumulative normalizer in fp32 |
| 224 | # this kernel is equivalent to |
| 225 | # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1) |
| 226 | chunk_local_cumsum_vector_kernel[grid]( |
| 227 | g_org, |
| 228 | g, |
| 229 | cu_seqlens, |
| 230 | chunk_indices, |
| 231 | T=T, |
| 232 | B=B, |
| 233 | H=H, |
| 234 | S=S, |
| 235 | BT=BT, |
| 236 | HEAD_FIRST=head_first, |
| 237 | REVERSE=reverse, |
| 238 | ) |
| 239 | return g |
| 240 | |
| 241 | |
| 242 | @input_guard |
| 243 | def chunk_local_cumsum( |
| 244 | g: torch.Tensor, |
| 245 | chunk_size: int, |
| 246 | reverse: bool = False, |
| 247 | cu_seqlens: torch.Tensor | None = None, |
| 248 | chunk_indices: torch.Tensor | None = None, |
| 249 | head_first: bool = False, |
| 250 | output_dtype: torch.dtype | None = torch.float, |
| 251 | **kwargs, |
| 252 | ) -> torch.Tensor: |
| 253 | if cu_seqlens is not None: |
| 254 | assert g.shape[0] == 1, ( |
| 255 | "Only batch size 1 is supported when cu_seqlens are provided" |
| 256 | ) |
| 257 | if len(g.shape) == 3: |
| 258 | return chunk_local_cumsum_scalar( |
| 259 | g, |
| 260 | chunk_size, |
| 261 | reverse, |
| 262 | cu_seqlens, |
| 263 | chunk_indices, |
| 264 | head_first, |
| 265 | output_dtype, |
| 266 | ) |
| 267 | elif len(g.shape) == 4: |
| 268 | return chunk_local_cumsum_vector( |
| 269 | g, |
| 270 | chunk_size, |
| 271 | reverse, |
| 272 | cu_seqlens, |
| 273 | chunk_indices, |
| 274 | head_first, |
| 275 | output_dtype, |
| 276 | ) |
| 277 | else: |
| 278 | raise ValueError( |
| 279 | f"Unsupported input shape {g.shape}. " |
| 280 | f"which should be (B, T, H, D) if `head_first=False` " |
| 281 | f"or (B, H, T, D) otherwise" |
| 282 | ) |