@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 525 lines 17.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Copyright (c) 2024, Tri Dao, Albert Gu.
5 # Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/ssd_chunk_scan.py
6
7 # ruff: noqa: E501,SIM102
8
9 from packaging import version
10
11 from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
12 from vllm.triton_utils import tl, triton
13
14 TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0")
15
16
17 @triton.autotune(
18 configs=[
19 # =================================================================
20 # Higher warp count configs for better latency hiding
21 # More warps = more instructions in flight = better memory latency hiding
22 # =================================================================
23 triton.Config(
24 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
25 num_stages=2,
26 num_warps=8,
27 ),
28 triton.Config(
29 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
30 num_stages=2,
31 num_warps=8,
32 ),
33 triton.Config(
34 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
35 num_stages=2,
36 num_warps=8,
37 ),
38 # Smaller tiles with more stages for software pipelining
39 triton.Config(
40 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
41 num_stages=3,
42 num_warps=4,
43 ),
44 triton.Config(
45 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64},
46 num_stages=2,
47 num_warps=4,
48 ),
49 # =================================================================
50 # Low register pressure configs (num_stages=1) for large dstate
51 # =================================================================
52 triton.Config(
53 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},
54 num_stages=1,
55 num_warps=4,
56 ),
57 triton.Config(
58 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
59 num_stages=1,
60 num_warps=4,
61 ),
62 triton.Config(
63 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
64 num_stages=1,
65 num_warps=4,
66 ),
67 triton.Config(
68 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
69 num_stages=1,
70 num_warps=4,
71 ),
72 # num_stages=2 configs - moderate register pressure
73 triton.Config(
74 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},
75 num_stages=2,
76 num_warps=4,
77 ),
78 triton.Config(
79 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
80 num_stages=2,
81 num_warps=4,
82 ),
83 triton.Config(
84 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
85 num_stages=2,
86 num_warps=4,
87 ),
88 # Original configs for larger dstate values
89 triton.Config(
90 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64},
91 num_stages=3,
92 num_warps=8,
93 ),
94 triton.Config(
95 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32},
96 num_stages=4,
97 num_warps=4,
98 ),
99 triton.Config(
100 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},
101 num_stages=4,
102 num_warps=4,
103 ),
104 triton.Config(
105 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
106 num_stages=4,
107 num_warps=4,
108 ),
109 triton.Config(
110 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},
111 num_stages=4,
112 num_warps=4,
113 ),
114 triton.Config(
115 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},
116 num_stages=4,
117 num_warps=4,
118 ),
119 triton.Config(
120 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64},
121 num_stages=4,
122 num_warps=4,
123 ),
124 triton.Config(
125 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
126 num_stages=4,
127 num_warps=4,
128 ),
129 triton.Config(
130 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
131 num_stages=5,
132 num_warps=2,
133 ),
134 triton.Config(
135 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
136 num_stages=5,
137 num_warps=2,
138 ),
139 triton.Config(
140 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
141 num_stages=4,
142 num_warps=2,
143 ),
144 ],
145 key=["chunk_size", "hdim", "dstate", "IS_CAUSAL"],
146 )
147 @triton.jit
148 def _chunk_scan_fwd_kernel(
149 # Pointers to matrices
150 cb_ptr,
151 x_ptr,
152 z_ptr,
153 out_ptr,
154 dt_ptr,
155 dA_cumsum_ptr,
156 seq_idx_ptr,
157 C_ptr,
158 states_ptr,
159 D_ptr,
160 initstates_ptr,
161 cu_chunk_seqlens_ptr,
162 # Matrix dimensions
163 chunk_size: tl.constexpr,
164 hdim: tl.constexpr,
165 dstate: tl.constexpr,
166 nheads_ngroups_ratio: tl.constexpr,
167 # Strides
168 stride_cb_chunk: tl.int64,
169 stride_cb_head: tl.int64,
170 stride_cb_csize_m: tl.int64,
171 stride_cb_csize_k: tl.constexpr,
172 stride_x_seqlen: tl.int64,
173 stride_x_head: tl.int64,
174 stride_x_hdim: tl.constexpr,
175 stride_z_seqlen: tl.int64,
176 stride_z_head: tl.int64,
177 stride_z_hdim: tl.constexpr,
178 stride_out_seqlen: tl.int64,
179 stride_out_head: tl.int64,
180 stride_out_hdim: tl.constexpr,
181 stride_dt_chunk: tl.int64,
182 stride_dt_head: tl.int64,
183 stride_dt_csize: tl.constexpr,
184 stride_dA_cs_chunk: tl.int64,
185 stride_dA_cs_head: tl.int64,
186 stride_dA_cs_csize: tl.constexpr,
187 stride_seq_idx_chunk: tl.constexpr,
188 stride_C_seqlen: tl.int64,
189 stride_C_head: tl.int64,
190 stride_C_dstate: tl.constexpr,
191 stride_states_chunk: tl.int64,
192 stride_states_head: tl.int64,
193 stride_states_hdim: tl.int64,
194 stride_states_dstate: tl.constexpr,
195 stride_init_states_batch: tl.int64,
196 stride_init_states_head: tl.int64,
197 stride_init_states_hdim: tl.int64,
198 stride_init_states_dstate: tl.constexpr,
199 stride_D_head: tl.constexpr,
200 # Meta-parameters
201 IS_CAUSAL: tl.constexpr,
202 HAS_D: tl.constexpr,
203 D_HAS_HDIM: tl.constexpr,
204 HAS_Z: tl.constexpr,
205 BLOCK_SIZE_M: tl.constexpr,
206 BLOCK_SIZE_N: tl.constexpr,
207 BLOCK_SIZE_K: tl.constexpr,
208 BLOCK_SIZE_DSTATE: tl.constexpr,
209 IS_TRITON_22: tl.constexpr,
210 HAS_INITSTATES: tl.constexpr,
211 ):
212 pid_c = tl.program_id(axis=1).to(tl.int64)
213 pid_h = tl.program_id(axis=2)
214 num_pid_n = tl.cdiv(hdim, BLOCK_SIZE_N)
215 pid_m = tl.program_id(axis=0) // num_pid_n
216 pid_n = tl.program_id(axis=0) % num_pid_n
217 cb_ptr += pid_c * stride_cb_chunk + (pid_h // nheads_ngroups_ratio) * stride_cb_head
218 chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c)
219 chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1)
220 x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head
221 dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head
222 dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head
223 C_ptr += (
224 chunk_seqlen_start * stride_C_seqlen
225 + (pid_h // nheads_ngroups_ratio) * stride_C_head
226 )
227
228 # M-block offsets and prev states
229 # - logic in next block may override these if there is an active offset
230 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
231
232 seq_idx_ptr += pid_c * stride_seq_idx_chunk
233 seq_idx = tl.load(seq_idx_ptr)
234 seq_idx_prev = tl.load(
235 seq_idx_ptr - stride_seq_idx_chunk, mask=pid_c >= 1, other=-1
236 )
237
238 if HAS_INITSTATES and (seq_idx != seq_idx_prev):
239 prev_states_ptr = (
240 initstates_ptr
241 + seq_idx * stride_init_states_batch
242 + pid_h * stride_init_states_head
243 )
244 prev_states_hdim = stride_init_states_hdim
245 prev_states_dstate = stride_init_states_dstate
246 else:
247 prev_states_ptr = (
248 states_ptr + (pid_c - 1) * stride_states_chunk + pid_h * stride_states_head
249 )
250 prev_states_hdim = stride_states_hdim
251 prev_states_dstate = stride_states_dstate
252
253 chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start
254
255 offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
256 dA_cs_m = tl.load(
257 dA_cumsum_ptr + offs_m * stride_dA_cs_csize, mask=offs_m < chunk_size, other=0.0
258 ).to(tl.float32)
259
260 acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
261
262 offs_out_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
263 offs_out_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
264
265 # Faster to just do 1 iteration with larger BLOCK_SIZE_K, up to block size 128
266 offs_k_dstate = tl.arange(
267 0, BLOCK_SIZE_DSTATE if BLOCK_SIZE_DSTATE <= 128 else BLOCK_SIZE_K
268 )
269 C_ptrs = C_ptr + (
270 offs_m[:, None] * stride_C_seqlen + offs_k_dstate[None, :] * stride_C_dstate
271 )
272
273 scale_m = fast_exp(dA_cs_m)
274 if BLOCK_SIZE_DSTATE <= 128:
275 C = tl.load(
276 C_ptrs,
277 mask=(offs_m[:, None] < chunk_size_limit)
278 & (offs_k_dstate[None, :] < dstate),
279 other=0.0,
280 )
281
282 if not HAS_INITSTATES and (seq_idx != seq_idx_prev):
283 # if no init states AND starting a new sequence, we need zeros
284 prev_states = tl.zeros(
285 (BLOCK_SIZE_DSTATE, BLOCK_SIZE_N), dtype=C_ptr.dtype.element_ty
286 )
287 else:
288 # otherwise read the previous state
289 prev_states_ptrs = (
290 prev_states_ptr
291 + offs_n[None, :] * prev_states_hdim
292 + offs_k_dstate[:, None] * prev_states_dstate
293 )
294 prev_states = tl.load(
295 prev_states_ptrs,
296 mask=(offs_k_dstate[:, None] < dstate) & (offs_n[None, :] < hdim),
297 other=0.0,
298 )
299 prev_states = prev_states.to(C_ptr.dtype.element_ty)
300
301 acc = tl.dot(C, prev_states) * scale_m[:, None]
302
303 else:
304 prev_states_ptrs = (
305 prev_states_ptr
306 + offs_n[None, :] * prev_states_hdim
307 + offs_k_dstate[:, None] * prev_states_dstate
308 )
309 for k in range(0, dstate, BLOCK_SIZE_K):
310 C = tl.load(
311 C_ptrs,
312 mask=(offs_m[:, None] < chunk_size_limit)
313 & (offs_k_dstate[None, :] < dstate - k),
314 other=0.0,
315 )
316 if not HAS_INITSTATES and (seq_idx != seq_idx_prev):
317 prev_states = tl.zeros(
318 (BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=C_ptr.dtype.element_ty
319 )
320 else:
321 prev_states = tl.load(
322 prev_states_ptrs,
323 mask=(offs_k_dstate[:, None] < dstate - k)
324 & (offs_n[None, :] < hdim),
325 other=0.0,
326 )
327 prev_states = prev_states.to(C_ptr.dtype.element_ty)
328 acc += tl.dot(C, prev_states)
329 C_ptrs += BLOCK_SIZE_K
330 prev_states_ptrs += BLOCK_SIZE_K
331 acc *= scale_m[:, None]
332
333 offs_k = tl.arange(0, BLOCK_SIZE_K)
334 cb_ptrs = cb_ptr + (
335 offs_m[:, None] * stride_cb_csize_m + offs_k[None, :] * stride_cb_csize_k
336 )
337 x_ptrs = x_ptr + (
338 offs_k[:, None] * stride_x_seqlen + offs_n[None, :] * stride_x_hdim
339 )
340 dt_ptrs = dt_ptr + offs_k * stride_dt_csize
341 dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize
342 K_MAX = (
343 chunk_size_limit
344 if not IS_CAUSAL
345 else min((pid_m + 1) * BLOCK_SIZE_M, chunk_size_limit)
346 )
347 for k in range(0, K_MAX, BLOCK_SIZE_K):
348 cb = tl.load(
349 cb_ptrs,
350 mask=(offs_m[:, None] < chunk_size) & (offs_k[None, :] < chunk_size - k),
351 other=0.0,
352 ).to(tl.float32)
353 dA_cs_k = tl.load(dA_cumsum_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(
354 tl.float32
355 )
356 # If there's seq_idx, we already set cb[i, j] = 0 for seq_idx[i] != seq_idx[j].
357 # So we don't need masking wrt seq_idx here.
358 cb *= fast_exp(tl.minimum(dA_cs_m[:, None] - dA_cs_k[None, :], 0.0))
359 dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32)
360 cb *= dt_k
361 if IS_CAUSAL:
362 mask = offs_m[:, None] >= k + offs_k[None, :]
363 cb = tl.where(mask, cb, 0.0)
364 cb = cb.to(x_ptr.dtype.element_ty)
365 x = tl.load(
366 x_ptrs,
367 mask=(offs_k[:, None] < chunk_size_limit - k) & (offs_n[None, :] < hdim),
368 other=0.0,
369 )
370 acc += tl.dot(cb, x)
371 cb_ptrs += BLOCK_SIZE_K * stride_cb_csize_k
372 x_ptrs += BLOCK_SIZE_K * stride_x_seqlen
373 dt_ptrs += BLOCK_SIZE_K * stride_dt_csize
374 dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize
375
376 offs_out_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
377 offs_out_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
378
379 if HAS_D:
380 if D_HAS_HDIM:
381 D = tl.load(
382 D_ptr + pid_h * stride_D_head + offs_n, mask=offs_n < hdim, other=0.0
383 ).to(tl.float32)
384 else:
385 D = tl.load(D_ptr + pid_h * stride_D_head).to(tl.float32)
386 x_residual = tl.load(
387 x_ptr
388 + (offs_m[:, None] * stride_x_seqlen + offs_n[None, :] * stride_x_hdim),
389 mask=(offs_m[:, None] < chunk_size_limit) & (offs_n[None, :] < hdim),
390 other=0.0,
391 ).to(tl.float32)
392 acc += x_residual * D
393
394 if HAS_Z:
395 z_ptr += chunk_seqlen_start * stride_z_seqlen + pid_h * stride_z_head
396 z_ptrs = z_ptr + (
397 stride_z_seqlen * offs_out_m[:, None] + stride_z_hdim * offs_out_n[None, :]
398 )
399 z = tl.load(
400 z_ptrs,
401 mask=(offs_out_m[:, None] < chunk_size_limit)
402 & (offs_out_n[None, :] < hdim),
403 other=0.0,
404 ).to(tl.float32)
405 acc *= z * tl.sigmoid(z)
406
407 out_ptr += chunk_seqlen_start * stride_out_seqlen + pid_h * stride_out_head
408 out_ptrs = out_ptr + (
409 stride_out_seqlen * offs_out_m[:, None] + offs_out_n[None, :] * stride_out_hdim
410 )
411 tl.store(
412 out_ptrs,
413 acc,
414 mask=(offs_out_m[:, None] < chunk_size_limit) & (offs_out_n[None, :] < hdim),
415 )
416
417
418 def _chunk_scan_fwd(
419 cb,
420 x,
421 dt,
422 dA_cumsum,
423 C,
424 states,
425 cu_chunk_seqlens,
426 out,
427 seq_idx,
428 D=None,
429 z=None,
430 initial_states=None,
431 ):
432 assert seq_idx is not None, "this implementation requires seq_idx"
433
434 seqlen, nheads, headdim = x.shape
435 _, nchunks, chunk_size = dt.shape
436 _, ngroups, dstate = C.shape
437 assert nheads % ngroups == 0
438 assert C.shape == (seqlen, ngroups, dstate)
439 assert cb.shape == (nchunks, ngroups, chunk_size, chunk_size)
440 if D is not None:
441 assert D.shape == (nheads, headdim) or D.shape == (nheads,)
442 if z is not None:
443 assert z.shape == x.shape
444 assert dt.shape == (nheads, nchunks, chunk_size)
445 assert dA_cumsum.shape == (nheads, nchunks, chunk_size)
446 assert states.shape == (nchunks, nheads, headdim, dstate)
447 assert seq_idx.shape == (nchunks,)
448
449 grid = lambda META: (
450 triton.cdiv(chunk_size, META["BLOCK_SIZE_M"])
451 * triton.cdiv(headdim, META["BLOCK_SIZE_N"]),
452 nchunks,
453 nheads,
454 )
455
456 z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0)
457 initial_states_strides = (
458 (
459 initial_states.stride(0),
460 initial_states.stride(1),
461 initial_states.stride(2),
462 initial_states.stride(3),
463 )
464 if initial_states is not None
465 else (0, 0, 0, 0)
466 )
467
468 _chunk_scan_fwd_kernel[grid](
469 cb_ptr=cb,
470 x_ptr=x,
471 z_ptr=z,
472 out_ptr=out,
473 dt_ptr=dt,
474 dA_cumsum_ptr=dA_cumsum,
475 seq_idx_ptr=seq_idx,
476 C_ptr=C,
477 states_ptr=states,
478 D_ptr=D,
479 initstates_ptr=initial_states,
480 cu_chunk_seqlens_ptr=cu_chunk_seqlens,
481 chunk_size=chunk_size,
482 hdim=headdim,
483 dstate=dstate,
484 nheads_ngroups_ratio=nheads // ngroups,
485 stride_cb_chunk=cb.stride(0),
486 stride_cb_head=cb.stride(1),
487 stride_cb_csize_m=cb.stride(2),
488 stride_cb_csize_k=cb.stride(3),
489 stride_x_seqlen=x.stride(0),
490 stride_x_head=x.stride(1),
491 stride_x_hdim=x.stride(2),
492 stride_z_seqlen=z_strides[0],
493 stride_z_head=z_strides[1],
494 stride_z_hdim=z_strides[2],
495 stride_out_seqlen=out.stride(0),
496 stride_out_head=out.stride(1),
497 stride_out_hdim=out.stride(2),
498 stride_dt_chunk=dt.stride(1),
499 stride_dt_head=dt.stride(0),
500 stride_dt_csize=dt.stride(2),
501 stride_dA_cs_chunk=dA_cumsum.stride(1),
502 stride_dA_cs_head=dA_cumsum.stride(0),
503 stride_dA_cs_csize=dA_cumsum.stride(2),
504 stride_seq_idx_chunk=seq_idx.stride(0),
505 stride_C_seqlen=C.stride(0),
506 stride_C_head=C.stride(1),
507 stride_C_dstate=C.stride(2),
508 stride_states_chunk=states.stride(0),
509 stride_states_head=states.stride(1),
510 stride_states_hdim=states.stride(2),
511 stride_states_dstate=states.stride(3),
512 stride_init_states_batch=initial_states_strides[0],
513 stride_init_states_head=initial_states_strides[1],
514 stride_init_states_hdim=initial_states_strides[2],
515 stride_init_states_dstate=initial_states_strides[3],
516 stride_D_head=D.stride(0) if D is not None else 0,
517 IS_CAUSAL=True,
518 HAS_D=D is not None,
519 D_HAS_HDIM=D.dim() == 2 if D is not None else True,
520 HAS_Z=z is not None,
521 BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
522 IS_TRITON_22=TRITON_22,
523 HAS_INITSTATES=initial_states is not None,
524 )
525 return