@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 701 lines 22.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Dispatch module for Mamba selective state update (SSU) backends.
4
5 Provides a unified `selective_state_update` function that dispatches to
6 the Triton, FlashInfer, or CPU backend based on the configured
7 `MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA)
8 the backend defaults to 'cpu'.
9 """
10
11 from abc import ABC, abstractmethod
12 from collections.abc import Callable
13 from functools import cache
14
15 import torch
16
17 from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
18 from vllm.logger import init_logger
19 from vllm.triton_utils import tl, triton
20 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
21 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
22 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
23
24 logger = init_logger(__name__)
25
26
27 @triton.jit(
28 do_not_specialize=["n_slots", "state_batch_indices_stride"],
29 do_not_specialize_on_alignment=["state_batch_indices"],
30 )
31 def _update_replayssm_ring_trackers_kernel(
32 ring_start,
33 prev_num_accepted,
34 prev_query_len,
35 state_batch_indices,
36 state_batch_indices_stride,
37 n_slots,
38 num_states,
39 logical_window: tl.constexpr,
40 ring_buffer_len: tl.constexpr,
41 pad_slot_id: tl.constexpr,
42 RESET: tl.constexpr,
43 BLOCK: tl.constexpr,
44 ) -> None:
45 offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
46 mask = offsets < n_slots
47 slots = tl.load(
48 state_batch_indices + offsets * state_batch_indices_stride,
49 mask=mask,
50 other=pad_slot_id,
51 )
52 valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states)
53 if RESET:
54 tl.store(ring_start + slots, 0, mask=valid)
55 tl.store(prev_num_accepted + slots, 0, mask=valid)
56 # MTP alone reads prev_query_len. Standard decode deliberately leaves
57 # it unchanged; reset clears the pending MTP window after prefill.
58 tl.store(prev_query_len + slots, 0, mask=valid)
59 else:
60 prev = tl.load(prev_num_accepted + slots, mask=valid, other=0)
61 start = tl.load(ring_start + slots, mask=valid, other=0)
62 must_checkpoint = prev + 1 > logical_window
63 next_start = tl.where(
64 must_checkpoint,
65 (start + prev) % ring_buffer_len,
66 start,
67 )
68 next_prev = tl.where(must_checkpoint, 1, prev + 1)
69 tl.store(ring_start + slots, next_start, mask=valid)
70 tl.store(prev_num_accepted + slots, next_prev, mask=valid)
71
72
73 @triton.jit(
74 do_not_specialize=["n_slots", "state_batch_indices_stride"],
75 do_not_specialize_on_alignment=["state_batch_indices"],
76 )
77 def _commit_replayssm_ring_trackers_kernel(
78 ring_start,
79 prev_num_accepted,
80 prev_query_len,
81 state_batch_indices,
82 num_accepted_tokens,
83 query_start_loc,
84 state_batch_indices_stride,
85 n_slots,
86 num_states,
87 logical_window: tl.constexpr,
88 ring_buffer_len: tl.constexpr,
89 pad_slot_id: tl.constexpr,
90 BLOCK: tl.constexpr,
91 ) -> None:
92 offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
93 mask = offsets < n_slots
94 slots = tl.load(
95 state_batch_indices + offsets * state_batch_indices_stride,
96 mask=mask,
97 other=pad_slot_id,
98 )
99 valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states)
100 accepted_since_checkpoint = tl.load(prev_num_accepted + slots, mask=valid, other=0)
101 current_ring_start = tl.load(ring_start + slots, mask=valid, other=0)
102 previous_speculative_query_len = tl.load(
103 prev_query_len + slots, mask=valid, other=0
104 )
105 accepted_from_previous_query = tl.load(
106 num_accepted_tokens + offsets, mask=mask, other=0
107 )
108 must_checkpoint = (previous_speculative_query_len > 0) & (
109 accepted_since_checkpoint + previous_speculative_query_len > logical_window
110 )
111 next_start = tl.where(
112 must_checkpoint,
113 (current_ring_start + accepted_since_checkpoint) % ring_buffer_len,
114 current_ring_start,
115 )
116 next_prev = tl.where(
117 previous_speculative_query_len == 0,
118 0,
119 tl.where(
120 must_checkpoint,
121 accepted_from_previous_query,
122 accepted_since_checkpoint + accepted_from_previous_query,
123 ),
124 )
125 current_speculative_query_len = tl.load(
126 query_start_loc + offsets + 1, mask=mask, other=0
127 ) - tl.load(query_start_loc + offsets, mask=mask, other=0)
128 tl.store(ring_start + slots, next_start, mask=valid)
129 tl.store(prev_num_accepted + slots, next_prev, mask=valid)
130 tl.store(prev_query_len + slots, current_speculative_query_len, mask=valid)
131
132
133 def update_replayssm_ring_trackers(
134 ring_start: torch.Tensor,
135 prev_num_accepted: torch.Tensor,
136 prev_query_len: torch.Tensor,
137 state_batch_indices: torch.Tensor,
138 logical_window: int | None = None,
139 ring_buffer_len: int | None = None,
140 pad_slot_id: int = NULL_BLOCK_ID,
141 ) -> None:
142 """Reset selected trackers, or advance them when a window is provided."""
143 if state_batch_indices.dim() > 1:
144 state_batch_indices = state_batch_indices[:, 0]
145 n_slots = state_batch_indices.numel()
146 if n_slots == 0:
147 return
148 reset = logical_window is None
149 if reset:
150 logical_window = 0
151 ring_buffer_len = 1
152 else:
153 assert ring_buffer_len is not None
154 block = 128
155 _update_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)](
156 ring_start,
157 prev_num_accepted,
158 prev_query_len,
159 state_batch_indices,
160 state_batch_indices.stride(0),
161 n_slots,
162 min(
163 ring_start.numel(),
164 prev_num_accepted.numel(),
165 prev_query_len.numel(),
166 ),
167 logical_window,
168 ring_buffer_len,
169 pad_slot_id,
170 RESET=reset,
171 BLOCK=block,
172 )
173
174
175 def reset_replayssm_ring_trackers(
176 ring_start: torch.Tensor,
177 prev_num_accepted: torch.Tensor,
178 prev_query_len: torch.Tensor,
179 state_batch_indices: torch.Tensor,
180 pad_slot_id: int = NULL_BLOCK_ID,
181 ) -> None:
182 """Reset selected ReplaySSM ring trackers."""
183 update_replayssm_ring_trackers(
184 ring_start,
185 prev_num_accepted,
186 prev_query_len,
187 state_batch_indices,
188 pad_slot_id=pad_slot_id,
189 )
190
191
192 def commit_replayssm_ring_trackers(
193 ring_start: torch.Tensor,
194 prev_num_accepted: torch.Tensor,
195 prev_query_len: torch.Tensor,
196 state_batch_indices: torch.Tensor,
197 num_accepted_tokens: torch.Tensor,
198 query_start_loc: torch.Tensor,
199 logical_window: int,
200 ring_buffer_len: int,
201 pad_slot_id: int = NULL_BLOCK_ID,
202 ) -> None:
203 """Commit the preceding speculative window and record the current one.
204
205 MTP evaluates a target token and its draft tokens together, but the number
206 accepted from that query is available only on the next forward pass. This
207 function then advances the per-request ring by the accepted prefix and
208 records the current query length for the following pass. A zero
209 ``prev_query_len`` means that reset/prefill left no prior MTP query to
210 commit; slot validity is handled independently by the kernel mask.
211
212 Standard single-token decode needs no delayed commit: its ReplaySSM kernel
213 advances the shared ring trackers directly after every token.
214 """
215 if state_batch_indices.dim() > 1:
216 state_batch_indices = state_batch_indices[:, 0]
217 n_slots = state_batch_indices.numel()
218 if n_slots == 0:
219 return
220 block = 128
221 _commit_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)](
222 ring_start,
223 prev_num_accepted,
224 prev_query_len,
225 state_batch_indices,
226 num_accepted_tokens,
227 query_start_loc,
228 state_batch_indices.stride(0),
229 n_slots,
230 min(
231 ring_start.numel(),
232 prev_num_accepted.numel(),
233 prev_query_len.numel(),
234 ),
235 logical_window,
236 ring_buffer_len,
237 pad_slot_id,
238 BLOCK=block,
239 )
240
241
242 class MambaSSUBackend(ABC):
243 """Abstract base class for Mamba SSU backends."""
244
245 def __init__(self, mamba_config: MambaConfig):
246 self._mamba_config = mamba_config
247
248 @property
249 @abstractmethod
250 def name(self) -> str: ...
251
252 @abstractmethod
253 def __call__(
254 self,
255 state: torch.Tensor,
256 x: torch.Tensor,
257 dt: torch.Tensor,
258 A: torch.Tensor,
259 B: torch.Tensor,
260 C: torch.Tensor,
261 D: torch.Tensor,
262 dt_bias: torch.Tensor,
263 z: torch.Tensor | None = None,
264 dt_softplus: bool = False,
265 state_batch_indices: torch.Tensor | None = None,
266 dst_state_batch_indices: torch.Tensor | None = None,
267 null_block_id: int = NULL_BLOCK_ID,
268 out: torch.Tensor | None = None,
269 num_accepted_tokens: torch.Tensor | None = None,
270 cu_seqlens: torch.Tensor | None = None,
271 is_blackwell: bool = False,
272 ) -> None: ...
273
274
275 class TritonSSUBackend(MambaSSUBackend):
276 """Triton-based SSU backend (vLLM's default)."""
277
278 def __init__(self, mamba_config: MambaConfig):
279 super().__init__(mamba_config)
280 from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
281 selective_state_update as _triton_selective_state_update,
282 )
283
284 self._kernel = _triton_selective_state_update
285
286 @property
287 def name(self) -> str:
288 return "triton"
289
290 def __call__(
291 self,
292 state: torch.Tensor,
293 x: torch.Tensor,
294 dt: torch.Tensor,
295 A: torch.Tensor,
296 B: torch.Tensor,
297 C: torch.Tensor,
298 D: torch.Tensor,
299 dt_bias: torch.Tensor,
300 z: torch.Tensor | None = None,
301 dt_softplus: bool = False,
302 state_batch_indices: torch.Tensor | None = None,
303 dst_state_batch_indices: torch.Tensor | None = None,
304 null_block_id: int = NULL_BLOCK_ID,
305 out: torch.Tensor | None = None,
306 num_accepted_tokens: torch.Tensor | None = None,
307 cu_seqlens: torch.Tensor | None = None,
308 is_blackwell: bool = False,
309 ) -> None:
310 self._kernel(
311 state,
312 x,
313 dt,
314 A,
315 B,
316 C,
317 D=D,
318 z=z,
319 dt_bias=dt_bias,
320 dt_softplus=dt_softplus,
321 state_batch_indices=state_batch_indices,
322 dst_state_batch_indices=dst_state_batch_indices,
323 null_block_id=null_block_id,
324 out=out,
325 num_accepted_tokens=num_accepted_tokens,
326 cu_seqlens=cu_seqlens,
327 is_blackwell=is_blackwell,
328 enable_stochastic_rounding=self._mamba_config.enable_stochastic_rounding,
329 cache_philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds,
330 )
331
332
333 class FlashInferSSUBackend(MambaSSUBackend):
334 """FlashInfer-based SSU backend."""
335
336 def __init__(self, mamba_config: MambaConfig):
337 super().__init__(mamba_config)
338 try:
339 from flashinfer.mamba import selective_state_update as _fi_ssu
340 except ImportError as e:
341 raise ImportError(
342 "FlashInfer is required for the flashinfer Mamba SSU backend. "
343 "Please install flashinfer (>= 0.6.4): "
344 "pip install flashinfer-python"
345 ) from e
346 logger.info_once("Using FlashInfer Mamba SSU algorithm: %s", self._algorithm)
347 self._kernel = _fi_ssu
348
349 @property
350 def _algorithm(self) -> MambaSSUAlgorithm:
351 return self._mamba_config.ssu_algorithm or "auto"
352
353 @property
354 def name(self) -> str:
355 return "flashinfer"
356
357 def __call__(
358 self,
359 state: torch.Tensor,
360 x: torch.Tensor,
361 dt: torch.Tensor,
362 A: torch.Tensor,
363 B: torch.Tensor,
364 C: torch.Tensor,
365 D: torch.Tensor,
366 dt_bias: torch.Tensor,
367 z: torch.Tensor | None = None,
368 dt_softplus: bool = False,
369 state_batch_indices: torch.Tensor | None = None,
370 dst_state_batch_indices: torch.Tensor | None = None,
371 null_block_id: int = NULL_BLOCK_ID,
372 out: torch.Tensor | None = None,
373 num_accepted_tokens: torch.Tensor | None = None,
374 cu_seqlens: torch.Tensor | None = None,
375 is_blackwell: bool = False,
376 ) -> None:
377 rand_seed = (
378 torch.randint(0, 2**32, (1,), device=state.device)
379 if self._mamba_config.enable_stochastic_rounding
380 else None
381 )
382 self._kernel(
383 state,
384 x,
385 dt,
386 A,
387 B,
388 C,
389 D=D,
390 z=z,
391 dt_bias=dt_bias,
392 dt_softplus=dt_softplus,
393 state_batch_indices=state_batch_indices,
394 dst_state_batch_indices=dst_state_batch_indices,
395 cu_seqlens=cu_seqlens,
396 num_accepted_tokens=num_accepted_tokens,
397 cache_steps=state_batch_indices.size(-1)
398 if cu_seqlens is not None and state_batch_indices is not None
399 else 0,
400 pad_slot_id=null_block_id,
401 out=out,
402 rand_seed=rand_seed,
403 philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10,
404 algorithm=self._algorithm,
405 )
406
407
408 class CPUSSUBackend(MambaSSUBackend):
409 """CPU SSU backend using the compiled C++ VSX/scalar kernel.
410
411 On CPU-only platforms (PowerPC, x86 without CUDA) this dispatches to
412 the vectorized C++ kernel registered as ``torch.ops._C.selective_state_update_cpu``.
413 That kernel uses vec_op SIMD intrinsics (VSX on ppc64le, AVX2 on x86,
414 scalar fallback elsewhere) and is parallelised with OpenMP across heads.
415
416 Falls back to the pure-PyTorch implementation only if the C++ op is
417 unavailable (e.g. a CPU-less build).
418 """
419
420 def __init__(self, mamba_config: MambaConfig):
421 super().__init__(mamba_config)
422 from vllm import _custom_ops as ops
423
424 self._cpp_kernel = ops.selective_state_update_cpu
425 logger.info("CPUSSUBackend: using compiled C++ selective_state_update kernel.")
426
427 @property
428 def name(self) -> str:
429 return "cpu"
430
431 def __call__(
432 self,
433 state: torch.Tensor,
434 x: torch.Tensor,
435 dt: torch.Tensor,
436 A: torch.Tensor,
437 B: torch.Tensor,
438 C: torch.Tensor,
439 D: torch.Tensor,
440 dt_bias: torch.Tensor,
441 z: torch.Tensor | None = None,
442 dt_softplus: bool = False,
443 state_batch_indices: torch.Tensor | None = None,
444 dst_state_batch_indices: torch.Tensor | None = None,
445 null_block_id: int = NULL_BLOCK_ID,
446 out: torch.Tensor | None = None,
447 num_accepted_tokens: torch.Tensor | None = None,
448 cu_seqlens: torch.Tensor | None = None,
449 is_blackwell: bool = False,
450 ) -> None:
451 # C++ kernel: state shape expected as (nstates, nheads, dim, dstate)
452 # The kernel writes in-place into `out` and updates `state`.
453 self._cpp_kernel(
454 state,
455 x,
456 dt,
457 A,
458 B,
459 C,
460 D,
461 z,
462 dt_bias,
463 dt_softplus,
464 state_batch_indices,
465 dst_state_batch_indices,
466 null_block_id,
467 out,
468 num_accepted_tokens,
469 cu_seqlens,
470 )
471
472
473 _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
474 MambaBackendEnum.TRITON: TritonSSUBackend,
475 MambaBackendEnum.FLASHINFER: FlashInferSSUBackend,
476 MambaBackendEnum.CPU: CPUSSUBackend,
477 }
478
479 _mamba_ssu_backend: MambaSSUBackend | None = None
480
481
482 _flashinfer_replayssm_kernel: Callable[..., torch.Tensor] | None = None
483
484
485 @cache
486 def flashinfer_replayssm_autotune_supported() -> bool:
487 """Return True when FlashInfer exposes ReplaySSM autotuning."""
488 try:
489 from flashinfer.mamba.checkpointing_ssu import ( # noqa: F401
490 CheckpointingSSURunner,
491 )
492 except ImportError:
493 return False
494 return True
495
496
497 def selective_state_update_replayssm_flashinfer(
498 state: torch.Tensor,
499 x: torch.Tensor,
500 dt: torch.Tensor,
501 A: torch.Tensor,
502 B: torch.Tensor,
503 C: torch.Tensor,
504 out: torch.Tensor,
505 x_cache: torch.Tensor,
506 B_cache: torch.Tensor,
507 dt_cache: torch.Tensor,
508 ring_start: torch.Tensor,
509 prev_num_accepted_tokens: torch.Tensor,
510 prev_query_len: torch.Tensor,
511 logical_window: int,
512 D: torch.Tensor | None = None,
513 dt_bias: torch.Tensor | None = None,
514 dt_softplus: bool = False,
515 state_batch_indices: torch.Tensor | None = None,
516 null_block_id: int = NULL_BLOCK_ID,
517 scratch: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
518 update_trackers: bool = True,
519 enable_stochastic_rounding: bool = False,
520 stochastic_rounding_philox_rounds: int = 0,
521 cu_seqlens: torch.Tensor | None = None,
522 max_seqlen: int | None = None,
523 enable_pdl: bool = False,
524 ) -> torch.Tensor:
525 """Run FlashInfer checkpointing SSU and optionally advance shared trackers."""
526 if _flashinfer_replayssm_kernel is None:
527 raise RuntimeError(
528 "FlashInfer ReplaySSM has not been initialized. "
529 "Call initialize_mamba_ssu_backend() with use_replayssm=True."
530 )
531
532 if x.dim() == 3:
533 dim = 0 if cu_seqlens is not None else 1
534 x = x.unsqueeze(dim)
535 dt = dt.unsqueeze(dim)
536 B = B.unsqueeze(dim)
537 C = C.unsqueeze(dim)
538 out = out.unsqueeze(dim)
539
540 indices = state_batch_indices
541 if indices is not None and indices.dim() > 1:
542 indices = indices[:, 0]
543
544 cb_scaled = cumAdt_vec = cb_old = None
545 if scratch is not None:
546 cb_scaled, cumAdt_vec, cb_old = scratch
547
548 rand_seed = (
549 torch.randint(0, 2**32, (1,), device=state.device, dtype=torch.int64)
550 if enable_stochastic_rounding
551 else None
552 )
553 result = _flashinfer_replayssm_kernel(
554 state,
555 x_cache,
556 B_cache,
557 dt_cache,
558 ring_start,
559 prev_num_accepted_tokens,
560 x,
561 dt,
562 A,
563 B,
564 C,
565 out,
566 D=D,
567 dt_bias=dt_bias,
568 dt_softplus=dt_softplus,
569 state_batch_indices=indices,
570 pad_slot_id=null_block_id,
571 rand_seed=rand_seed,
572 philox_rounds=stochastic_rounding_philox_rounds or 10,
573 cu_seqlens=cu_seqlens,
574 max_seqlen=max_seqlen,
575 enable_pdl=enable_pdl,
576 cb_scaled=cb_scaled,
577 cumAdt_vec=cumAdt_vec,
578 cb_old=cb_old,
579 )
580 if update_trackers and indices is not None:
581 update_replayssm_ring_trackers(
582 ring_start,
583 prev_num_accepted_tokens,
584 prev_query_len,
585 indices,
586 logical_window=logical_window,
587 ring_buffer_len=x_cache.size(2),
588 pad_slot_id=null_block_id,
589 )
590 return result
591
592
593 def initialize_mamba_ssu_backend(
594 mamba_config: MambaConfig,
595 kv_cache_config: KVCacheConfig,
596 *,
597 use_replayssm: bool = False,
598 ) -> None:
599 """Initialize the Mamba SSU backend and optional FlashInfer ReplaySSM."""
600 if not any(
601 isinstance(g.kv_cache_spec, MambaSpec)
602 and g.kv_cache_spec.mamba_type
603 in (MambaAttentionBackendEnum.MAMBA1, MambaAttentionBackendEnum.MAMBA2)
604 for g in kv_cache_config.kv_cache_groups
605 ):
606 return
607
608 global _flashinfer_replayssm_kernel, _mamba_ssu_backend
609 backend = mamba_config.backend
610
611 if backend == MambaBackendEnum.TRITON:
612 from vllm.platforms import current_platform
613
614 if current_platform.is_cpu():
615 logger.info(
616 "CPU platform detected: overriding Mamba SSU backend "
617 "from 'triton' to 'cpu'."
618 )
619 backend = MambaBackendEnum.CPU
620
621 if backend not in _BACKEND_REGISTRY:
622 raise ValueError(
623 f"Unknown Mamba SSU backend: {backend}. "
624 f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
625 )
626 if use_replayssm and backend not in (
627 MambaBackendEnum.TRITON,
628 MambaBackendEnum.FLASHINFER,
629 ):
630 raise ValueError(f"ReplaySSM does not support mamba backend {backend.value!r}")
631
632 backend_cls = _BACKEND_REGISTRY[backend]
633 if not isinstance(_mamba_ssu_backend, backend_cls):
634 _mamba_ssu_backend = backend_cls(mamba_config)
635 logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
636
637 _flashinfer_replayssm_kernel = None
638 if use_replayssm and backend == MambaBackendEnum.FLASHINFER:
639 try:
640 from flashinfer.mamba.checkpointing_ssu import checkpointing_ssu
641 except ImportError as e:
642 raise ImportError(
643 "FlashInfer ReplaySSM requires a compatible flashinfer-python package"
644 ) from e
645 _flashinfer_replayssm_kernel = checkpointing_ssu
646 if use_replayssm:
647 logger.info("Using %s ReplaySSM backend.", backend.value)
648
649
650 def get_mamba_ssu_backend() -> MambaSSUBackend:
651 """Get the current Mamba SSU backend. Raises if not initialized."""
652 if _mamba_ssu_backend is None:
653 raise RuntimeError(
654 "Mamba SSU backend has not been initialized. "
655 "Call initialize_mamba_ssu_backend() first."
656 )
657 return _mamba_ssu_backend
658
659
660 def selective_state_update(
661 state: torch.Tensor,
662 x: torch.Tensor,
663 dt: torch.Tensor,
664 A: torch.Tensor,
665 B: torch.Tensor,
666 C: torch.Tensor,
667 D: torch.Tensor,
668 dt_bias: torch.Tensor,
669 z: torch.Tensor | None = None,
670 dt_softplus: bool = False,
671 state_batch_indices: torch.Tensor | None = None,
672 dst_state_batch_indices: torch.Tensor | None = None,
673 null_block_id: int = NULL_BLOCK_ID,
674 out: torch.Tensor | None = None,
675 num_accepted_tokens: torch.Tensor | None = None,
676 cu_seqlens: torch.Tensor | None = None,
677 is_blackwell: bool = False,
678 ) -> None:
679 """Unified dispatch for Mamba selective state update.
680
681 Delegates to the initialized backend (Triton or FlashInfer).
682 """
683 get_mamba_ssu_backend()(
684 state,
685 x,
686 dt,
687 A,
688 B,
689 C,
690 D,
691 dt_bias,
692 z=z,
693 dt_softplus=dt_softplus,
694 state_batch_indices=state_batch_indices,
695 dst_state_batch_indices=dst_state_batch_indices,
696 null_block_id=null_block_id,
697 out=out,
698 num_accepted_tokens=num_accepted_tokens,
699 cu_seqlens=cu_seqlens,
700 is_blackwell=is_blackwell,
701 )