| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from enum import Enum, EnumMeta |
| 5 | from typing import Any, Literal, get_args |
| 6 | |
| 7 | from pydantic import field_validator |
| 8 | |
| 9 | from vllm.config.utils import config |
| 10 | |
| 11 | |
| 12 | class _MambaBackendEnumMeta(EnumMeta): |
| 13 | """Metaclass for MambaBackendEnum to provide better error messages.""" |
| 14 | |
| 15 | def __getitem__(cls, name: str): |
| 16 | try: |
| 17 | return super().__getitem__(name) |
| 18 | except KeyError: |
| 19 | valid = ", ".join(cls.__members__.keys()) |
| 20 | raise ValueError( |
| 21 | f"Unknown Mamba SSU backend: '{name}'. Valid options are: {valid}" |
| 22 | ) from None |
| 23 | |
| 24 | |
| 25 | class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): |
| 26 | """Enumeration of supported Mamba SSU (selective state update) backends.""" |
| 27 | |
| 28 | TRITON = "triton" |
| 29 | FLASHINFER = "flashinfer" |
| 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.""" |
| 39 | |
| 40 | backend: MambaBackendEnum = MambaBackendEnum.TRITON |
| 41 | """Mamba SSU backend to use.""" |
| 42 | |
| 43 | enable_stochastic_rounding: bool = False |
| 44 | """Enable stochastic rounding when writing SSM state to fp16 cache. |
| 45 | Uses random bits to unbias the rounding error, which can improve |
| 46 | numerical stability for long sequences.""" |
| 47 | stochastic_rounding_philox_rounds: int = 0 |
| 48 | """Number of Philox PRNG rounds for stochastic rounding random number |
| 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: |
| 61 | """Enable parsing of the `backend` enum type from string.""" |
| 62 | if isinstance(value, str): |
| 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 | |
| 88 | if not current_platform.is_cuda(): |
| 89 | raise ValueError( |
| 90 | "Stochastic rounding for Mamba cache is only supported " |
| 91 | "on NVIDIA CUDA platforms. Please do not specify " |
| 92 | "`--enable-mamba-cache-stochastic-rounding`." |
| 93 | ) |
| 94 | if ( |
| 95 | self.backend == MambaBackendEnum.TRITON |
| 96 | and not current_platform.is_device_capability_family(100) |
| 97 | ): |
| 98 | raise ValueError( |
| 99 | "Stochastic rounding for Mamba cache with triton backend requires " |
| 100 | "compute capability 10.0 (data center Blackwell). The `cvt.rs` " |
| 101 | "PTX instruction is not supported on your GPU. Please do not " |
| 102 | "specify `--enable-mamba-cache-stochastic-rounding`, " |
| 103 | "or set `--mamba-backend flashinfer`." |
| 104 | ) |