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, 02:15 UTC
Gate SSU dispatch setup (#40039)
Signed-off-by: Roi Koren <roik@nvidia.com>
roikoren755 committed
Apr 16, 2026 at 23:06 UTC
b897f00c9c35dcc7b229973cf665a49d7082b8bf
4 files changed
+73
-10
tests/kernels/mamba/test_ssu_dispatch.py
+48
-4
@@ -13,6 +13,11 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
13
selective_state_update,
14
)
15
from vllm.utils.torch_utils import set_random_seed
16
+from vllm.v1.kv_cache_interface import (
17
+ KVCacheConfig,
18
+ KVCacheGroupSpec,
19
+ MambaSpec,
20
+)
21
22
try:
23
import flashinfer.mamba # noqa: F401
@@ -22,22 +27,40 @@ except ImportError:
27
HAS_FLASHINFER = False
28
29
30
+def _kv_cache_config_with_ssu(mamba_type: str = "mamba2") -> KVCacheConfig:
31
+ spec = MambaSpec(
32
+ block_size=16,
33
+ shapes=((16, 64),),
34
+ dtypes=(torch.float16,),
35
+ mamba_type=mamba_type,
36
+ )
37
+ return KVCacheConfig(
38
+ num_blocks=1,
39
+ kv_cache_tensors=[],
40
+ kv_cache_groups=[KVCacheGroupSpec(layer_names=["l0"], kv_cache_spec=spec)],
41
+ )
42
+
43
+
44
def test_default_backend_is_triton():
26
- initialize_mamba_ssu_backend(MambaConfig())
45
+ initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu())
46
backend = get_mamba_ssu_backend()
47
assert isinstance(backend, TritonSSUBackend)
48
assert backend.name == "triton"
49
50
51
def test_explicit_triton_backend():
33
- initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
52
+ initialize_mamba_ssu_backend(
53
+ MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
54
+ )
55
backend = get_mamba_ssu_backend()
56
assert isinstance(backend, TritonSSUBackend)
57
58
59
@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
60
def test_flashinfer_backend_init():
40
- initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
61
+ initialize_mamba_ssu_backend(
62
+ MambaConfig(backend=MambaBackendEnum.FLASHINFER), _kv_cache_config_with_ssu()
63
+ )
64
backend = get_mamba_ssu_backend()
65
assert isinstance(backend, FlashInferSSUBackend)
66
assert backend.name == "flashinfer"
@@ -53,6 +76,25 @@ def test_uninitialized_backend_raises():
76
mod._mamba_ssu_backend = old
77
78
79
+@pytest.mark.parametrize(
80
+ "mamba_type", ["linear_attention", "gdn_attention", "short_conv"]
81
+)
82
+def test_init_is_noop_for_non_ssu_mamba_type(mamba_type):
83
+ import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
84
+
85
+ old = mod._mamba_ssu_backend
86
+ mod._mamba_ssu_backend = None
87
+ try:
88
+ initialize_mamba_ssu_backend(
89
+ MambaConfig(), _kv_cache_config_with_ssu(mamba_type)
90
+ )
91
+ assert mod._mamba_ssu_backend is None
92
+ with pytest.raises(RuntimeError, match="not been initialized"):
93
+ get_mamba_ssu_backend()
94
+ finally:
95
+ mod._mamba_ssu_backend = old
96
+
97
+
98
@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
99
def test_flashinfer_import_error():
100
with pytest.raises(ImportError, match="FlashInfer is required"):
@@ -61,7 +103,9 @@ def test_flashinfer_import_error():
103
104
def test_triton_basic_call():
105
set_random_seed(0)
64
- initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
106
+ initialize_mamba_ssu_backend(
107
+ MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
108
+ )
109
device = "cuda"
110
batch_size = 2
111
dim = 64
vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
+19
-4
@@ -15,6 +15,7 @@ import torch
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
+from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
19
20
logger = init_logger(__name__)
21
@@ -188,12 +189,22 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
189
_mamba_ssu_backend: MambaSSUBackend | None = None
190
191
191
-def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
192
+def initialize_mamba_ssu_backend(
193
+ mamba_config: MambaConfig,
194
+ kv_cache_config: KVCacheConfig,
195
+) -> None:
196
"""Initialize the global Mamba SSU backend.
197
194
- Args:
195
- mamba_config: Mamba configuration.
198
+ No-op if `kv_cache_config` contains no specs that call
199
+ selective_state_update.
200
"""
201
+ if not any(
202
+ isinstance(g.kv_cache_spec, MambaSpec)
203
+ and g.kv_cache_spec.mamba_type in ("mamba1", "mamba2")
204
+ for g in kv_cache_config.kv_cache_groups
205
+ ):
206
+ return
207
+
208
global _mamba_ssu_backend
209
210
backend = mamba_config.backend
@@ -203,7 +214,11 @@ def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
214
f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
215
)
216
206
- _mamba_ssu_backend = _BACKEND_REGISTRY[backend](mamba_config)
217
+ backend_cls = _BACKEND_REGISTRY[backend]
218
+ if isinstance(_mamba_ssu_backend, backend_cls):
219
+ return
220
+
221
+ _mamba_ssu_backend = backend_cls(mamba_config)
222
logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
223
224
vllm/v1/worker/gpu/model_runner.py
+3
-1
@@ -363,7 +363,9 @@ 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)
366
+ initialize_mamba_ssu_backend(
367
+ self.vllm_config.mamba_config, self.kv_cache_config
368
+ )
369
cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
370
attn_cg_support.min_cg_support,
371
attn_cg_support.min_cg_attn_backend,
vllm/v1/worker/gpu_model_runner.py
+3
-1
@@ -6738,7 +6738,9 @@ class GPUModelRunner(
6738
self.may_add_encoder_only_layers_to_kv_cache_config()
6739
self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
6740
self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
6741
- initialize_mamba_ssu_backend(self.vllm_config.mamba_config)
6741
+ initialize_mamba_ssu_backend(
6742
+ self.vllm_config.mamba_config, self.kv_cache_config
6743
+ )
6744
# The kernel block size for all KV cache groups. For example, if
6745
# kv_cache_manager uses block_size 256 for a given group, but the attention
6746
# backends for that group only supports block_size 64, we will return