| 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 os |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.triton_utils import tl, triton |
| 16 | |
| 17 | from .index import prepare_chunk_indices |
| 18 | from .op import make_tensor_descriptor |
| 19 | from .utils import input_guard, is_amd, is_tma_supported |
| 20 | |
| 21 | FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "ieee") |
| 22 | ALLOWED_TRIL_PRECISIONS = ["ieee", "tf32"] if is_amd else ["ieee", "tf32", "tf32x3"] |
| 23 | assert FLA_TRIL_PRECISION in ALLOWED_TRIL_PRECISIONS, ( |
| 24 | f"FLA_TRIL_PRECISION must be one of {ALLOWED_TRIL_PRECISIONS}, but got {FLA_TRIL_PRECISION}" |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) |
| 29 | @triton.autotune( |
| 30 | configs=[ |
| 31 | triton.Config({}, num_warps=num_warps, num_stages=num_stages) |
| 32 | for num_warps in [1, 2, 4, 8] |
| 33 | for num_stages in [2, 3, 4, 5] |
| 34 | ], |
| 35 | key=["BT"], |
| 36 | ) |
| 37 | @triton.jit(do_not_specialize=["T"]) |
| 38 | def solve_tril_16x16_kernel( |
| 39 | A, |
| 40 | Ai, |
| 41 | cu_seqlens, |
| 42 | chunk_indices, |
| 43 | T, |
| 44 | H: tl.constexpr, |
| 45 | BT: tl.constexpr, |
| 46 | USE_TMA: tl.constexpr, |
| 47 | IS_VARLEN: tl.constexpr, |
| 48 | DOT_PRECISION: tl.constexpr, |
| 49 | ): |
| 50 | i_t, i_bh = tl.program_id(0), tl.program_id(1) |
| 51 | i_b, i_h = i_bh // H, i_bh % H |
| 52 | if IS_VARLEN: |
| 53 | i_n, i_t = ( |
| 54 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 55 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 56 | ) |
| 57 | bos, eos = ( |
| 58 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 59 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 60 | ) |
| 61 | T = eos - bos |
| 62 | else: |
| 63 | bos, eos = i_b * T, i_b * T + T |
| 64 | o_i = tl.arange(0, 16) |
| 65 | m_A = o_i[:, None] > o_i[None, :] |
| 66 | m_I = o_i[:, None] == o_i[None, :] |
| 67 | |
| 68 | A = A + (bos * H + i_h) * BT |
| 69 | Ai = Ai + (bos * H + i_h) * 16 |
| 70 | |
| 71 | offset = (i_t * 16) % BT |
| 72 | if not USE_TMA: |
| 73 | p_A = tl.make_block_ptr( |
| 74 | A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) |
| 75 | ) |
| 76 | # [16, 16] |
| 77 | b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) |
| 78 | else: |
| 79 | desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) |
| 80 | desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16]) |
| 81 | b_A = desc.load([i_t * 16, offset]).to(tl.float32) |
| 82 | b_A = -tl.where(m_A, b_A, 0) |
| 83 | |
| 84 | for i in range(2, min(16, T - i_t * 16)): |
| 85 | # [16] |
| 86 | b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset) |
| 87 | b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) |
| 88 | b_A = tl.where((o_i == i)[:, None], b_a, b_A) |
| 89 | b_A += m_I |
| 90 | if not USE_TMA: |
| 91 | p_Ai = tl.make_block_ptr( |
| 92 | Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0) |
| 93 | ) |
| 94 | tl.store( |
| 95 | p_Ai, |
| 96 | b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 97 | boundary_check=(0, 1), |
| 98 | ) |
| 99 | else: |
| 100 | desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) |
| 101 | |
| 102 | |
| 103 | @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) |
| 104 | @triton.autotune( |
| 105 | configs=[ |
| 106 | triton.Config({}, num_warps=num_warps, num_stages=num_stages) |
| 107 | for num_warps in [1, 2, 4, 8] |
| 108 | for num_stages in [2, 3, 4, 5] |
| 109 | ], |
| 110 | key=["H", "BT", "IS_VARLEN"], |
| 111 | ) |
| 112 | @triton.jit(do_not_specialize=["T"]) |
| 113 | def merge_16x16_to_32x32_inverse_kernel( |
| 114 | A, |
| 115 | Ai, |
| 116 | cu_seqlens, |
| 117 | chunk_indices, |
| 118 | T, |
| 119 | H: tl.constexpr, |
| 120 | BT: tl.constexpr, |
| 121 | USE_TMA: tl.constexpr, |
| 122 | IS_VARLEN: tl.constexpr, |
| 123 | DOT_PRECISION: tl.constexpr, |
| 124 | ): |
| 125 | i_t, i_bh = tl.program_id(0), tl.program_id(1) |
| 126 | i_b, i_h = i_bh // H, i_bh % H |
| 127 | if IS_VARLEN: |
| 128 | i_n, i_t = ( |
| 129 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 130 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 131 | ) |
| 132 | bos, eos = ( |
| 133 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 134 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 135 | ) |
| 136 | T = eos - bos |
| 137 | else: |
| 138 | bos, eos = i_b * T, i_b * T + T |
| 139 | |
| 140 | o_i = tl.arange(0, 16) |
| 141 | m_A = o_i[:, None] > o_i[None, :] |
| 142 | m_I = o_i[:, None] == o_i[None, :] |
| 143 | A += (bos * H + i_h) * BT |
| 144 | Ai += (bos * H + i_h) * BT |
| 145 | |
| 146 | if not USE_TMA: |
| 147 | p_A_11 = tl.make_block_ptr( |
| 148 | A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) |
| 149 | ) |
| 150 | p_A_22 = tl.make_block_ptr( |
| 151 | A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) |
| 152 | ) |
| 153 | b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) |
| 154 | b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) |
| 155 | else: |
| 156 | desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) |
| 157 | desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) |
| 158 | b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) |
| 159 | b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) |
| 160 | |
| 161 | # [16, 16] |
| 162 | b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) |
| 163 | b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) |
| 164 | |
| 165 | for i in range(2, min(16, T - i_t * BT)): |
| 166 | b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) |
| 167 | b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) |
| 168 | b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) |
| 169 | for i in range(16 + 2, min(32, T - i_t * BT)): |
| 170 | b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) |
| 171 | b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) |
| 172 | b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) |
| 173 | |
| 174 | b_Ai_11 += m_I |
| 175 | b_Ai_22 += m_I |
| 176 | |
| 177 | if not USE_TMA: |
| 178 | p_A_21 = tl.make_block_ptr( |
| 179 | A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) |
| 180 | ) |
| 181 | b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) |
| 182 | else: |
| 183 | b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) |
| 184 | |
| 185 | b_Ai_21 = -tl.dot( |
| 186 | tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), |
| 187 | b_Ai_11, |
| 188 | input_precision=DOT_PRECISION, |
| 189 | ) |
| 190 | |
| 191 | if not USE_TMA: |
| 192 | p_Ai_11 = tl.make_block_ptr( |
| 193 | Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) |
| 194 | ) |
| 195 | p_Ai_21 = tl.make_block_ptr( |
| 196 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) |
| 197 | ) |
| 198 | p_Ai_22 = tl.make_block_ptr( |
| 199 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) |
| 200 | ) |
| 201 | tl.store( |
| 202 | p_Ai_11, |
| 203 | b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 204 | boundary_check=(0, 1), |
| 205 | ) |
| 206 | tl.store( |
| 207 | p_Ai_22, |
| 208 | b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 209 | boundary_check=(0, 1), |
| 210 | ) |
| 211 | tl.store( |
| 212 | p_Ai_21, |
| 213 | b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 214 | boundary_check=(0, 1), |
| 215 | ) |
| 216 | else: |
| 217 | desc_o.store( |
| 218 | [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 219 | ) |
| 220 | desc_o.store( |
| 221 | [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 222 | ) |
| 223 | desc_o.store( |
| 224 | [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 225 | ) |
| 226 | |
| 227 | |
| 228 | @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) |
| 229 | @triton.autotune( |
| 230 | configs=[ |
| 231 | triton.Config({}, num_warps=num_warps, num_stages=num_stages) |
| 232 | for num_warps in [2, 4, 8] |
| 233 | for num_stages in [2, 3, 4, 5] |
| 234 | ], |
| 235 | key=["H", "BT", "IS_VARLEN"], |
| 236 | ) |
| 237 | @triton.jit(do_not_specialize=["T"]) |
| 238 | def merge_16x16_to_64x64_inverse_kernel( |
| 239 | A, |
| 240 | Ai, |
| 241 | cu_seqlens, |
| 242 | chunk_indices, |
| 243 | T, |
| 244 | H: tl.constexpr, |
| 245 | BT: tl.constexpr, |
| 246 | USE_TMA: tl.constexpr, |
| 247 | IS_VARLEN: tl.constexpr, |
| 248 | DOT_PRECISION: tl.constexpr, |
| 249 | ): |
| 250 | i_t, i_bh = tl.program_id(0), tl.program_id(1) |
| 251 | i_b, i_h = i_bh // H, i_bh % H |
| 252 | if IS_VARLEN: |
| 253 | i_n, i_t = ( |
| 254 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 255 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 256 | ) |
| 257 | bos, eos = ( |
| 258 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 259 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 260 | ) |
| 261 | T = eos - bos |
| 262 | else: |
| 263 | bos, eos = i_b * T, i_b * T + T |
| 264 | |
| 265 | o_i = tl.arange(0, 16) |
| 266 | m_A = o_i[:, None] > o_i[None, :] |
| 267 | m_I = o_i[:, None] == o_i[None, :] |
| 268 | A += (bos * H + i_h) * BT |
| 269 | Ai += (bos * H + i_h) * BT |
| 270 | |
| 271 | if not USE_TMA: |
| 272 | p_A_11 = tl.make_block_ptr( |
| 273 | A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) |
| 274 | ) |
| 275 | p_A_22 = tl.make_block_ptr( |
| 276 | A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) |
| 277 | ) |
| 278 | p_A_33 = tl.make_block_ptr( |
| 279 | A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) |
| 280 | ) |
| 281 | p_A_44 = tl.make_block_ptr( |
| 282 | A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) |
| 283 | ) |
| 284 | b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) |
| 285 | b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) |
| 286 | b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) |
| 287 | b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) |
| 288 | else: |
| 289 | desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) |
| 290 | desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) |
| 291 | b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) |
| 292 | b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) |
| 293 | b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) |
| 294 | b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) |
| 295 | |
| 296 | # [16, 16] |
| 297 | b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) |
| 298 | b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) |
| 299 | b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) |
| 300 | b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) |
| 301 | |
| 302 | for i in range(2, min(16, T - i_t * BT)): |
| 303 | b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) |
| 304 | b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) |
| 305 | b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) |
| 306 | for i in range(16 + 2, min(32, T - i_t * BT)): |
| 307 | b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) |
| 308 | b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) |
| 309 | b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) |
| 310 | for i in range(32 + 2, min(48, T - i_t * BT)): |
| 311 | b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) |
| 312 | b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) |
| 313 | b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) |
| 314 | for i in range(48 + 2, min(64, T - i_t * BT)): |
| 315 | b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) |
| 316 | b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) |
| 317 | b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) |
| 318 | b_Ai_11 += m_I |
| 319 | b_Ai_22 += m_I |
| 320 | b_Ai_33 += m_I |
| 321 | b_Ai_44 += m_I |
| 322 | |
| 323 | if not USE_TMA: |
| 324 | p_A_21 = tl.make_block_ptr( |
| 325 | A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) |
| 326 | ) |
| 327 | p_A_31 = tl.make_block_ptr( |
| 328 | A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) |
| 329 | ) |
| 330 | p_A_32 = tl.make_block_ptr( |
| 331 | A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) |
| 332 | ) |
| 333 | p_A_41 = tl.make_block_ptr( |
| 334 | A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) |
| 335 | ) |
| 336 | p_A_42 = tl.make_block_ptr( |
| 337 | A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) |
| 338 | ) |
| 339 | p_A_43 = tl.make_block_ptr( |
| 340 | A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) |
| 341 | ) |
| 342 | b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) |
| 343 | b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) |
| 344 | b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) |
| 345 | b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) |
| 346 | b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) |
| 347 | b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) |
| 348 | else: |
| 349 | b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) |
| 350 | b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) |
| 351 | b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32) |
| 352 | b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32) |
| 353 | b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32) |
| 354 | b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32) |
| 355 | |
| 356 | b_Ai_21 = -tl.dot( |
| 357 | tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), |
| 358 | b_Ai_11, |
| 359 | input_precision=DOT_PRECISION, |
| 360 | ) |
| 361 | b_Ai_32 = -tl.dot( |
| 362 | tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), |
| 363 | b_Ai_22, |
| 364 | input_precision=DOT_PRECISION, |
| 365 | ) |
| 366 | b_Ai_43 = -tl.dot( |
| 367 | tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), |
| 368 | b_Ai_33, |
| 369 | input_precision=DOT_PRECISION, |
| 370 | ) |
| 371 | |
| 372 | b_Ai_31 = -tl.dot( |
| 373 | b_Ai_33, |
| 374 | tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) |
| 375 | + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), |
| 376 | input_precision=DOT_PRECISION, |
| 377 | ) |
| 378 | b_Ai_42 = -tl.dot( |
| 379 | b_Ai_44, |
| 380 | tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) |
| 381 | + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), |
| 382 | input_precision=DOT_PRECISION, |
| 383 | ) |
| 384 | b_Ai_41 = -tl.dot( |
| 385 | b_Ai_44, |
| 386 | tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) |
| 387 | + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) |
| 388 | + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), |
| 389 | input_precision=DOT_PRECISION, |
| 390 | ) |
| 391 | |
| 392 | if not USE_TMA: |
| 393 | p_Ai_11 = tl.make_block_ptr( |
| 394 | Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) |
| 395 | ) |
| 396 | p_Ai_22 = tl.make_block_ptr( |
| 397 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) |
| 398 | ) |
| 399 | p_Ai_33 = tl.make_block_ptr( |
| 400 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) |
| 401 | ) |
| 402 | p_Ai_44 = tl.make_block_ptr( |
| 403 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) |
| 404 | ) |
| 405 | p_Ai_21 = tl.make_block_ptr( |
| 406 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) |
| 407 | ) |
| 408 | p_Ai_31 = tl.make_block_ptr( |
| 409 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) |
| 410 | ) |
| 411 | p_Ai_32 = tl.make_block_ptr( |
| 412 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) |
| 413 | ) |
| 414 | p_Ai_41 = tl.make_block_ptr( |
| 415 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) |
| 416 | ) |
| 417 | p_Ai_42 = tl.make_block_ptr( |
| 418 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) |
| 419 | ) |
| 420 | p_Ai_43 = tl.make_block_ptr( |
| 421 | Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) |
| 422 | ) |
| 423 | tl.store( |
| 424 | p_Ai_11, |
| 425 | b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 426 | boundary_check=(0, 1), |
| 427 | ) |
| 428 | tl.store( |
| 429 | p_Ai_22, |
| 430 | b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 431 | boundary_check=(0, 1), |
| 432 | ) |
| 433 | tl.store( |
| 434 | p_Ai_33, |
| 435 | b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 436 | boundary_check=(0, 1), |
| 437 | ) |
| 438 | tl.store( |
| 439 | p_Ai_44, |
| 440 | b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 441 | boundary_check=(0, 1), |
| 442 | ) |
| 443 | tl.store( |
| 444 | p_Ai_21, |
| 445 | b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 446 | boundary_check=(0, 1), |
| 447 | ) |
| 448 | tl.store( |
| 449 | p_Ai_31, |
| 450 | b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 451 | boundary_check=(0, 1), |
| 452 | ) |
| 453 | tl.store( |
| 454 | p_Ai_32, |
| 455 | b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 456 | boundary_check=(0, 1), |
| 457 | ) |
| 458 | tl.store( |
| 459 | p_Ai_41, |
| 460 | b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 461 | boundary_check=(0, 1), |
| 462 | ) |
| 463 | tl.store( |
| 464 | p_Ai_42, |
| 465 | b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 466 | boundary_check=(0, 1), |
| 467 | ) |
| 468 | tl.store( |
| 469 | p_Ai_43, |
| 470 | b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), |
| 471 | boundary_check=(0, 1), |
| 472 | ) |
| 473 | else: |
| 474 | desc_o.store( |
| 475 | [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 476 | ) |
| 477 | desc_o.store( |
| 478 | [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 479 | ) |
| 480 | desc_o.store( |
| 481 | [i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 482 | ) |
| 483 | desc_o.store( |
| 484 | [i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 485 | ) |
| 486 | desc_o.store( |
| 487 | [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 488 | ) |
| 489 | desc_o.store( |
| 490 | [i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 491 | ) |
| 492 | desc_o.store( |
| 493 | [i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 494 | ) |
| 495 | desc_o.store( |
| 496 | [i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 497 | ) |
| 498 | desc_o.store( |
| 499 | [i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 500 | ) |
| 501 | desc_o.store( |
| 502 | [i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne") |
| 503 | ) |
| 504 | |
| 505 | |
| 506 | @input_guard |
| 507 | def solve_tril( |
| 508 | A: torch.Tensor, |
| 509 | cu_seqlens: torch.Tensor | None = None, |
| 510 | chunk_indices: torch.Tensor | None = None, |
| 511 | output_dtype: torch.dtype = torch.float, |
| 512 | ) -> torch.Tensor: |
| 513 | """ |
| 514 | Compute the inverse of the matrix I + A |
| 515 | A should be strictly lower triangular, i.e., A.triu() == 0. |
| 516 | |
| 517 | Args: |
| 518 | A (torch.Tensor): |
| 519 | [B, T, H, BT], where BT should only be 16, 32, or 64. |
| 520 | cu_seqlens (torch.Tensor): |
| 521 | The cumulative sequence lengths of the input tensor. Default: `None`. |
| 522 | chunk_indices (torch.Tensor): |
| 523 | Pre-computed chunk indices. Default: `None`. |
| 524 | output_dtype (torch.dtype): |
| 525 | The dtype of the output tensor. Default: `torch.float`. |
| 526 | If `None`, the output dtype will be the same as the input dtype. |
| 527 | |
| 528 | Returns: |
| 529 | (I + A)^-1 with the same shape as A |
| 530 | """ |
| 531 | assert A.shape[-1] in [16, 32, 64] |
| 532 | output_dtype = A.dtype if output_dtype is None else output_dtype |
| 533 | |
| 534 | B, T, H, BT = A.shape |
| 535 | if chunk_indices is None and cu_seqlens is not None: |
| 536 | chunk_indices = prepare_chunk_indices(cu_seqlens, BT) |
| 537 | NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) |
| 538 | |
| 539 | Ai = torch.zeros_like(A, dtype=output_dtype) |
| 540 | if BT == 16: |
| 541 | merge_fn = solve_tril_16x16_kernel |
| 542 | elif BT == 32: |
| 543 | merge_fn = merge_16x16_to_32x32_inverse_kernel |
| 544 | elif BT == 64: |
| 545 | merge_fn = merge_16x16_to_64x64_inverse_kernel |
| 546 | |
| 547 | merge_fn[NT, B * H]( |
| 548 | A=A, |
| 549 | Ai=Ai, |
| 550 | cu_seqlens=cu_seqlens, |
| 551 | chunk_indices=chunk_indices, |
| 552 | T=T, |
| 553 | H=H, |
| 554 | BT=BT, |
| 555 | USE_TMA=is_tma_supported, |
| 556 | DOT_PRECISION=FLA_TRIL_PRECISION, |
| 557 | ) |
| 558 | return Ai |