[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>

tests/kernels/mamba/test_ssu_dispatch.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from unittest.mock import Mock
+
 import pytest
 import torch
 
-from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     FlashInferSSUBackend,
     TritonSSUBackend,
@@ -69,6 +71,48 @@ def test_flashinfer_backend_init():
     assert backend.name == "flashinfer"
 
 
+@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
+@pytest.mark.parametrize(
+    ("algorithm", "expected"),
+    [
+        (None, "auto"),
+        ("auto", "auto"),
+        ("simple", "simple"),
+        ("vertical", "vertical"),
+        ("horizontal", "horizontal"),
+    ],
+)
+def test_flashinfer_forwards_ssu_algorithm(
+    algorithm: MambaSSUAlgorithm | None,
+    expected: MambaSSUAlgorithm,
+    monkeypatch,
+):
+    import flashinfer.mamba
+
+    kernel = Mock()
+    monkeypatch.setattr(flashinfer.mamba, "selective_state_update", kernel)
+    backend = FlashInferSSUBackend(
+        MambaConfig(
+            backend=MambaBackendEnum.FLASHINFER,
+            ssu_algorithm=algorithm,
+        )
+    )
+
+    tensor = torch.empty(1)
+    backend(
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+    )
+
+    assert kernel.call_args.kwargs["algorithm"] == expected
+
+
 def test_uninitialized_backend_raises():
     import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
 

vllm/config/mamba.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from enum import Enum, EnumMeta
-from typing import Any
+from typing import Any, Literal, get_args
 
 from pydantic import field_validator
 
@@ -30,6 +30,9 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):
     CPU = "cpu"
 
 
+MambaSSUAlgorithm = Literal["auto", "simple", "vertical", "horizontal"]
+
+
 @config
 class MambaConfig:
     """Configuration for Mamba SSM backends."""
@@ -46,6 +49,12 @@ class MambaConfig:
     generation. 0 uses the Triton default. Higher values improve randomness
     quality at the cost of compute."""
 
+    ssu_algorithm: MambaSSUAlgorithm | None = None
+    """Selective state update algorithm to use with the FlashInfer backend.
+    None defaults to FlashInfer's "auto" algorithm. Forced algorithms must
+    be supported by FlashInfer for the active GPU, state dtype, and decoding
+    mode."""
+
     @field_validator("backend", mode="before")
     @classmethod
     def validate_backend_before(cls, value: Any) -> Any:
@@ -54,7 +63,25 @@ class MambaConfig:
             return MambaBackendEnum[value.upper()]
         return value
 
+    def validate_ssu_algorithm(self) -> None:
+        if self.ssu_algorithm is None:
+            return
+        valid_algorithms = get_args(MambaSSUAlgorithm)
+        if self.ssu_algorithm not in valid_algorithms:
+            valid = ", ".join(valid_algorithms)
+            raise ValueError(
+                f"Unknown Mamba SSU algorithm: '{self.ssu_algorithm}'. "
+                f"Valid options are: {valid}"
+            )
+        if self.backend != MambaBackendEnum.FLASHINFER:
+            raise ValueError(
+                "Mamba SSU algorithm selection is only supported with the "
+                "FlashInfer backend. Please set `--mamba-backend flashinfer`, "
+                "or omit `--mamba-ssu-algorithm`."
+            )
+
     def __post_init__(self):
+        self.validate_ssu_algorithm()
         if self.enable_stochastic_rounding:
             from vllm.platforms import current_platform
 

vllm/engine/arg_utils.py

@@ -77,7 +77,7 @@ from vllm.config.device import Device
 from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend
 from vllm.config.load import SafetensorsLoadStrategy
 from vllm.config.lora import MaxLoRARanks
-from vllm.config.mamba import MambaBackendEnum
+from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm
 from vllm.config.model import (
     ConvertOption,
     HfOverrides,
@@ -707,6 +707,7 @@ class EngineArgs:
     use_replayssm: bool = CacheConfig.use_replayssm
 
     mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON
+    mamba_ssu_algorithm: MambaSSUAlgorithm | None = None
     enable_mamba_cache_stochastic_rounding: bool = (
         MambaConfig.enable_stochastic_rounding
     )
@@ -961,6 +962,9 @@ class EngineArgs:
             description=MambaConfig.__doc__,
         )
         mamba_group.add_argument("--mamba-backend", **mamba_kwargs["backend"])
+        mamba_group.add_argument(
+            "--mamba-ssu-algorithm", **mamba_kwargs["ssu_algorithm"]
+        )
         mamba_group.add_argument(
             "--enable-mamba-cache-stochastic-rounding",
             **mamba_kwargs["enable_stochastic_rounding"],
@@ -2356,6 +2360,8 @@ class EngineArgs:
             mamba_config.backend = MambaBackendEnum[self.mamba_backend.upper()]
         else:
             mamba_config.backend = self.mamba_backend
+        if self.mamba_ssu_algorithm is not None:
+            mamba_config.ssu_algorithm = self.mamba_ssu_algorithm
         if self.enable_mamba_cache_stochastic_rounding:
             mamba_config.enable_stochastic_rounding = (
                 self.enable_mamba_cache_stochastic_rounding
@@ -2364,6 +2370,7 @@ class EngineArgs:
             mamba_config.stochastic_rounding_philox_rounds = (
                 self.mamba_cache_philox_rounds
             )
+        mamba_config.validate_ssu_algorithm()
 
         # Kernel config overrides
         kernel_config = copy.deepcopy(self.kernel_config)

vllm/model_executor/layers/mamba/ops/ssu_dispatch.py

@@ -13,7 +13,7 @@ from abc import ABC, abstractmethod
 
 import torch
 
-from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
 from vllm.logger import init_logger
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
@@ -126,8 +126,13 @@ class FlashInferSSUBackend(MambaSSUBackend):
                 "Please install flashinfer (>= 0.6.4): "
                 "pip install flashinfer-python"
             ) from e
+        logger.info_once("Using FlashInfer Mamba SSU algorithm: %s", self._algorithm)
         self._kernel = _fi_ssu
 
+    @property
+    def _algorithm(self) -> MambaSSUAlgorithm:
+        return self._mamba_config.ssu_algorithm or "auto"
+
     @property
     def name(self) -> str:
         return "flashinfer"
@@ -157,7 +162,6 @@ class FlashInferSSUBackend(MambaSSUBackend):
             if self._mamba_config.enable_stochastic_rounding
             else None
         )
-
         self._kernel(
             state,
             x,
@@ -180,6 +184,7 @@ class FlashInferSSUBackend(MambaSSUBackend):
             out=out,
             rand_seed=rand_seed,
             philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10,
+            algorithm=self._algorithm,
         )