@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ecf5ff7ce3
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 04:45 UTC

[Mamba] Flashinfer selective_state_update (#36162)

Signed-off-by: Roi Koren <roik@nvidia.com>

roikoren755 committed Apr 14, 2026 at 22:10 UTC ecf5ff7ce3ac2a8db51654712cbcea3540faaef8
15 files changed +516 -71
tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml
+1
@@ -8,4 +8,5 @@ server_args: >-
8 --max-model-len 4096
9 --tensor-parallel-size 8
10 --enable-expert-parallel
11 + --mamba-backend flashinfer
12 --speculative-config '{"method":"mtp","num_speculative_tokens":5}'
tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml
+1
@@ -8,4 +8,5 @@ server_args: >-
8 --max-model-len 4096
9 --tensor-parallel-size 2
10 --enable-expert-parallel
11 + --mamba-backend flashinfer
12 --speculative-config '{"method":"mtp","num_speculative_tokens":5}'
tests/kernels/mamba/test_ssu_dispatch.py new
+92
@@ -0,0 +1,92 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.config.mamba import MambaBackendEnum, MambaConfig
8 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
9 + FlashInferSSUBackend,
10 + TritonSSUBackend,
11 + get_mamba_ssu_backend,
12 + initialize_mamba_ssu_backend,
13 + selective_state_update,
14 +)
15 +from vllm.utils.torch_utils import set_random_seed
16 +
17 +try:
18 + import flashinfer.mamba # noqa: F401
19 +
20 + HAS_FLASHINFER = True
21 +except ImportError:
22 + HAS_FLASHINFER = False
23 +
24 +
25 +def test_default_backend_is_triton():
26 + initialize_mamba_ssu_backend(MambaConfig())
27 + backend = get_mamba_ssu_backend()
28 + assert isinstance(backend, TritonSSUBackend)
29 + assert backend.name == "triton"
30 +
31 +
32 +def test_explicit_triton_backend():
33 + initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
34 + backend = get_mamba_ssu_backend()
35 + assert isinstance(backend, TritonSSUBackend)
36 +
37 +
38 +@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
39 +def test_flashinfer_backend_init():
40 + initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
41 + backend = get_mamba_ssu_backend()
42 + assert isinstance(backend, FlashInferSSUBackend)
43 + assert backend.name == "flashinfer"
44 +
45 +
46 +def test_uninitialized_backend_raises():
47 + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
48 +
49 + old = mod._mamba_ssu_backend
50 + mod._mamba_ssu_backend = None
51 + with pytest.raises(RuntimeError, match="not been initialized"):
52 + get_mamba_ssu_backend()
53 + mod._mamba_ssu_backend = old
54 +
55 +
56 +@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
57 +def test_flashinfer_import_error():
58 + with pytest.raises(ImportError, match="FlashInfer is required"):
59 + FlashInferSSUBackend(MambaConfig())
60 +
61 +
62 +def test_triton_basic_call():
63 + set_random_seed(0)
64 + initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
65 + device = "cuda"
66 + batch_size = 2
67 + dim = 64
68 + dstate = 16
69 +
70 + state = torch.randn(batch_size, dim, dstate, device=device)
71 + x = torch.randn(batch_size, dim, device=device)
72 + out = torch.empty_like(x)
73 + dt = torch.randn(batch_size, dim, device=device)
74 + dt_bias = torch.rand(dim, device=device) - 4.0
75 + A = -torch.rand(dim, dstate, device=device)
76 + B = torch.randn(batch_size, dstate, device=device)
77 + C = torch.randn(batch_size, dstate, device=device)
78 + D = torch.randn(dim, device=device)
79 +
80 + selective_state_update(
81 + state,
82 + x,
83 + dt,
84 + A,
85 + B,
86 + C,
87 + D=D,
88 + dt_bias=dt_bias,
89 + dt_softplus=True,
90 + out=out,
91 + )
92 + assert not torch.isnan(out).any()
vllm/config/__init__.py
+3
@@ -16,6 +16,7 @@ from vllm.config.kv_events import KVEventsConfig
16 from vllm.config.kv_transfer import KVTransferConfig
17 from vllm.config.load import LoadConfig
18 from vllm.config.lora import LoRAConfig
19 +from vllm.config.mamba import MambaConfig
20 from vllm.config.model import (
21 ModelConfig,
22 iter_architecture_defaults,
@@ -83,6 +84,8 @@ __all__ = [
84 "LoadConfig",
85 # From vllm.config.lora
86 "LoRAConfig",
87 + # From vllm.config.mamba
88 + "MambaConfig",
89 # From vllm.config.model
90 "ModelConfig",
91 "iter_architecture_defaults",
vllm/config/cache.py
-34
@@ -123,14 +123,6 @@ class CacheConfig:
123 - "align": only cache the mamba state of the last token of each scheduler step and
124 when the token is at position i * block_size.
125 """
126 - enable_mamba_cache_stochastic_rounding: bool = False
127 - """Enable stochastic rounding when writing SSM state to fp16 cache.
128 - Uses random bits to unbias the rounding error, which can improve
129 - numerical stability for long sequences."""
130 - mamba_cache_philox_rounds: int = 0
131 - """Number of Philox PRNG rounds for stochastic rounding random number
132 - generation. 0 uses the Triton default. Higher values improve randomness
133 - quality at the cost of compute."""
126
127 # Will be set after profiling.
128 num_gpu_blocks: int | None = field(default=None, init=False)
@@ -258,29 +250,3 @@ class CacheConfig:
250 str(cache_dtype),
251 )
252 return cache_dtype
261 -
262 - def __post_init__(self):
263 - if self.enable_mamba_cache_stochastic_rounding:
264 - from vllm.platforms import current_platform
265 -
266 - if not current_platform.is_cuda():
267 - raise ValueError(
268 - "Stochastic rounding for Mamba cache is only supported "
269 - "on NVIDIA CUDA platforms. Please do not specify "
270 - "`--enable-mamba-cache-stochastic-rounding`."
271 - )
272 - if not current_platform.is_device_capability_family(100):
273 - raise ValueError(
274 - "Stochastic rounding for Mamba cache requires compute "
275 - "capability 10.0 (data center Blackwell). The `cvt.rs` PTX "
276 - "instruction is not supported on your GPU. Please do not specify "
277 - "`--enable-mamba-cache-stochastic-rounding`."
278 - )
279 - if self.mamba_ssm_cache_dtype != "float16":
280 - raise ValueError(
281 - "Stochastic rounding for Mamba cache requires "
282 - "the SSM cache to be float16. Please set it explicitly, "
283 - "by specifying `--mamba-ssm-cache-dtype float16`, or disable "
284 - "stochastic rounding by not specifying "
285 - "`--enable-mamba-cache-stochastic-rounding`."
286 - )
vllm/config/mamba.py new
+76
@@ -0,0 +1,76 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from enum import Enum, EnumMeta
5 +from typing import Any
6 +
7 +from pydantic import field_validator
8 +
9 +from vllm.config.utils import config
10 +
11 +
12 +class _MambaBackendEnumMeta(EnumMeta):
13 + """Metaclass for MambaBackendEnum to provide better error messages."""
14 +
15 + def __getitem__(cls, name: str):
16 + try:
17 + return super().__getitem__(name)
18 + except KeyError:
19 + valid = ", ".join(cls.__members__.keys())
20 + raise ValueError(
21 + f"Unknown Mamba SSU backend: '{name}'. Valid options are: {valid}"
22 + ) from None
23 +
24 +
25 +class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):
26 + """Enumeration of supported Mamba SSU (selective state update) backends."""
27 +
28 + TRITON = "triton"
29 + FLASHINFER = "flashinfer"
30 +
31 +
32 +@config
33 +class MambaConfig:
34 + """Configuration for Mamba SSM backends."""
35 +
36 + backend: MambaBackendEnum = MambaBackendEnum.TRITON
37 + """Mamba SSU backend to use."""
38 +
39 + enable_stochastic_rounding: bool = False
40 + """Enable stochastic rounding when writing SSM state to fp16 cache.
41 + Uses random bits to unbias the rounding error, which can improve
42 + numerical stability for long sequences."""
43 + stochastic_rounding_philox_rounds: int = 0
44 + """Number of Philox PRNG rounds for stochastic rounding random number
45 + generation. 0 uses the Triton default. Higher values improve randomness
46 + quality at the cost of compute."""
47 +
48 + @field_validator("backend", mode="before")
49 + @classmethod
50 + def validate_backend_before(cls, value: Any) -> Any:
51 + """Enable parsing of the `backend` enum type from string."""
52 + if isinstance(value, str):
53 + return MambaBackendEnum[value.upper()]
54 + return value
55 +
56 + def __post_init__(self):
57 + if self.enable_stochastic_rounding:
58 + from vllm.platforms import current_platform
59 +
60 + if not current_platform.is_cuda():
61 + raise ValueError(
62 + "Stochastic rounding for Mamba cache is only supported "
63 + "on NVIDIA CUDA platforms. Please do not specify "
64 + "`--enable-mamba-cache-stochastic-rounding`."
65 + )
66 + if (
67 + self.backend == MambaBackendEnum.TRITON
68 + and not current_platform.is_device_capability_family(100)
69 + ):
70 + raise ValueError(
71 + "Stochastic rounding for Mamba cache with triton backend requires "
72 + "compute capability 10.0 (data center Blackwell). The `cvt.rs` "
73 + "PTX instruction is not supported on your GPU. Please do not "
74 + "specify `--enable-mamba-cache-stochastic-rounding`, "
75 + "or set `--mamba-backend flashinfer`."
76 + )
vllm/config/vllm.py
+15
@@ -37,6 +37,7 @@ from .kv_events import KVEventsConfig
37 from .kv_transfer import KVTransferConfig
38 from .load import LoadConfig
39 from .lora import LoRAConfig
40 +from .mamba import MambaConfig
41 from .model import ModelConfig
42 from .observability import ObservabilityConfig
43 from .offload import OffloadConfig
@@ -275,6 +276,8 @@ class VllmConfig:
276 """Model weight offloading configuration."""
277 attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
278 """Attention configuration."""
279 + mamba_config: MambaConfig = Field(default_factory=MambaConfig)
280 + """Mamba configuration."""
281 kernel_config: KernelConfig = Field(default_factory=KernelConfig)
282 """Kernel configuration."""
283 lora_config: LoRAConfig | None = None
@@ -717,6 +720,18 @@ class VllmConfig:
720 if self.lora_config is not None:
721 self.lora_config.verify_with_model_config(self.model_config)
722
723 + if (
724 + self.mamba_config.enable_stochastic_rounding
725 + and self.cache_config.mamba_ssm_cache_dtype != "float16"
726 + ):
727 + raise ValueError(
728 + "Stochastic rounding for Mamba cache requires "
729 + "the SSM cache to be float16. Please set it explicitly, "
730 + "by specifying `--mamba-ssm-cache-dtype float16`, or disable "
731 + "stochastic rounding by not specifying "
732 + "`--enable-mamba-cache-stochastic-rounding`."
733 + )
734 +
735 if self.quant_config is None and self.model_config is not None:
736 self.quant_config = VllmConfig._get_quantization_config(
737 self.model_config, self.load_config
vllm/engine/arg_utils.py
+42 -11
@@ -45,6 +45,7 @@ from vllm.config import (
45 KVTransferConfig,
46 LoadConfig,
47 LoRAConfig,
48 + MambaConfig,
49 ModelConfig,
50 MultiModalConfig,
51 ObservabilityConfig,
@@ -72,6 +73,7 @@ from vllm.config.cache import (
73 from vllm.config.device import Device
74 from vllm.config.kernel import IrOpPriorityConfig, MoEBackend
75 from vllm.config.lora import MaxLoRARanks
76 +from vllm.config.mamba import MambaBackendEnum
77 from vllm.config.model import (
78 ConvertOption,
79 HfOverrides,
@@ -578,6 +580,7 @@ class EngineArgs:
580 pooler_config: PoolerConfig | None = ModelConfig.pooler_config
581 compilation_config: CompilationConfig = get_field(VllmConfig, "compilation_config")
582 attention_config: AttentionConfig = get_field(VllmConfig, "attention_config")
583 + mamba_config: MambaConfig = get_field(VllmConfig, "mamba_config")
584 kernel_config: KernelConfig = get_field(VllmConfig, "kernel_config")
585 enable_flashinfer_autotune: bool = get_field(
586 KernelConfig, "enable_flashinfer_autotune"
@@ -610,10 +613,12 @@ class EngineArgs:
613 mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype
614 mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
615 mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode
616 +
617 + mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON
618 enable_mamba_cache_stochastic_rounding: bool = (
614 - CacheConfig.enable_mamba_cache_stochastic_rounding
619 + MambaConfig.enable_stochastic_rounding
620 )
616 - mamba_cache_philox_rounds: int = CacheConfig.mamba_cache_philox_rounds
621 + mamba_cache_philox_rounds: int = MambaConfig.stochastic_rounding_philox_rounds
622
623 additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config")
624
@@ -655,6 +660,8 @@ class EngineArgs:
660 self.compilation_config = CompilationConfig(**self.compilation_config)
661 if isinstance(self.attention_config, dict):
662 self.attention_config = AttentionConfig(**self.attention_config)
663 + if isinstance(self.mamba_config, dict):
664 + self.mamba_config = MambaConfig(**self.mamba_config)
665 if isinstance(self.kernel_config, dict):
666 self.kernel_config = KernelConfig(**self.kernel_config)
667 if isinstance(self.eplb_config, dict):
@@ -825,6 +832,22 @@ class EngineArgs:
832 "--attention-backend", **attention_kwargs["backend"]
833 )
834
835 + # Mamba arguments
836 + mamba_kwargs = get_kwargs(MambaConfig)
837 + mamba_group = parser.add_argument_group(
838 + title="MambaConfig",
839 + description=MambaConfig.__doc__,
840 + )
841 + mamba_group.add_argument("--mamba-backend", **mamba_kwargs["backend"])
842 + mamba_group.add_argument(
843 + "--enable-mamba-cache-stochastic-rounding",
844 + **mamba_kwargs["enable_stochastic_rounding"],
845 + )
846 + mamba_group.add_argument(
847 + "--mamba-cache-philox-rounds",
848 + **mamba_kwargs["stochastic_rounding_philox_rounds"],
849 + )
850 +
851 # Structured outputs arguments
852 structured_outputs_kwargs = get_kwargs(StructuredOutputsConfig)
853 structured_outputs_group = parser.add_argument_group(
@@ -1050,13 +1073,6 @@ class EngineArgs:
1073 cache_group.add_argument(
1074 "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]
1075 )
1053 - cache_group.add_argument(
1054 - "--enable-mamba-cache-stochastic-rounding",
1055 - **cache_kwargs["enable_mamba_cache_stochastic_rounding"],
1056 - )
1057 - cache_group.add_argument(
1058 - "--mamba-cache-philox-rounds", **cache_kwargs["mamba_cache_philox_rounds"]
1059 - )
1076 cache_group.add_argument(
1077 "--kv-offloading-size", **cache_kwargs["kv_offloading_size"]
1078 )
@@ -1622,8 +1638,6 @@ class EngineArgs:
1638 mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,
1639 mamba_block_size=self.mamba_block_size,
1640 mamba_cache_mode=self.mamba_cache_mode,
1625 - enable_mamba_cache_stochastic_rounding=self.enable_mamba_cache_stochastic_rounding,
1626 - mamba_cache_philox_rounds=self.mamba_cache_philox_rounds,
1641 kv_offloading_size=self.kv_offloading_size,
1642 kv_offloading_backend=self.kv_offloading_backend,
1643 )
@@ -1934,6 +1948,22 @@ class EngineArgs:
1948 self.attention_backend
1949 )
1950
1951 + # Mamba config overrides
1952 + mamba_config = copy.deepcopy(self.mamba_config)
1953 + # Convert string to enum if needed (CLI parsing returns a string)
1954 + if isinstance(self.mamba_backend, str):
1955 + mamba_config.backend = MambaBackendEnum[self.mamba_backend.upper()]
1956 + else:
1957 + mamba_config.backend = self.mamba_backend
1958 + if self.enable_mamba_cache_stochastic_rounding:
1959 + mamba_config.enable_stochastic_rounding = (
1960 + self.enable_mamba_cache_stochastic_rounding
1961 + )
1962 + if self.mamba_cache_philox_rounds:
1963 + mamba_config.stochastic_rounding_philox_rounds = (
1964 + self.mamba_cache_philox_rounds
1965 + )
1966 +
1967 # Kernel config overrides
1968 kernel_config = copy.deepcopy(self.kernel_config)
1969 if self.enable_flashinfer_autotune is not None:
@@ -2032,6 +2062,7 @@ class EngineArgs:
2062 load_config=load_config,
2063 offload_config=offload_config,
2064 attention_config=attention_config,
2065 + mamba_config=mamba_config,
2066 kernel_config=kernel_config,
2067 lora_config=lora_config,
2068 speculative_config=speculative_config,
vllm/model_executor/layers/mamba/mamba_mixer.py
+3 -7
@@ -30,10 +30,8 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
30 causal_conv1d_fn,
31 causal_conv1d_update,
32 )
33 -from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
34 - selective_scan_fn,
35 - selective_state_update,
36 -)
33 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_scan_fn
34 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update
35 from vllm.model_executor.utils import set_weight_attrs
36 from vllm.platforms import current_platform
37 from vllm.utils.torch_utils import (
@@ -431,14 +429,12 @@ class MambaMixer(MambaBase, PluggableLayer):
429 B_d,
430 C_d,
431 self.D,
434 - gate_d.transpose(0, 1),
432 time_proj_bias,
433 + z=gate_d.transpose(0, 1),
434 dt_softplus=True,
435 state_batch_indices=state_indices_tensor_d_input,
436 dst_state_batch_indices=state_indices_tensor_d_output,
437 out=scan_outputs_d,
440 - enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding,
441 - cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds,
438 )
439 scan_outputs_d = scan_outputs_d.transpose(0, 1)
440
vllm/model_executor/layers/mamba/mamba_mixer2.py
+2 -5
@@ -31,10 +31,10 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
31 causal_conv1d_update,
32 )
33 from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated
34 -from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update
34 from vllm.model_executor.layers.mamba.ops.ssd_combined import (
35 mamba_chunk_scan_combined_varlen,
36 )
37 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update
38 from vllm.model_executor.layers.quantization import QuantizationConfig
39 from vllm.model_executor.model_loader.weight_utils import (
40 LoaderFunction,
@@ -890,8 +890,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
890 B_d,
891 C_d,
892 D_d,
893 - z=None,
894 - dt_bias=dt_bias,
893 + dt_bias,
894 dt_softplus=True,
895 state_batch_indices=state_indices_tensor_d_input,
896 dst_state_batch_indices=state_indices_tensor_d_output,
@@ -899,8 +898,6 @@ class MambaMixer2(MambaBase, PluggableLayer):
898 num_accepted_tokens=num_accepted_tokens,
899 cu_seqlens=query_start_loc_d,
900 is_blackwell=self.is_blackwell,
902 - enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding,
903 - cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds,
901 )
902
903 def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+10 -10
@@ -323,9 +323,9 @@ def selective_state_update(
323 A,
324 B,
325 C,
326 - D=None,
326 + D,
327 + dt_bias,
328 z=None,
328 - dt_bias=None,
329 dt_softplus=False,
330 state_batch_indices=None,
331 dst_state_batch_indices=None,
@@ -374,11 +374,11 @@ def selective_state_update(
374 B = B.unsqueeze(1)
375 if C.dim() == 2:
376 C = C.unsqueeze(1)
377 - if D is not None and D.dim() == 1:
377 + if D.dim() == 1:
378 D = D.unsqueeze(0)
379 if z is not None and z.dim() == 2:
380 z = z.unsqueeze(1)
381 - if dt_bias is not None and dt_bias.dim() == 1:
381 + if dt_bias.dim() == 1:
382 dt_bias = dt_bias.unsqueeze(0)
383 if out.dim() == 2:
384 out = out.unsqueeze(1)
@@ -410,12 +410,10 @@ def selective_state_update(
410 assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
411 assert B.shape == (batch, ngroups, dstate)
412 assert C.shape == B.shape
413 - if D is not None:
414 - assert D.shape == (nheads, dim)
413 + assert D.shape == (nheads, dim)
414 if z is not None:
415 assert z.shape == x.shape
417 - if dt_bias is not None:
418 - assert dt_bias.shape == (nheads, dim)
416 + assert dt_bias.shape == (nheads, dim)
417 if state_batch_indices is not None:
418 assert state_batch_indices.shape[0] >= N
419 assert state_batch_indices.shape[1] >= max_seqlen
@@ -506,7 +504,8 @@ def selective_state_update(
504 dt.stride(0),
505 dt.stride(1),
506 dt.stride(2),
509 - *(dt_bias.stride(0), dt_bias.stride(1)) if dt_bias is not None else 0,
507 + dt_bias.stride(0),
508 + dt_bias.stride(1),
509 A.stride(0),
510 A.stride(1),
511 A.stride(2),
@@ -516,7 +515,8 @@ def selective_state_update(
515 C.stride(0),
516 C.stride(1),
517 C.stride(2),
519 - *(D.stride(0), D.stride(1)) if D is not None else 0,
518 + D.stride(0),
519 + D.stride(1),
520 z_strides[0],
521 z_strides[1],
522 z_strides[2],
vllm/model_executor/layers/mamba/ops/ssu_dispatch.py new
+261
@@ -0,0 +1,261 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Dispatch module for Mamba selective state update (SSU) backends.
5 +
6 +Provides a unified `selective_state_update` function that dispatches to
7 +either the Triton or FlashInfer backend based on the configured
8 +`MambaBackendEnum`. Follows SGLang's dispatch pattern adapted for vLLM.
9 +"""
10 +
11 +from abc import ABC, abstractmethod
12 +
13 +import torch
14 +
15 +from vllm.config.mamba import MambaBackendEnum, MambaConfig
16 +from vllm.logger import init_logger
17 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +class MambaSSUBackend(ABC):
23 + """Abstract base class for Mamba SSU backends."""
24 +
25 + def __init__(self, mamba_config: MambaConfig):
26 + self._mamba_config = mamba_config
27 +
28 + @property
29 + @abstractmethod
30 + def name(self) -> str: ...
31 +
32 + @abstractmethod
33 + def __call__(
34 + self,
35 + state: torch.Tensor,
36 + x: torch.Tensor,
37 + dt: torch.Tensor,
38 + A: torch.Tensor,
39 + B: torch.Tensor,
40 + C: torch.Tensor,
41 + D: torch.Tensor,
42 + dt_bias: torch.Tensor,
43 + z: torch.Tensor | None = None,
44 + dt_softplus: bool = False,
45 + state_batch_indices: torch.Tensor | None = None,
46 + dst_state_batch_indices: torch.Tensor | None = None,
47 + null_block_id: int = NULL_BLOCK_ID,
48 + out: torch.Tensor | None = None,
49 + num_accepted_tokens: torch.Tensor | None = None,
50 + cu_seqlens: torch.Tensor | None = None,
51 + is_blackwell: bool = False,
52 + ) -> None: ...
53 +
54 +
55 +class TritonSSUBackend(MambaSSUBackend):
56 + """Triton-based SSU backend (vLLM's default)."""
57 +
58 + def __init__(self, mamba_config: MambaConfig):
59 + super().__init__(mamba_config)
60 + from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
61 + selective_state_update as _triton_selective_state_update,
62 + )
63 +
64 + self._kernel = _triton_selective_state_update
65 +
66 + @property
67 + def name(self) -> str:
68 + return "triton"
69 +
70 + def __call__(
71 + self,
72 + state: torch.Tensor,
73 + x: torch.Tensor,
74 + dt: torch.Tensor,
75 + A: torch.Tensor,
76 + B: torch.Tensor,
77 + C: torch.Tensor,
78 + D: torch.Tensor,
79 + dt_bias: torch.Tensor,
80 + z: torch.Tensor | None = None,
81 + dt_softplus: bool = False,
82 + state_batch_indices: torch.Tensor | None = None,
83 + dst_state_batch_indices: torch.Tensor | None = None,
84 + null_block_id: int = NULL_BLOCK_ID,
85 + out: torch.Tensor | None = None,
86 + num_accepted_tokens: torch.Tensor | None = None,
87 + cu_seqlens: torch.Tensor | None = None,
88 + is_blackwell: bool = False,
89 + ) -> None:
90 + self._kernel(
91 + state,
92 + x,
93 + dt,
94 + A,
95 + B,
96 + C,
97 + D=D,
98 + z=z,
99 + dt_bias=dt_bias,
100 + dt_softplus=dt_softplus,
101 + state_batch_indices=state_batch_indices,
102 + dst_state_batch_indices=dst_state_batch_indices,
103 + null_block_id=null_block_id,
104 + out=out,
105 + num_accepted_tokens=num_accepted_tokens,
106 + cu_seqlens=cu_seqlens,
107 + is_blackwell=is_blackwell,
108 + enable_stochastic_rounding=self._mamba_config.enable_stochastic_rounding,
109 + cache_philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds,
110 + )
111 +
112 +
113 +class FlashInferSSUBackend(MambaSSUBackend):
114 + """FlashInfer-based SSU backend."""
115 +
116 + def __init__(self, mamba_config: MambaConfig):
117 + super().__init__(mamba_config)
118 + try:
119 + from flashinfer.mamba import selective_state_update as _fi_ssu
120 + except ImportError as e:
121 + raise ImportError(
122 + "FlashInfer is required for the flashinfer Mamba SSU backend. "
123 + "Please install flashinfer (>= 0.6.4): "
124 + "pip install flashinfer-python"
125 + ) from e
126 + self._kernel = _fi_ssu
127 +
128 + @property
129 + def name(self) -> str:
130 + return "flashinfer"
131 +
132 + def __call__(
133 + self,
134 + state: torch.Tensor,
135 + x: torch.Tensor,
136 + dt: torch.Tensor,
137 + A: torch.Tensor,
138 + B: torch.Tensor,
139 + C: torch.Tensor,
140 + D: torch.Tensor,
141 + dt_bias: torch.Tensor,
142 + z: torch.Tensor | None = None,
143 + dt_softplus: bool = False,
144 + state_batch_indices: torch.Tensor | None = None,
145 + dst_state_batch_indices: torch.Tensor | None = None,
146 + null_block_id: int = NULL_BLOCK_ID,
147 + out: torch.Tensor | None = None,
148 + num_accepted_tokens: torch.Tensor | None = None,
149 + cu_seqlens: torch.Tensor | None = None,
150 + is_blackwell: bool = False,
151 + ) -> None:
152 + rand_seed = (
153 + torch.randint(0, 2**32, (1,), device=state.device)
154 + if self._mamba_config.enable_stochastic_rounding
155 + else None
156 + )
157 +
158 + self._kernel(
159 + state,
160 + x,
161 + dt,
162 + A,
163 + B,
164 + C,
165 + D=D,
166 + z=z,
167 + dt_bias=dt_bias,
168 + dt_softplus=dt_softplus,
169 + state_batch_indices=state_batch_indices,
170 + dst_state_batch_indices=dst_state_batch_indices,
171 + cu_seqlens=cu_seqlens,
172 + num_accepted_tokens=num_accepted_tokens,
173 + cache_steps=state_batch_indices.size(-1)
174 + if cu_seqlens is not None and state_batch_indices is not None
175 + else 0,
176 + pad_slot_id=null_block_id,
177 + out=out,
178 + rand_seed=rand_seed,
179 + philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10,
180 + )
181 +
182 +
183 +_BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
184 + MambaBackendEnum.TRITON: TritonSSUBackend,
185 + MambaBackendEnum.FLASHINFER: FlashInferSSUBackend,
186 +}
187 +
188 +_mamba_ssu_backend: MambaSSUBackend | None = None
189 +
190 +
191 +def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
192 + """Initialize the global Mamba SSU backend.
193 +
194 + Args:
195 + mamba_config: Mamba configuration.
196 + """
197 + global _mamba_ssu_backend
198 +
199 + backend = mamba_config.backend
200 + if backend not in _BACKEND_REGISTRY:
201 + raise ValueError(
202 + f"Unknown Mamba SSU backend: {backend}. "
203 + f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
204 + )
205 +
206 + _mamba_ssu_backend = _BACKEND_REGISTRY[backend](mamba_config)
207 + logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
208 +
209 +
210 +def get_mamba_ssu_backend() -> MambaSSUBackend:
211 + """Get the current Mamba SSU backend. Raises if not initialized."""
212 + if _mamba_ssu_backend is None:
213 + raise RuntimeError(
214 + "Mamba SSU backend has not been initialized. "
215 + "Call initialize_mamba_ssu_backend() first."
216 + )
217 + return _mamba_ssu_backend
218 +
219 +
220 +def selective_state_update(
221 + state: torch.Tensor,
222 + x: torch.Tensor,
223 + dt: torch.Tensor,
224 + A: torch.Tensor,
225 + B: torch.Tensor,
226 + C: torch.Tensor,
227 + D: torch.Tensor,
228 + dt_bias: torch.Tensor,
229 + z: torch.Tensor | None = None,
230 + dt_softplus: bool = False,
231 + state_batch_indices: torch.Tensor | None = None,
232 + dst_state_batch_indices: torch.Tensor | None = None,
233 + null_block_id: int = NULL_BLOCK_ID,
234 + out: torch.Tensor | None = None,
235 + num_accepted_tokens: torch.Tensor | None = None,
236 + cu_seqlens: torch.Tensor | None = None,
237 + is_blackwell: bool = False,
238 +) -> None:
239 + """Unified dispatch for Mamba selective state update.
240 +
241 + Delegates to the initialized backend (Triton or FlashInfer).
242 + """
243 + get_mamba_ssu_backend()(
244 + state,
245 + x,
246 + dt,
247 + A,
248 + B,
249 + C,
250 + D,
251 + dt_bias,
252 + z=z,
253 + dt_softplus=dt_softplus,
254 + state_batch_indices=state_batch_indices,
255 + dst_state_batch_indices=dst_state_batch_indices,
256 + null_block_id=null_block_id,
257 + out=out,
258 + num_accepted_tokens=num_accepted_tokens,
259 + cu_seqlens=cu_seqlens,
260 + is_blackwell=is_blackwell,
261 + )
vllm/model_executor/models/plamo2.py
+2 -4
@@ -38,10 +38,10 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
38 causal_conv1d_fn,
39 causal_conv1d_update,
40 )
41 -from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update
41 from vllm.model_executor.layers.mamba.ops.ssd_combined import (
42 mamba_chunk_scan_combined_varlen,
43 )
44 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update
45 from vllm.model_executor.layers.quantization import QuantizationConfig
46 from vllm.model_executor.layers.rotary_embedding import get_rope
47 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -447,13 +447,11 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):
447 B,
448 C,
449 D,
450 + dt_bias,
451 z=gate_d.reshape(num_decodes, -1, self.head_dim),
451 - dt_bias=dt_bias,
452 dt_softplus=True,
453 state_batch_indices=state_indices_tensor_d,
454 out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),
455 - enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding,
456 - cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds,
455 )
456
457 # 4. Final linear projection
vllm/v1/worker/gpu/model_runner.py
+4
@@ -36,6 +36,9 @@ from vllm.distributed.parallel_state import (
36 )
37 from vllm.forward_context import BatchDescriptor, set_forward_context
38 from vllm.logger import init_logger
39 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
40 + initialize_mamba_ssu_backend,
41 +)
42 from vllm.model_executor.model_loader import get_model_loader
43 from vllm.multimodal import MULTIMODAL_REGISTRY
44 from vllm.sequence import IntermediateTensors
@@ -360,6 +363,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
363 self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend(
364 self.kv_cache_config, self.vllm_config, self.device
365 )
366 + initialize_mamba_ssu_backend(self.vllm_config.mamba_config)
367 cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
368 attn_cg_support.min_cg_support,
369 attn_cg_support.min_cg_attn_backend,
vllm/v1/worker/gpu_model_runner.py
+4
@@ -56,6 +56,9 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
56 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
57 RoutedExpertsCapturer,
58 )
59 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
60 + initialize_mamba_ssu_backend,
61 +)
62 from vllm.model_executor.layers.rotary_embedding import (
63 MRotaryEmbedding,
64 XDRotaryEmbedding,
@@ -6750,6 +6753,7 @@ class GPUModelRunner(
6753 self.may_add_encoder_only_layers_to_kv_cache_config()
6754 self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
6755 self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
6756 + initialize_mamba_ssu_backend(self.vllm_config.mamba_config)
6757 # The kernel block size for all KV cache groups. For example, if
6758 # kv_cache_manager uses block_size 256 for a given group, but the attention
6759 # backends for that group only supports block_size 64, we will return