vllm.model_executor.layers.mamba.checkpoint

Classes:

• MambaPrefillCheckpointBuilder –

  Build backend-neutral prefill checkpoint locations.

• MambaPrefillCheckpointExporter –

  Store each checkpoint's pre-convolution window and recurrent state.

Functions:

• compute_mamba_prefill_checkpoints –

  Per-row internal prefill checkpoint offsets and cache block columns.

MambaPrefillCheckpointBuilder

Build backend-neutral prefill checkpoint locations.

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

class MambaPrefillCheckpointBuilder:
    """Build backend-neutral prefill checkpoint locations."""

    def __init__(self, vllm_config: VllmConfig, kv_cache_spec: MambaSpec) -> None:
        self.vllm_config = vllm_config
        self.kv_cache_spec = kv_cache_spec
        speculative_config = vllm_config.speculative_config
        self.drop_eagle_block = (
            speculative_config is not None
            and speculative_config.use_eagle_block_drop()
            and not vllm_config.num_hash_lookahead_tokens
        )

    def build(
        self,
        m: CommonAttentionMetadata,
        request_rows: list[int],
    ) -> MambaPrefillCheckpointMetadata | None:
        if self.vllm_config.cache_config.mamba_cache_mode != "align":
            return None
        if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0:
            return None
        cache_config = self.vllm_config.cache_config
        hash_block_size = cache_config.hash_block_size
        cache_hit_alignment_tokens = cache_config.cache_hit_alignment_tokens
        if hash_block_size is None or cache_hit_alignment_tokens is None:
            # Unresolved only for the minimal KV cache built for profiling.
            return None
        assert m.seq_lens_cpu_upper_bound is not None
        all_query_lens = m.query_start_loc_cpu.diff().tolist()
        query_lens = [all_query_lens[row] for row in request_rows]
        seq_lens = m.seq_lens_cpu_upper_bound.tolist()
        block_size = self.kv_cache_spec.block_size
        checkpoint_offsets, checkpoint_cols = compute_mamba_prefill_checkpoints(
            [seq_lens[row] for row in request_rows],
            query_lens,
            hash_block_size=hash_block_size,
            mamba_block_size=block_size,
            checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment,
            drop_eagle_block=self.drop_eagle_block,
            cache_hit_alignment_tokens=cache_hit_alignment_tokens,
            num_hash_lookahead_tokens=self.vllm_config.num_hash_lookahead_tokens,
        )
        if not any(checkpoint_offsets):
            return None
        checkpoint_offsets_tensor = async_tensor_h2d(
            checkpoint_offsets,
            dtype=torch.int32,
            device=m.query_start_loc.device,
        )
        request_rows_tensor = async_tensor_h2d(
            request_rows, dtype=torch.int64, device=m.query_start_loc.device
        )
        checkpoint_cols_tensor = async_tensor_h2d(
            checkpoint_cols, dtype=torch.int64, device=m.query_start_loc.device
        )
        return MambaPrefillCheckpointMetadata(
            checkpoint_offsets_tensor,
            _gather_checkpoint_state_indices(
                m.block_table_tensor, request_rows_tensor, checkpoint_cols_tensor
            ),
            offsets=checkpoint_offsets,
            query_lens=query_lens,
            request_rows=request_rows_tensor,
            block_cols=checkpoint_cols_tensor,
        )

MambaPrefillCheckpointExporter dataclass

Store each checkpoint's pre-convolution window and recurrent state.

The recurrent states are rows of a buffer the prefill kernel emitted; recurrent_rows picks each checkpoint row's state out of it, and by default row i holds checkpoint i's. Layers with only a conv state pass no recurrent buffers. token_indices maps the rows' tokens, laid out by cu_seqlens, to their rows of raw_qkv.

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

