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