« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15# -*- coding: utf-8 -*-
16# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
18# Keep the validated kernel adapter close to its upstream implementation.
19# pylint: disable=line-too-long,missing-public-type-hints,missing-public-docstring
20# pylint: disable=non-google-docstring,disallowed-name,unused-argument,invalid-name
21# pylint: disable=missing-module-docstring,missing-function-docstring
22# pylint: disable=abstract-method,arguments-differ
24import warnings
25from typing import Optional
27import torch
29from .triton.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
30from .triton.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o
31from .triton.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
32from .triton.cumsum import chunk_local_cumsum
33from .triton.solve_tril import solve_tril
34from .triton.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
35from .triton.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd
38def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> tuple[torch.Tensor, torch.Tensor]:
39 inv_norm = torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + eps)
40 return (x * inv_norm).to(x.dtype), inv_norm
43def chunk_gated_delta_rule_fwd_prepare(
44 k: torch.Tensor,
45 v: torch.Tensor,
46 g: torch.Tensor,
47 beta: torch.Tensor,
48 cu_seqlens: Optional[torch.LongTensor] = None,
49 chunk_size: int = 64,
50):
51 """Compute forward intermediates that do not depend on the initial state."""
52 g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, head_first=False)
53 A = chunk_scaled_dot_kkt_fwd(
54 k=k,
55 g=g,
56 beta=beta,
57 cu_seqlens=cu_seqlens,
58 chunk_size=chunk_size,
59 output_dtype=torch.float32,
60 )
61 A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
62 w, u = recompute_w_u_fwd(
63 k=k,
64 v=v,
65 beta=beta,
66 A=A,
67 g=g,
68 cu_seqlens=cu_seqlens,
69 )
70 return g, A, w, u
73def chunk_gated_delta_rule_fwd_apply_state(
74 k: torch.Tensor,
75 g: torch.Tensor,
76 w: torch.Tensor,
77 u: torch.Tensor,
78 initial_state: Optional[torch.Tensor],
79 output_final_state: bool,
80 cu_seqlens: Optional[torch.LongTensor] = None,
81 chunk_size: int = 64,
82):
83 """Apply the recurrent initial state and return local state intermediates."""
84 return chunk_gated_delta_rule_fwd_h(
85 k=k,
86 w=w,
87 u=u,
88 g=g,
89 initial_state=initial_state,
90 output_final_state=output_final_state,
91 chunk_size=chunk_size,
92 cu_seqlens=cu_seqlens,
93 )
96def chunk_gated_delta_rule_fwd_output(
97 q: torch.Tensor,
98 k: torch.Tensor,
99 v_new: torch.Tensor,
100 h: torch.Tensor,
101 g: torch.Tensor,
102 scale: float,
103 cu_seqlens: Optional[torch.LongTensor] = None,
104 chunk_size: int = 64,
105):
106 """Compute local outputs after recurrent states have been applied."""
107 return chunk_fwd_o(
108 q=q,
109 k=k,
110 v=v_new,
111 h=h,
112 g=g,
113 scale=scale,
114 cu_seqlens=cu_seqlens,
115 chunk_size=chunk_size,
116 )
119def chunk_gated_delta_rule_fwd(
120 q: torch.Tensor,
121 k: torch.Tensor,
122 v: torch.Tensor,
123 g: torch.Tensor,
124 beta: torch.Tensor,
125 scale: float,
126 initial_state: torch.Tensor,
127 output_final_state: bool,
128 cu_seqlens: Optional[torch.LongTensor] = None,
129 chunk_size: int = 64,
130):
131 g, A, w, u = chunk_gated_delta_rule_fwd_prepare(
132 k=k,
133 v=v,
134 g=g,
135 beta=beta,
136 cu_seqlens=cu_seqlens,
137 chunk_size=chunk_size,
138 )
139 h, v_new, final_state = chunk_gated_delta_rule_fwd_apply_state(
140 k=k,
141 g=g,
142 w=w,
143 u=u,
144 initial_state=initial_state,
145 output_final_state=output_final_state,
146 cu_seqlens=cu_seqlens,
147 chunk_size=chunk_size,
148 )
149 o = chunk_gated_delta_rule_fwd_output(
150 q=q,
151 k=k,
152 v_new=v_new,
153 h=h,
154 g=g,
155 scale=scale,
156 cu_seqlens=cu_seqlens,
157 chunk_size=chunk_size,
158 )
159 return g, o, A, final_state
162def chunk_gated_delta_rule_bwd_prepare(
163 q: torch.Tensor,
164 k: torch.Tensor,
165 v: torch.Tensor,
166 g: torch.Tensor,
167 beta: torch.Tensor,
168 A: torch.Tensor,
169 scale: float,
170 initial_state: Optional[torch.Tensor],
171 do: torch.Tensor,
172 cu_seqlens: Optional[torch.LongTensor] = None,
173 chunk_size: int = 64,
174):
175 """Compute backward intermediates that do not depend on final-state grad."""
176 w, u = recompute_w_u_fwd(
177 k=k,
178 v=v,
179 beta=beta,
180 A=A,
181 g=g,
182 cu_seqlens=cu_seqlens,
183 )
184 h, v_new, _ = chunk_gated_delta_rule_fwd_apply_state(
185 k=k,
186 g=g,
187 w=w,
188 u=u,
189 initial_state=initial_state,
190 output_final_state=False,
191 cu_seqlens=cu_seqlens,
192 chunk_size=chunk_size,
193 )
194 dv = chunk_bwd_dv_local(
195 q=q,
196 k=k,
197 g=g,
198 do=do,
199 scale=scale,
200 cu_seqlens=cu_seqlens,
201 chunk_size=chunk_size,
202 )
203 return w, h, v_new, dv
206def chunk_gated_delta_rule_bwd_state(
207 q: torch.Tensor,
208 k: torch.Tensor,
209 w: torch.Tensor,
210 g: torch.Tensor,
211 initial_state: Optional[torch.Tensor],
212 dht: Optional[torch.Tensor],
213 do: torch.Tensor,
214 dv: torch.Tensor,
215 scale: float,
216 cu_seqlens: Optional[torch.LongTensor] = None,
217 chunk_size: int = 64,
218):
219 """Apply the final-state gradient and produce the initial-state gradient."""
220 return chunk_gated_delta_rule_bwd_dhu(
221 q=q,
222 k=k,
223 w=w,
224 g=g,
225 h0=initial_state,
226 dht=dht,
227 do=do,
228 dv=dv,
229 scale=scale,
230 cu_seqlens=cu_seqlens,
231 chunk_size=chunk_size,
232 )
235def chunk_gated_delta_rule_bwd_finish(
236 q: torch.Tensor,
237 k: torch.Tensor,
238 v: torch.Tensor,
239 g: torch.Tensor,
240 beta: torch.Tensor,
241 A: torch.Tensor,
242 w: torch.Tensor,
243 h: torch.Tensor,
244 v_new: torch.Tensor,
245 dv: torch.Tensor,
246 do: torch.Tensor,
247 dh: torch.Tensor,
248 scale: float,
249 cu_seqlens: Optional[torch.LongTensor] = None,
250 chunk_size: int = 64,
251):
252 """Finish local tensor gradients after the state-gradient handoff."""
253 dq, dk, dw, dg = chunk_bwd_dqkwg(
254 q=q,
255 k=k,
256 v=v_new,
257 w=w,
258 g=g,
259 h=h,
260 dv=dv,
261 do=do,
262 dh=dh,
263 chunk_size=chunk_size,
264 scale=scale,
265 cu_seqlens=cu_seqlens,
266 )
267 dk2, dv, db, dg2 = prepare_wy_repr_bwd(
268 k=k,
269 v=v,
270 beta=beta,
271 g=g,
272 A=A,
273 dw=dw,
274 du=dv,
275 cu_seqlens=cu_seqlens,
276 chunk_size=chunk_size,
277 )
278 dk.add_(dk2)
279 dg.add_(dg2)
280 if dg.dtype != torch.float32:
281 raise ValueError(f"dg current type is {dg.dtype} , should be float32")
282 dg = chunk_local_cumsum(
283 dg,
284 chunk_size=chunk_size,
285 reverse=True,
286 cu_seqlens=cu_seqlens,
287 head_first=False,
288 )
289 return dq, dk, dv, db, dg
292def chunk_gated_delta_rule_bwd(
293 q: torch.Tensor,
294 k: torch.Tensor,
295 v: torch.Tensor,
296 g: torch.Tensor,
297 beta: torch.Tensor,
298 A: torch.Tensor,
299 scale: float,
300 initial_state: torch.Tensor,
301 do: torch.Tensor,
302 dht: torch.Tensor,
303 cu_seqlens: Optional[torch.LongTensor] = None,
304 chunk_size: int = 64,
305):
306 w, h, v_new, dv = chunk_gated_delta_rule_bwd_prepare(
307 q=q,
308 k=k,
309 v=v,
310 g=g,
311 beta=beta,
312 A=A,
313 scale=scale,
314 initial_state=initial_state,
315 do=do,
316 cu_seqlens=cu_seqlens,
317 chunk_size=chunk_size,
318 )
319 dh, dh0, dv = chunk_gated_delta_rule_bwd_state(
320 q=q,
321 k=k,
322 w=w,
323 g=g,
324 initial_state=initial_state,
325 dht=dht,
326 do=do,
327 dv=dv,
328 scale=scale,
329 cu_seqlens=cu_seqlens,
330 chunk_size=chunk_size,
331 )
332 dq, dk, dv, db, dg = chunk_gated_delta_rule_bwd_finish(
333 q=q,
334 k=k,
335 v=v,
336 g=g,
337 beta=beta,
338 A=A,
339 w=w,
340 h=h,
341 v_new=v_new,
342 dv=dv,
343 do=do,
344 dh=dh,
345 scale=scale,
346 cu_seqlens=cu_seqlens,
347 chunk_size=chunk_size,
348 )
349 return dq, dk, dv, db, dg, dh0
352@torch.compiler.disable
353@input_guard
354def chunk_gated_delta_rule_fwd_prepare_saved(
355 q: torch.Tensor,
356 k: torch.Tensor,
357 v: torch.Tensor,
358 g: torch.Tensor,
359 beta: torch.Tensor,
360 scale: float = None,
361 use_qk_l2norm_in_kernel: bool = False,
362 chunk_size: int = 64,
363):
364 """Prepare fused GDN forward tensors without consuming the initial state."""
365 if scale is None:
366 scale = k.shape[-1] ** -0.5
368 q_norm, q_inv_norm = q, q.new_empty(0)
369 k_norm, k_inv_norm = k, k.new_empty(0)
370 if use_qk_l2norm_in_kernel:
371 q_norm, q_inv_norm = _l2norm(q)
372 k_norm, k_inv_norm = _l2norm(k)
374 g_cumsum, A, w, u = chunk_gated_delta_rule_fwd_prepare(
375 k=k_norm,
376 v=v,
377 g=g,
378 beta=beta,
379 chunk_size=chunk_size,
380 )
381 return (
382 q_norm,
383 k_norm,
384 q_inv_norm,
385 k_inv_norm,
386 g_cumsum,
387 A,
388 w,
389 u,
390 scale,
391 )
394@torch.compiler.disable
395@input_guard
396def chunk_gated_delta_rule_fwd_apply_state_saved(
397 k_norm: torch.Tensor,
398 g_cumsum: torch.Tensor,
399 w: torch.Tensor,
400 u: torch.Tensor,
401 initial_state: torch.Tensor = None,
402 output_final_state: bool = True,
403 chunk_size: int = 64,
404):
405 """Apply an initial state to fused prepared forward tensors."""
406 return chunk_gated_delta_rule_fwd_apply_state(
407 k=k_norm,
408 g=g_cumsum,
409 w=w,
410 u=u,
411 initial_state=initial_state,
412 output_final_state=output_final_state,
413 chunk_size=chunk_size,
414 )
417@torch.compiler.disable
418@input_guard
419def chunk_gated_delta_rule_fwd_output_saved(
420 q_norm: torch.Tensor,
421 k_norm: torch.Tensor,
422 g_cumsum: torch.Tensor,
423 h: torch.Tensor,
424 v_new: torch.Tensor,
425 scale: float,
426 chunk_size: int = 64,
427):
428 """Compute fused GDN output from prepared, state-applied tensors."""
429 return chunk_gated_delta_rule_fwd_output(
430 q=q_norm,
431 k=k_norm,
432 v_new=v_new,
433 h=h,
434 g=g_cumsum,
435 scale=scale,
436 chunk_size=chunk_size,
437 )
440@torch.compiler.disable
441@input_guard
442def chunk_gated_delta_rule_fwd_saved(
443 q: torch.Tensor,
444 k: torch.Tensor,
445 v: torch.Tensor,
446 g: torch.Tensor,
447 beta: torch.Tensor,
448 scale: float = None,
449 initial_state: torch.Tensor = None,
450 use_qk_l2norm_in_kernel: bool = False,
451 chunk_size: int = 64,
452):
453 """Run fused GDN forward and return the tensors required by its backward."""
454 (
455 q_norm,
456 k_norm,
457 q_inv_norm,
458 k_inv_norm,
459 g_cumsum,
460 A,
461 w,
462 u,
463 scale,
464 ) = chunk_gated_delta_rule_fwd_prepare_saved(
465 q,
466 k,
467 v,
468 g,
469 beta,
470 scale=scale,
471 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
472 chunk_size=chunk_size,
473 )
474 h, v_new, final_state = chunk_gated_delta_rule_fwd_apply_state_saved(
475 k_norm,
476 g_cumsum,
477 w,
478 u,
479 initial_state=initial_state,
480 output_final_state=True,
481 chunk_size=chunk_size,
482 )
483 output = chunk_gated_delta_rule_fwd_output_saved(
484 q_norm,
485 k_norm,
486 g_cumsum,
487 h,
488 v_new,
489 scale,
490 chunk_size=chunk_size,
491 )
492 return (
493 output.to(q.dtype),
494 final_state,
495 q_norm,
496 k_norm,
497 q_inv_norm,
498 k_inv_norm,
499 g_cumsum,
500 A,
501 scale,
502 )
505@torch.compiler.disable
506@input_guard
507def chunk_gated_delta_rule_bwd_prepare_saved(
508 q_norm: torch.Tensor,
509 k_norm: torch.Tensor,
510 v: torch.Tensor,
511 g_cumsum: torch.Tensor,
512 beta: torch.Tensor,
513 A: torch.Tensor,
514 initial_state: torch.Tensor,
515 grad_output: torch.Tensor,
516 scale: float,
517 chunk_size: int = 64,
518):
519 """Prepare fused backward tensors before the final-state grad arrives."""
520 return chunk_gated_delta_rule_bwd_prepare(
521 q=q_norm,
522 k=k_norm,
523 v=v,
524 g=g_cumsum,
525 beta=beta,
526 A=A,
527 scale=scale,
528 initial_state=initial_state,
529 do=grad_output,
530 chunk_size=chunk_size,
531 )
534@torch.compiler.disable
535@input_guard
536def chunk_gated_delta_rule_bwd_state_saved(
537 q_norm: torch.Tensor,
538 k_norm: torch.Tensor,
539 g_cumsum: torch.Tensor,
540 w: torch.Tensor,
541 initial_state: torch.Tensor,
542 grad_final_state: torch.Tensor,
543 grad_output: torch.Tensor,
544 dv_local: torch.Tensor,
545 scale: float,
546 chunk_size: int = 64,
547):
548 """Consume the final-state grad and produce the initial-state grad."""
549 return chunk_gated_delta_rule_bwd_state(
550 q=q_norm,
551 k=k_norm,
552 w=w,
553 g=g_cumsum,
554 initial_state=initial_state,
555 dht=grad_final_state,
556 do=grad_output,
557 dv=dv_local,
558 scale=scale,
559 chunk_size=chunk_size,
560 )
563@torch.compiler.disable
564@input_guard
565def chunk_gated_delta_rule_bwd_finish_saved(
566 q: torch.Tensor,
567 k: torch.Tensor,
568 q_norm: torch.Tensor,
569 k_norm: torch.Tensor,
570 v: torch.Tensor,
571 g_cumsum: torch.Tensor,
572 beta: torch.Tensor,
573 A: torch.Tensor,
574 w: torch.Tensor,
575 h: torch.Tensor,
576 v_new: torch.Tensor,
577 dv: torch.Tensor,
578 grad_output: torch.Tensor,
579 dh: torch.Tensor,
580 q_inv_norm: torch.Tensor,
581 k_inv_norm: torch.Tensor,
582 scale: float,
583 use_qk_l2norm_in_kernel: bool = False,
584 chunk_size: int = 64,
585):
586 """Finish local fused gradients after the P2P state-gradient handoff."""
587 dq, dk, dv, dbeta, dg = chunk_gated_delta_rule_bwd_finish(
588 q=q_norm,
589 k=k_norm,
590 v=v,
591 g=g_cumsum,
592 beta=beta,
593 A=A,
594 w=w,
595 h=h,
596 v_new=v_new,
597 dv=dv,
598 do=grad_output,
599 dh=dh,
600 scale=scale,
601 chunk_size=chunk_size,
602 )
603 if use_qk_l2norm_in_kernel:
604 with torch.enable_grad():
605 q_leaf = q.detach().requires_grad_(True)
606 k_leaf = k.detach().requires_grad_(True)
607 q_recomputed, _ = _l2norm(q_leaf)
608 k_recomputed, _ = _l2norm(k_leaf)
609 dq, dk = torch.autograd.grad(
610 (q_recomputed, k_recomputed),
611 (q_leaf, k_leaf),
612 grad_outputs=(dq, dk),
613 )
614 del q_inv_norm, k_inv_norm
615 return dq, dk, dv, dg, dbeta
618@torch.compiler.disable
619@input_guard
620def chunk_gated_delta_rule_bwd_saved(
621 q: torch.Tensor,
622 k: torch.Tensor,
623 q_norm: torch.Tensor,
624 k_norm: torch.Tensor,
625 v: torch.Tensor,
626 g_cumsum: torch.Tensor,
627 beta: torch.Tensor,
628 A: torch.Tensor,
629 initial_state: torch.Tensor,
630 grad_output: torch.Tensor,
631 grad_final_state: torch.Tensor,
632 q_inv_norm: torch.Tensor,
633 k_inv_norm: torch.Tensor,
634 scale: float,
635 use_qk_l2norm_in_kernel: bool = False,
636 chunk_size: int = 64,
637):
638 """Run fused GDN backward from a context saved by the forward helper."""
639 w, h, v_new, dv = chunk_gated_delta_rule_bwd_prepare_saved(
640 q_norm,
641 k_norm,
642 v,
643 g_cumsum,
644 beta,
645 A,
646 initial_state,
647 grad_output,
648 scale,
649 chunk_size=chunk_size,
650 )
651 dh, dh0, dv = chunk_gated_delta_rule_bwd_state_saved(
652 q_norm,
653 k_norm,
654 g_cumsum,
655 w,
656 initial_state,
657 grad_final_state,
658 grad_output,
659 dv,
660 scale,
661 chunk_size=chunk_size,
662 )
663 dq, dk, dv, dg, dbeta = chunk_gated_delta_rule_bwd_finish_saved(
664 q,
665 k,
666 q_norm,
667 k_norm,
668 v,
669 g_cumsum,
670 beta,
671 A,
672 w,
673 h,
674 v_new,
675 dv,
676 grad_output,
677 dh,
678 q_inv_norm,
679 k_inv_norm,
680 scale,
681 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
682 chunk_size=chunk_size,
683 )
684 return dq, dk, dv, dg, dbeta, dh0
687class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
688 """Autograd wrapper for the Triton-Ascend chunk Gated Delta Rule."""
690 @staticmethod
691 @input_guard
692 @autocast_custom_fwd
693 def forward(
694 ctx,
695 q: torch.Tensor,
696 k: torch.Tensor,
697 v: torch.Tensor,
698 g: torch.Tensor,
699 beta: torch.Tensor,
700 scale: float,
701 initial_state: torch.Tensor,
702 output_final_state: bool,
703 cu_seqlens: Optional[torch.LongTensor] = None,
704 use_qk_l2norm_in_kernel: bool = False,
705 chunk_size: int = 64,
706 ):
707 g, o, A, final_state = chunk_gated_delta_rule_fwd(
708 q=q,
709 k=k,
710 v=v,
711 g=g,
712 beta=beta,
713 scale=scale,
714 initial_state=initial_state,
715 output_final_state=output_final_state,
716 cu_seqlens=cu_seqlens,
717 chunk_size=chunk_size,
718 )
720 saved_initial_state = initial_state if initial_state is not None else q.new_empty(0)
721 saved_cu_seqlens = cu_seqlens if cu_seqlens is not None else q.new_empty(0, dtype=torch.long)
722 ctx.save_for_backward(q, k, v, g, beta, A, saved_initial_state, saved_cu_seqlens)
723 ctx.has_initial_state = initial_state is not None
724 ctx.has_cu_seqlens = cu_seqlens is not None
725 ctx.scale = scale
726 ctx.chunk_size = chunk_size
727 return o.to(q.dtype), final_state
729 @staticmethod
730 @input_guard
731 @autocast_custom_bwd
732 def backward(
733 ctx,
734 do: torch.Tensor,
735 dht: torch.Tensor
736 ):
737 q, k, v, g, beta, A, initial_state, cu_seqlens = ctx.saved_tensors
738 if not ctx.has_initial_state:
739 initial_state = None
740 if not ctx.has_cu_seqlens:
741 cu_seqlens = None
742 dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
743 q=q,
744 k=k,
745 v=v,
746 g=g,
747 beta=beta,
748 A=A,
749 scale=ctx.scale,
750 initial_state=initial_state,
751 do=do,
752 dht=dht,
753 cu_seqlens=cu_seqlens,
754 chunk_size=ctx.chunk_size,
755 )
756 return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), None, dh0, None, None, None, None
759@torch.compiler.disable
760def chunk_gated_delta_rule(
761 q: torch.Tensor,
762 k: torch.Tensor,
763 v: torch.Tensor,
764 g: torch.Tensor,
765 beta: torch.Tensor,
766 scale: float = None,
767 initial_state: torch.Tensor = None,
768 output_final_state: bool = False,
769 use_qk_l2norm_in_kernel: bool = False,
770 cu_seqlens: Optional[torch.LongTensor] = None,
771 chunk_size: int = 64,
772 head_first: bool = False,
773):
774 r"""
775 Args:
776 q (torch.Tensor):
777 queries of shape `[B, T, H, K]`.
778 k (torch.Tensor):
779 keys of shape `[B, T, H, K]`.
780 v (torch.Tensor):
781 values of shape `[B, T, H, V]`.
782 g (torch.Tensor):
783 (forget) gating tensor (in log space!) of shape `[B, T, H]`.
784 beta (torch.Tensor):
785 betas of shape `[B, T, H]`.
786 scale (Optional[float]):
787 Scale factor for the RetNet attention scores.
788 If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
789 initial_state (Optional[torch.Tensor]):
790 Initial state of shape `[N, H, K, V]` for `N` input sequences.
791 For equal-length input sequences, `N` equals the batch size `B`.
792 Default: `None`.
793 output_final_state (Optional[bool]):
794 Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.
795 use_qk_l2norm_in_kernel (bool):
796 Whether to apply L2norm to the q/k tensor internally. Default: `False`.
797 cu_seqlens (torch.LongTensor):
798 Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
799 consistent with the FlashAttention API.
800 head_first (Optional[bool]):
801 Whether the inputs are in the head-first format. Default: `False`.
802 This argument has been deprecated.
803 Returns:
804 o (torch.Tensor):
805 Outputs of shape `[B, T, H, V]`.
806 final_state (torch.Tensor):
807 Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.
809 Examples::
810 >>> import torch
811 >>> import torch.nn.functional as F
812 >>> from einops import rearrange
813 >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
814 # inputs with equal lengths
815 >>> B, T, H, K, V = 4, 2048, 4, 512, 512
816 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
817 >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
818 >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
819 >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()
820 >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))
821 >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')
822 >>> o, ht = chunk_gated_delta_rule(
823 q, k, v, g, beta,
824 initial_state=h0,
825 output_final_state=True
826 )
827 # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
828 >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
829 # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
830 >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
831 >>> o, ht = chunk_gated_delta_rule(
832 q, k, v, g, beta,
833 initial_state=h0,
834 output_final_state=True,
835 cu_seqlens=cu_seqlens
836 )
837 """
838 if q.dtype != k.dtype or k.dtype != v.dtype:
839 raise ValueError(
840 f"q current type is {q.dtype}, k current type is {k.dtype}, "
841 f"v current type is {v.dtype}, they should be equal"
842 )
843 if q.dtype == torch.float32:
844 raise ValueError(
845 "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
846 )
847 if len(beta.shape) != 3:
848 raise ValueError(
849 f"beta current shape len is {len(beta.shape)}, beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
850 )
852 if head_first:
853 warnings.warn(
854 "head_first is deprecated and will be removed in a future version. "
855 "Please use head_first=False for now instead."
856 )
857 if not head_first and q.shape[1] < q.shape[2]:
858 warnings.warn(
859 f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
860 "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
861 "when head_first=False was specified. "
862 "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
863 )
864 if cu_seqlens is not None:
865 if q.shape[0] != 1:
866 raise ValueError(
867 f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
868 f"Please flatten variable-length inputs before processing."
869 )
870 if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
871 raise ValueError(
872 f"The number of initial states is expected to be equal to the number of input sequences, "
873 f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}."
874 )
875 if scale is None:
876 scale = k.shape[-1] ** -0.5
878 if use_qk_l2norm_in_kernel:
879 q, _ = _l2norm(q)
880 k, _ = _l2norm(k)
882 o, final_state = ChunkGatedDeltaRuleFunction.apply(
883 q,
884 k,
885 v,
886 g,
887 beta,
888 scale,
889 initial_state,
890 output_final_state,
891 cu_seqlens,
892 False,
893 chunk_size,
894 )
895 return o, final_state