Skip to content

vllm.model_executor.layers.mamba.ops.ssd_checkpoint

Functions:

store_prefill_checkpoint(checkpoint, chunk_idx, x, conv_state, chunk_states, ssm_state, query_start_loc)

Write each prefill row's mid-prefill checkpoint into its paged block.

Rows without a checkpoint (offset 0) write nothing.

Parameters:

  • checkpoint

    (MambaPrefillCheckpointMetadata) –

    Per-row offsets and destination blocks.

  • chunk_idx

    (Tensor) –

    Per-row index into chunk_states of the checkpoint.

  • x

    (Tensor) –

    Pre-conv activations, (dim, num_prefill_tokens).

  • conv_state

    (Tensor) –

    Paged conv state, (num_blocks, dim, state_len).

  • chunk_states

    (Tensor) –

    SSD states at every chunk end.

  • ssm_state

    (Tensor) –

    Paged SSM state.

  • query_start_loc

    (Tensor) –

    Prefill query start locations.

Source code in vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py 
 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
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def store_prefill_checkpoint(
    checkpoint: MambaPrefillCheckpointMetadata,
    chunk_idx: torch.Tensor,
    x: torch.Tensor,
    conv_state: torch.Tensor,
    chunk_states: torch.Tensor,
    ssm_state: torch.Tensor,
    query_start_loc: torch.Tensor,
) -> None:
    """Write each prefill row's mid-prefill checkpoint into its paged block.

    Rows without a checkpoint (offset 0) write nothing.

    Args:
        checkpoint: Per-row offsets and destination blocks.
        chunk_idx: Per-row index into ``chunk_states`` of the checkpoint.
        x: Pre-conv activations, ``(dim, num_prefill_tokens)``.
        conv_state: Paged conv state, ``(num_blocks, dim, state_len)``.
        chunk_states: SSD states at every chunk end.
        ssm_state: Paged SSM state.
        query_start_loc: Prefill query start locations.

    """
    dim = x.size(0)
    state_len = conv_state.size(-1)
    ssm_row_size = chunk_states[0].numel()
    block_size = 1024
    grid = (
        chunk_idx.numel(),
        triton.cdiv(max(dim * state_len, ssm_row_size), block_size),
    )
    _store_prefill_checkpoint_kernel[grid](
        x,
        conv_state,
        chunk_states,
        ssm_state,
        query_start_loc,
        checkpoint.checkpoint_offsets,
        checkpoint.state_indices,
        chunk_idx,
        x.stride(0),
        x.stride(1),
        conv_state.stride(0),
        conv_state.stride(1),
        conv_state.stride(2),
        chunk_states.stride(0),
        ssm_state.stride(0),
        DIM=dim,
        STATE_LEN=state_len,
        SSM_ROW_SIZE=ssm_row_size,
        BLOCK_SIZE=block_size,
    )