Skip to content

vllm.model_executor.layers.mamba.abstract

Classes:

  • MambaBase –

    Base class for Mamba-like layers which support the v1 engine.

Functions:

MambaBase

Bases: AttentionLayerBase

Base class for Mamba-like layers which support the v1 engine. Inherit from this class if you implement a custom layer.

Methods:

Source code in vllm/model_executor/layers/mamba/abstract.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
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
class MambaBase(AttentionLayerBase):
    """Base class for Mamba-like layers which support the v1 engine.
    Inherit from this class if you implement a custom layer.
    """

    # Contains the KV cache (mamba state) for the layer
    # in the shape specified by `self.get_state_shape`.
    kv_cache: tuple[torch.Tensor, ...]
    supports_dcp: bool = False

    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
        """Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views.

        Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm)
        packed contiguously; slice them out and reinterpret per dtype/shape.
        """
        pages = kv_cache.squeeze(dim=(1, 2))
        states: list[torch.Tensor] = []
        offset = 0
        for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()):
            nbytes = prod(shape) * get_dtype_size(dtype)
            state = pages[:, offset : offset + nbytes].view(dtype)
            states.append(state.view(-1, *shape))
            offset += nbytes
        self.kv_cache = tuple(states)

    @abstractmethod
    def get_state_shape(self) -> Iterable[tuple[int, ...]]:
        """Defines the shape of the state.
        For mamba layers this is usually a (conv_state, ssm_state) tuple.
        In this case, returns (conv_state_shape, ssm_state_shape).
        """
        pass

    @property
    @abstractmethod
    def mamba_type(self) -> MambaAttentionBackendEnum:
        pass

    @property
    def is_kv_cache_tp_replicated(self) -> bool:
        return False

    @abstractmethod
    def get_state_dtype(self) -> tuple[torch.dtype, ...]:
        pass

    def get_prefill_checkpoint_alignment(self) -> int | None:
        """Token alignment of the checkpoint offsets this layer can export from
        inside a prefill, if any; in align mode, the platform makes Mamba blocks
        multiples of it.
        """
        return None

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
        mamba_block_size = vllm_config.cache_config.mamba_block_size
        assert mamba_block_size is not None
        page_size_padded = vllm_config.cache_config.mamba_page_size_padded
        return MambaSpec(
            shapes=tuple(self.get_state_shape()),
            dtypes=self.get_state_dtype(),
            block_size=mamba_block_size,
            page_size_padded=page_size_padded,
            mamba_type=self.mamba_type,
            tp_replicated=self.is_kv_cache_tp_replicated,
            mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode,
            # ReplaySSM verifies the whole window off one checkpoint, so it
            # never writes the baseline's per-draft-token state slots.
            num_speculative_blocks=(
                0
                if vllm_config.cache_config.use_replayssm
                else vllm_config.num_speculative_tokens
            ),
        )

    def get_attn_backend(self) -> type[AttentionBackend]:
        """Get the attention backend class for this Mamba layer."""
        return get_mamba_attn_backend(self.mamba_type)

bind_kv_cache(kv_cache)

Unpack a raw [B, 1, 1, C] int8 page view into per-state views.

Each block's C bytes hold the layer's states (e.g. conv, ssm) packed contiguously; slice them out and reinterpret per dtype/shape.

Source code in vllm/model_executor/layers/mamba/abstract.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
    """Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views.

    Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm)
    packed contiguously; slice them out and reinterpret per dtype/shape.
    """
    pages = kv_cache.squeeze(dim=(1, 2))
    states: list[torch.Tensor] = []
    offset = 0
    for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()):
        nbytes = prod(shape) * get_dtype_size(dtype)
        state = pages[:, offset : offset + nbytes].view(dtype)
        states.append(state.view(-1, *shape))
        offset += nbytes
    self.kv_cache = tuple(states)

get_attn_backend()

Get the attention backend class for this Mamba layer.

Source code in vllm/model_executor/layers/mamba/abstract.py 
93
94
95
def get_attn_backend(self) -> type[AttentionBackend]:
    """Get the attention backend class for this Mamba layer."""
    return get_mamba_attn_backend(self.mamba_type)

get_prefill_checkpoint_alignment()

Token alignment of the checkpoint offsets this layer can export from inside a prefill, if any; in align mode, the platform makes Mamba blocks multiples of it.

Source code in vllm/model_executor/layers/mamba/abstract.py 
65
66
67
68
69
70
def get_prefill_checkpoint_alignment(self) -> int | None:
    """Token alignment of the checkpoint offsets this layer can export from
    inside a prefill, if any; in align mode, the platform makes Mamba blocks
    multiples of it.
    """
    return None

get_state_shape() abstractmethod

Defines the shape of the state. For mamba layers this is usually a (conv_state, ssm_state) tuple. In this case, returns (conv_state_shape, ssm_state_shape).

Source code in vllm/model_executor/layers/mamba/abstract.py 
44
45
46
47
48
49
50
@abstractmethod
def get_state_shape(self) -> Iterable[tuple[int, ...]]:
    """Defines the shape of the state.
    For mamba layers this is usually a (conv_state, ssm_state) tuple.
    In this case, returns (conv_state_shape, ssm_state_shape).
    """
    pass

get_mamba_prefill_checkpoint_alignment(vllm_config)

LCM of the Mamba layers' prefill checkpoint alignments.

None unless every Mamba layer declares one: the scheduler takes internal checkpoints only when every Mamba group can export them.

Source code in vllm/model_executor/layers/mamba/abstract.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def get_mamba_prefill_checkpoint_alignment(vllm_config: VllmConfig) -> int | None:
    """LCM of the Mamba layers' prefill checkpoint alignments.

    None unless every Mamba layer declares one: the scheduler takes internal
    checkpoints only when every Mamba group can export them.
    """
    alignment = 1
    for layer in get_layers_from_vllm_config(
        vllm_config,
        MambaBase,  # type: ignore[type-abstract]
    ).values():
        layer_alignment = layer.get_prefill_checkpoint_alignment()
        if layer_alignment is None:
            return None
        alignment = lcm(alignment, layer_alignment)
    return alignment