| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li |
| 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-2026, Songlin Yang, Yu Zhang, Zhiyuan Li |
| 9 | # ruff: noqa: E501 |
| 10 | |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import ( |
| 15 | chunk_gated_delta_rule_fwd_h, |
| 16 | ) |
| 17 | from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum |
| 18 | from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices |
| 19 | from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd |
| 20 | from vllm.third_party.flash_linear_attention.ops.op import ( |
| 21 | exp2, |
| 22 | log, |
| 23 | make_tensor_descriptor, |
| 24 | ) |
| 25 | from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd |
| 26 | from vllm.triton_utils import tl, triton |
| 27 | from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 |
| 28 | |
| 29 | from .chunk_intra import chunk_kda_fwd_intra |
| 30 | |
| 31 | BT_LIST_AUTOTUNE = [32, 64, 128] |
| 32 | NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] |
| 33 | |
| 34 | |
| 35 | # num_stages=4 is excluded. The `u` loop has a trip count of 2 (V / BV), and a |
| 36 | # 4-stage pipeline over it emits a third, consumer-less async copy into LDS that |
| 37 | # the `w` loop then reuses; at num_warps=4 that races on gfx950 and `u` comes |
| 38 | # back with non-deterministic O(1e38) garbage once a batch reaches 4096 tokens. |
| 39 | _RECOMPUTE_W_U_NUM_STAGES = [2, 3] |
| 40 | |
| 41 | |
| 42 | @triton.heuristics( |
| 43 | { |
| 44 | "STORE_QG": lambda args: args["qg"] is not None, |
| 45 | "STORE_KG": lambda args: args["kg"] is not None, |
| 46 | "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, |
| 47 | } |
| 48 | ) |
| 49 | @triton.autotune( |
| 50 | configs=[ |
| 51 | triton.Config({}, num_warps=num_warps, num_stages=num_stages) |
| 52 | for num_warps in [2, 4, 8] |
| 53 | for num_stages in _RECOMPUTE_W_U_NUM_STAGES |
| 54 | ], |
| 55 | key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], |
| 56 | ) |
| 57 | @triton.jit(do_not_specialize=["T"]) |
| 58 | def recompute_w_u_fwd_kernel( |
| 59 | q, |
| 60 | k, |
| 61 | qg, |
| 62 | kg, |
| 63 | v, |
| 64 | beta, |
| 65 | w, |
| 66 | u, |
| 67 | A, |
| 68 | gk, |
| 69 | cu_seqlens, |
| 70 | chunk_indices, |
| 71 | T, |
| 72 | H: tl.constexpr, |
| 73 | K: tl.constexpr, |
| 74 | V: tl.constexpr, |
| 75 | BT: tl.constexpr, |
| 76 | BK: tl.constexpr, |
| 77 | BV: tl.constexpr, |
| 78 | STORE_QG: tl.constexpr, |
| 79 | STORE_KG: tl.constexpr, |
| 80 | IS_VARLEN: tl.constexpr, |
| 81 | DOT_PRECISION: tl.constexpr, |
| 82 | ): |
| 83 | i_t, i_bh = tl.program_id(0), tl.program_id(1) |
| 84 | i_b, i_h = i_bh // H, i_bh % H |
| 85 | if IS_VARLEN: |
| 86 | i_n, i_t = ( |
| 87 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 88 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 89 | ) |
| 90 | bos, eos = ( |
| 91 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 92 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 93 | ) |
| 94 | T = eos - bos |
| 95 | else: |
| 96 | bos, eos = i_b * T, i_b * T + T |
| 97 | o_b = i_t * BT + tl.arange(0, BT) |
| 98 | b_b = tl.load(beta + (bos + o_b) * H + i_h, mask=o_b < T, other=0).to(tl.float32) |
| 99 | |
| 100 | desc_A = make_tensor_descriptor(A + (bos * H + i_h) * BT, [T, BT], [H * BT, 1], [BT, BT]) |
| 101 | b_A = desc_A.load([i_t * BT, 0]) |
| 102 | |
| 103 | desc_v = make_tensor_descriptor(v + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV]) |
| 104 | desc_u = make_tensor_descriptor(u + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV]) |
| 105 | for i_v in range(tl.cdiv(V, BV)): |
| 106 | b_v = desc_v.load([i_t * BT, i_v * BV]) |
| 107 | b_vb = (b_v * b_b[:, None]).to(b_v.dtype) |
| 108 | b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) |
| 109 | desc_u.store([i_t * BT, i_v * BV], b_u.to(desc_u.dtype)) |
| 110 | |
| 111 | desc_w = make_tensor_descriptor(w + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 112 | desc_k = make_tensor_descriptor(k + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 113 | desc_gk = make_tensor_descriptor(gk + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 114 | for i_k in range(tl.cdiv(K, BK)): |
| 115 | b_k = desc_k.load([i_t * BT, i_k * BK]) |
| 116 | b_kb = b_k * b_b[:, None] |
| 117 | |
| 118 | b_gk = desc_gk.load([i_t * BT, i_k * BK]) |
| 119 | b_kb *= exp2(b_gk) |
| 120 | if STORE_QG: |
| 121 | desc_q = make_tensor_descriptor(q + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 122 | desc_qg = make_tensor_descriptor(qg + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 123 | b_q = desc_q.load([i_t * BT, i_k * BK]) |
| 124 | b_qg = b_q * exp2(b_gk) |
| 125 | desc_qg.store([i_t * BT, i_k * BK], b_qg.to(desc_qg.dtype)) |
| 126 | if STORE_KG: |
| 127 | last_idx = min(i_t * BT + BT, T) - 1 |
| 128 | |
| 129 | o_k = i_k * BK + tl.arange(0, BK) |
| 130 | m_k = o_k < K |
| 131 | b_gn = tl.load( |
| 132 | gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 |
| 133 | ) |
| 134 | b_kg = b_k * exp2(b_gn - b_gk) |
| 135 | |
| 136 | desc_kg = make_tensor_descriptor(kg + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 137 | desc_kg.store([i_t * BT, i_k * BK], b_kg.to(desc_kg.dtype)) |
| 138 | |
| 139 | b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) |
| 140 | desc_w.store([i_t * BT, i_k * BK], b_w.to(desc_w.dtype)) |
| 141 | |
| 142 | |
| 143 | def recompute_w_u_fwd( |
| 144 | k: torch.Tensor, |
| 145 | v: torch.Tensor, |
| 146 | beta: torch.Tensor, |
| 147 | A: torch.Tensor, |
| 148 | q: torch.Tensor | None = None, |
| 149 | gk: torch.Tensor | None = None, |
| 150 | cu_seqlens: torch.Tensor | None = None, |
| 151 | chunk_indices: torch.Tensor | None = None, |
| 152 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 153 | B, T, H, K, V = *k.shape, v.shape[-1] |
| 154 | BT = A.shape[-1] |
| 155 | BK = 64 |
| 156 | BV = 64 |
| 157 | |
| 158 | if chunk_indices is None and cu_seqlens is not None: |
| 159 | chunk_indices = prepare_chunk_indices(cu_seqlens, BT) |
| 160 | NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) |
| 161 | |
| 162 | w = torch.empty_like(k) |
| 163 | u = torch.empty_like(v) |
| 164 | kg = torch.empty_like(k) if gk is not None else None |
| 165 | recompute_w_u_fwd_kernel[(NT, B * H)]( |
| 166 | q=q, |
| 167 | k=k, |
| 168 | qg=None, |
| 169 | kg=kg, |
| 170 | v=v, |
| 171 | beta=beta, |
| 172 | w=w, |
| 173 | u=u, |
| 174 | A=A, |
| 175 | gk=gk, |
| 176 | cu_seqlens=cu_seqlens, |
| 177 | chunk_indices=chunk_indices, |
| 178 | T=T, |
| 179 | H=H, |
| 180 | K=K, |
| 181 | V=V, |
| 182 | BT=BT, |
| 183 | BK=BK, |
| 184 | BV=BV, |
| 185 | DOT_PRECISION="ieee", |
| 186 | ) |
| 187 | return w, u, None, kg |
| 188 | |
| 189 | |
| 190 | @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) |
| 191 | @triton.autotune( |
| 192 | configs=[ |
| 193 | triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) |
| 194 | for BK in [32, 64] |
| 195 | for BV in [64, 128] |
| 196 | for num_warps in [2, 4, 8] |
| 197 | for num_stages in [2, 3, 4] |
| 198 | ], |
| 199 | key=["BT"], |
| 200 | ) |
| 201 | @triton.jit(do_not_specialize=["T"]) |
| 202 | def chunk_gla_fwd_kernel_o( |
| 203 | q, |
| 204 | v, |
| 205 | g, |
| 206 | h, |
| 207 | o, |
| 208 | A, |
| 209 | cu_seqlens, |
| 210 | chunk_indices, |
| 211 | scale, |
| 212 | T, |
| 213 | H: tl.constexpr, |
| 214 | K: tl.constexpr, |
| 215 | V: tl.constexpr, |
| 216 | BT: tl.constexpr, |
| 217 | BK: tl.constexpr, |
| 218 | BV: tl.constexpr, |
| 219 | IS_VARLEN: tl.constexpr, |
| 220 | ): |
| 221 | i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| 222 | i_b, i_h = i_bh // H, i_bh % H |
| 223 | if IS_VARLEN: |
| 224 | i_tg = i_t |
| 225 | i_n, i_t = ( |
| 226 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 227 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 228 | ) |
| 229 | bos, eos = ( |
| 230 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 231 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 232 | ) |
| 233 | T = eos - bos |
| 234 | NT = tl.cdiv(T, BT) |
| 235 | else: |
| 236 | NT = tl.cdiv(T, BT) |
| 237 | i_tg = i_b * NT + i_t |
| 238 | bos, eos = i_b * T, i_b * T + T |
| 239 | |
| 240 | m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] |
| 241 | |
| 242 | b_o = tl.zeros([BT, BV], dtype=tl.float32) |
| 243 | desc_q = make_tensor_descriptor(q + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 244 | desc_g = make_tensor_descriptor(g + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK]) |
| 245 | desc_h = make_tensor_descriptor(h + (i_tg * H + i_h) * K * V, [V, K], [K, 1], [BV, BK]) |
| 246 | for i_k in range(tl.cdiv(K, BK)): |
| 247 | # [BT, BK] |
| 248 | b_q = desc_q.load([i_t * BT, i_k * BK]) |
| 249 | b_q = (b_q * scale).to(b_q.dtype) |
| 250 | # [BT, BK] |
| 251 | b_g = desc_g.load([i_t * BT, i_k * BK]) |
| 252 | # [BT, BK] |
| 253 | b_qg = (b_q * exp2(b_g)).to(b_q.dtype) |
| 254 | # [BV, BK] |
| 255 | b_h = desc_h.load([i_v * BV, i_k * BK]) |
| 256 | # [BT, BV] |
| 257 | if i_k >= 0: |
| 258 | b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) |
| 259 | desc_v = make_tensor_descriptor(v + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV]) |
| 260 | desc_o = make_tensor_descriptor(o + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV]) |
| 261 | desc_A = make_tensor_descriptor(A + (bos * H + i_h) * BT, [T, BT], [H * BT, 1], [BT, BT]) |
| 262 | # [BT, BV] |
| 263 | b_v = desc_v.load([i_t * BT, i_v * BV]) |
| 264 | # [BT, BT] |
| 265 | b_A = desc_A.load([i_t * BT, 0]) |
| 266 | b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) |
| 267 | b_o += tl.dot(b_A, b_v, allow_tf32=False) |
| 268 | desc_o.store([i_t * BT, i_v * BV], b_o.to(desc_o.dtype)) |
| 269 | |
| 270 | |
| 271 | def chunk_gla_fwd_o_gk( |
| 272 | q: torch.Tensor, |
| 273 | v: torch.Tensor, |
| 274 | g: torch.Tensor, |
| 275 | A: torch.Tensor, |
| 276 | h: torch.Tensor, |
| 277 | o: torch.Tensor, |
| 278 | scale: float, |
| 279 | cu_seqlens: torch.Tensor | None = None, |
| 280 | chunk_indices: torch.Tensor | None = None, |
| 281 | chunk_size: int = FLA_CHUNK_SIZE, |
| 282 | ): |
| 283 | B, T, H, K, V = *q.shape, v.shape[-1] |
| 284 | BT = chunk_size |
| 285 | |
| 286 | if chunk_indices is None and cu_seqlens is not None: |
| 287 | chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) |
| 288 | NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) |
| 289 | |
| 290 | def grid(meta): |
| 291 | return (cdiv(V, meta["BV"]), NT, B * H) |
| 292 | |
| 293 | chunk_gla_fwd_kernel_o[grid]( |
| 294 | q=q, |
| 295 | v=v, |
| 296 | g=g, |
| 297 | h=h, |
| 298 | o=o, |
| 299 | A=A, |
| 300 | cu_seqlens=cu_seqlens, |
| 301 | chunk_indices=chunk_indices, |
| 302 | scale=scale, |
| 303 | T=T, |
| 304 | H=H, |
| 305 | K=K, |
| 306 | V=V, |
| 307 | BT=BT, |
| 308 | ) |
| 309 | return o |
| 310 | |
| 311 | |
| 312 | @triton.heuristics( |
| 313 | { |
| 314 | "HAS_BIAS": lambda args: args["g_bias"] is not None, |
| 315 | "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, |
| 316 | } |
| 317 | ) |
| 318 | @triton.autotune( |
| 319 | configs=[ |
| 320 | triton.Config({"BS": BS}, num_warps=num_warps) |
| 321 | for BS in [32, 64] |
| 322 | for num_warps in [2, 4, 8] |
| 323 | ], |
| 324 | key=["H", "S", "BT", "IS_VARLEN"], |
| 325 | ) |
| 326 | @triton.jit(do_not_specialize=["T"]) |
| 327 | def kda_gate_chunk_cumsum_vector_kernel( |
| 328 | s, |
| 329 | raw_beta, |
| 330 | A_log, |
| 331 | g_bias, |
| 332 | o, |
| 333 | beta_out, |
| 334 | cu_seqlens, |
| 335 | chunk_indices, |
| 336 | cumsum_scale, |
| 337 | lower_bound, |
| 338 | beta, |
| 339 | threshold, |
| 340 | T, |
| 341 | stride_beta_batch, |
| 342 | stride_beta_token, |
| 343 | stride_beta_head, |
| 344 | H: tl.constexpr, |
| 345 | S: tl.constexpr, |
| 346 | BT: tl.constexpr, |
| 347 | BS: tl.constexpr, |
| 348 | HAS_BIAS: tl.constexpr, |
| 349 | IS_VARLEN: tl.constexpr, |
| 350 | USE_LOWER_BOUND: tl.constexpr, |
| 351 | ): |
| 352 | i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| 353 | i_b, i_h = i_bh // H, i_bh % H |
| 354 | if IS_VARLEN: |
| 355 | i_n, i_t = ( |
| 356 | tl.load(chunk_indices + i_t * 2).to(tl.int32), |
| 357 | tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), |
| 358 | ) |
| 359 | bos, eos = ( |
| 360 | tl.load(cu_seqlens + i_n).to(tl.int32), |
| 361 | tl.load(cu_seqlens + i_n + 1).to(tl.int32), |
| 362 | ) |
| 363 | T = eos - bos |
| 364 | else: |
| 365 | bos = i_b * T |
| 366 | |
| 367 | if i_s == 0: |
| 368 | o_beta_t = tl.arange(0, BT) |
| 369 | m_beta = i_t * BT + o_beta_t < T |
| 370 | if IS_VARLEN: |
| 371 | p_beta = ( |
| 372 | raw_beta |
| 373 | + (bos + i_t * BT + o_beta_t) * stride_beta_token |
| 374 | + i_h * stride_beta_head |
| 375 | ) |
| 376 | else: |
| 377 | p_beta = ( |
| 378 | raw_beta |
| 379 | + i_b * stride_beta_batch |
| 380 | + (i_t * BT + o_beta_t) * stride_beta_token |
| 381 | + i_h * stride_beta_head |
| 382 | ) |
| 383 | b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32) |
| 384 | p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h |
| 385 | tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta) |
| 386 | return |
| 387 | |
| 388 | i_s -= 1 |
| 389 | |
| 390 | desc_s = make_tensor_descriptor(s + (bos * H + i_h) * S, [T, S], [H * S, 1], [BT, BS]) |
| 391 | desc_o = make_tensor_descriptor(o + (bos * H + i_h) * S, [T, S], [H * S, 1], [BT, BS]) |
| 392 | |
| 393 | b_s = desc_s.load([i_t * BT, i_s * BS]).to(tl.float32) |
| 394 | if HAS_BIAS: |
| 395 | desc_bias = make_tensor_descriptor(g_bias + i_h * S, [S], [1], [BS]) |
| 396 | b_bias = desc_bias.load([i_s * BS]).to(tl.float32) |
| 397 | b_s += b_bias[None, :] |
| 398 | |
| 399 | b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) |
| 400 | if USE_LOWER_BOUND: |
| 401 | b_gate = lower_bound * tl.sigmoid(b_a * b_s) |
| 402 | else: |
| 403 | b_g_scaled = b_s * beta |
| 404 | b_softplus = tl.where( |
| 405 | b_g_scaled > threshold, |
| 406 | b_s, |
| 407 | (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), |
| 408 | ) |
| 409 | b_gate = -b_a * b_softplus |
| 410 | |
| 411 | # Boundary loads return zero, but bias and gate activation can make padded |
| 412 | # rows nonzero. Padding trails valid rows, so it only affects masked stores. |
| 413 | b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale |
| 414 | desc_o.store([i_t * BT, i_s * BS], b_o.to(desc_o.dtype)) |
| 415 | |
| 416 | |
| 417 | def fused_kda_gate_chunk_cumsum( |
| 418 | raw_g: torch.Tensor, |
| 419 | raw_beta: torch.Tensor, |
| 420 | A_log: torch.Tensor, |
| 421 | g_bias: torch.Tensor | None = None, |
| 422 | beta: float = 1.0, |
| 423 | threshold: float = 20.0, |
| 424 | lower_bound: float | None = None, |
| 425 | cu_seqlens: torch.Tensor | None = None, |
| 426 | chunk_indices: torch.Tensor | None = None, |
| 427 | chunk_size: int = FLA_CHUNK_SIZE, |
| 428 | output_dtype: torch.dtype | None = torch.float, |
| 429 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 430 | if cu_seqlens is not None: |
| 431 | assert raw_g.shape[0] == 1, ( |
| 432 | "Only batch size 1 is supported when cu_seqlens are provided" |
| 433 | ) |
| 434 | B, T, H, D = raw_g.shape |
| 435 | if raw_beta.shape != (B, T, H): |
| 436 | raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}") |
| 437 | if chunk_indices is None and cu_seqlens is not None: |
| 438 | chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) |
| 439 | NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) |
| 440 | |
| 441 | A_log = A_log.reshape(-1) |
| 442 | if g_bias is not None: |
| 443 | g_bias = g_bias.reshape(-1) |
| 444 | y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) |
| 445 | beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32) |
| 446 | |
| 447 | def grid(meta): |
| 448 | # For each (chunk, head), program 0 computes beta without extending a |
| 449 | # gate tile's critical path. The remaining programs cover the gate dim. |
| 450 | return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H) |
| 451 | |
| 452 | kda_gate_chunk_cumsum_vector_kernel[grid]( |
| 453 | s=raw_g, |
| 454 | raw_beta=raw_beta, |
| 455 | A_log=A_log, |
| 456 | g_bias=g_bias, |
| 457 | o=y, |
| 458 | beta_out=beta_out, |
| 459 | cu_seqlens=cu_seqlens, |
| 460 | chunk_indices=chunk_indices, |
| 461 | # RCP_LN2 folds in the natural-log -> log2 conversion so downstream |
| 462 | # exp2-based kernels reproduce exp(g). Keep this in sync with the |
| 463 | # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. |
| 464 | cumsum_scale=RCP_LN2, |
| 465 | lower_bound=lower_bound or 0.0, |
| 466 | beta=beta, |
| 467 | threshold=threshold, |
| 468 | T=T, |
| 469 | stride_beta_batch=raw_beta.stride(0), |
| 470 | stride_beta_token=raw_beta.stride(1), |
| 471 | stride_beta_head=raw_beta.stride(2), |
| 472 | H=H, |
| 473 | S=D, |
| 474 | BT=chunk_size, |
| 475 | USE_LOWER_BOUND=lower_bound is not None, |
| 476 | ) |
| 477 | return y, beta_out |
| 478 | |
| 479 | |
| 480 | def _chunk_kda_fwd_with_cumulative_g( |
| 481 | q: torch.Tensor, |
| 482 | k: torch.Tensor, |
| 483 | v: torch.Tensor, |
| 484 | g: torch.Tensor, |
| 485 | beta: torch.Tensor, |
| 486 | scale: float, |
| 487 | initial_state: torch.Tensor, |
| 488 | output_final_state: bool, |
| 489 | cu_seqlens: torch.Tensor | None = None, |
| 490 | chunk_indices: torch.Tensor | None = None, |
| 491 | chunk_offsets: torch.Tensor | None = None, |
| 492 | chunk_size: int = FLA_CHUNK_SIZE, |
| 493 | safe_gate: bool = False, |
| 494 | ): |
| 495 | # `g` must already be chunk-local cumulatively-summed AND scaled by |
| 496 | # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). |
| 497 | # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of |
| 498 | # calling this helper directly unless that invariant is upheld. |
| 499 | Aqk, A = chunk_kda_fwd_intra( |
| 500 | q=q, |
| 501 | k=k, |
| 502 | gk=g, |
| 503 | beta=beta, |
| 504 | scale=scale, |
| 505 | cu_seqlens=cu_seqlens, |
| 506 | chunk_indices=chunk_indices, |
| 507 | chunk_size=chunk_size, |
| 508 | safe_gate=safe_gate, |
| 509 | ) |
| 510 | w, u, _, kg = recompute_w_u_fwd( |
| 511 | k=k, |
| 512 | v=v, |
| 513 | beta=beta, |
| 514 | A=A, |
| 515 | gk=g, |
| 516 | cu_seqlens=cu_seqlens, |
| 517 | chunk_indices=chunk_indices, |
| 518 | ) |
| 519 | del A |
| 520 | h, v_new, final_state = chunk_gated_delta_rule_fwd_h( |
| 521 | k=kg, |
| 522 | w=w, |
| 523 | u=u, |
| 524 | gk=g, |
| 525 | initial_state=initial_state, |
| 526 | output_final_state=output_final_state, |
| 527 | cu_seqlens=cu_seqlens, |
| 528 | chunk_indices=chunk_indices, |
| 529 | chunk_offsets=chunk_offsets, |
| 530 | use_exp2=True, |
| 531 | ) |
| 532 | del w, u, kg |
| 533 | o = chunk_gla_fwd_o_gk( |
| 534 | q=q, |
| 535 | v=v_new, |
| 536 | g=g, |
| 537 | A=Aqk, |
| 538 | h=h, |
| 539 | o=v, |
| 540 | scale=scale, |
| 541 | cu_seqlens=cu_seqlens, |
| 542 | chunk_indices=chunk_indices, |
| 543 | chunk_size=chunk_size, |
| 544 | ) |
| 545 | del Aqk, v_new, h |
| 546 | return o, final_state |
| 547 | |
| 548 | |
| 549 | def chunk_kda_fwd( |
| 550 | q: torch.Tensor, |
| 551 | k: torch.Tensor, |
| 552 | v: torch.Tensor, |
| 553 | g: torch.Tensor, |
| 554 | beta: torch.Tensor, |
| 555 | scale: float, |
| 556 | initial_state: torch.Tensor, |
| 557 | output_final_state: bool, |
| 558 | cu_seqlens: torch.Tensor | None = None, |
| 559 | ): |
| 560 | chunk_size = FLA_CHUNK_SIZE |
| 561 | chunk_indices = ( |
| 562 | prepare_chunk_indices(cu_seqlens, chunk_size) |
| 563 | if cu_seqlens is not None |
| 564 | else None |
| 565 | ) |
| 566 | g = chunk_local_cumsum( |
| 567 | g, |
| 568 | chunk_size=chunk_size, |
| 569 | cu_seqlens=cu_seqlens, |
| 570 | chunk_indices=chunk_indices, |
| 571 | ) |
| 572 | # KDA evaluates cumulative gate decays with exp2. Convert from natural-log |
| 573 | # space so exp(x) is preserved as exp2(x / ln(2)). |
| 574 | g = g * RCP_LN2 |
| 575 | return _chunk_kda_fwd_with_cumulative_g( |
| 576 | q=q, |
| 577 | k=k, |
| 578 | v=v, |
| 579 | g=g, |
| 580 | beta=beta, |
| 581 | scale=scale, |
| 582 | initial_state=initial_state, |
| 583 | output_final_state=output_final_state, |
| 584 | cu_seqlens=cu_seqlens, |
| 585 | chunk_indices=chunk_indices, |
| 586 | chunk_size=chunk_size, |
| 587 | ) |
| 588 | |
| 589 | |
| 590 | def chunk_kda_with_fused_gate_fwd( |
| 591 | q: torch.Tensor, |
| 592 | k: torch.Tensor, |
| 593 | v: torch.Tensor, |
| 594 | raw_g: torch.Tensor, |
| 595 | raw_beta: torch.Tensor, |
| 596 | A_log: torch.Tensor, |
| 597 | g_bias: torch.Tensor | None, |
| 598 | scale: float, |
| 599 | initial_state: torch.Tensor, |
| 600 | output_final_state: bool, |
| 601 | lower_bound: float | None = None, |
| 602 | cu_seqlens: torch.Tensor | None = None, |
| 603 | chunk_indices: torch.Tensor | None = None, |
| 604 | chunk_offsets: torch.Tensor | None = None, |
| 605 | ): |
| 606 | chunk_size = FLA_CHUNK_SIZE |
| 607 | # Building chunk metadata from a device `cu_seqlens` costs a blocking |
| 608 | # device->host copy. Callers that already hold it, i.e. the attention |
| 609 | # metadata builder, should pass it in. |
| 610 | if chunk_indices is None and cu_seqlens is not None: |
| 611 | chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) |
| 612 | g, beta = fused_kda_gate_chunk_cumsum( |
| 613 | raw_g, |
| 614 | raw_beta=raw_beta, |
| 615 | A_log=A_log, |
| 616 | g_bias=g_bias, |
| 617 | cu_seqlens=cu_seqlens, |
| 618 | chunk_indices=chunk_indices, |
| 619 | chunk_size=chunk_size, |
| 620 | lower_bound=lower_bound, |
| 621 | ) |
| 622 | return _chunk_kda_fwd_with_cumulative_g( |
| 623 | q=q, |
| 624 | k=k, |
| 625 | v=v, |
| 626 | g=g, |
| 627 | beta=beta, |
| 628 | scale=scale, |
| 629 | initial_state=initial_state, |
| 630 | output_final_state=output_final_state, |
| 631 | cu_seqlens=cu_seqlens, |
| 632 | chunk_indices=chunk_indices, |
| 633 | chunk_offsets=chunk_offsets, |
| 634 | chunk_size=chunk_size, |
| 635 | safe_gate=lower_bound is not None, |
| 636 | ) |
| 637 | |
| 638 | |
| 639 | def chunk_kda( |
| 640 | q: torch.Tensor, |
| 641 | k: torch.Tensor, |
| 642 | v: torch.Tensor, |
| 643 | g: torch.Tensor, |
| 644 | beta: torch.Tensor, |
| 645 | scale: float = None, |
| 646 | initial_state: torch.Tensor = None, |
| 647 | output_final_state: bool = False, |
| 648 | use_qk_l2norm_in_kernel: bool = False, |
| 649 | cu_seqlens: torch.Tensor | None = None, |
| 650 | **kwargs, |
| 651 | ): |
| 652 | if scale is None: |
| 653 | scale = k.shape[-1] ** -0.5 |
| 654 | |
| 655 | if use_qk_l2norm_in_kernel: |
| 656 | q = l2norm_fwd(q.contiguous()) |
| 657 | k = l2norm_fwd(k.contiguous()) |
| 658 | |
| 659 | o, final_state = chunk_kda_fwd( |
| 660 | q=q, |
| 661 | k=k, |
| 662 | v=v.contiguous(), |
| 663 | g=g.contiguous(), |
| 664 | beta=beta.contiguous(), |
| 665 | scale=scale, |
| 666 | initial_state=initial_state.contiguous(), |
| 667 | output_final_state=output_final_state, |
| 668 | cu_seqlens=cu_seqlens, |
| 669 | ) |
| 670 | return o, final_state |
| 671 | |
| 672 | |
| 673 | def chunk_kda_with_fused_gate( |
| 674 | q: torch.Tensor, |
| 675 | k: torch.Tensor, |
| 676 | v: torch.Tensor, |
| 677 | raw_g: torch.Tensor, |
| 678 | raw_beta: torch.Tensor, |
| 679 | A_log: torch.Tensor, |
| 680 | g_bias: torch.Tensor | None, |
| 681 | scale: float | None = None, |
| 682 | initial_state: torch.Tensor | None = None, |
| 683 | output_final_state: bool = False, |
| 684 | lower_bound: float | None = None, |
| 685 | use_qk_l2norm_in_kernel: bool = False, |
| 686 | cu_seqlens: torch.Tensor | None = None, |
| 687 | chunk_indices: torch.Tensor | None = None, |
| 688 | chunk_offsets: torch.Tensor | None = None, |
| 689 | **kwargs, |
| 690 | ): |
| 691 | """Run chunk KDA from raw gate and beta projections.""" |
| 692 | if scale is None: |
| 693 | scale = k.shape[-1] ** -0.5 |
| 694 | |
| 695 | if use_qk_l2norm_in_kernel: |
| 696 | q = l2norm_fwd(q.contiguous()) |
| 697 | k = l2norm_fwd(k.contiguous()) |
| 698 | |
| 699 | o, final_state = chunk_kda_with_fused_gate_fwd( |
| 700 | q=q, |
| 701 | k=k, |
| 702 | v=v.contiguous(), |
| 703 | raw_g=raw_g.contiguous(), |
| 704 | raw_beta=raw_beta, |
| 705 | A_log=A_log, |
| 706 | g_bias=g_bias, |
| 707 | scale=scale, |
| 708 | initial_state=initial_state.contiguous() if initial_state is not None else None, |
| 709 | output_final_state=output_final_state, |
| 710 | lower_bound=lower_bound, |
| 711 | cu_seqlens=cu_seqlens, |
| 712 | chunk_indices=chunk_indices, |
| 713 | chunk_offsets=chunk_offsets, |
| 714 | ) |
| 715 | return o, final_state |
| 716 | |
| 717 | |
| 718 | @triton.autotune( |
| 719 | configs=[ |
| 720 | triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) |
| 721 | for bt in BT_LIST_AUTOTUNE |
| 722 | for nw in NUM_WARPS_AUTOTUNE |
| 723 | for ns in [2, 3] |
| 724 | ], |
| 725 | key=["H", "D"], |
| 726 | ) |
| 727 | @triton.jit |
| 728 | def kda_gate_fwd_kernel( |
| 729 | g, |
| 730 | A, |
| 731 | y, |
| 732 | g_bias, |
| 733 | lower_bound, |
| 734 | beta: tl.constexpr, |
| 735 | threshold: tl.constexpr, |
| 736 | T, |
| 737 | H, |
| 738 | D: tl.constexpr, |
| 739 | BT: tl.constexpr, |
| 740 | BD: tl.constexpr, |
| 741 | HAS_BIAS: tl.constexpr, |
| 742 | USE_LOWER_BOUND: tl.constexpr, |
| 743 | ): |
| 744 | i_t, i_h = tl.program_id(0), tl.program_id(1) |
| 745 | n_t = i_t * BT |
| 746 | |
| 747 | b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) |
| 748 | |
| 749 | stride_row = H * D |
| 750 | stride_col = 1 |
| 751 | |
| 752 | desc_g = make_tensor_descriptor(g + i_h * D, [T, D], [stride_row, stride_col], [BT, BD]) |
| 753 | desc_y = make_tensor_descriptor(y + i_h * D, [T, D], [stride_row, stride_col], [BT, BD]) |
| 754 | |
| 755 | b_g = desc_g.load([n_t, 0]).to(tl.float32) |
| 756 | |
| 757 | if HAS_BIAS: |
| 758 | n_d = tl.arange(0, BD) |
| 759 | bias_mask = n_d < D |
| 760 | b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( |
| 761 | tl.float32 |
| 762 | ) |
| 763 | b_g = b_g + b_bias[None, :] |
| 764 | |
| 765 | if USE_LOWER_BOUND: |
| 766 | b_y = lower_bound * tl.sigmoid(b_a * b_g) |
| 767 | else: |
| 768 | g_scaled = b_g * beta |
| 769 | use_linear = g_scaled > threshold |
| 770 | sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) |
| 771 | b_y = -b_a * sp |
| 772 | |
| 773 | desc_y.store([n_t, 0], b_y.to(desc_y.dtype)) |
| 774 | |
| 775 | |
| 776 | def fused_kda_gate( |
| 777 | g: torch.Tensor, |
| 778 | A: torch.Tensor, |
| 779 | head_k_dim: int, |
| 780 | g_bias: torch.Tensor | None = None, |
| 781 | beta: float = 1.0, |
| 782 | threshold: float = 20.0, |
| 783 | lower_bound: float | None = None, |
| 784 | ) -> torch.Tensor: |
| 785 | """Forward pass for KDA gate: |
| 786 | input g: [..., H*D] |
| 787 | param A: [H] or [1, 1, H, 1] |
| 788 | beta: softplus beta parameter |
| 789 | threshold: softplus threshold parameter |
| 790 | return : [..., H, D] |
| 791 | """ |
| 792 | orig_shape = g.shape[:-1] |
| 793 | |
| 794 | g = g.view(-1, g.shape[-1]) |
| 795 | T = g.shape[0] |
| 796 | HD = g.shape[1] |
| 797 | H = A.numel() |
| 798 | assert H * head_k_dim == HD |
| 799 | |
| 800 | y = torch.empty_like(g, dtype=torch.float32) |
| 801 | |
| 802 | def grid(meta): |
| 803 | return (cdiv(T, meta["BT"]), H) |
| 804 | |
| 805 | kda_gate_fwd_kernel[grid]( |
| 806 | g, |
| 807 | A, |
| 808 | y, |
| 809 | g_bias, |
| 810 | lower_bound or 0.0, |
| 811 | beta, |
| 812 | threshold, |
| 813 | T, |
| 814 | H, |
| 815 | head_k_dim, |
| 816 | BD=next_power_of_2(head_k_dim), |
| 817 | HAS_BIAS=g_bias is not None, |
| 818 | USE_LOWER_BOUND=lower_bound is not None, |
| 819 | ) |
| 820 | |
| 821 | y = y.view(*orig_shape, H, head_k_dim) |
| 822 | return y |