@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 382 lines 12.9 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 vllm.triton_utils import tl, triton
14
15 from .index import prepare_chunk_indices, prepare_chunk_offsets
16 from .op import exp, exp2
17 from .utils import FLA_CHUNK_SIZE, use_cuda_graph
18
19 NUM_WARPS = [2, 4, 8, 16]
20 # Triton's AMD backend fails to lower this kernel with num_stages=4.
21 _CHUNK_DELTA_H_NUM_STAGES = [2, 3] if torch.version.hip else [2, 3, 4]
22
23
24 @triton.heuristics(
25 {
26 "USE_G": lambda args: args["g"] is not None,
27 "USE_GK": lambda args: args["gk"] is not None,
28 "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
29 "STORE_FINAL_STATE": lambda args: args["ht"] is not None,
30 "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None,
31 "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
32 }
33 )
34 @triton.autotune(
35 configs=[
36 triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages)
37 for num_warps in [2, 4]
38 for num_stages in _CHUNK_DELTA_H_NUM_STAGES
39 for BV in [32, 64]
40 ],
41 key=["H", "K", "V", "BT"],
42 use_cuda_graph=use_cuda_graph,
43 )
44 @triton.jit(do_not_specialize=["T"])
45 def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
46 k,
47 v,
48 w,
49 v_new,
50 g,
51 gk,
52 h,
53 h0,
54 ht,
55 cu_seqlens,
56 chunk_offsets,
57 T,
58 H: tl.constexpr,
59 Hg: tl.constexpr,
60 K: tl.constexpr,
61 V: tl.constexpr,
62 BT: tl.constexpr,
63 BV: tl.constexpr,
64 USE_G: tl.constexpr,
65 USE_GK: tl.constexpr,
66 USE_INITIAL_STATE: tl.constexpr,
67 STORE_FINAL_STATE: tl.constexpr,
68 SAVE_NEW_VALUE: tl.constexpr,
69 IS_VARLEN: tl.constexpr,
70 USE_EXP2: tl.constexpr,
71 ):
72 i_v, i_nh = tl.program_id(0), tl.program_id(1)
73 i_n, i_h = i_nh // H, i_nh % H
74 if IS_VARLEN:
75 bos, eos = (
76 tl.load(cu_seqlens + i_n).to(tl.int32),
77 tl.load(cu_seqlens + i_n + 1).to(tl.int32),
78 )
79 T = eos - bos
80 NT = tl.cdiv(T, BT)
81 boh = tl.load(chunk_offsets + i_n).to(tl.int32)
82 else:
83 bos, eos = i_n * T, i_n * T + T
84 NT = tl.cdiv(T, BT)
85 boh = i_n * NT
86
87 # [BV, BK]
88 b_h1 = tl.zeros([BV, 64], dtype=tl.float32)
89 if K > 64:
90 b_h2 = tl.zeros([BV, 64], dtype=tl.float32)
91 if K > 128:
92 b_h3 = tl.zeros([BV, 64], dtype=tl.float32)
93 if K > 192:
94 b_h4 = tl.zeros([BV, 64], dtype=tl.float32)
95
96 # calculate offset
97 h += ((boh * H + i_h) * V * K).to(tl.int64)
98 v += ((bos * H + i_h) * V).to(tl.int64)
99 k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64)
100 w += ((bos * H + i_h) * K).to(tl.int64)
101 if SAVE_NEW_VALUE:
102 v_new += ((bos * H + i_h) * V).to(tl.int64)
103 stride_v = H * V
104 stride_h = H * V * K
105 stride_k = Hg * K
106 stride_w = H * K
107 if USE_INITIAL_STATE:
108 h0 = h0 + i_nh * V * K
109 if STORE_FINAL_STATE:
110 ht = ht + i_nh * V * K
111
112 # load initial state
113 if USE_INITIAL_STATE:
114 p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
115 b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32)
116 if K > 64:
117 p_h0_2 = tl.make_block_ptr(
118 h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
119 )
120 b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32)
121 if K > 128:
122 p_h0_3 = tl.make_block_ptr(
123 h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
124 )
125 b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32)
126 if K > 192:
127 p_h0_4 = tl.make_block_ptr(
128 h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
129 )
130 b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32)
131
132 # main recurrence
133 for i_t in range(NT):
134 p_h1 = tl.make_block_ptr(
135 h + i_t.to(tl.int64) * stride_h,
136 (V, K),
137 (K, 1),
138 (i_v * BV, 0),
139 (BV, 64),
140 (1, 0),
141 )
142 tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1))
143 if K > 64:
144 p_h2 = tl.make_block_ptr(
145 h + i_t.to(tl.int64) * stride_h,
146 (V, K),
147 (K, 1),
148 (i_v * BV, 64),
149 (BV, 64),
150 (1, 0),
151 )
152 tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1))
153 if K > 128:
154 p_h3 = tl.make_block_ptr(
155 h + i_t.to(tl.int64) * stride_h,
156 (V, K),
157 (K, 1),
158 (i_v * BV, 128),
159 (BV, 64),
160 (1, 0),
161 )
162 tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1))
163 if K > 192:
164 p_h4 = tl.make_block_ptr(
165 h + i_t.to(tl.int64) * stride_h,
166 (V, K),
167 (K, 1),
168 (i_v * BV, 192),
169 (BV, 64),
170 (1, 0),
171 )
172 tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))
173
174 p_w = tl.make_block_ptr(
175 w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)
176 )
177 b_w = tl.load(p_w, boundary_check=(0, 1))
178 b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype))
179 if K > 64:
180 p_w = tl.make_block_ptr(
181 w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0)
182 )
183 b_w = tl.load(p_w, boundary_check=(0, 1))
184 b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype))
185 if K > 128:
186 p_w = tl.make_block_ptr(
187 w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0)
188 )
189 b_w = tl.load(p_w, boundary_check=(0, 1))
190 b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype))
191 if K > 192:
192 p_w = tl.make_block_ptr(
193 w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0)
194 )
195 b_w = tl.load(p_w, boundary_check=(0, 1))
196 b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype))
197 p_v = tl.make_block_ptr(
198 v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
199 )
200 b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v
201
202 if SAVE_NEW_VALUE:
203 p_v = tl.make_block_ptr(
204 v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
205 )
206 tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1))
207
208 last_idx = min((i_t.to(tl.int64) + 1) * BT, T) - 1
209 if USE_G:
210 m_t = (i_t.to(tl.int64) * BT + tl.arange(0, BT)) < T
211 b_g_last = tl.load(g + bos * H + last_idx * H + i_h)
212 p_g = tl.make_block_ptr(
213 g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)
214 )
215 b_g = tl.load(p_g, boundary_check=(0,))
216 if USE_EXP2:
217 b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None]
218 b_g_last = exp2(b_g_last)
219 else:
220 b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None]
221 b_g_last = exp(b_g_last)
222 b_h1 *= b_g_last
223 if K > 64:
224 b_h2 *= b_g_last
225 if K > 128:
226 b_h3 *= b_g_last
227 if K > 192:
228 b_h4 *= b_g_last
229
230 if USE_GK:
231 o_k1 = tl.arange(0, 64)
232 b_gk_last1 = tl.load(
233 gk + (bos + last_idx) * H * K + i_h * K + o_k1,
234 mask=(o_k1 < K),
235 other=0.0,
236 )
237 if USE_EXP2:
238 b_h1 *= exp2(b_gk_last1)[None, :]
239 else:
240 b_h1 *= exp(b_gk_last1)[None, :]
241 if K > 64:
242 o_k2 = 64 + o_k1
243 b_gk_last2 = tl.load(
244 gk + (bos + last_idx) * H * K + i_h * K + o_k2,
245 mask=(o_k2 < K),
246 other=0.0,
247 )
248 if USE_EXP2:
249 b_h2 *= exp2(b_gk_last2)[None, :]
250 else:
251 b_h2 *= exp(b_gk_last2)[None, :]
252 if K > 128:
253 o_k3 = 128 + o_k1
254 b_gk_last3 = tl.load(
255 gk + (bos + last_idx) * H * K + i_h * K + o_k3,
256 mask=(o_k3 < K),
257 other=0.0,
258 )
259 if USE_EXP2:
260 b_h3 *= exp2(b_gk_last3)[None, :]
261 else:
262 b_h3 *= exp(b_gk_last3)[None, :]
263 if K > 192:
264 o_k4 = 192 + o_k1
265 b_gk_last4 = tl.load(
266 gk + (bos + last_idx) * H * K + i_h * K + o_k4,
267 mask=(o_k4 < K),
268 other=0.0,
269 )
270 if USE_EXP2:
271 b_h4 *= exp2(b_gk_last4)[None, :]
272 else:
273 b_h4 *= exp(b_gk_last4)[None, :]
274 b_v = b_v.to(k.dtype.element_ty)
275
276 p_k = tl.make_block_ptr(
277 k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)
278 )
279 b_k = tl.load(p_k, boundary_check=(0, 1))
280 b_h1 += tl.trans(tl.dot(b_k, b_v))
281 if K > 64:
282 p_k = tl.make_block_ptr(
283 k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
284 )
285 b_k = tl.load(p_k, boundary_check=(0, 1))
286 b_h2 += tl.trans(tl.dot(b_k, b_v))
287 if K > 128:
288 p_k = tl.make_block_ptr(
289 k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)
290 )
291 b_k = tl.load(p_k, boundary_check=(0, 1))
292 b_h3 += tl.trans(tl.dot(b_k, b_v))
293 if K > 192:
294 p_k = tl.make_block_ptr(
295 k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)
296 )
297 b_k = tl.load(p_k, boundary_check=(0, 1))
298 b_h4 += tl.trans(tl.dot(b_k, b_v))
299 # epilogue
300 if STORE_FINAL_STATE:
301 p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
302 tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
303 if K > 64:
304 p_ht = tl.make_block_ptr(
305 ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
306 )
307 tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
308 if K > 128:
309 p_ht = tl.make_block_ptr(
310 ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
311 )
312 tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
313 if K > 192:
314 p_ht = tl.make_block_ptr(
315 ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
316 )
317 tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
318
319
320 def chunk_gated_delta_rule_fwd_h(
321 k: torch.Tensor,
322 w: torch.Tensor,
323 u: torch.Tensor,
324 g: torch.Tensor | None = None,
325 gk: torch.Tensor | None = None,
326 initial_state: torch.Tensor | None = None,
327 output_final_state: bool = False,
328 chunk_size: int = FLA_CHUNK_SIZE,
329 save_new_value: bool = True,
330 cu_seqlens: torch.Tensor | None = None,
331 chunk_indices: torch.Tensor | None = None,
332 chunk_offsets: torch.Tensor | None = None,
333 use_exp2: bool = False,
334 ) -> tuple[torch.Tensor, torch.Tensor]:
335 # This kernel is slightly different from fla to support Q/K with different head numbers.
336 # In fla, Q/K always have the same head number, so Hg is always equal to H.
337 B, T, Hg, K, V = *k.shape, u.shape[-1]
338 H = u.shape[-2]
339 BT = chunk_size
340
341 if chunk_indices is None and cu_seqlens is not None:
342 chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
343 # N: the actual number of sequences in the batch with either equal or variable lengths
344 if cu_seqlens is None:
345 N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
346 else:
347 N, NT = len(cu_seqlens) - 1, len(chunk_indices)
348 if chunk_offsets is None:
349 chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)
350 assert K <= 256, "current kernel does not support head dimension larger than 256."
351
352 h = k.new_empty(B, NT, H, V, K)
353 final_state = (
354 k.new_empty(N, H, V, K, dtype=torch.float32) if output_final_state else None
355 )
356
357 v_new = torch.empty_like(u) if save_new_value else None
358
359 def grid(meta):
360 return (triton.cdiv(V, meta["BV"]), N * H)
361
362 chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid](
363 k=k,
364 v=u,
365 w=w,
366 v_new=v_new,
367 g=g,
368 gk=gk,
369 h=h,
370 h0=initial_state,
371 ht=final_state,
372 cu_seqlens=cu_seqlens,
373 chunk_offsets=chunk_offsets,
374 T=T,
375 H=H,
376 Hg=Hg,
377 K=K,
378 V=V,
379 BT=BT,
380 USE_EXP2=use_exp2,
381 )
382 return h, v_new, final_state