vllm.model_executor.layers.mamba.mamba_utils

Classes:

• MambaCopySpec –

  Data class specifying the memory-copy parameters for Mamba states used for

• MambaStateDtypeCalculator –
• MambaStateShapeCalculator –

Functions:

• get_conv_copy_spec –

  Return a MambaCopySpec for copying a convolutional state slice.

• get_conv_state_layout –

  Return the SSM conv state layout.

• get_temporal_copy_spec –

  Return a MambaCopySpec for copying a temporal state slice.

• is_conv_state_dim_first –

  True when the conv state is stored as (dim, state_len) per block.

Attributes:

• MambaStateCopyFunc (TypeAlias) –

  Type alias for a function that computes a MambaCopySpec for copying state slices.

MambaStateCopyFunc = Callable[[torch.Tensor, list[int], int, int], MambaCopySpec] module-attribute

Type alias for a function that computes a MambaCopySpec for copying state slices. Parameters: state: torch.Tensor - the Mamba state tensor (e.g., conv or temporal states). block_ids: list[int] - the list of block indices for the state to copy. cur_block_idx: int - current block index within block_ids to copy from. num_accepted_tokens: int - number of accepted tokens used to compute the copy offset. Range: 1 .. 1 + num_speculative_tokens (inclusive).

MambaCopySpec dataclass

Data class specifying the memory-copy parameters for Mamba states used for prefix caching in align mode.

Attributes:

• start_addr (int) –

  Starting address for the memory copy operation.

• num_elements (int) –

  Number of elements to copy from the starting address.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@dataclass
class MambaCopySpec:
    """Data class specifying the memory-copy parameters for Mamba states used for
    prefix caching in align mode.

    Attributes:
        start_addr (int): Starting address for the memory copy operation.
        num_elements (int): Number of elements to copy from the starting address.

    """

    start_addr: int
    num_elements: int

MambaStateDtypeCalculator

Methods:

• append_replayssm_ring –

  Append the ReplaySSM ring dtypes to a base (conv, ssm) tuple:

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

class MambaStateDtypeCalculator:
    @classmethod
    def linear_attention_state_dtype(
        cls,
        model_dtype: ModelDType | torch.dtype,
        mamba_cache_dtype: MambaDType,
    ) -> tuple[torch.dtype, ...]:
        state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
        return (state_dtype,)

    @classmethod
    def mamba1_state_dtype(
        cls,
        model_dtype: ModelDType | torch.dtype,
        mamba_cache_dtype: MambaDType,
        mamba_ssm_cache_dtype: MambaDType,
    ) -> tuple[torch.dtype, ...]:
        return cls._mamba_state_dtype(
            model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype
        )

    @classmethod
    def mamba2_state_dtype(
        cls,
        model_dtype: ModelDType | torch.dtype,
        mamba_cache_dtype: MambaDType,
        mamba_ssm_cache_dtype: MambaDType,
    ) -> tuple[torch.dtype, ...]:
        return cls._mamba_state_dtype(
            model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype
        )

    @classmethod
    def append_replayssm_ring(
        cls,
        base_dtypes: tuple[torch.dtype, ...],
        model_dtype: ModelDType | torch.dtype,
    ) -> tuple[torch.dtype, ...]:
        """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple:
        ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``.
        """
        activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype)
        return (*base_dtypes, activation_dtype, torch.float32, activation_dtype)

    @classmethod
    def _mamba_state_dtype(
        cls,
        model_dtype: ModelDType | torch.dtype,
        mamba_cache_dtype: MambaDType,
        mamba_ssm_cache_dtype: MambaDType,
    ) -> tuple[torch.dtype, ...]:
        conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
        if mamba_ssm_cache_dtype == "auto":
            temporal_state_dtype = conv_state_dtype
        else:
            temporal_state_dtype = STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype]

        return (conv_state_dtype, temporal_state_dtype)

    @classmethod
    def short_conv_state_dtype(
        cls,
        model_dtype: ModelDType | torch.dtype,
        mamba_cache_dtype: MambaDType,
    ) -> tuple[torch.dtype, ...]:
        conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
        return (conv_state_dtype,)

    @classmethod
    def gated_delta_net_state_dtype(
        cls,
        model_dtype: ModelDType | torch.dtype,
        mamba_cache_dtype: MambaDType,
        mamba_ssm_cache_dtype: MambaDType = "auto",
    ) -> tuple[torch.dtype, torch.dtype]:
        return cls._mamba_state_dtype(
            model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype
        )

    @classmethod
    def kda_state_dtype(
        cls,
        model_dtype: ModelDType | torch.dtype,
        mamba_cache_dtype: MambaDType,
        mamba_ssm_cache_dtype: MambaDType = "auto",
    ) -> tuple[torch.dtype, torch.dtype]:
        conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
        if mamba_ssm_cache_dtype == "auto":
            recurrent_state_dtype = torch.float32
        elif mamba_ssm_cache_dtype in ("float32", "bfloat16"):
            recurrent_state_dtype = STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype]
        else:
            raise ValueError(
                "KDA recurrent state supports only auto, float32, and bfloat16; "
                f"got {mamba_ssm_cache_dtype}"
            )
        return (conv_state_dtype, recurrent_state_dtype)

