Skip to content

vllm.model_executor.layers.mamba.ops.replayssm_quantization

Prefill conversion for FlashInfer shared quantized checkpoints.

Functions:

dequantize_replayssm_state(state, state_scale, state_indices, has_initial)

Gather existing checkpoints into the floating-point prefill input.

Source code in vllm/model_executor/layers/mamba/ops/replayssm_quantization.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def dequantize_replayssm_state(
    state: torch.Tensor,
    state_scale: torch.Tensor,
    state_indices: torch.Tensor,
    has_initial: torch.Tensor,
) -> torch.Tensor:
    """Gather existing checkpoints into the floating-point prefill input."""
    output = torch.empty(
        (state_indices.numel(), *state.shape[1:]),
        dtype=torch.float32,
        device=state.device,
    )
    rows = state.shape[1] * state.shape[2]
    _dequantize_checkpoint[(state_indices.numel(), rows)](
        state,
        state_scale,
        state_indices.contiguous(),
        has_initial,
        output,
        state.stride(0),
        state_scale.stride(0),
        ROWS=rows,
        N=state.shape[3],
    )
    return output

quantize_replayssm_state(states, state, state_scale, state_indices, source_indices=None)

Store prefill states with per-row decode scales and nearest rounding.

Source code in vllm/model_executor/layers/mamba/ops/replayssm_quantization.py 
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
def quantize_replayssm_state(
    states: torch.Tensor,
    state: torch.Tensor,
    state_scale: torch.Tensor,
    state_indices: torch.Tensor,
    source_indices: torch.Tensor | None = None,
) -> None:
    """Store prefill states with per-row decode scales and nearest rounding."""
    if state.dtype not in (torch.float8_e4m3fn, torch.int8):
        raise ValueError("Shared ReplaySSM requires FP8 or INT8 state")
    rows = state.shape[1] * state.shape[2]
    _quantize_checkpoint[(state_indices.numel(), rows)](
        states,
        state,
        state_scale,
        state_indices.contiguous(),
        source_indices,
        states.stride(0),
        state.stride(0),
        state_scale.stride(0),
        ROWS=rows,
        N=state.shape[3],
        QMAX=127.0 if state.dtype == torch.int8 else 448.0,
        HAS_SOURCES=source_indices is not None,
        INT8=state.dtype == torch.int8,
    )