diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 703a5df163..c37f74eb50 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/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 diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index a68842f4bf..e8988aab19 100644 --- a/vllm/config/mamba.py +++ b/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 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 96c8da1f29..41b9d3c6e4 100644 --- a/vllm/engine/arg_utils.py +++ b/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) diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 1795a036b1..24bdba688e 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/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, )