| 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 | import torch |
| 11 | |
| 12 | from vllm.triton_utils import tl, triton |
| 13 | |
| 14 | |
| 15 | @triton.heuristics( |
| 16 | { |
| 17 | "USE_INITIAL_STATE": lambda args: args["h0"] is not None, |
| 18 | "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, |
| 19 | "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, |
| 20 | "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, |
| 21 | } |
| 22 | ) |
| 23 | @triton.jit(do_not_specialize=["N", "T"]) |
| 24 | def fused_sigmoid_gating_delta_rule_update_kernel( |
| 25 | A_log, |
| 26 | a, |
| 27 | b, |
| 28 | dt_bias, |
| 29 | beta, |
| 30 | threshold, |
| 31 | q, |
| 32 | k, |
| 33 | v, |
| 34 | o, |
| 35 | h0, |
| 36 | ht, |
| 37 | cu_seqlens, |
| 38 | ssm_state_indices, |
| 39 | num_accepted_tokens, |
| 40 | scale, |
| 41 | N: tl.int64, # num of sequences |
| 42 | T: tl.int64, # num of tokens |
| 43 | B: tl.constexpr, |
| 44 | H: tl.constexpr, |
| 45 | HV: tl.constexpr, |
| 46 | K: tl.constexpr, |
| 47 | V: tl.constexpr, |
| 48 | BK: tl.constexpr, |
| 49 | BV: tl.constexpr, |
| 50 | stride_init_state_token: tl.constexpr, |
| 51 | stride_final_state_token: tl.constexpr, |
| 52 | stride_indices_seq: tl.constexpr, |
| 53 | stride_indices_tok: tl.constexpr, |
| 54 | USE_INITIAL_STATE: tl.constexpr, # whether to use initial state |
| 55 | INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace |
| 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 | |
| 87 | p_A_log = A_log + i_hv |
| 88 | if not IS_KDA: |
| 89 | p_a = a + bos * HV + i_hv |
| 90 | p_dt_bias = dt_bias + i_hv |
| 91 | else: |
| 92 | p_a = a + (bos * HV + i_hv) * K + o_k |
| 93 | p_dt_bias = dt_bias + i_hv * K + o_k |
| 94 | |
| 95 | p_b = b + bos * HV + i_hv |
| 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 | b_b = tl.load(p_b).to(tl.float32) |
| 127 | |
| 128 | # If the model is loaded in fp16, without the .float() here, A might be -inf |
| 129 | x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32) |
| 130 | softplus_x = tl.where( |
| 131 | beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x |
| 132 | ) |
| 133 | b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x |
| 134 | |
| 135 | # compute beta_output = sigmoid(b) |
| 136 | b_beta = tl.sigmoid(b_b.to(tl.float32)) |
| 137 | |
| 138 | if USE_QK_L2NORM_IN_KERNEL: |
| 139 | b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)) |
| 140 | b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)) |
| 141 | b_q = b_q * scale |
| 142 | # [BV, BK] |
| 143 | if not IS_KDA: |
| 144 | b_h *= tl.exp(b_g) |
| 145 | else: |
| 146 | b_h *= tl.exp(b_g[None, :]) |
| 147 | # [BV] |
| 148 | b_v -= tl.sum(b_h * b_k[None, :], 1) |
| 149 | b_v *= b_beta |
| 150 | # [BV, BK] |
| 151 | b_h += b_v[:, None] * b_k[None, :] |
| 152 | # [BV] |
| 153 | b_o = tl.sum(b_h * b_q[None, :], 1) |
| 154 | tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) |
| 155 | |
| 156 | # keep the states for multi-query tokens |
| 157 | if INPLACE_FINAL_STATE: |
| 158 | # Load state index and check for invalid entries |
| 159 | final_state_idx = tl.load( |
| 160 | ssm_state_indices + i_n * stride_indices_seq + i_t |
| 161 | ).to(tl.int64) |
| 162 | # Only store if state index is valid (not NULL_BLOCK_ID=0) |
| 163 | if final_state_idx > 0: |
| 164 | p_ht = ht + final_state_idx * stride_final_state_token |
| 165 | p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] |
| 166 | tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) |
| 167 | else: |
| 168 | p_ht = ht + (bos + i_t) * stride_final_state_token |
| 169 | p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] |
| 170 | tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) |
| 171 | |
| 172 | # Update pointers for next timestep |
| 173 | p_q += H * K |
| 174 | p_k += H * K |
| 175 | p_o += HV * V |
| 176 | p_v += HV * V |
| 177 | p_b += HV |
| 178 | p_a += HV |
| 179 | |
| 180 | |
| 181 | def fused_sigmoid_gating_delta_rule_update( |
| 182 | A_log: torch.Tensor, |
| 183 | a: torch.Tensor, |
| 184 | b: torch.Tensor, |
| 185 | dt_bias: torch.Tensor, |
| 186 | q: torch.Tensor, |
| 187 | k: torch.Tensor, |
| 188 | v: torch.Tensor, |
| 189 | beta: float = 1.0, |
| 190 | threshold: float = 20.0, |
| 191 | scale: float = None, |
| 192 | initial_state: torch.Tensor = None, |
| 193 | inplace_final_state: bool = True, |
| 194 | cu_seqlens: torch.Tensor | None = None, |
| 195 | ssm_state_indices: torch.Tensor | None = None, |
| 196 | num_accepted_tokens: torch.Tensor | None = None, |
| 197 | use_qk_l2norm_in_kernel: bool = False, |
| 198 | is_kda: bool = False, |
| 199 | ): |
| 200 | """ |
| 201 | Fused triton implementation of sigmoid gating delta rule update. |
| 202 | This function uses a single fused kernel that combines both sigmoid gating |
| 203 | computation and the recurrent delta rule update for better performance. |
| 204 | """ |
| 205 | B, T, H, K, V = *k.shape, v.shape[-1] |
| 206 | HV = v.shape[2] |
| 207 | N = B if cu_seqlens is None else len(cu_seqlens) - 1 |
| 208 | BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) |
| 209 | NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) |
| 210 | assert NK == 1, "NK > 1 is not supported yet" |
| 211 | num_stages = 3 |
| 212 | num_warps = 4 |
| 213 | |
| 214 | if cu_seqlens is not None and q.shape[0] != 1: |
| 215 | raise ValueError( |
| 216 | f"The batch size is expected to be 1 rather than {q.shape[0]}" |
| 217 | f" when using `cu_seqlens`. Please flatten variable-length" |
| 218 | f" inputs before processing." |
| 219 | ) |
| 220 | if scale is None: |
| 221 | scale = k.shape[-1] ** -0.5 |
| 222 | else: |
| 223 | assert scale > 0, "scale must be positive" |
| 224 | |
| 225 | o = q.new_empty(NK, *v.shape) |
| 226 | if inplace_final_state: |
| 227 | final_state = initial_state |
| 228 | else: |
| 229 | final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype) |
| 230 | |
| 231 | stride_init_state_token = initial_state.stride(0) |
| 232 | stride_final_state_token = final_state.stride(0) |
| 233 | |
| 234 | if ssm_state_indices is None: |
| 235 | stride_indices_seq, stride_indices_tok = 1, 1 |
| 236 | elif ssm_state_indices.ndim == 1: |
| 237 | stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1 |
| 238 | else: |
| 239 | stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() |
| 240 | |
| 241 | grid = (NK, NV, N * HV) |
| 242 | fused_sigmoid_gating_delta_rule_update_kernel[grid]( |
| 243 | A_log=A_log, |
| 244 | a=a.contiguous(), |
| 245 | b=b.contiguous(), |
| 246 | dt_bias=dt_bias, |
| 247 | beta=beta, |
| 248 | threshold=threshold, |
| 249 | q=q.contiguous(), |
| 250 | k=k.contiguous(), |
| 251 | v=v.contiguous(), |
| 252 | o=o, |
| 253 | h0=initial_state, |
| 254 | ht=final_state, |
| 255 | cu_seqlens=cu_seqlens, |
| 256 | ssm_state_indices=ssm_state_indices, |
| 257 | num_accepted_tokens=num_accepted_tokens, |
| 258 | scale=scale, |
| 259 | N=N, |
| 260 | T=T, |
| 261 | B=B, |
| 262 | H=H, |
| 263 | HV=HV, |
| 264 | K=K, |
| 265 | V=V, |
| 266 | BK=BK, |
| 267 | BV=BV, |
| 268 | stride_init_state_token=stride_init_state_token, |
| 269 | stride_final_state_token=stride_final_state_token, |
| 270 | stride_indices_seq=stride_indices_seq, |
| 271 | stride_indices_tok=stride_indices_tok, |
| 272 | INPLACE_FINAL_STATE=inplace_final_state, |
| 273 | USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, |
| 274 | IS_KDA=is_kda, |
| 275 | num_warps=num_warps, |
| 276 | num_stages=num_stages, |
| 277 | ) |
| 278 | o = o.squeeze(0) |
| 279 | return o, final_state |