@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 245 lines 8.51 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 #
5 # This file contains code copied from the flash-linear-attention project.
6 # The original source code was licensed under the MIT license and included
7 # the following copyright notice:
8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
9 # ruff: noqa: E501
10
11 import torch
12
13 from .chunk_delta_h import chunk_gated_delta_rule_fwd_h
14 from .chunk_o import chunk_fwd_o
15 from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
16 from .cumsum import chunk_local_cumsum
17 from .l2norm import l2norm_fwd
18 from .solve_tril import solve_tril
19 from .utils import FLA_CHUNK_SIZE, SUPPRESS_LEVEL, input_guard
20 from .wy_fast import recompute_w_u_fwd
21
22
23 def chunk_gated_delta_rule_fwd(
24 q: torch.Tensor,
25 k: torch.Tensor,
26 v: torch.Tensor,
27 g: torch.Tensor,
28 beta: torch.Tensor,
29 scale: float,
30 initial_state: torch.Tensor,
31 output_final_state: bool,
32 cu_seqlens: torch.Tensor | None = None,
33 chunk_indices: torch.Tensor | None = None,
34 chunk_offsets: torch.Tensor | None = None,
35 core_attn_out: torch.Tensor | None = None,
36 ):
37 g = chunk_local_cumsum(
38 g, chunk_size=FLA_CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
39 )
40 # obtain WY representation. u is actually the new v.
41 A = chunk_scaled_dot_kkt_fwd(
42 k=k,
43 beta=beta,
44 g=g,
45 cu_seqlens=cu_seqlens,
46 chunk_indices=chunk_indices,
47 output_dtype=torch.float32,
48 )
49 A = solve_tril(
50 A=A, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, output_dtype=k.dtype
51 )
52 w, u = recompute_w_u_fwd(
53 k=k,
54 v=v,
55 beta=beta,
56 A=A,
57 g_cumsum=g,
58 cu_seqlens=cu_seqlens,
59 chunk_indices=chunk_indices,
60 )
61 h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
62 k=k,
63 w=w,
64 u=u,
65 g=g,
66 initial_state=initial_state,
67 output_final_state=output_final_state,
68 cu_seqlens=cu_seqlens,
69 chunk_indices=chunk_indices,
70 chunk_offsets=chunk_offsets,
71 )
72 o = chunk_fwd_o(
73 q=q,
74 k=k,
75 v=v_new,
76 h=h,
77 g=g,
78 scale=scale,
79 cu_seqlens=cu_seqlens,
80 chunk_indices=chunk_indices,
81 core_attn_out=core_attn_out,
82 )
83 if SUPPRESS_LEVEL < 3:
84 return g, o, A, final_state, None, None, None
85 elif SUPPRESS_LEVEL >= 3:
86 return g, o, A, final_state, w, h, v_new
87
88
89 class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
90 @staticmethod
91 @input_guard
92 @torch.amp.custom_fwd(device_type="cuda")
93 def forward(
94 ctx,
95 q: torch.Tensor,
96 k: torch.Tensor,
97 v: torch.Tensor,
98 g: torch.Tensor,
99 beta: torch.Tensor,
100 scale: float,
101 initial_state: torch.Tensor,
102 output_final_state: bool,
103 cu_seqlens: torch.Tensor | None = None,
104 chunk_indices: torch.Tensor | None = None,
105 chunk_offsets: torch.Tensor | None = None,
106 use_qk_l2norm_in_kernel: bool = False,
107 core_attn_out: torch.Tensor | None = None,
108 ):
109 if use_qk_l2norm_in_kernel:
110 q = l2norm_fwd(q)
111 k = l2norm_fwd(k)
112
113 g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd(
114 q=q,
115 k=k,
116 v=v,
117 g=g,
118 beta=beta,
119 scale=scale,
120 initial_state=initial_state,
121 output_final_state=output_final_state,
122 cu_seqlens=cu_seqlens,
123 chunk_indices=chunk_indices,
124 chunk_offsets=chunk_offsets,
125 core_attn_out=core_attn_out,
126 )
127 ctx.scale = scale
128 ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
129 if core_attn_out is not None:
130 assert not torch.is_grad_enabled(), (
131 "core_attn_out buffer reuse is only supported for inference"
132 )
133 assert q.dtype == o.dtype, "Incompatible dtype for inplace computation"
134 return o.to(q.dtype), final_state
135
136
137 @torch.compiler.disable
138 def chunk_gated_delta_rule(
139 q: torch.Tensor,
140 k: torch.Tensor,
141 v: torch.Tensor,
142 g: torch.Tensor,
143 beta: torch.Tensor,
144 scale: float = None,
145 initial_state: torch.Tensor = None,
146 output_final_state: bool = False,
147 cu_seqlens: torch.Tensor | None = None,
148 chunk_indices: torch.Tensor | None = None,
149 chunk_offsets: torch.Tensor | None = None,
150 use_qk_l2norm_in_kernel: bool = False,
151 core_attn_out: torch.Tensor | None = None,
152 ):
153 r"""
154 Args:
155 q (torch.Tensor):
156 Queries of shape `[B, T, H, K]`.
157 k (torch.Tensor):
158 Keys of shape `[B, T, H, K]`.
159 v (torch.Tensor):
160 Values of shape `[B, T, H, V]`.
161 g (torch.Tensor):
162 (forget) Gating tensor (in log space!) of shape `[B, T, H]`.
163 beta (torch.Tensor):
164 Betas of shape `[B, T, H]`.
165 scale (Optional[int]):
166 Scale factor for the RetNet attention scores.
167 If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
168 initial_state (Optional[torch.Tensor]):
169 Initial state of shape `[N, H, V, K]` for `N` input sequences.
170 For equal-length input sequences, `N` equals the batch size `B`.
171 Default: `None`.
172 output_final_state (Optional[bool]):
173 Whether to output the final state of shape `[N, H, V, K]`. Default: `False`.
174 cu_seqlens (torch.Tensor):
175 Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
176 consistent with the FlashAttention API.
177 Returns:
178 o (torch.Tensor):
179 Outputs of shape `[B, T, H, V]`.
180 final_state (torch.Tensor):
181 Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`.
182
183 Examples::
184 >>> import torch
185 >>> import torch.nn.functional as F
186 >>> from einops import rearrange
187 >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
188 # inputs with equal lengths
189 >>> B, T, H, K, V = 4, 2048, 4, 512, 512
190 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
191 >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
192 >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
193 >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()
194 >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))
195 >>> h0 = torch.randn(B, H, V, K, dtype=torch.bfloat16, device='cuda')
196 >>> o, ht = chunk_gated_delta_rule(
197 q, k, v, g, beta,
198 initial_state=h0,
199 output_final_state=True
200 )
201 # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
202 >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
203 # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
204 >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32)
205 >>> o_var, ht_var = chunk_gated_delta_rule(
206 q, k, v, g, beta,
207 initial_state=h0,
208 output_final_state=True,
209 cu_seqlens=cu_seqlens
210 )
211 """
212 assert q.dtype == k.dtype == v.dtype
213 assert q.dtype != torch.float32, (
214 "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
215 )
216 assert len(beta.shape) == 3, "beta must be of shape [B, T, H]."
217 if cu_seqlens is not None:
218 if q.shape[0] != 1:
219 raise ValueError(
220 f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
221 f"Please flatten variable-length inputs before processing."
222 )
223 if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
224 raise ValueError(
225 f"The number of initial states is expected to be equal to the number of input sequences, "
226 f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}."
227 )
228 if scale is None:
229 scale = k.shape[-1] ** -0.5
230 o, final_state = ChunkGatedDeltaRuleFunction.apply(
231 q,
232 k,
233 v,
234 g,
235 beta,
236 scale,
237 initial_state,
238 output_final_state,
239 cu_seqlens,
240 chunk_indices,
241 chunk_offsets,
242 use_qk_l2norm_in_kernel,
243 core_attn_out,
244 )
245 return o, final_state