vllm.model_executor.layers.mamba.checkpoint

Classes:

• MambaPrefillCheckpointBuilder –

  Build backend-neutral prefill checkpoint locations.

• MambaPrefillCheckpointExporter –

  Export a mid-prefill checkpoint into backend-specific paged states.

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

    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
        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
        hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size
        speculative_config = self.vllm_config.speculative_config
        drop_eagle_block = (
            speculative_config is not None and speculative_config.use_eagle_block_drop()
        )
        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=drop_eagle_block,
        )
        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
            ),
            request_rows=request_rows_tensor,
            block_cols=checkpoint_cols_tensor,
        )

MambaPrefillCheckpointExporter

Bases: ABC

Export a mid-prefill checkpoint into backend-specific paged states.

Methods:

• export –

  Write checkpoint state into the paged cache.

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

class MambaPrefillCheckpointExporter(ABC):
    """Export a mid-prefill checkpoint into backend-specific paged states."""

    @abstractmethod
    def export(
        self,
        checkpoint: MambaPrefillCheckpointMetadata,
        *args,
        **kwargs,
    ) -> None:
        """Write checkpoint state into the paged cache."""
        raise NotImplementedError

export(checkpoint, *args, **kwargs) abstractmethod

Write checkpoint state into the paged cache.

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

@abstractmethod
def export(
    self,
    checkpoint: MambaPrefillCheckpointMetadata,
    *args,
    **kwargs,
) -> None:
    """Write checkpoint state into the paged cache."""
    raise NotImplementedError

compute_mamba_prefill_checkpoints(seq_lens, query_lens, hash_block_size, mamba_block_size, checkpoint_alignment, drop_eagle_block)

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,
) -> 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, hash_block_size, drop_eagle_block=drop_eagle_block
        )
        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