Skip to content

vllm.config.mamba

Classes:

  • MambaBackendEnum –

    Enumeration of supported Mamba SSU (selective state update) backends.

  • MambaConfig –

    Configuration for Mamba SSM backends.

MambaBackendEnum

Bases: Enum

Enumeration of supported Mamba SSU (selective state update) backends.

Source code in vllm/config/mamba.py 
25
26
27
28
29
30
class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):
    """Enumeration of supported Mamba SSU (selective state update) backends."""

    TRITON = "triton"
    FLASHINFER = "flashinfer"
    CPU = "cpu"

MambaConfig

Configuration for Mamba SSM backends.

Methods:

Attributes:

Source code in vllm/config/mamba.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
@config
class MambaConfig:
    """Configuration for Mamba SSM backends."""

    backend: MambaBackendEnum = MambaBackendEnum.TRITON
    """Mamba SSU backend to use."""

    enable_stochastic_rounding: bool = False
    """Enable stochastic rounding when writing SSM state to fp16 cache.
    Uses random bits to unbias the rounding error, which can improve
    numerical stability for long sequences. The Triton backend requires
    compute capability 10.0 on CUDA or gfx950 on ROCm."""
    stochastic_rounding_philox_rounds: int = 0
    """Number of Philox PRNG rounds for stochastic rounding random number
    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."""

    prefill_backend: MambaBackendEnum = MambaBackendEnum.TRITON
    """Kernel backend for the SSD scan of Mamba2 prefills. "flashinfer" runs
    FlashInfer's fused SSD scan (SM100/SM103 only); prefills that store Mamba
    checkpoints always use "triton"."""

    @field_validator("backend", "prefill_backend", mode="before")
    @classmethod
    def validate_backend_before(cls, value: Any) -> Any:
        """Enable parsing of the backend enum fields from strings."""
        if isinstance(value, str):
            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 validate_prefill_backend(self) -> None:
        if self.prefill_backend == MambaBackendEnum.CPU:
            raise ValueError(
                "Unknown Mamba prefill backend: 'cpu'. "
                "Valid options are: triton, flashinfer"
            )
        if self.prefill_backend == MambaBackendEnum.FLASHINFER:
            from vllm.platforms import current_platform

            if not current_platform.is_device_capability_family(100):
                raise ValueError(
                    "--mamba-prefill-backend flashinfer requires SM100/SM103."
                )

    def __post_init__(self):
        self.validate_ssu_algorithm()
        self.validate_prefill_backend()
        if self.enable_stochastic_rounding:
            from vllm.platforms import current_platform

            if current_platform.is_rocm():
                from vllm.platforms.rocm import on_gfx950

                if not on_gfx950():
                    raise ValueError(
                        "Stochastic rounding for Mamba cache on ROCm requires "
                        "gfx950 (MI350/MI355). Please do not specify "
                        "`--enable-mamba-cache-stochastic-rounding`."
                    )
            elif not current_platform.is_cuda():
                raise ValueError(
                    "Stochastic rounding for Mamba cache is only supported "
                    "on NVIDIA CUDA and AMD gfx950 platforms. Please do not specify "
                    "`--enable-mamba-cache-stochastic-rounding`."
                )
            elif (
                self.backend == MambaBackendEnum.TRITON
                and not current_platform.is_device_capability_family(100)
            ):
                raise ValueError(
                    "Stochastic rounding for Mamba cache with triton backend requires "
                    "compute capability 10.0 (data center Blackwell). The `cvt.rs` "
                    "PTX instruction is not supported on your GPU. Please do not "
                    "specify `--enable-mamba-cache-stochastic-rounding`, "
                    "or set `--mamba-backend flashinfer`."
                )

backend = MambaBackendEnum.TRITON class-attribute instance-attribute

Mamba SSU backend to use.

enable_stochastic_rounding = False class-attribute instance-attribute

Enable stochastic rounding when writing SSM state to fp16 cache. Uses random bits to unbias the rounding error, which can improve numerical stability for long sequences. The Triton backend requires compute capability 10.0 on CUDA or gfx950 on ROCm.

prefill_backend = MambaBackendEnum.TRITON class-attribute instance-attribute

Kernel backend for the SSD scan of Mamba2 prefills. "flashinfer" runs FlashInfer's fused SSD scan (SM100/SM103 only); prefills that store Mamba checkpoints always use "triton".

ssu_algorithm = None class-attribute instance-attribute

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.

stochastic_rounding_philox_rounds = 0 class-attribute instance-attribute

Number of Philox PRNG rounds for stochastic rounding random number generation. 0 uses the Triton default. Higher values improve randomness quality at the cost of compute.

validate_backend_before(value) classmethod

Enable parsing of the backend enum fields from strings.

Source code in vllm/config/mamba.py 
64
65
66
67
68
69
70
@field_validator("backend", "prefill_backend", mode="before")
@classmethod
def validate_backend_before(cls, value: Any) -> Any:
    """Enable parsing of the backend enum fields from strings."""
    if isinstance(value, str):
        return MambaBackendEnum[value.upper()]
    return value

_MambaBackendEnumMeta

Bases: EnumMeta

Metaclass for MambaBackendEnum to provide better error messages.

Source code in vllm/config/mamba.py 
12
13
14
15
16
17
18
19
20
21
22
class _MambaBackendEnumMeta(EnumMeta):
    """Metaclass for MambaBackendEnum to provide better error messages."""

    def __getitem__(cls, name: str):
        try:
            return super().__getitem__(name)
        except KeyError:
            valid = ", ".join(cls.__members__.keys())
            raise ValueError(
                f"Unknown Mamba SSU backend: '{name}'. Valid options are: {valid}"
            ) from None