| 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 .op import exp, log |
| 16 | |
| 17 | |
| 18 | @triton.heuristics( |
| 19 | { |
| 20 | "USE_INITIAL_STATE": lambda args: args["h0"] is not None, |
| 21 | "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, |
| 22 | "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, |
| 23 | "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, |
| 24 | } |
| 25 | ) |
| 26 | @triton.jit(do_not_specialize=["N", "T"]) |
| 27 | def fused_recurrent_gated_delta_rule_fwd_kernel( |
| 28 | q, |
| 29 | k, |
| 30 | v, |
| 31 | g, |
| 32 | beta, |
| 33 | o, |
| 34 | h0, |
| 35 | ht, |
| 36 | cu_seqlens, |
| 37 | ssm_state_indices, |
| 38 | num_accepted_tokens, |
| 39 | scale, |
| 40 | N: tl.int64, # num of sequences |
| 41 | T: tl.int64, # num of tokens |
| 42 | B: tl.constexpr, |
| 43 | H: tl.constexpr, |
| 44 | HV: tl.constexpr, |
| 45 | K: tl.constexpr, |
| 46 | V: tl.constexpr, |
| 47 | BK: tl.constexpr, |
| 48 | BV: tl.constexpr, |
| 49 | stride_init_state_token: tl.constexpr, |
| 50 | stride_final_state_token: tl.constexpr, |
| 51 | stride_indices_seq: tl.constexpr, |
| 52 | stride_indices_tok: tl.constexpr, |
| 53 | USE_INITIAL_STATE: tl.constexpr, # whether to use initial state |
| 54 | INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace |
| 55 | IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, |
| 56 | USE_QK_L2NORM_IN_KERNEL: tl.constexpr, |
| 57 | IS_VARLEN: tl.constexpr, |
| 58 | IS_CONTINUOUS_BATCHING: tl.constexpr, |
| 59 | IS_SPEC_DECODING: tl.constexpr, |
| 60 | IS_KDA: tl.constexpr, |
| 61 | ): |
| 62 | i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| 63 | i_n, i_hv = i_nh // HV, i_nh % HV |
| 64 | i_h = i_hv // (HV // H) |
| 65 | if IS_VARLEN: |
| 66 | bos, eos = ( |
| 67 | tl.load(cu_seqlens + i_n).to(tl.int64), |
| 68 | tl.load(cu_seqlens + i_n + 1).to(tl.int64), |
| 69 | ) |
| 70 | all = T |
| 71 | T = eos - bos |
| 72 | else: |
| 73 | bos, eos = i_n * T, i_n * T + T |
| 74 | all = B * T |
| 75 | |
| 76 | if T == 0: |
| 77 | # no tokens to process for this sequence |
| 78 | return |
| 79 | |
| 80 | o_k = i_k * BK + tl.arange(0, BK) |
| 81 | o_v = i_v * BV + tl.arange(0, BV) |
| 82 | |
| 83 | p_q = q + (bos * H + i_h) * K + o_k |
| 84 | p_k = k + (bos * H + i_h) * K + o_k |
| 85 | p_v = v + (bos * HV + i_hv) * V + o_v |
| 86 | if IS_BETA_HEADWISE: |
| 87 | p_beta = beta + (bos * HV + i_hv) * V + o_v |
| 88 | else: |
| 89 | p_beta = beta + bos * HV + i_hv |
| 90 | |
| 91 | if not IS_KDA: |
| 92 | p_g = g + bos * HV + i_hv |
| 93 | else: |
| 94 | p_gk = g + (bos * HV + i_hv) * K + o_k |
| 95 | |
| 96 | p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v |
| 97 | |
| 98 | mask_k = o_k < K |
| 99 | mask_v = o_v < V |
| 100 | mask_h = mask_v[:, None] & mask_k[None, :] |
| 101 | |
| 102 | b_h = tl.zeros([BV, BK], dtype=tl.float32) |
| 103 | if USE_INITIAL_STATE: |
| 104 | if IS_CONTINUOUS_BATCHING: |
| 105 | if IS_SPEC_DECODING: |
| 106 | i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 |
| 107 | else: |
| 108 | i_t = 0 |
| 109 | # Load state index and check for invalid entries |
| 110 | state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( |
| 111 | tl.int64 |
| 112 | ) |
| 113 | # Skip if state index is invalid (NULL_BLOCK_ID=0) |
| 114 | if state_idx <= 0: |
| 115 | return |
| 116 | p_h0 = h0 + state_idx * stride_init_state_token |
| 117 | else: |
| 118 | p_h0 = h0 + bos * HV * V * K |
| 119 | p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :] |
| 120 | b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) |
| 121 | |
| 122 | for i_t in range(0, T): |
| 123 | b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) |
| 124 | b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) |
| 125 | b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) |
| 126 | |
| 127 | if USE_QK_L2NORM_IN_KERNEL: |
| 128 | b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) |
| 129 | b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) |
| 130 | b_q = b_q * scale |
| 131 | # [BV, BK] |
| 132 | if not IS_KDA: |
| 133 | b_g = tl.load(p_g).to(tl.float32) |
| 134 | b_h *= exp(b_g) |
| 135 | else: |
| 136 | b_gk = tl.load(p_gk).to(tl.float32) |
| 137 | b_h *= exp(b_gk[None, :]) |
| 138 | # [BV] |
| 139 | b_v -= tl.sum(b_h * b_k[None, :], 1) |
| 140 | if IS_BETA_HEADWISE: |
| 141 | b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) |
| 142 | else: |
| 143 | b_beta = tl.load(p_beta).to(tl.float32) |
| 144 | b_v *= b_beta |
| 145 | # [BV, BK] |
| 146 | b_h += b_v[:, None] * b_k[None, :] |
| 147 | # [BV] |
| 148 | b_o = tl.sum(b_h * b_q[None, :], 1) |
| 149 | tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) |
| 150 | |
| 151 | # keep the states for multi-query tokens |
| 152 | if INPLACE_FINAL_STATE: |
| 153 | # Load state index and check for invalid entries |
| 154 | final_state_idx = tl.load( |
| 155 | ssm_state_indices + i_n * stride_indices_seq + i_t |
| 156 | ).to(tl.int64) |
| 157 | # Only store if state index is valid (not NULL_BLOCK_ID=0) |
| 158 | if final_state_idx > 0: |
| 159 | p_ht = ht + final_state_idx * stride_final_state_token |
| 160 | p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] |
| 161 | tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) |
| 162 | else: |
| 163 | p_ht = ht + (bos + i_t) * stride_final_state_token |
| 164 | p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] |
| 165 | tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) |
| 166 | |
| 167 | p_q += H * K |
| 168 | p_k += H * K |
| 169 | p_o += HV * V |
| 170 | p_v += HV * V |
| 171 | if not IS_KDA: |
| 172 | p_g += HV |
| 173 | else: |
| 174 | p_gk += HV * K |
| 175 | p_beta += HV * (V if IS_BETA_HEADWISE else 1) |
| 176 | |
| 177 | |
| 178 | def fused_recurrent_gated_delta_rule_fwd( |
| 179 | q: torch.Tensor, |
| 180 | k: torch.Tensor, |
| 181 | v: torch.Tensor, |
| 182 | g: torch.Tensor, |
| 183 | beta: torch.Tensor, |
| 184 | scale: float, |
| 185 | initial_state: torch.Tensor, |
| 186 | inplace_final_state: bool = True, |
| 187 | cu_seqlens: torch.Tensor | None = None, |
| 188 | ssm_state_indices: torch.Tensor | None = None, |
| 189 | num_accepted_tokens: torch.Tensor | None = None, |
| 190 | use_qk_l2norm_in_kernel: bool = False, |
| 191 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 192 | B, T, H, K, V = *k.shape, v.shape[-1] |
| 193 | HV = v.shape[2] |
| 194 | N = B if cu_seqlens is None else len(cu_seqlens) - 1 |
| 195 | BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) |
| 196 | NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) |
| 197 | assert NK == 1, "NK > 1 is not supported yet" |
| 198 | num_stages = 3 |
| 199 | num_warps = 1 |
| 200 | |
| 201 | o = q.new_empty(NK, *v.shape) |
| 202 | if inplace_final_state: |
| 203 | final_state = initial_state |
| 204 | else: |
| 205 | final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype) |
| 206 | |
| 207 | stride_init_state_token = initial_state.stride(0) |
| 208 | stride_final_state_token = final_state.stride(0) |
| 209 | |
| 210 | if ssm_state_indices is None: |
| 211 | stride_indices_seq, stride_indices_tok = 1, 1 |
| 212 | elif ssm_state_indices.ndim == 1: |
| 213 | stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1 |
| 214 | else: |
| 215 | stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() |
| 216 | |
| 217 | grid = (NK, NV, N * HV) |
| 218 | fused_recurrent_gated_delta_rule_fwd_kernel[grid]( |
| 219 | q=q, |
| 220 | k=k, |
| 221 | v=v, |
| 222 | g=g, |
| 223 | beta=beta, |
| 224 | o=o, |
| 225 | h0=initial_state, |
| 226 | ht=final_state, |
| 227 | cu_seqlens=cu_seqlens, |
| 228 | ssm_state_indices=ssm_state_indices, |
| 229 | num_accepted_tokens=num_accepted_tokens, |
| 230 | scale=scale, |
| 231 | N=N, |
| 232 | T=T, |
| 233 | B=B, |
| 234 | H=H, |
| 235 | HV=HV, |
| 236 | K=K, |
| 237 | V=V, |
| 238 | BK=BK, |
| 239 | BV=BV, |
| 240 | stride_init_state_token=stride_init_state_token, |
| 241 | stride_final_state_token=stride_final_state_token, |
| 242 | stride_indices_seq=stride_indices_seq, |
| 243 | stride_indices_tok=stride_indices_tok, |
| 244 | IS_BETA_HEADWISE=beta.ndim == v.ndim, |
| 245 | USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, |
| 246 | INPLACE_FINAL_STATE=inplace_final_state, |
| 247 | IS_KDA=False, |
| 248 | num_warps=num_warps, |
| 249 | num_stages=num_stages, |
| 250 | ) |
| 251 | o = o.squeeze(0) |
| 252 | return o, final_state |
| 253 | |
| 254 | |
| 255 | @triton.jit |
| 256 | def fused_recurrent_gated_delta_rule_packed_decode_kernel( |
| 257 | mixed_qkv, |
| 258 | a, |
| 259 | b, |
| 260 | A_log, |
| 261 | dt_bias, |
| 262 | o, |
| 263 | h0, |
| 264 | ht, |
| 265 | ssm_state_indices, |
| 266 | scale, |
| 267 | stride_mixed_qkv_tok: tl.constexpr, |
| 268 | stride_a_tok: tl.constexpr, |
| 269 | stride_b_tok: tl.constexpr, |
| 270 | stride_init_state_token: tl.constexpr, |
| 271 | stride_final_state_token: tl.constexpr, |
| 272 | stride_indices_seq: tl.constexpr, |
| 273 | H: tl.constexpr, |
| 274 | HV: tl.constexpr, |
| 275 | K: tl.constexpr, |
| 276 | V: tl.constexpr, |
| 277 | BK: tl.constexpr, |
| 278 | BV: tl.constexpr, |
| 279 | SOFTPLUS_THRESHOLD: tl.constexpr, |
| 280 | USE_QK_L2NORM_IN_KERNEL: tl.constexpr, |
| 281 | SPLIT_BATCH_HEAD_GRID: tl.constexpr, |
| 282 | ): |
| 283 | if SPLIT_BATCH_HEAD_GRID: |
| 284 | i_v, i_hv, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2) |
| 285 | else: |
| 286 | i_v, i_nh = tl.program_id(0), tl.program_id(1) |
| 287 | i_n, i_hv = i_nh // HV, i_nh % HV |
| 288 | i_h = i_hv // (HV // H) |
| 289 | |
| 290 | o_k = tl.arange(0, BK) |
| 291 | o_v = i_v * BV + tl.arange(0, BV) |
| 292 | mask_k = o_k < K |
| 293 | mask_v = o_v < V |
| 294 | mask_h = mask_v[:, None] & mask_k[None, :] |
| 295 | |
| 296 | state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64) |
| 297 | p_o = o + (i_n * HV + i_hv) * V + o_v |
| 298 | |
| 299 | # Skip if state index is invalid (NULL_BLOCK_ID=0) |
| 300 | if state_idx <= 0: |
| 301 | zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty) |
| 302 | tl.store(p_o, zero, mask=mask_v) |
| 303 | return |
| 304 | |
| 305 | p_h0 = h0 + state_idx * stride_init_state_token |
| 306 | p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :] |
| 307 | b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) |
| 308 | |
| 309 | p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok |
| 310 | q_off = i_h * K + o_k |
| 311 | k_off = (H * K) + i_h * K + o_k |
| 312 | v_off = (2 * H * K) + i_hv * V + o_v |
| 313 | b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32) |
| 314 | b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32) |
| 315 | b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32) |
| 316 | |
| 317 | if USE_QK_L2NORM_IN_KERNEL: |
| 318 | b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) |
| 319 | b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) |
| 320 | b_q = b_q * scale |
| 321 | |
| 322 | a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32) |
| 323 | b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32) |
| 324 | A_log_val = tl.load(A_log + i_hv).to(tl.float32) |
| 325 | dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32) |
| 326 | x = a_val + dt_bias_val |
| 327 | softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) |
| 328 | g_val = -tl.exp(A_log_val) * softplus_x |
| 329 | beta_val = tl.sigmoid(b_val) |
| 330 | |
| 331 | b_h *= exp(g_val) |
| 332 | b_v -= tl.sum(b_h * b_k[None, :], 1) |
| 333 | b_v *= beta_val |
| 334 | b_h += b_v[:, None] * b_k[None, :] |
| 335 | b_o = tl.sum(b_h * b_q[None, :], 1) |
| 336 | tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) |
| 337 | |
| 338 | p_ht = ht + state_idx * stride_final_state_token |
| 339 | p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] |
| 340 | tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) |
| 341 | |
| 342 | |
| 343 | def fused_recurrent_gated_delta_rule_packed_decode( |
| 344 | mixed_qkv: torch.Tensor, |
| 345 | a: torch.Tensor, |
| 346 | b: torch.Tensor, |
| 347 | A_log: torch.Tensor, |
| 348 | dt_bias: torch.Tensor, |
| 349 | scale: float, |
| 350 | initial_state: torch.Tensor, |
| 351 | out: torch.Tensor, |
| 352 | ssm_state_indices: torch.Tensor, |
| 353 | use_qk_l2norm_in_kernel: bool = False, |
| 354 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 355 | if mixed_qkv.ndim != 2: |
| 356 | raise ValueError( |
| 357 | f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})." |
| 358 | ) |
| 359 | if mixed_qkv.stride(-1) != 1: |
| 360 | raise ValueError("`mixed_qkv` must be contiguous in the last dim.") |
| 361 | if a.ndim != 2 or b.ndim != 2: |
| 362 | raise ValueError( |
| 363 | f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})." |
| 364 | ) |
| 365 | if a.stride(-1) != 1 or b.stride(-1) != 1: |
| 366 | raise ValueError("`a`/`b` must be contiguous in the last dim.") |
| 367 | if A_log.ndim != 1 or dt_bias.ndim != 1: |
| 368 | raise ValueError("`A_log`/`dt_bias` must be 1D tensors.") |
| 369 | if A_log.stride(0) != 1 or dt_bias.stride(0) != 1: |
| 370 | raise ValueError("`A_log`/`dt_bias` must be contiguous.") |
| 371 | if ssm_state_indices.ndim != 1: |
| 372 | raise ValueError( |
| 373 | f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})." |
| 374 | ) |
| 375 | if not out.is_contiguous(): |
| 376 | raise ValueError("`out` must be contiguous.") |
| 377 | |
| 378 | dev = mixed_qkv.device |
| 379 | if ( |
| 380 | a.device != dev |
| 381 | or b.device != dev |
| 382 | or A_log.device != dev |
| 383 | or dt_bias.device != dev |
| 384 | or initial_state.device != dev |
| 385 | or out.device != dev |
| 386 | or ssm_state_indices.device != dev |
| 387 | ): |
| 388 | raise ValueError("All inputs must be on the same device.") |
| 389 | |
| 390 | B = mixed_qkv.shape[0] |
| 391 | if a.shape[0] != B or b.shape[0] != B: |
| 392 | raise ValueError( |
| 393 | "Mismatched batch sizes: " |
| 394 | f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}." |
| 395 | ) |
| 396 | if ssm_state_indices.shape[0] != B: |
| 397 | raise ValueError( |
| 398 | f"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},))." |
| 399 | ) |
| 400 | |
| 401 | if initial_state.ndim != 4: |
| 402 | raise ValueError( |
| 403 | f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})." |
| 404 | ) |
| 405 | if initial_state.stride(-1) != 1: |
| 406 | raise ValueError("`initial_state` must be contiguous in the last dim.") |
| 407 | HV, V, K = initial_state.shape[-3:] |
| 408 | if a.shape[1] != HV or b.shape[1] != HV: |
| 409 | raise ValueError( |
| 410 | f"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)})." |
| 411 | ) |
| 412 | if A_log.numel() != HV or dt_bias.numel() != HV: |
| 413 | raise ValueError( |
| 414 | f"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()})." |
| 415 | ) |
| 416 | if out.shape != (B, 1, HV, V): |
| 417 | raise ValueError( |
| 418 | f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})." |
| 419 | ) |
| 420 | |
| 421 | qkv_dim = mixed_qkv.shape[1] |
| 422 | qk_dim = qkv_dim - HV * V |
| 423 | if qk_dim <= 0 or qk_dim % 2 != 0: |
| 424 | raise ValueError( |
| 425 | f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}." |
| 426 | ) |
| 427 | q_dim = qk_dim // 2 |
| 428 | if q_dim % K != 0: |
| 429 | raise ValueError(f"Invalid packed Q size {q_dim}: must be divisible by K={K}.") |
| 430 | H = q_dim // K |
| 431 | if H <= 0 or HV % H != 0: |
| 432 | raise ValueError( |
| 433 | f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}." |
| 434 | ) |
| 435 | |
| 436 | BK = triton.next_power_of_2(K) |
| 437 | if triton.cdiv(K, BK) != 1: |
| 438 | raise ValueError( |
| 439 | f"Packed decode kernel only supports NK=1 (got K={K}, BK={BK})." |
| 440 | ) |
| 441 | BV = min(triton.next_power_of_2(V), 32) |
| 442 | num_stages = 3 |
| 443 | num_warps = 1 |
| 444 | |
| 445 | stride_mixed_qkv_tok = mixed_qkv.stride(0) |
| 446 | stride_a_tok = a.stride(0) |
| 447 | stride_b_tok = b.stride(0) |
| 448 | stride_init_state_token = initial_state.stride(0) |
| 449 | stride_final_state_token = initial_state.stride(0) |
| 450 | stride_indices_seq = ssm_state_indices.stride(0) |
| 451 | |
| 452 | NV = triton.cdiv(V, BV) |
| 453 | # CUDA limits grid Y/Z dimensions to 65535. |
| 454 | split_batch_head_grid = B * HV > 65535 |
| 455 | grid = (NV, HV, B) if split_batch_head_grid else (NV, B * HV) |
| 456 | fused_recurrent_gated_delta_rule_packed_decode_kernel[grid]( |
| 457 | mixed_qkv=mixed_qkv, |
| 458 | a=a, |
| 459 | b=b, |
| 460 | A_log=A_log, |
| 461 | dt_bias=dt_bias, |
| 462 | o=out, |
| 463 | h0=initial_state, |
| 464 | ht=initial_state, |
| 465 | ssm_state_indices=ssm_state_indices, |
| 466 | scale=scale, |
| 467 | stride_mixed_qkv_tok=stride_mixed_qkv_tok, |
| 468 | stride_a_tok=stride_a_tok, |
| 469 | stride_b_tok=stride_b_tok, |
| 470 | stride_init_state_token=stride_init_state_token, |
| 471 | stride_final_state_token=stride_final_state_token, |
| 472 | stride_indices_seq=stride_indices_seq, |
| 473 | H=H, |
| 474 | HV=HV, |
| 475 | K=K, |
| 476 | V=V, |
| 477 | BK=BK, |
| 478 | BV=BV, |
| 479 | SOFTPLUS_THRESHOLD=20.0, |
| 480 | USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, |
| 481 | SPLIT_BATCH_HEAD_GRID=split_batch_head_grid, |
| 482 | num_warps=num_warps, |
| 483 | num_stages=num_stages, |
| 484 | ) |
| 485 | return out, initial_state |
| 486 | |
| 487 | |
| 488 | class FusedRecurrentFunction(torch.autograd.Function): |
| 489 | @staticmethod |
| 490 | def forward( |
| 491 | ctx, |
| 492 | q: torch.Tensor, |
| 493 | k: torch.Tensor, |
| 494 | v: torch.Tensor, |
| 495 | g: torch.Tensor, |
| 496 | beta: torch.Tensor, |
| 497 | scale: float, |
| 498 | initial_state: torch.Tensor, |
| 499 | inplace_final_state: bool = True, |
| 500 | cu_seqlens: torch.Tensor | None = None, |
| 501 | ssm_state_indices: torch.Tensor | None = None, |
| 502 | num_accepted_tokens: torch.Tensor | None = None, |
| 503 | use_qk_l2norm_in_kernel: bool = False, |
| 504 | ): |
| 505 | o, final_state = fused_recurrent_gated_delta_rule_fwd( |
| 506 | q=q.contiguous(), |
| 507 | k=k.contiguous(), |
| 508 | v=v.contiguous(), |
| 509 | g=g.contiguous(), |
| 510 | beta=beta.contiguous(), |
| 511 | scale=scale, |
| 512 | initial_state=initial_state, |
| 513 | inplace_final_state=inplace_final_state, |
| 514 | cu_seqlens=cu_seqlens, |
| 515 | ssm_state_indices=ssm_state_indices, |
| 516 | num_accepted_tokens=num_accepted_tokens, |
| 517 | use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, |
| 518 | ) |
| 519 | |
| 520 | return o, final_state |
| 521 | |
| 522 | |
| 523 | def fused_recurrent_gated_delta_rule( |
| 524 | q: torch.Tensor, |
| 525 | k: torch.Tensor, |
| 526 | v: torch.Tensor, |
| 527 | g: torch.Tensor, |
| 528 | beta: torch.Tensor = None, |
| 529 | scale: float = None, |
| 530 | initial_state: torch.Tensor = None, |
| 531 | inplace_final_state: bool = True, |
| 532 | cu_seqlens: torch.Tensor | None = None, |
| 533 | ssm_state_indices: torch.Tensor | None = None, |
| 534 | num_accepted_tokens: torch.Tensor | None = None, |
| 535 | use_qk_l2norm_in_kernel: bool = False, |
| 536 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 537 | r""" |
| 538 | Args: |
| 539 | q (torch.Tensor): |
| 540 | queries of shape `[B, T, H, K]`. |
| 541 | k (torch.Tensor): |
| 542 | keys of shape `[B, T, H, K]`. |
| 543 | v (torch.Tensor): |
| 544 | values of shape `[B, T, HV, V]`. |
| 545 | GVA is applied if `HV > H`. |
| 546 | g (torch.Tensor): |
| 547 | g (decays) of shape `[B, T, HV]`. |
| 548 | beta (torch.Tensor): |
| 549 | betas of shape `[B, T, HV]`. |
| 550 | scale (Optional[int]): |
| 551 | Scale factor for the RetNet attention scores. |
| 552 | If not provided, it will default to `1 / sqrt(K)`. Default: `None`. |
| 553 | initial_state (Optional[torch.Tensor]): |
| 554 | Initial state of shape `[N, HV, V, K]` for `N` input sequences. |
| 555 | For equal-length input sequences, `N` equals the batch size `B`. |
| 556 | Default: `None`. |
| 557 | inplace_final_state: bool: |
| 558 | Whether to store the final state in-place to save memory. |
| 559 | Default: `True`. |
| 560 | cu_seqlens (torch.Tensor): |
| 561 | Cumulative sequence lengths of shape `[N+1]` used for variable-length training, |
| 562 | consistent with the FlashAttention API. |
| 563 | ssm_state_indices (Optional[torch.Tensor]): |
| 564 | Indices to map the input sequences to the initial/final states. |
| 565 | num_accepted_tokens (Optional[torch.Tensor]): |
| 566 | Number of accepted tokens for each sequence during decoding. |
| 567 | |
| 568 | Returns: |
| 569 | o (torch.Tensor): |
| 570 | Outputs of shape `[B, T, HV, V]`. |
| 571 | final_state (torch.Tensor): |
| 572 | Final state of shape `[N, HV, V, K]`. |
| 573 | |
| 574 | Examples:: |
| 575 | >>> import torch |
| 576 | >>> import torch.nn.functional as F |
| 577 | >>> from einops import rearrange |
| 578 | >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule |
| 579 | # inputs with equal lengths |
| 580 | >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 |
| 581 | >>> q = torch.randn(B, T, H, K, device='cuda') |
| 582 | >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) |
| 583 | >>> v = torch.randn(B, T, HV, V, device='cuda') |
| 584 | >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) |
| 585 | >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() |
| 586 | >>> h0 = torch.randn(B, HV, V, K, device='cuda') |
| 587 | >>> o, ht = fused_gated_recurrent_delta_rule( |
| 588 | q, k, v, g, beta, |
| 589 | initial_state=h0, |
| 590 | ) |
| 591 | # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required |
| 592 | >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) |
| 593 | # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected |
| 594 | >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32) |
| 595 | >>> o_var, ht_var = fused_gated_recurrent_delta_rule( |
| 596 | q, k, v, g, beta, |
| 597 | initial_state=h0, |
| 598 | cu_seqlens=cu_seqlens |
| 599 | ) |
| 600 | """ |
| 601 | if cu_seqlens is not None and q.shape[0] != 1: |
| 602 | raise ValueError( |
| 603 | f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." |
| 604 | f"Please flatten variable-length inputs before processing." |
| 605 | ) |
| 606 | if scale is None: |
| 607 | scale = k.shape[-1] ** -0.5 |
| 608 | else: |
| 609 | assert scale > 0, "scale must be positive" |
| 610 | if beta is None: |
| 611 | beta = torch.ones_like(q[..., 0]) |
| 612 | o, final_state = FusedRecurrentFunction.apply( |
| 613 | q, |
| 614 | k, |
| 615 | v, |
| 616 | g, |
| 617 | beta, |
| 618 | scale, |
| 619 | initial_state, |
| 620 | inplace_final_state, |
| 621 | cu_seqlens, |
| 622 | ssm_state_indices, |
| 623 | num_accepted_tokens, |
| 624 | use_qk_l2norm_in_kernel, |
| 625 | ) |
| 626 | return o, final_state |