@dataclass(frozen=True)
class MambaPrefillCheckpointExporter:
    """Store each checkpoint's pre-convolution window and recurrent state.

    The recurrent states are rows of a buffer the prefill kernel emitted;
    ``recurrent_rows`` picks each checkpoint row's state out of it, and by
    default row ``i`` holds checkpoint ``i``'s. Layers with only a conv state
    pass no recurrent buffers. ``token_indices`` maps the rows' tokens, laid out
    by ``cu_seqlens``, to their rows of ``raw_qkv``.
    """

    state_len: int | None = None

    def export(
        self,
        checkpoint: MambaPrefillCheckpointMetadata,
        *,
        raw_qkv: torch.Tensor,
        conv_state: torch.Tensor,
        cu_seqlens: torch.Tensor,
        recurrent_checkpoint: torch.Tensor | None = None,
        recurrent_state: torch.Tensor | None = None,
        recurrent_rows: torch.Tensor | None = None,
        token_indices: torch.Tensor | None = None,
    ) -> None:
        state_len = (
            self.state_len if self.state_len is not None else conv_state.shape[-1]
        )
        # The block is cached as a checkpoint, so it needs its whole conv window.
        assert checkpoint.offsets is None or all(
            offset == 0 or offset >= state_len for offset in checkpoint.offsets
        )
        assert (
            recurrent_rows is None
            or recurrent_rows.numel() == checkpoint.checkpoint_offsets.numel()
        )
        width = raw_qkv.shape[-1]
        has_recurrent = recurrent_checkpoint is not None
        if recurrent_checkpoint is None:
            # Unread placeholders for the recurrent pointers.
            recurrent_checkpoint = recurrent_state = conv_state
            recurrent_row_size = 0
        else:
            assert recurrent_state is not None
            recurrent_row_size = recurrent_checkpoint[0].numel()
        block_size = 512
        store_cache_checkpoints_kernel[
            (
                checkpoint.checkpoint_offsets.numel(),
                triton.cdiv(max(width * state_len, recurrent_row_size), block_size),
            )
        ](
            raw_qkv,
            conv_state,
            recurrent_checkpoint,
            recurrent_state,
            cu_seqlens,
            checkpoint.checkpoint_offsets,
            checkpoint.state_indices,
            recurrent_rows if recurrent_rows is not None else checkpoint.state_indices,
            token_indices if token_indices is not None else cu_seqlens,
            raw_qkv.stride(0),
            raw_qkv.stride(1),
            conv_state.stride(0),
            conv_state.stride(1),
            conv_state.stride(2),
            recurrent_checkpoint.stride(0),
            recurrent_state.stride(0),
            checkpoint.checkpoint_offsets.stride(0),
            state_len,
            width,
            recurrent_row_size,
            NULL_BLOCK_ID,
            has_recurrent,
            recurrent_rows is not None,
            token_indices is not None,
            block_size,
        )

compute_mamba_prefill_checkpoints(seq_lens, query_lens, hash_block_size, mamba_block_size, checkpoint_alignment, drop_eagle_block, cache_hit_alignment_tokens, num_hash_lookahead_tokens=0)

Per-row internal prefill checkpoint offsets and cache block columns.

Backends call this instead of re-deriving the rules, so they decline in lockstep with the scheduler and MambaManager: allocating a checkpoint block without writing it leaves the prefix cache serving uninitialized state.

Returns:

• list[int] –

  (offsets, cols): the checkpoint's offset into each row's query and

• list[int] –

  its block-table column. 0 and -1 mean the row has none.

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

def compute_mamba_prefill_checkpoints(
    seq_lens: list[int],
    query_lens: list[int],
    hash_block_size: int,
    mamba_block_size: int,
    checkpoint_alignment: int | None,
    drop_eagle_block: bool,
    cache_hit_alignment_tokens: int,
    num_hash_lookahead_tokens: int = 0,
) -> tuple[list[int], list[int]]:
    """Per-row internal prefill checkpoint offsets and cache block columns.

    Backends call this instead of re-deriving the rules, so they decline in
    lockstep with the scheduler and ``MambaManager``: allocating a checkpoint
    block without writing it leaves the prefix cache serving uninitialized
    state.

    Returns:
        ``(offsets, cols)``: the checkpoint's offset into each row's query and
        its block-table column. ``0`` and ``-1`` mean the row has none.

    """
    offsets: list[int] = []
    cols: list[int] = []
    for seq_len, query_len in zip(seq_lens, query_lens):
        query_start = seq_len - query_len
        position = get_mamba_prefill_checkpoint_position(
            seq_len,
            cache_hit_alignment_tokens,
            drop_eagle_block=drop_eagle_block,
            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
        )
        valid = is_mamba_prefill_checkpoint_valid(
            query_start=query_start,
            query_end=seq_len,
            checkpoint_position=position,
            hash_block_size=hash_block_size,
            mamba_block_size=mamba_block_size,
            checkpoint_alignment=checkpoint_alignment,
        )
        offsets.append(position - query_start if valid else 0)
        cols.append(cdiv(seq_len, mamba_block_size) - 2 if valid else -1)
    return offsets, cols