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