| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import abstractmethod |
| 4 | from collections.abc import Iterable |
| 5 | from math import prod |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 11 | from vllm.utils.torch_utils import get_dtype_size |
| 12 | from vllm.v1.attention.backend import AttentionBackend |
| 13 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 14 | from vllm.v1.attention.selector import get_mamba_attn_backend |
| 15 | from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec |
| 16 | |
| 17 | |
| 18 | class MambaBase(AttentionLayerBase): |
| 19 | """Base class for Mamba-like layers which support the v1 engine. |
| 20 | Inherit from this class if you implement a custom layer. |
| 21 | """ |
| 22 | |
| 23 | # Contains the KV cache (mamba state) for the layer |
| 24 | # in the shape specified by `self.get_state_shape`. |
| 25 | kv_cache: tuple[torch.Tensor, ...] |
| 26 | supports_dcp: bool = False |
| 27 | |
| 28 | def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: |
| 29 | """Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views. |
| 30 | |
| 31 | Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm) |
| 32 | packed contiguously; slice them out and reinterpret per dtype/shape. |
| 33 | """ |
| 34 | pages = kv_cache.squeeze(dim=(1, 2)) |
| 35 | states: list[torch.Tensor] = [] |
| 36 | offset = 0 |
| 37 | for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()): |
| 38 | nbytes = prod(shape) * get_dtype_size(dtype) |
| 39 | state = pages[:, offset : offset + nbytes].view(dtype) |
| 40 | states.append(state.view(-1, *shape)) |
| 41 | offset += nbytes |
| 42 | self.kv_cache = tuple(states) |
| 43 | |
| 44 | @abstractmethod |
| 45 | def get_state_shape(self) -> Iterable[tuple[int, ...]]: |
| 46 | """Defines the shape of the state. |
| 47 | For mamba layers this is usually a (conv_state, ssm_state) tuple. |
| 48 | In this case, returns (conv_state_shape, ssm_state_shape). |
| 49 | """ |
| 50 | pass |
| 51 | |
| 52 | @property |
| 53 | @abstractmethod |
| 54 | def mamba_type(self) -> MambaAttentionBackendEnum: |
| 55 | pass |
| 56 | |
| 57 | @property |
| 58 | def is_kv_cache_tp_replicated(self) -> bool: |
| 59 | return False |
| 60 | |
| 61 | @abstractmethod |
| 62 | def get_state_dtype(self) -> tuple[torch.dtype, ...]: |
| 63 | pass |
| 64 | |
| 65 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: |
| 66 | mamba_block_size = vllm_config.cache_config.mamba_block_size |
| 67 | assert mamba_block_size is not None |
| 68 | page_size_padded = vllm_config.cache_config.mamba_page_size_padded |
| 69 | return MambaSpec( |
| 70 | shapes=tuple(self.get_state_shape()), |
| 71 | dtypes=self.get_state_dtype(), |
| 72 | block_size=mamba_block_size, |
| 73 | page_size_padded=page_size_padded, |
| 74 | mamba_type=self.mamba_type, |
| 75 | tp_replicated=self.is_kv_cache_tp_replicated, |
| 76 | mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, |
| 77 | # ReplaySSM verifies the whole window off one checkpoint, so it |
| 78 | # never writes the baseline's per-draft-token state slots. |
| 79 | num_speculative_blocks=( |
| 80 | 0 |
| 81 | if vllm_config.cache_config.use_replayssm |
| 82 | else vllm_config.num_speculative_tokens |
| 83 | ), |
| 84 | ) |
| 85 | |
| 86 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 87 | """Get the attention backend class for this Mamba layer.""" |
| 88 | return get_mamba_attn_backend(self.mamba_type) |