append_replayssm_ring(base_dtypes, model_dtype) classmethod

Append the ReplaySSM ring dtypes to a base (conv, ssm) tuple: (x_cache, dt_cache, B_cache) = (activation, fp32, activation).

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def append_replayssm_ring(
    cls,
    base_dtypes: tuple[torch.dtype, ...],
    model_dtype: ModelDType | torch.dtype,
) -> tuple[torch.dtype, ...]:
    """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple:
    ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``.
    """
    activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype)
    return (*base_dtypes, activation_dtype, torch.float32, activation_dtype)

MambaStateShapeCalculator

Methods:

• append_replayssm_ring –

  Append the physical ReplaySSM ring shapes.

• extra_groups_for_head_shards –

  Compute the increase in group numbers to account for

• gated_delta_net_state_splits –

  How TP splits gated_delta_net_state_shape's states: conv

• kda_state_splits –

  How TP splits kda_state_shape's states: conv [Q, K, V], and the

• linear_attention_state_splits –

  How TP splits linear_attention_state_shape's state: by head.

• mamba1_state_splits –

  How TP splits mamba1_state_shape's states: conv [x], and the

• mamba2_state_splits –

  How TP splits mamba2_state_shape's states: conv [x, B, C], and

• short_conv_state_splits –

  How TP splits short_conv_state_shape's conv state: by channel.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

