[Kernel] Add support for Flashinfer Mamba SSU algorithm selection (#50157)
Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
amitz-nv committed
Aug 4, 2026 at 19:40 UTC
52c0e3cb08b8178829e1f2db4ac90e9bb98c8a5f
4 files changed
+88
-5
tests/kernels/mamba/test_ssu_dispatch.py
+45
-1
@@ -1,10 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from unittest.mock import Mock
5
+
6
import pytest
7
import torch
8
7
-from vllm.config.mamba import MambaBackendEnum, MambaConfig
9
+from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
10
from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
11
FlashInferSSUBackend,
12
TritonSSUBackend,
@@ -69,6 +71,48 @@ def test_flashinfer_backend_init():
71
assert backend.name == "flashinfer"
72
73
74
+@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
75
+@pytest.mark.parametrize(
76
+ ("algorithm", "expected"),
77
+ [
78
+ (None, "auto"),
79
+ ("auto", "auto"),
80
+ ("simple", "simple"),
81
+ ("vertical", "vertical"),
82
+ ("horizontal", "horizontal"),
83
+ ],
84
+)
85
+def test_flashinfer_forwards_ssu_algorithm(
86
+ algorithm: MambaSSUAlgorithm | None,
87
+ expected: MambaSSUAlgorithm,
88
+ monkeypatch,
89
+):
90
+ import flashinfer.mamba
91
+
92
+ kernel = Mock()
93
+ monkeypatch.setattr(flashinfer.mamba, "selective_state_update", kernel)
94
+ backend = FlashInferSSUBackend(
95
+ MambaConfig(
96
+ backend=MambaBackendEnum.FLASHINFER,
97
+ ssu_algorithm=algorithm,
98
+ )
99
+ )
100
+
101
+ tensor = torch.empty(1)
102
+ backend(
103
+ tensor,
104
+ tensor,
105
+ tensor,
106
+ tensor,
107
+ tensor,
108
+ tensor,
109
+ tensor,
110
+ tensor,
111
+ )
112
+
113
+ assert kernel.call_args.kwargs["algorithm"] == expected
114
+
115
+
116
def test_uninitialized_backend_raises():
117
import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
118
vllm/config/mamba.py
+28
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from enum import Enum, EnumMeta
5
-from typing import Any
5
+from typing import Any, Literal, get_args
6
7
from pydantic import field_validator
8
@@ -30,6 +30,9 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):
30
CPU = "cpu"
31
32
33
+MambaSSUAlgorithm = Literal["auto", "simple", "vertical", "horizontal"]
34
+
35
+
36
@config
37
class MambaConfig:
38
"""Configuration for Mamba SSM backends."""
@@ -46,6 +49,12 @@ class MambaConfig:
49
generation. 0 uses the Triton default. Higher values improve randomness
50
quality at the cost of compute."""
51
52
+ ssu_algorithm: MambaSSUAlgorithm | None = None
53
+ """Selective state update algorithm to use with the FlashInfer backend.
54
+ None defaults to FlashInfer's "auto" algorithm. Forced algorithms must
55
+ be supported by FlashInfer for the active GPU, state dtype, and decoding
56
+ mode."""
57
+
58
@field_validator("backend", mode="before")
59
@classmethod
60
def validate_backend_before(cls, value: Any) -> Any:
@@ -54,7 +63,25 @@ class MambaConfig:
63
return MambaBackendEnum[value.upper()]
64
return value
65
66
+ def validate_ssu_algorithm(self) -> None:
67
+ if self.ssu_algorithm is None:
68
+ return
69
+ valid_algorithms = get_args(MambaSSUAlgorithm)
70
+ if self.ssu_algorithm not in valid_algorithms:
71
+ valid = ", ".join(valid_algorithms)
72
+ raise ValueError(
73
+ f"Unknown Mamba SSU algorithm: '{self.ssu_algorithm}'. "
74
+ f"Valid options are: {valid}"
75
+ )
76
+ if self.backend != MambaBackendEnum.FLASHINFER:
77
+ raise ValueError(
78
+ "Mamba SSU algorithm selection is only supported with the "
79
+ "FlashInfer backend. Please set `--mamba-backend flashinfer`, "
80
+ "or omit `--mamba-ssu-algorithm`."
81
+ )
82
+
83
def __post_init__(self):
84
+ self.validate_ssu_algorithm()
85
if self.enable_stochastic_rounding:
86
from vllm.platforms import current_platform
87
vllm/engine/arg_utils.py
+8
-1
@@ -77,7 +77,7 @@ from vllm.config.device import Device
77
from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend
78
from vllm.config.load import SafetensorsLoadStrategy
79
from vllm.config.lora import MaxLoRARanks
80
-from vllm.config.mamba import MambaBackendEnum
80
+from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm
81
from vllm.config.model import (
82
ConvertOption,
83
HfOverrides,
@@ -707,6 +707,7 @@ class EngineArgs:
707
use_replayssm: bool = CacheConfig.use_replayssm
708
709
mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON
710
+ mamba_ssu_algorithm: MambaSSUAlgorithm | None = None
711
enable_mamba_cache_stochastic_rounding: bool = (
712
MambaConfig.enable_stochastic_rounding
713
)
@@ -961,6 +962,9 @@ class EngineArgs:
962
description=MambaConfig.__doc__,
963
)
964
mamba_group.add_argument("--mamba-backend", **mamba_kwargs["backend"])
965
+ mamba_group.add_argument(
966
+ "--mamba-ssu-algorithm", **mamba_kwargs["ssu_algorithm"]
967
+ )
968
mamba_group.add_argument(
969
"--enable-mamba-cache-stochastic-rounding",
970
**mamba_kwargs["enable_stochastic_rounding"],
@@ -2356,6 +2360,8 @@ class EngineArgs:
2360
mamba_config.backend = MambaBackendEnum[self.mamba_backend.upper()]
2361
else:
2362
mamba_config.backend = self.mamba_backend
2363
+ if self.mamba_ssu_algorithm is not None:
2364
+ mamba_config.ssu_algorithm = self.mamba_ssu_algorithm
2365
if self.enable_mamba_cache_stochastic_rounding:
2366
mamba_config.enable_stochastic_rounding = (
2367
self.enable_mamba_cache_stochastic_rounding
@@ -2364,6 +2370,7 @@ class EngineArgs:
2370
mamba_config.stochastic_rounding_philox_rounds = (
2371
self.mamba_cache_philox_rounds
2372
)
2373
+ mamba_config.validate_ssu_algorithm()
2374
2375
# Kernel config overrides
2376
kernel_config = copy.deepcopy(self.kernel_config)
vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
+7
-2
@@ -13,7 +13,7 @@ from abc import ABC, abstractmethod
13
14
import torch
15
16
-from vllm.config.mamba import MambaBackendEnum, MambaConfig
16
+from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
17
from vllm.logger import init_logger
18
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
19
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
@@ -126,8 +126,13 @@ class FlashInferSSUBackend(MambaSSUBackend):
126
"Please install flashinfer (>= 0.6.4): "
127
"pip install flashinfer-python"
128
) from e
129
+ logger.info_once("Using FlashInfer Mamba SSU algorithm: %s", self._algorithm)
130
self._kernel = _fi_ssu
131
132
+ @property
133
+ def _algorithm(self) -> MambaSSUAlgorithm:
134
+ return self._mamba_config.ssu_algorithm or "auto"
135
+
136
@property
137
def name(self) -> str:
138
return "flashinfer"
@@ -157,7 +162,6 @@ class FlashInferSSUBackend(MambaSSUBackend):
162
if self._mamba_config.enable_stochastic_rounding
163
else None
164
)
160
-
165
self._kernel(
166
state,
167
x,
@@ -180,6 +184,7 @@ class FlashInferSSUBackend(MambaSSUBackend):
184
out=out,
185
rand_seed=rand_seed,
186
philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10,
187
+ algorithm=self._algorithm,
188
)
189
190