@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 857 lines 27.7 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/selective_state_update.py
6
7 import functools
8 import json
9 import os
10 from contextlib import contextmanager
11 from typing import Any
12
13 import torch
14 from packaging import version
15
16 import vllm.envs as envs
17 from vllm import _custom_ops as ops
18 from vllm.logger import init_logger
19 from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
20 from vllm.platforms import current_platform
21 from vllm.triton_utils import HAS_TRITON, tl, triton
22 from vllm.utils.platform_utils import get_device_name_as_file_name
23 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
24
25 if current_platform.is_xpu():
26 from vllm._xpu_ops import xpu_ops
27
28 logger = init_logger(__name__)
29
30 TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0"))
31
32
33 # ---------------------------------------------------------------------------
34 # JSON config loading
35 # ---------------------------------------------------------------------------
36
37 _CONFIGS_DIR = os.path.join(
38 os.path.dirname(os.path.realpath(__file__)), "configs", "selective_state_update"
39 )
40
41
42 def get_ssm_config_file_name(
43 headdim: int, dstate: int, cache_dtype: str, device_name: str
44 ) -> str:
45 """Return the JSON filename for the given kernel shape.
46
47 Layout: ``configs/selective_state_update/
48 headdim=<H>,dstate=<D>,device_name=<dev>,cache_dtype=<dt>.json``.
49 """
50 return (
51 f"headdim={headdim},dstate={dstate},"
52 f"device_name={device_name},cache_dtype={cache_dtype}.json"
53 )
54
55
56 def get_ssm_device_name() -> str:
57 return get_device_name_as_file_name()
58
59
60 def _canonical_cache_dtype(cache_dtype: str) -> str:
61 """Canonical key for config lookup. bf16 and fp16 share the same tuned
62 configs because the kernel only sees bit width when accessing state."""
63 return "float16" if cache_dtype == "bfloat16" else cache_dtype
64
65
66 @functools.cache
67 def get_ssm_configs(
68 headdim: int, dstate: int, cache_dtype: str
69 ) -> dict[int, Any] | None:
70 """Return tuned (BLOCK_SIZE_M, num_warps) configs for *selective_state_update*
71 keyed by ``effective_batch = batch * nheads``, or ``None`` if no config
72 file is found for the (headdim, dstate, cache_dtype, device) combination.
73
74 They can be generated with:
75 benchmarks/kernels/benchmark_selective_state_update.py --save-configs
76 """
77 cache_dtype = _canonical_cache_dtype(cache_dtype)
78 device_name = get_ssm_device_name()
79 json_file_name = get_ssm_config_file_name(headdim, dstate, cache_dtype, device_name)
80
81 config_file_paths: list[str] = []
82
83 # User-supplied override
84 user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER
85 if user_defined_config_folder is not None:
86 config_file_paths.append(
87 os.path.join(user_defined_config_folder, json_file_name)
88 )
89
90 # Bundled default
91 config_file_paths.append(os.path.join(_CONFIGS_DIR, json_file_name))
92
93 for path in config_file_paths:
94 if os.path.exists(path):
95 with open(path) as f:
96 logger.info_once(
97 "Using SSM config from %s for selective_state_update.",
98 path,
99 scope="global",
100 )
101 raw = json.load(f)
102 if isinstance(raw, dict):
103 # triton_version included in the config file only for reference
104 raw.pop("triton_version", None)
105 return {int(k): v for k, v in raw.items() if k.isdigit()}
106
107 logger.warning_once(
108 "Using default Mamba SSU config. Performance might be sub-optimal! "
109 "Config file not found at %s",
110 ", ".join(config_file_paths),
111 )
112 return None
113
114
115 def _get_default_ssm_launch_config(
116 dstate: int,
117 is_blackwell: bool,
118 ) -> tuple[int, int]:
119 """Hard-coded fallback heuristic used when no tuned config is available."""
120 BLOCK_SIZE_M, num_warps = 4, 8
121 if dstate <= 16:
122 BLOCK_SIZE_M, num_warps = 32, 4
123 elif dstate <= 32:
124 BLOCK_SIZE_M, num_warps = 16, 4
125 elif dstate <= 64:
126 BLOCK_SIZE_M, num_warps = 8, 4
127 else:
128 if is_blackwell:
129 BLOCK_SIZE_M, num_warps = 32, 8
130 elif dstate <= 128:
131 BLOCK_SIZE_M, num_warps = 4, 4
132 return BLOCK_SIZE_M, num_warps
133
134
135 @functools.cache
136 def _try_get_optimal_ssm_config_cached(
137 headdim: int,
138 dstate: int,
139 batch: int,
140 nheads: int,
141 cache_dtype: str,
142 is_blackwell: bool,
143 ) -> tuple[int, int]:
144 """Cached resolution. See :func:`try_get_optimal_ssm_config`."""
145 effective_batch = batch * nheads
146 configs = get_ssm_configs(headdim, dstate, cache_dtype)
147 if configs:
148 # Pick the closest effective_batch in the tuned grid (MoE strategy).
149 closest = min(configs.keys(), key=lambda x: abs(x - effective_batch))
150 cfg = configs[closest]
151 return cfg["BLOCK_SIZE_M"], cfg["num_warps"]
152
153 return _get_default_ssm_launch_config(dstate, is_blackwell)
154
155
156 # Override hook for benchmarks/tests, see `override_ssm_config`.
157 _ssm_config_override: tuple[int, int] | None = None
158
159
160 @contextmanager
161 def override_ssm_config(config: tuple[int, int]):
162 """Pin ``try_get_optimal_ssm_config`` to ``config`` for the duration of
163 the context. Used by the tuning benchmark to time specific configs."""
164 global _ssm_config_override
165 prev = _ssm_config_override
166 _ssm_config_override = config
167 try:
168 yield
169 finally:
170 _ssm_config_override = prev
171
172
173 def try_get_optimal_ssm_config(
174 headdim: int,
175 dstate: int,
176 batch: int,
177 nheads: int,
178 cache_dtype: str,
179 is_blackwell: bool,
180 ) -> tuple[int, int]:
181 """Return (BLOCK_SIZE_M, num_warps) for the given kernel shape.
182
183 Tuning is keyed on ``effective_batch = batch * nheads`` (the kernel grid
184 scales with the product), so configs transfer across (model, TP) combos
185 sharing ``(headdim, dstate, cache_dtype)``.
186 """
187 if _ssm_config_override is not None:
188 return _ssm_config_override
189 return _try_get_optimal_ssm_config_cached(
190 headdim, dstate, batch, nheads, cache_dtype, is_blackwell
191 )
192
193
194 if TRITON3:
195
196 @triton.jit
197 def softplus(dt):
198 dt = tl.where(dt <= 20.0, tl.math.log(tl.math.exp(dt) + 1), dt)
199 return dt
200 else:
201
202 @triton.jit
203 def softplus(dt):
204 dt = tl.where(dt <= 20.0, tl.math.log1p(tl.exp(dt)), dt)
205 return dt
206
207
208 @triton.jit
209 def convert_rs_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor:
210 y = tl.inline_asm_elementwise(
211 asm="""{
212 cvt.rs.f16x2.f32 $0, $2, $1, $3;
213 }""",
214 constraints="=r,r,r,r,r",
215 args=(x, rand),
216 dtype=tl.float16,
217 is_pure=True,
218 pack=2,
219 )
220 return y
221
222
223 @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None})
224 @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None})
225 @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})
226 @triton.heuristics(
227 {
228 "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]
229 is not None
230 }
231 )
232 @triton.heuristics(
233 {"IS_SPEC_DECODING": lambda args: args["num_accepted_tokens_ptr"] is not None}
234 )
235 @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens_ptr"] is not None})
236 @triton.heuristics(
237 {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}
238 )
239 @triton.jit(do_not_specialize=["N"])
240 def _selective_scan_update_kernel(
241 # Pointers to matrices
242 state_ptr,
243 rand_seed_ptr,
244 x_ptr,
245 dt_ptr,
246 dt_bias_ptr,
247 A_ptr,
248 B_ptr,
249 C_ptr,
250 D_ptr,
251 z_ptr,
252 out_ptr,
253 state_batch_indices_ptr,
254 dst_state_batch_indices_ptr,
255 null_block_id,
256 num_accepted_tokens_ptr,
257 cu_seqlens_ptr,
258 # Matrix dimensions
259 N,
260 nheads,
261 dim,
262 dstate,
263 nheads_ngroups_ratio,
264 # Strides
265 stride_state_batch,
266 stride_state_head,
267 stride_state_dim,
268 stride_state_dstate,
269 stride_x_batch,
270 stride_x_head,
271 stride_x_dim,
272 stride_dt_batch,
273 stride_dt_head,
274 stride_dt_dim,
275 stride_dt_bias_head,
276 stride_dt_bias_dim,
277 stride_A_head,
278 stride_A_dim,
279 stride_A_dstate,
280 stride_B_batch,
281 stride_B_group,
282 stride_B_dstate,
283 stride_C_batch,
284 stride_C_group,
285 stride_C_dstate,
286 stride_D_head,
287 stride_D_dim,
288 stride_z_batch,
289 stride_z_head,
290 stride_z_dim,
291 stride_out_batch,
292 stride_out_head,
293 stride_out_dim,
294 stride_state_indices_batch,
295 stride_state_indices_T,
296 stride_dst_state_indices_batch,
297 stride_dst_state_indices_T,
298 # Meta-parameters
299 DT_SOFTPLUS: tl.constexpr,
300 TIE_HDIM: tl.constexpr,
301 BLOCK_SIZE_M: tl.constexpr,
302 HAS_DT_BIAS: tl.constexpr,
303 HAS_D: tl.constexpr,
304 HAS_Z: tl.constexpr,
305 HAS_STATE_BATCH_INDICES: tl.constexpr,
306 IS_SPEC_DECODING: tl.constexpr,
307 IS_VARLEN: tl.constexpr,
308 BLOCK_SIZE_DSTATE: tl.constexpr,
309 USE_RS_ROUNDING: tl.constexpr,
310 PHILOX_ROUNDS: tl.constexpr,
311 ):
312 pid_m = tl.program_id(axis=0)
313 pid_b = tl.program_id(axis=1)
314 pid_h = tl.program_id(axis=2)
315
316 if IS_VARLEN:
317 bos = tl.load(cu_seqlens_ptr + pid_b).to(tl.int64)
318 eos = tl.load(cu_seqlens_ptr + pid_b + 1).to(tl.int64)
319 seq_len = eos - bos
320
321 if seq_len == 0:
322 return
323 else:
324 bos = pid_b
325 seq_len = 1
326
327 state_ptr_base = state_ptr
328
329 # If HAS_STATE_BATCH_INDICES is true, then the ssm state's batch coordinate
330 # is taken from the state_batch_indices_ptr Otherwise, the state coordinate
331 # is the same as the batch id.
332 if HAS_STATE_BATCH_INDICES:
333 if IS_SPEC_DECODING:
334 num_accepted = tl.load(num_accepted_tokens_ptr + pid_b).to(tl.int64)
335 init_token_idx = tl.maximum(num_accepted - 1, 0)
336 else:
337 init_token_idx = 0
338
339 dst_state_batch_indices_ptr += pid_b * stride_dst_state_indices_batch
340 if not IS_SPEC_DECODING:
341 dst_state_batch_idx = tl.load(
342 dst_state_batch_indices_ptr
343 + init_token_idx * stride_dst_state_indices_T
344 ).to(tl.int64)
345 dst_state_ptr = state_ptr + (
346 dst_state_batch_idx * stride_state_batch + pid_h * stride_state_head
347 )
348
349 state_batch_indices_ptr += (
350 pid_b * stride_state_indices_batch + init_token_idx * stride_state_indices_T
351 )
352 state_batch_idx = tl.load(state_batch_indices_ptr).to(tl.int64)
353 state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head
354 else:
355 dst_state_ptr = (
356 state_ptr + pid_b * stride_state_batch + pid_h * stride_state_head
357 )
358 state_ptr += pid_b * stride_state_batch + pid_h * stride_state_head
359
360 x_ptr += bos * stride_x_batch + pid_h * stride_x_head
361 dt_ptr += bos * stride_dt_batch + pid_h * stride_dt_head
362 if HAS_DT_BIAS:
363 dt_bias_ptr += pid_h * stride_dt_bias_head
364 A_ptr += pid_h * stride_A_head
365 B_ptr += bos * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group
366 C_ptr += bos * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group
367 if HAS_Z:
368 z_ptr += bos * stride_z_batch + pid_h * stride_z_head
369 out_ptr += bos * stride_out_batch + pid_h * stride_out_head
370
371 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
372 offs_n = tl.arange(0, BLOCK_SIZE_DSTATE)
373 state_ptrs = state_ptr + (
374 offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate
375 )
376 if not IS_SPEC_DECODING:
377 dst_state_ptrs = dst_state_ptr + (
378 offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate
379 )
380
381 mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)
382 if HAS_STATE_BATCH_INDICES:
383 mask &= state_batch_idx != null_block_id
384 state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32)
385
386 if HAS_DT_BIAS:
387 dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim
388 if HAS_D:
389 D_ptr += pid_h * stride_D_head
390 D_ptrs = D_ptr + offs_m * stride_D_dim
391 A_ptrs = A_ptr + offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate
392
393 for i_t in range(seq_len):
394 x_ptrs = x_ptr + offs_m * stride_x_dim
395 dt_ptrs = dt_ptr + offs_m * stride_dt_dim
396 B_ptrs = B_ptr + offs_n * stride_B_dstate
397 C_ptrs = C_ptr + offs_n * stride_C_dstate
398 if HAS_Z:
399 z_ptrs = z_ptr + offs_m * stride_z_dim
400 out_ptrs = out_ptr + offs_m * stride_out_dim
401
402 x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
403 if not TIE_HDIM:
404 dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
405 if HAS_DT_BIAS:
406 dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
407 if DT_SOFTPLUS:
408 dt = softplus(dt)
409 A = tl.load(
410 A_ptrs,
411 mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate),
412 other=0.0,
413 ).to(tl.float32)
414 dA = fast_exp(A * dt[:, None])
415 else:
416 dt = tl.load(dt_ptr).to(tl.float32)
417 if HAS_DT_BIAS:
418 dt += tl.load(dt_bias_ptr).to(tl.float32)
419 if DT_SOFTPLUS:
420 dt = softplus(dt)
421 A = tl.load(A_ptr).to(tl.float32)
422 dA = fast_exp(A * dt) # scalar, not a matrix
423
424 B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)
425 C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)
426 if HAS_D:
427 D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
428 if HAS_Z:
429 z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
430
431 dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt
432 state = state * dA + dB * x[:, None]
433
434 if IS_SPEC_DECODING:
435 dst_idx_ptr = dst_state_batch_indices_ptr + i_t * stride_dst_state_indices_T
436 token_dst_idx = tl.load(dst_idx_ptr).to(tl.int64)
437 if token_dst_idx != null_block_id:
438 token_dst_ptrs = (
439 state_ptr_base
440 + token_dst_idx * stride_state_batch
441 + pid_h * stride_state_head
442 + offs_m[:, None] * stride_state_dim
443 + offs_n[None, :] * stride_state_dstate
444 )
445 tl.store(
446 token_dst_ptrs, state.to(token_dst_ptrs.dtype.element_ty), mask=mask
447 )
448
449 out = tl.sum(state * C[None, :], axis=1)
450 if HAS_D:
451 out += x * D
452 if HAS_Z:
453 out *= z * tl.sigmoid(z)
454 tl.store(out_ptrs, out, mask=offs_m < dim)
455
456 x_ptr += stride_x_batch
457 dt_ptr += stride_dt_batch
458 B_ptr += stride_B_batch
459 C_ptr += stride_C_batch
460 out_ptr += stride_out_batch
461 if HAS_Z:
462 z_ptr += stride_z_batch
463
464 if not IS_SPEC_DECODING:
465 if USE_RS_ROUNDING:
466 # Load random seed
467 rand_seed = tl.load(rand_seed_ptr)
468 # Generate random offsets for each element in state
469 if HAS_STATE_BATCH_INDICES:
470 rand_offsets = (
471 state_batch_idx * stride_state_batch + pid_h * stride_state_head
472 )
473 else:
474 rand_offsets = pid_b * stride_state_batch + pid_h * stride_state_head
475 rand_offsets += (
476 offs_m[:, None] * stride_state_dim
477 + offs_n[None, :] * stride_state_dstate
478 )
479 # Generate random 32-bits for each element in state
480 if PHILOX_ROUNDS > 0:
481 rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS)
482 else:
483 rand = tl.randint(rand_seed, rand_offsets)
484 # Convert state to fp16 with RS rounding
485 state = convert_rs_fp16x2(state, rand)
486 tl.static_assert(state.dtype == tl.float16, "state must be fp16")
487 tl.static_assert(
488 dst_state_ptrs.dtype.element_ty == tl.float16,
489 "dst_state_ptrs must be fp16",
490 )
491 else:
492 state = state.to(dst_state_ptrs.dtype.element_ty)
493 tl.store(dst_state_ptrs, state, mask=mask)
494
495
496 def selective_state_update(
497 state,
498 x,
499 dt,
500 A,
501 B,
502 C,
503 D,
504 dt_bias,
505 z=None,
506 dt_softplus=False,
507 state_batch_indices=None,
508 dst_state_batch_indices=None,
509 null_block_id=NULL_BLOCK_ID,
510 out=None,
511 num_accepted_tokens=None,
512 cu_seqlens=None,
513 is_blackwell=False,
514 enable_stochastic_rounding=False,
515 cache_philox_rounds=0,
516 ):
517 """Argument:
518 state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
519 x: (batch, dim) or (batch, nheads, dim)
520 dt: (batch, dim) or (batch, nheads, dim)
521 A: (dim, dstate) or (nheads, dim, dstate)
522 B: (batch, dstate) or (batch, ngroups, dstate)
523 C: (batch, dstate) or (batch, ngroups, dstate)
524 D: (dim,) or (nheads, dim)
525 z: (batch, dim) or (batch, nheads, dim)
526 dt_bias: (dim,) or (nheads, dim)
527 null_block_id: int
528 if state_batch_indices is passed, lets the kernel identify
529 padded entries that will not be processed,
530 for example: state_batch_indices = [null_block_id, 1, 20,
531 null_block_id] in this case, the kernel will not process
532 entries at indices 0 and 3
533 out: Preallocated ssm output tensor. Assume same shape as x.
534 In-place updated.
535 num_accepted_tokens: (batch,)
536 number of accepted tokens from previous verification step,
537 tells the kernel which initial state to use
538 cu_seqlens: (batch,)
539 length per sequence, for variable length in speculative decoding cases
540 """
541 if state.dim() == 3:
542 state = state.unsqueeze(1)
543 if x.dim() == 2:
544 x = x.unsqueeze(1)
545 if dt.dim() == 2:
546 dt = dt.unsqueeze(1)
547 if A.dim() == 2:
548 A = A.unsqueeze(0)
549 if B.dim() == 2:
550 B = B.unsqueeze(1)
551 if C.dim() == 2:
552 C = C.unsqueeze(1)
553 if D.dim() == 1:
554 D = D.unsqueeze(0)
555 if z is not None and z.dim() == 2:
556 z = z.unsqueeze(1)
557 if dt_bias.dim() == 1:
558 dt_bias = dt_bias.unsqueeze(0)
559 if out.dim() == 2:
560 out = out.unsqueeze(1)
561 if state_batch_indices is not None and state_batch_indices.dim() == 1:
562 state_batch_indices = state_batch_indices.unsqueeze(1)
563 if dst_state_batch_indices is not None and dst_state_batch_indices.dim() == 1:
564 dst_state_batch_indices = dst_state_batch_indices.unsqueeze(1)
565 if num_accepted_tokens is not None:
566 assert state_batch_indices is not None and state_batch_indices.dim() == 2
567 assert dst_state_batch_indices is None or dst_state_batch_indices.dim() == 2
568
569 _, nheads, dim, dstate = state.shape
570 batch = x.shape[0]
571 if cu_seqlens is not None:
572 N = len(cu_seqlens) - 1
573 # Only used to verify the shape of
574 # state_batch_indices and dst_state_batch_indices
575 max_seqlen = (
576 state_batch_indices.size(-1) if state_batch_indices is not None else 1
577 )
578 else:
579 N = batch
580 max_seqlen = 1
581
582 assert x.shape == (batch, nheads, dim)
583 assert dt.shape == x.shape
584 assert A.shape == (nheads, dim, dstate)
585 ngroups = B.shape[1]
586 assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
587 assert B.shape == (batch, ngroups, dstate)
588 assert C.shape == B.shape
589 assert D.shape == (nheads, dim)
590 if z is not None:
591 assert z.shape == x.shape
592 assert dt_bias.shape == (nheads, dim)
593 if state_batch_indices is not None:
594 assert state_batch_indices.shape[0] >= N
595 assert state_batch_indices.shape[1] >= max_seqlen
596 if dst_state_batch_indices is not None:
597 assert dst_state_batch_indices.shape[0] >= N
598 assert dst_state_batch_indices.shape[1] >= max_seqlen
599 else:
600 # revert to the default behavior of in-place state updates
601 dst_state_batch_indices = state_batch_indices
602 assert out.shape == x.shape
603 if num_accepted_tokens is not None:
604 assert num_accepted_tokens.shape == (N,)
605
606 grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), N, nheads)
607 z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0)
608 state_batch_indices_strides = (
609 (state_batch_indices.stride(0), state_batch_indices.stride(1))
610 if state_batch_indices is not None
611 else (0, 0)
612 )
613 dst_state_batch_indices_strides = (
614 (dst_state_batch_indices.stride(0), dst_state_batch_indices.stride(1))
615 if dst_state_batch_indices is not None
616 else (0, 0)
617 )
618 # We don't want autotune since it will overwrite the state.
619 # Load from JSON config if available, otherwise fall back to heuristic.
620 cache_dtype = str(state.dtype).removeprefix("torch.")
621 BLOCK_SIZE_M, num_warps = try_get_optimal_ssm_config(
622 dim, dstate, N, nheads, cache_dtype, is_blackwell
623 )
624
625 tie_hdim = (
626 A.stride(-1) == 0
627 and A.stride(-2) == 0
628 and dt.stride(-1) == 0
629 and dt_bias.stride(-1) == 0
630 )
631 rand_seed = (
632 torch.randint(0, 2**32, (1,), device=state.device)
633 if enable_stochastic_rounding
634 else None
635 )
636
637 with torch.accelerator.device_index(x.device.index):
638 _selective_scan_update_kernel[grid](
639 state,
640 rand_seed,
641 x,
642 dt,
643 dt_bias,
644 A,
645 B,
646 C,
647 D,
648 z,
649 out,
650 state_batch_indices,
651 dst_state_batch_indices,
652 null_block_id,
653 num_accepted_tokens,
654 cu_seqlens,
655 N,
656 nheads,
657 dim,
658 dstate,
659 nheads // ngroups,
660 state.stride(0),
661 state.stride(1),
662 state.stride(2),
663 state.stride(3),
664 x.stride(0),
665 x.stride(1),
666 x.stride(2),
667 dt.stride(0),
668 dt.stride(1),
669 dt.stride(2),
670 dt_bias.stride(0),
671 dt_bias.stride(1),
672 A.stride(0),
673 A.stride(1),
674 A.stride(2),
675 B.stride(0),
676 B.stride(1),
677 B.stride(2),
678 C.stride(0),
679 C.stride(1),
680 C.stride(2),
681 D.stride(0),
682 D.stride(1),
683 z_strides[0],
684 z_strides[1],
685 z_strides[2],
686 out.stride(0),
687 out.stride(1),
688 out.stride(2),
689 state_batch_indices_strides[0],
690 state_batch_indices_strides[1],
691 dst_state_batch_indices_strides[0],
692 dst_state_batch_indices_strides[1],
693 dt_softplus,
694 tie_hdim,
695 BLOCK_SIZE_M,
696 num_warps=num_warps,
697 USE_RS_ROUNDING=enable_stochastic_rounding,
698 PHILOX_ROUNDS=cache_philox_rounds,
699 )
700
701
702 def selective_scan_fn(
703 u,
704 ssm_states,
705 delta,
706 A,
707 B,
708 C,
709 D=None,
710 z=None,
711 delta_bias=None,
712 delta_softplus=False,
713 query_start_loc=None,
714 cache_indices=None,
715 has_initial_state=None,
716 null_block_id=NULL_BLOCK_ID,
717 block_size=1024,
718 block_idx_first_scheduled_token=None,
719 block_idx_last_scheduled_token=None,
720 initial_state_idx=None,
721 cu_chunk_seqlen=None,
722 last_chunk_indices=None,
723 ) -> torch.Tensor:
724 """u: (dim, total_length) for varlen or (batch, dim, seqlen)
725 applies changes in place.
726 ssm_states: (batch, dim, dstate) or (batch, nheads, dim, dstate)
727 applies changes in place.
728 delta: (dim, total_length) for varlen or (batch, dim, seqlen)
729 A: (dim, dstate)
730 B: (ngroups, dstate, total_length) for varlen or
731 (batch,ngroups,dstate,seqlen)
732 C: (ngroups, dstate, total_length) for varlen or
733 (batch,ngroups,dstate,seqlen)
734 D: (dim,)
735 z: (dim, total_length) for varlen or (batch, dim, seqlen)
736 dt_bias: (dim,) or (dim)
737 query_start_loc: (batch + 1) int32
738 The cumulative sequence lengths of the sequences in
739 the batch, used to index into sequence. prepended with 0.
740 for example: query_start_loc = torch.Tensor([0,10,16,17]),
741 x.shape=(dim,17)
742 cache_indices: (batch) int32
743 A tensor with each cell is a correspondent
744 input and output ssm_state indices
745 - Without APC: (batch,) - single state index per batch item
746 - With APC: (batch, max_positions) - cache block indices for read/write
747 Each non-zero value indicates a cache block to load from and/or write to.
748 has_initial_state: (batch) bool
749 A tensor populated with ones and zeros,
750 indicate if the ssm_state at the corresponding index should be
751 used as initial state. Not providing argument assumes
752 there's no initial state
753 null_block_id: int
754 if cache_indices is passed, lets the kernel identify padding entries
755 that will not be processed,
756 for example: cache_indices = [null_block_id, 1 ,20 ,null_block_id]
757 in this case, the kernel will not process entries at indices 0 and 3
758 block_size: int
759 The block size to align the cached states to
760 block_idx_first_scheduled_token: (batch,), dtype int32
761 The pointer into cache_indices, where the first
762 cache block to be filled is located.
763 block_idx_last_scheduled_token: (batch,), dtype int32
764 The pointer into cache_indices, where the last cache block
765 to be filled is located.
766 initial_state_idx: (batch,), dtype int32
767 The pointer into cache_indices, where the cache block
768 containing the initial state is located.
769
770 Returns
771 -------
772 output: (dim, total_length) for varlen or (batch, dim, seqlen)
773 supports inplace replacement
774
775 """
776 if u.stride(-1) != 1:
777 u = u.contiguous()
778 if delta.stride(-1) != 1:
779 delta = delta.contiguous()
780 if D is not None:
781 D = D.contiguous()
782 if B.stride(-1) != 1:
783 B = B.contiguous()
784 if C.stride(-1) != 1:
785 C = C.contiguous()
786 if z is not None and z.stride(-1) != 1:
787 z = z.contiguous()
788 if B.dim() == 3 and query_start_loc is None:
789 B = B.unsqueeze(1)
790 if B.dim() == 2 and query_start_loc is not None:
791 B = B.unsqueeze(0)
792 if C.dim() == 3 and query_start_loc is None:
793 C = C.unsqueeze(1)
794 if C.dim() == 2 and query_start_loc is not None:
795 C = C.unsqueeze(0)
796
797 if current_platform.is_xpu():
798 xpu_ops.selective_scan_fwd(
799 u,
800 delta,
801 A,
802 B,
803 C,
804 D,
805 z,
806 delta_bias,
807 delta_softplus,
808 query_start_loc,
809 cache_indices,
810 has_initial_state,
811 ssm_states,
812 null_block_id,
813 block_size,
814 block_idx_first_scheduled_token,
815 block_idx_last_scheduled_token,
816 initial_state_idx,
817 cu_chunk_seqlen,
818 last_chunk_indices,
819 )
820 else:
821 ops.selective_scan_fwd(
822 u,
823 delta,
824 A,
825 B,
826 C,
827 D,
828 z,
829 delta_bias,
830 delta_softplus,
831 query_start_loc,
832 cache_indices,
833 has_initial_state,
834 ssm_states,
835 null_block_id,
836 block_size,
837 block_idx_first_scheduled_token,
838 block_idx_last_scheduled_token,
839 initial_state_idx,
840 cu_chunk_seqlen,
841 last_chunk_indices,
842 )
843
844 if z is None:
845 return delta # output written inplace to delta
846 else:
847 return z # output written inplace to z
848
849
850 from vllm.platforms import current_platform # noqa: E402
851
852 if current_platform.is_cpu():
853 from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import (
854 selective_state_update as selective_state_update_cpu,
855 )
856
857 selective_state_update = selective_state_update_cpu # type: ignore