class MambaStateShapeCalculator:
    @classmethod
    def linear_attention_state_shape(
        cls,
        num_heads: int,
        tp_size: int,
        head_dim: int,
    ) -> tuple[tuple[int, int, int]]:
        state_shape = (num_heads // tp_size, head_dim, head_dim)
        return (state_shape,)

    @staticmethod
    def _orient_conv_shape(dim: int, state_len: int) -> tuple[int, int]:
        """Return (dim, state_len) for DS layout, (state_len, dim) for SD."""
        if is_conv_state_dim_first():
            return (dim, state_len)
        return (state_len, dim)

    @classmethod
    def mamba1_state_shape(
        cls,
        tp_world_size: int,
        intermediate_size: int,
        state_size: int,
        conv_kernel: int,
    ) -> tuple[tuple[int, int], tuple[int, int]]:
        conv_dim = divide(intermediate_size, tp_world_size)
        conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1)

        temporal_state_shape = (divide(intermediate_size, tp_world_size), state_size)

        return conv_state_shape, temporal_state_shape

    @classmethod
    def mamba2_state_shape(
        cls,
        tp_world_size: int,
        intermediate_size: int,
        n_groups: int,
        num_heads: int,
        head_dim: int,
        state_size: int,
        conv_kernel: int,
        num_spec: int = 0,
    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
        # if n_groups is not divisible by world_size, need to extend the shards
        # to ensure all groups needed by a head is sharded along with it
        n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size)
        # heads and n_groups are TP-ed
        conv_dim = intermediate_size + 2 * n_groups * state_size

        conv_state_shape = cls._orient_conv_shape(
            divide(conv_dim, tp_world_size), conv_kernel - 1 + num_spec
        )

        # These are not TP-ed as they depend on A, dt_bias, D
        # - they are typically small
        #   e.g., (h_heads, head_dim, state_size) = (128, 64, 128)
        temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size)
        return conv_state_shape, temporal_state_shape

    @classmethod
    def append_replayssm_ring(
        cls,
        base_shapes: tuple[tuple[int, int], tuple[int, int, int]],
        n_groups: int,
        tp_world_size: int,
        logical_window: int,
        backend: MambaBackendEnum,
        num_speculative_tokens: int = 0,
    ) -> tuple[
        tuple[int, int],
        tuple[int, int, int],
        tuple[int, int, int],
        tuple[int, int],
        tuple[int, int, int],
    ]:
        """Append the physical ReplaySSM ring shapes.

        ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``;
        B_cache extends ``n_groups`` for head shards as in ``mamba2_state_shape``.
        """
        ring_buffer_len = logical_window
        if backend == MambaBackendEnum.FLASHINFER:
            # FlashInfer keeps the live window and current verify window together.
            ring_buffer_len += 1 + num_speculative_tokens
        local_nheads, head_dim, state_size = base_shapes[1]
        n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size)
        local_ngroups = divide(n_groups, tp_world_size)
        return (
            *base_shapes,
            (local_nheads, ring_buffer_len, head_dim),
            (local_nheads, ring_buffer_len),
            (local_ngroups, ring_buffer_len, state_size),
        )

    @classmethod
    def short_conv_state_shape(
        cls,
        tp_world_size: int,
        intermediate_size: int,
        conv_kernel: int,
        num_spec: int = 0,
    ) -> tuple[tuple[int, int]]:
        conv_dim = divide(intermediate_size, tp_world_size)
        conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1 + num_spec)
        return (conv_state_shape,)

    @classmethod
    def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int):
        """Compute the increase in group numbers to account for
        replication in order to accompany the head shards."""
        # in the case ngoups % tp_size == 0, this will be zero
        if ngroups % tp_size == 0:
            return 0

        # for n_groups == 1, this is exactly tp_size - n_groups
        return tp_size - ngroups

    @classmethod
    def gated_delta_net_state_shape(
        cls,
        tp_world_size: int,
        num_k_heads: int,
        num_v_heads: int,
        head_k_dim: int,
        head_v_dim: int,
        conv_kernel_size: int,
        num_spec: int = 0,
    ):
        conv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads
        conv_state_shape = cls._orient_conv_shape(
            divide(conv_dim, tp_world_size),
            conv_kernel_size - 1 + num_spec,
        )

        temporal_state_shape = (
            divide(num_v_heads, tp_world_size),
            head_v_dim,
            head_k_dim,
        )
        return conv_state_shape, temporal_state_shape

    @classmethod
    def _conv_split(cls, *segments: tuple[int, int]) -> MambaStateSplit:
        return MambaStateSplit(0 if is_conv_state_dim_first() else 1, segments)

    @classmethod
    def linear_attention_state_splits(
        cls, tp_size: int, num_heads: int
    ) -> tuple[MambaStateSplit, ...] | None:
        """How TP splits ``linear_attention_state_shape``'s state: by head."""
        if num_heads % tp_size:
            return None
        return (MambaStateSplit(0, ((num_heads // tp_size, 1),)),)

    @classmethod
    def short_conv_state_splits(
        cls, tp_world_size: int, intermediate_size: int
    ) -> tuple[MambaStateSplit, ...] | None:
        """How TP splits ``short_conv_state_shape``'s conv state: by channel."""
        if intermediate_size % tp_world_size:
            return None
        return (cls._conv_split((intermediate_size // tp_world_size, 1)),)

    @classmethod
    def mamba1_state_splits(
        cls, tp_world_size: int, intermediate_size: int
    ) -> tuple[MambaStateSplit, ...] | None:
        """How TP splits ``mamba1_state_shape``'s states: conv [x], and the
        temporal state by channel."""
        if intermediate_size % tp_world_size:
            return None
        channels = (intermediate_size // tp_world_size, 1)
        return (cls._conv_split(channels), MambaStateSplit(0, (channels,)))

    @classmethod
    def mamba2_state_splits(
        cls,
        tp_world_size: int,
        intermediate_size: int,
        n_groups: int,
        num_heads: int,
        state_size: int,
    ) -> tuple[MambaStateSplit, ...] | None:
        """How TP splits ``mamba2_state_shape``'s states: conv [x, B, C], and
        the temporal state by head. A single group that TP does not divide is
        held whole by every rank (see ``extra_groups_for_head_shards``)."""
        if intermediate_size % tp_world_size or num_heads % tp_world_size:
            return None
        if n_groups % tp_world_size == 0:
            b_and_c = (n_groups * state_size // tp_world_size, 1)
        elif n_groups == 1:
            b_and_c = (state_size, tp_world_size)
        else:
            return None
        x = (intermediate_size // tp_world_size, 1)
        heads = (num_heads // tp_world_size, 1)
        return (cls._conv_split(x, b_and_c, b_and_c), MambaStateSplit(0, (heads,)))

    @classmethod
    def gated_delta_net_state_splits(
        cls,
        tp_world_size: int,
        num_k_heads: int,
        num_v_heads: int,
        head_k_dim: int,
        head_v_dim: int,
    ) -> tuple[MambaStateSplit, ...] | None:
        """How TP splits ``gated_delta_net_state_shape``'s states: conv
        [Q, K, V], and the temporal state by value head."""
        if num_k_heads % tp_world_size or num_v_heads % tp_world_size:
            return None
        qk = (num_k_heads // tp_world_size * head_k_dim, 1)
        v_heads = num_v_heads // tp_world_size
        conv = cls._conv_split(qk, qk, (v_heads * head_v_dim, 1))
        return (conv, MambaStateSplit(0, ((v_heads, 1),)))

    @classmethod
    def kda_state_shape(
        cls,
        tp_world_size: int,
        num_heads: int,
        head_dim: int,
        num_k_heads: int | None = None,
        head_k_dim: int | None = None,
        conv_kernel_size: int = 4,
        num_spec: int = 0,
    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
        if num_k_heads is None:
            num_k_heads = num_heads
        if head_k_dim is None:
            head_k_dim = head_dim

        proj_size = num_heads * head_dim
        proj_k_size = num_k_heads * head_k_dim

        conv_dim = proj_size + 2 * proj_k_size
        conv_state_shape = cls._orient_conv_shape(
            divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec
        )
        recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
        return (conv_state_shape, recurrent_state_shape)

    @classmethod
    def kda_state_splits(
        cls,
        tp_world_size: int,
        num_heads: int,
        head_dim: int,
        num_k_heads: int | None = None,
        head_k_dim: int | None = None,
    ) -> tuple[MambaStateSplit, ...] | None:
        """How TP splits ``kda_state_shape``'s states: conv [Q, K, V], and the
        recurrent state by head. The same layout as gated delta net's."""
        return cls.gated_delta_net_state_splits(
            tp_world_size,
            num_heads if num_k_heads is None else num_k_heads,
            num_heads,
            head_dim if head_k_dim is None else head_k_dim,
            head_dim,
        )

_orient_conv_shape(dim, state_len) staticmethod

Return (dim, state_len) for DS layout, (state_len, dim) for SD.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@staticmethod
def _orient_conv_shape(dim: int, state_len: int) -> tuple[int, int]:
    """Return (dim, state_len) for DS layout, (state_len, dim) for SD."""
    if is_conv_state_dim_first():
        return (dim, state_len)
    return (state_len, dim)

append_replayssm_ring(base_shapes, n_groups, tp_world_size, logical_window, backend, num_speculative_tokens=0) classmethod

Append the physical ReplaySSM ring shapes.

base_shapes[1] is (nheads // tp, head_dim, state_size); B_cache extends n_groups for head shards as in mamba2_state_shape.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def append_replayssm_ring(
    cls,
    base_shapes: tuple[tuple[int, int], tuple[int, int, int]],
    n_groups: int,
    tp_world_size: int,
    logical_window: int,
    backend: MambaBackendEnum,
    num_speculative_tokens: int = 0,
) -> tuple[
    tuple[int, int],
    tuple[int, int, int],
    tuple[int, int, int],
    tuple[int, int],
    tuple[int, int, int],
]:
    """Append the physical ReplaySSM ring shapes.

    ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``;
    B_cache extends ``n_groups`` for head shards as in ``mamba2_state_shape``.
    """
    ring_buffer_len = logical_window
    if backend == MambaBackendEnum.FLASHINFER:
        # FlashInfer keeps the live window and current verify window together.
        ring_buffer_len += 1 + num_speculative_tokens
    local_nheads, head_dim, state_size = base_shapes[1]
    n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size)
    local_ngroups = divide(n_groups, tp_world_size)
    return (
        *base_shapes,
        (local_nheads, ring_buffer_len, head_dim),
        (local_nheads, ring_buffer_len),
        (local_ngroups, ring_buffer_len, state_size),
    )

extra_groups_for_head_shards(ngroups, tp_size) classmethod

Compute the increase in group numbers to account for replication in order to accompany the head shards.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int):
    """Compute the increase in group numbers to account for
    replication in order to accompany the head shards."""
    # in the case ngoups % tp_size == 0, this will be zero
    if ngroups % tp_size == 0:
        return 0

    # for n_groups == 1, this is exactly tp_size - n_groups
    return tp_size - ngroups

gated_delta_net_state_splits(tp_world_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim) classmethod

How TP splits gated_delta_net_state_shape's states: conv [Q, K, V], and the temporal state by value head.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def gated_delta_net_state_splits(
    cls,
    tp_world_size: int,
    num_k_heads: int,
    num_v_heads: int,
    head_k_dim: int,
    head_v_dim: int,
) -> tuple[MambaStateSplit, ...] | None:
    """How TP splits ``gated_delta_net_state_shape``'s states: conv
    [Q, K, V], and the temporal state by value head."""
    if num_k_heads % tp_world_size or num_v_heads % tp_world_size:
        return None
    qk = (num_k_heads // tp_world_size * head_k_dim, 1)
    v_heads = num_v_heads // tp_world_size
    conv = cls._conv_split(qk, qk, (v_heads * head_v_dim, 1))
    return (conv, MambaStateSplit(0, ((v_heads, 1),)))

kda_state_splits(tp_world_size, num_heads, head_dim, num_k_heads=None, head_k_dim=None) classmethod

How TP splits kda_state_shape's states: conv [Q, K, V], and the recurrent state by head. The same layout as gated delta net's.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def kda_state_splits(
    cls,
    tp_world_size: int,
    num_heads: int,
    head_dim: int,
    num_k_heads: int | None = None,
    head_k_dim: int | None = None,
) -> tuple[MambaStateSplit, ...] | None:
    """How TP splits ``kda_state_shape``'s states: conv [Q, K, V], and the
    recurrent state by head. The same layout as gated delta net's."""
    return cls.gated_delta_net_state_splits(
        tp_world_size,
        num_heads if num_k_heads is None else num_k_heads,
        num_heads,
        head_dim if head_k_dim is None else head_k_dim,
        head_dim,
    )

linear_attention_state_splits(tp_size, num_heads) classmethod

How TP splits linear_attention_state_shape's state: by head.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def linear_attention_state_splits(
    cls, tp_size: int, num_heads: int
) -> tuple[MambaStateSplit, ...] | None:
    """How TP splits ``linear_attention_state_shape``'s state: by head."""
    if num_heads % tp_size:
        return None
    return (MambaStateSplit(0, ((num_heads // tp_size, 1),)),)

mamba1_state_splits(tp_world_size, intermediate_size) classmethod

How TP splits mamba1_state_shape's states: conv [x], and the temporal state by channel.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def mamba1_state_splits(
    cls, tp_world_size: int, intermediate_size: int
) -> tuple[MambaStateSplit, ...] | None:
    """How TP splits ``mamba1_state_shape``'s states: conv [x], and the
    temporal state by channel."""
    if intermediate_size % tp_world_size:
        return None
    channels = (intermediate_size // tp_world_size, 1)
    return (cls._conv_split(channels), MambaStateSplit(0, (channels,)))

mamba2_state_splits(tp_world_size, intermediate_size, n_groups, num_heads, state_size) classmethod

How TP splits mamba2_state_shape's states: conv [x, B, C], and the temporal state by head. A single group that TP does not divide is held whole by every rank (see extra_groups_for_head_shards).

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def mamba2_state_splits(
    cls,
    tp_world_size: int,
    intermediate_size: int,
    n_groups: int,
    num_heads: int,
    state_size: int,
) -> tuple[MambaStateSplit, ...] | None:
    """How TP splits ``mamba2_state_shape``'s states: conv [x, B, C], and
    the temporal state by head. A single group that TP does not divide is
    held whole by every rank (see ``extra_groups_for_head_shards``)."""
    if intermediate_size % tp_world_size or num_heads % tp_world_size:
        return None
    if n_groups % tp_world_size == 0:
        b_and_c = (n_groups * state_size // tp_world_size, 1)
    elif n_groups == 1:
        b_and_c = (state_size, tp_world_size)
    else:
        return None
    x = (intermediate_size // tp_world_size, 1)
    heads = (num_heads // tp_world_size, 1)
    return (cls._conv_split(x, b_and_c, b_and_c), MambaStateSplit(0, (heads,)))

short_conv_state_splits(tp_world_size, intermediate_size) classmethod

How TP splits short_conv_state_shape's conv state: by channel.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@classmethod
def short_conv_state_splits(
    cls, tp_world_size: int, intermediate_size: int
) -> tuple[MambaStateSplit, ...] | None:
    """How TP splits ``short_conv_state_shape``'s conv state: by channel."""
    if intermediate_size % tp_world_size:
        return None
    return (cls._conv_split((intermediate_size // tp_world_size, 1)),)

get_conv_copy_spec(state, block_ids, cur_block_idx, num_accepted_tokens)

Return a MambaCopySpec for copying a convolutional state slice.

Works for both SD layout (num_blocks, state_len, dim) and DS layout (num_blocks, dim, state_len).

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

def get_conv_copy_spec(
    state: torch.Tensor,
    block_ids: list[int],
    cur_block_idx: int,
    num_accepted_tokens: int,
) -> MambaCopySpec:
    """Return a MambaCopySpec for copying a convolutional state slice.

    Works for both SD layout ``(num_blocks, state_len, dim)`` and
    DS layout ``(num_blocks, dim, state_len)``.
    """
    src_block_id = block_ids[cur_block_idx]
    offset = num_accepted_tokens - 1
    if is_conv_state_dim_first():
        # DS offset > 0 is handled by the fused postprocess kernel.
        assert offset == 0, (
            "DS conv state with num_accepted_tokens > 1 must be handled by "
            "the fused postprocess kernel, not get_conv_copy_spec"
        )
        src_state = state[src_block_id]
    else:
        # SD layout: (num_blocks, state_len, dim), with dim contiguous.
        src_state = state[src_block_id, offset:]
    return MambaCopySpec(
        start_addr=src_state.data_ptr(), num_elements=src_state.numel()
    )

get_conv_state_layout() cached

Return the SSM conv state layout.

SD = (state_len, dim) — dim is the innermost contiguous dimension. DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HND for KV cache), consistent with SSM temporal state layout.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

@functools.lru_cache
def get_conv_state_layout() -> ConvStateLayoutType:
    """Return the SSM conv state layout.

    SD = (state_len, dim) — dim is the innermost contiguous dimension.
    DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HND for KV
         cache), consistent with SSM temporal state layout.
    """
    layout: ConvStateLayoutType | None = envs.VLLM_SSM_CONV_STATE_LAYOUT
    if layout is not None:
        logger.info_once(
            "VLLM_SSM_CONV_STATE_LAYOUT env detected. "
            "Setting SSM conv state layout to %s.",
            layout,
        )
        return layout

    return "SD"

get_temporal_copy_spec(state, block_ids, cur_block_idx, num_accepted_tokens)

Return a MambaCopySpec for copying a temporal state slice.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

def get_temporal_copy_spec(
    state: torch.Tensor,
    block_ids: list[int],
    cur_block_idx: int,
    num_accepted_tokens: int,
) -> MambaCopySpec:
    """Return a MambaCopySpec for copying a temporal state slice."""
    src_block_id = block_ids[cur_block_idx + num_accepted_tokens - 1]
    src_state = state[src_block_id]
    return MambaCopySpec(
        start_addr=src_state.data_ptr(), num_elements=src_state.numel()
    )

is_conv_state_dim_first()

True when the conv state is stored as (dim, state_len) per block.

Source code in vllm/model_executor/layers/mamba/mamba_utils.py

def is_conv_state_dim_first() -> bool:
    """True when the conv state is stored as (dim, state_len) per block."""
    return get_conv_state_layout() == "DS"