Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / gated_delta_net.py: 0%
153 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
17# Keep the validated kernel adapter close to its upstream implementation.
18# pylint: disable=line-too-long,missing-public-type-hints,missing-public-docstring
19# pylint: disable=non-google-docstring,disallowed-name,unused-argument,invalid-name
20# pylint: disable=missing-module-docstring,missing-function-docstring
21# pylint: disable=abstract-method,arguments-differ
22# pylint: disable=forbidden-backend-import
24import warnings
25from typing import Optional
27import torch
29from ._triton.gated_delta_net.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
30from ._triton.gated_delta_net.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o
31from ._triton.gated_delta_net.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
32from ._triton.gated_delta_net.cumsum import chunk_local_cumsum
33from ._triton.gated_delta_net.solve_tril import solve_tril
34from ._triton.gated_delta_net.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
35from ._triton.gated_delta_net.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 matrix_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 matrix_a = solve_tril(A=matrix_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=matrix_a,
67 g=g,
68 cu_seqlens=cu_seqlens,
69 )
70 return g, matrix_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, matrix_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 output = 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, output, matrix_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 cu_seqlens: torch.Tensor = None,
364):
365 """Prepare fused GDN forward tensors without consuming the initial state."""
366 if scale is None:
367 scale = k.shape[-1] ** -0.5
369 q_norm, q_inv_norm = q, q.new_empty(0)
370 k_norm, k_inv_norm = k, k.new_empty(0)
371 if use_qk_l2norm_in_kernel:
372 q_norm, q_inv_norm = _l2norm(q)
373 k_norm, k_inv_norm = _l2norm(k)
375 g_cumsum, matrix_a, w, u = chunk_gated_delta_rule_fwd_prepare(
376 k=k_norm,
377 v=v,
378 g=g,
379 beta=beta,
380 chunk_size=chunk_size,
381 cu_seqlens=cu_seqlens,
382 )
383 return (
384 q_norm,
385 k_norm,
386 q_inv_norm,
387 k_inv_norm,
388 g_cumsum,
389 matrix_a,
390 w,
391 u,
392 scale,
393 )
396@torch.compiler.disable
397@input_guard
398def chunk_gated_delta_rule_fwd_apply_state_saved(
399 k_norm: torch.Tensor,
400 g_cumsum: torch.Tensor,
401 w: torch.Tensor,
402 u: torch.Tensor,
403 initial_state: torch.Tensor = None,
404 output_final_state: bool = True,
405 chunk_size: int = 64,
406 cu_seqlens: torch.Tensor = None,
407):
408 """Apply an initial state to fused prepared forward tensors."""
409 return chunk_gated_delta_rule_fwd_apply_state(
410 k=k_norm,
411 g=g_cumsum,
412 w=w,
413 u=u,
414 initial_state=initial_state,
415 output_final_state=output_final_state,
416 chunk_size=chunk_size,
417 cu_seqlens=cu_seqlens,
418 )
421@torch.compiler.disable
422@input_guard
423def chunk_gated_delta_rule_fwd_output_saved(
424 q_norm: torch.Tensor,
425 k_norm: torch.Tensor,
426 g_cumsum: torch.Tensor,
427 h: torch.Tensor,
428 v_new: torch.Tensor,
429 scale: float,
430 chunk_size: int = 64,
431 cu_seqlens: torch.Tensor = None,
432):
433 """Compute fused GDN output from prepared, state-applied tensors."""
434 return chunk_gated_delta_rule_fwd_output(
435 q=q_norm,
436 k=k_norm,
437 v_new=v_new,
438 h=h,
439 g=g_cumsum,
440 scale=scale,
441 chunk_size=chunk_size,
442 cu_seqlens=cu_seqlens,
443 )
446@torch.compiler.disable
447@input_guard
448def chunk_gated_delta_rule_fwd_saved(
449 q: torch.Tensor,
450 k: torch.Tensor,
451 v: torch.Tensor,
452 g: torch.Tensor,
453 beta: torch.Tensor,
454 scale: float = None,
455 initial_state: torch.Tensor = None,
456 use_qk_l2norm_in_kernel: bool = False,
457 chunk_size: int = 64,
458 cu_seqlens: torch.Tensor = None,
459):
460 """Run fused GDN forward and return the tensors required by its backward."""
461 (
462 q_norm,
463 k_norm,
464 q_inv_norm,
465 k_inv_norm,
466 g_cumsum,
467 matrix_a,
468 w,
469 u,
470 scale,
471 ) = chunk_gated_delta_rule_fwd_prepare_saved(
472 q,
473 k,
474 v,
475 g,
476 beta,
477 scale=scale,
478 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
479 chunk_size=chunk_size,
480 cu_seqlens=cu_seqlens,
481 )
482 h, v_new, final_state = chunk_gated_delta_rule_fwd_apply_state_saved(
483 k_norm,
484 g_cumsum,
485 w,
486 u,
487 initial_state=initial_state,
488 output_final_state=True,
489 chunk_size=chunk_size,
490 cu_seqlens=cu_seqlens,
491 )
492 output = chunk_gated_delta_rule_fwd_output_saved(
493 q_norm,
494 k_norm,
495 g_cumsum,
496 h,
497 v_new,
498 scale,
499 chunk_size=chunk_size,
500 cu_seqlens=cu_seqlens,
501 )
502 return (
503 output.to(q.dtype),
504 final_state,
505 q_norm,
506 k_norm,
507 q_inv_norm,
508 k_inv_norm,
509 g_cumsum,
510 matrix_a,
511 scale,
512 )
515@torch.compiler.disable
516@input_guard
517def chunk_gated_delta_rule_bwd_prepare_saved(
518 q_norm: torch.Tensor,
519 k_norm: torch.Tensor,
520 v: torch.Tensor,
521 g_cumsum: torch.Tensor,
522 beta: torch.Tensor,
523 A: torch.Tensor,
524 initial_state: torch.Tensor,
525 grad_output: torch.Tensor,
526 scale: float,
527 chunk_size: int = 64,
528 cu_seqlens: torch.Tensor = None,
529):
530 """Prepare fused backward tensors before the final-state grad arrives."""
531 return chunk_gated_delta_rule_bwd_prepare(
532 q=q_norm,
533 k=k_norm,
534 v=v,
535 g=g_cumsum,
536 beta=beta,
537 A=A,
538 scale=scale,
539 initial_state=initial_state,
540 do=grad_output,
541 chunk_size=chunk_size,
542 cu_seqlens=cu_seqlens,
543 )
546@torch.compiler.disable
547@input_guard
548def chunk_gated_delta_rule_bwd_state_saved(
549 q_norm: torch.Tensor,
550 k_norm: torch.Tensor,
551 g_cumsum: torch.Tensor,
552 w: torch.Tensor,
553 initial_state: torch.Tensor,
554 grad_final_state: torch.Tensor,
555 grad_output: torch.Tensor,
556 dv_local: torch.Tensor,
557 scale: float,
558 chunk_size: int = 64,
559 cu_seqlens: torch.Tensor = None,
560):
561 """Consume the final-state grad and produce the initial-state grad."""
562 return chunk_gated_delta_rule_bwd_state(
563 q=q_norm,
564 k=k_norm,
565 w=w,
566 g=g_cumsum,
567 initial_state=initial_state,
568 dht=grad_final_state,
569 do=grad_output,
570 dv=dv_local,
571 scale=scale,
572 chunk_size=chunk_size,
573 cu_seqlens=cu_seqlens,
574 )
577@torch.compiler.disable
578@input_guard
579def chunk_gated_delta_rule_bwd_finish_saved(
580 q: torch.Tensor,
581 k: torch.Tensor,
582 q_norm: torch.Tensor,
583 k_norm: torch.Tensor,
584 v: torch.Tensor,
585 g_cumsum: torch.Tensor,
586 beta: torch.Tensor,
587 A: torch.Tensor,
588 w: torch.Tensor,
589 h: torch.Tensor,
590 v_new: torch.Tensor,
591 dv: torch.Tensor,
592 grad_output: torch.Tensor,
593 dh: torch.Tensor,
594 q_inv_norm: torch.Tensor,
595 k_inv_norm: torch.Tensor,
596 scale: float,
597 use_qk_l2norm_in_kernel: bool = False,
598 chunk_size: int = 64,
599 cu_seqlens: torch.Tensor = None,
600):
601 """Finish local fused gradients after the P2P state-gradient handoff."""
602 dq, dk, dv, dbeta, dg = chunk_gated_delta_rule_bwd_finish(
603 q=q_norm,
604 k=k_norm,
605 v=v,
606 g=g_cumsum,
607 beta=beta,
608 A=A,
609 w=w,
610 h=h,
611 v_new=v_new,
612 dv=dv,
613 do=grad_output,
614 dh=dh,
615 scale=scale,
616 chunk_size=chunk_size,
617 cu_seqlens=cu_seqlens,
618 )
619 if use_qk_l2norm_in_kernel:
620 with torch.enable_grad():
621 q_leaf = q.detach().requires_grad_(True)
622 k_leaf = k.detach().requires_grad_(True)
623 q_recomputed, _ = _l2norm(q_leaf)
624 k_recomputed, _ = _l2norm(k_leaf)
625 dq, dk = torch.autograd.grad(
626 (q_recomputed, k_recomputed),
627 (q_leaf, k_leaf),
628 grad_outputs=(dq, dk),
629 )
630 del q_inv_norm, k_inv_norm
631 return dq, dk, dv, dg, dbeta
634@torch.compiler.disable
635@input_guard
636def chunk_gated_delta_rule_bwd_saved(
637 q: torch.Tensor,
638 k: torch.Tensor,
639 q_norm: torch.Tensor,
640 k_norm: torch.Tensor,
641 v: torch.Tensor,
642 g_cumsum: torch.Tensor,
643 beta: torch.Tensor,
644 A: torch.Tensor,
645 initial_state: torch.Tensor,
646 grad_output: torch.Tensor,
647 grad_final_state: torch.Tensor,
648 q_inv_norm: torch.Tensor,
649 k_inv_norm: torch.Tensor,
650 scale: float,
651 use_qk_l2norm_in_kernel: bool = False,
652 chunk_size: int = 64,
653 cu_seqlens: torch.Tensor = None,
654):
655 """Run fused GDN backward from a context saved by the forward helper."""
656 w, h, v_new, dv = chunk_gated_delta_rule_bwd_prepare_saved(
657 q_norm,
658 k_norm,
659 v,
660 g_cumsum,
661 beta,
662 A,
663 initial_state,
664 grad_output,
665 scale,
666 chunk_size=chunk_size,
667 cu_seqlens=cu_seqlens,
668 )
669 dh, dh0, dv = chunk_gated_delta_rule_bwd_state_saved(
670 q_norm,
671 k_norm,
672 g_cumsum,
673 w,
674 initial_state,
675 grad_final_state,
676 grad_output,
677 dv,
678 scale,
679 chunk_size=chunk_size,
680 cu_seqlens=cu_seqlens,
681 )
682 dq, dk, dv, dg, dbeta = chunk_gated_delta_rule_bwd_finish_saved(
683 q,
684 k,
685 q_norm,
686 k_norm,
687 v,
688 g_cumsum,
689 beta,
690 A,
691 w,
692 h,
693 v_new,
694 dv,
695 grad_output,
696 dh,
697 q_inv_norm,
698 k_inv_norm,
699 scale,
700 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
701 chunk_size=chunk_size,
702 cu_seqlens=cu_seqlens,
703 )
704 return dq, dk, dv, dg, dbeta, dh0
707class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
708 """Autograd wrapper for the Triton-Ascend chunk Gated Delta Rule."""
710 @staticmethod
711 @input_guard
712 @autocast_custom_fwd
713 def forward(
714 ctx,
715 q: torch.Tensor,
716 k: torch.Tensor,
717 v: torch.Tensor,
718 g: torch.Tensor,
719 beta: torch.Tensor,
720 scale: float,
721 initial_state: torch.Tensor,
722 output_final_state: bool,
723 cu_seqlens: Optional[torch.LongTensor] = None,
724 use_qk_l2norm_in_kernel: bool = False,
725 chunk_size: int = 64,
726 ):
727 g, output, matrix_a, final_state = chunk_gated_delta_rule_fwd(
728 q=q,
729 k=k,
730 v=v,
731 g=g,
732 beta=beta,
733 scale=scale,
734 initial_state=initial_state,
735 output_final_state=output_final_state,
736 cu_seqlens=cu_seqlens,
737 chunk_size=chunk_size,
738 )
740 saved_initial_state = initial_state if initial_state is not None else q.new_empty(0)
741 saved_cu_seqlens = cu_seqlens if cu_seqlens is not None else q.new_empty(0, dtype=torch.long)
742 ctx.save_for_backward(q, k, v, g, beta, matrix_a, saved_initial_state, saved_cu_seqlens)
743 ctx.has_initial_state = initial_state is not None
744 ctx.has_cu_seqlens = cu_seqlens is not None
745 ctx.scale = scale
746 ctx.chunk_size = chunk_size
747 return output.to(q.dtype), final_state
749 @staticmethod
750 @input_guard
751 @autocast_custom_bwd
752 def backward(
753 ctx,
754 do: torch.Tensor,
755 dht: torch.Tensor
756 ):
757 q, k, v, g, beta, matrix_a, initial_state, cu_seqlens = ctx.saved_tensors
758 if not ctx.has_initial_state:
759 initial_state = None
760 if not ctx.has_cu_seqlens:
761 cu_seqlens = None
762 dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
763 q=q,
764 k=k,
765 v=v,
766 g=g,
767 beta=beta,
768 A=matrix_a,
769 scale=ctx.scale,
770 initial_state=initial_state,
771 do=do,
772 dht=dht,
773 cu_seqlens=cu_seqlens,
774 chunk_size=ctx.chunk_size,
775 )
776 return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), None, dh0, None, None, None, None
779@torch.compiler.disable
780def chunk_gated_delta_rule(
781 q: torch.Tensor,
782 k: torch.Tensor,
783 v: torch.Tensor,
784 g: torch.Tensor,
785 beta: torch.Tensor,
786 scale: float = None,
787 initial_state: torch.Tensor = None,
788 output_final_state: bool = False,
789 use_qk_l2norm_in_kernel: bool = False,
790 cu_seqlens: Optional[torch.LongTensor] = None,
791 chunk_size: int = 64,
792 head_first: bool = False,
793):
794 r"""
795 Args:
796 q (torch.Tensor):
797 queries of shape `[B, T, H, K]`.
798 k (torch.Tensor):
799 keys of shape `[B, T, H, K]`.
800 v (torch.Tensor):
801 values of shape `[B, T, H, V]`.
802 g (torch.Tensor):
803 (forget) gating tensor (in log space!) of shape `[B, T, H]`.
804 beta (torch.Tensor):
805 betas of shape `[B, T, H]`.
806 scale (Optional[float]):
807 Scale factor for the RetNet attention scores.
808 If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
809 initial_state (Optional[torch.Tensor]):
810 Initial state of shape `[N, H, K, V]` for `N` input sequences.
811 For equal-length input sequences, `N` equals the batch size `B`.
812 Default: `None`.
813 output_final_state (Optional[bool]):
814 Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.
815 use_qk_l2norm_in_kernel (bool):
816 Whether to apply L2norm to the q/k tensor internally. Default: `False`.
817 cu_seqlens (torch.LongTensor):
818 Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
819 consistent with the FlashAttention API.
820 head_first (Optional[bool]):
821 Whether the inputs are in the head-first format. Default: `False`.
822 This argument has been deprecated.
823 Returns:
824 o (torch.Tensor):
825 Outputs of shape `[B, T, H, V]`.
826 final_state (torch.Tensor):
827 Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.
829 Examples::
830 >>> import torch
831 >>> import torch.nn.functional as F
832 >>> from einops import rearrange
833 >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
834 # inputs with equal lengths
835 >>> B, T, H, K, V = 4, 2048, 4, 512, 512
836 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
837 >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
838 >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
839 >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()
840 >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))
841 >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')
842 >>> o, ht = chunk_gated_delta_rule(
843 q, k, v, g, beta,
844 initial_state=h0,
845 output_final_state=True
846 )
847 # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
848 >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
849 # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
850 >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
851 >>> o, ht = chunk_gated_delta_rule(
852 q, k, v, g, beta,
853 initial_state=h0,
854 output_final_state=True,
855 cu_seqlens=cu_seqlens
856 )
857 """
858 if q.dtype != k.dtype or k.dtype != v.dtype:
859 raise ValueError(
860 f"q current type is {q.dtype}, k current type is {k.dtype}, "
861 f"v current type is {v.dtype}, they should be equal"
862 )
863 if q.dtype == torch.float32:
864 raise ValueError(
865 "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
866 )
867 if len(beta.shape) != 3:
868 raise ValueError(
869 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."
870 )
872 if head_first:
873 warnings.warn(
874 "head_first is deprecated and will be removed in a future version. "
875 "Please use head_first=False for now instead."
876 )
877 if not head_first and q.shape[1] < q.shape[2]:
878 warnings.warn(
879 f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
880 "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
881 "when head_first=False was specified. "
882 "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
883 )
884 if cu_seqlens is not None:
885 if q.shape[0] != 1:
886 raise ValueError(
887 f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
888 f"Please flatten variable-length inputs before processing."
889 )
890 if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
891 raise ValueError(
892 f"The number of initial states is expected to be equal to the number of input sequences, "
893 f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}."
894 )
895 if scale is None:
896 scale = k.shape[-1] ** -0.5
898 if use_qk_l2norm_in_kernel:
899 q, _ = _l2norm(q)
900 k, _ = _l2norm(k)
902 output, final_state = ChunkGatedDeltaRuleFunction.apply(
903 q,
904 k,
905 v,
906 g,
907 beta,
908 scale,
909 initial_state,
910 output_final_state,
911 cu_seqlens,
912 False,
913 chunk_size,
914 )
915 return output, final_state