Skip to content

vllm.model_executor.layers.mamba.ops.ssd_flashinfer

Mamba2 prefill SSD scan through FlashInfer's fused SSD (Cake backend).

Runs the chunked scan of mamba_chunk_scan_combined_varlen in two kernel launches instead of five, with the packed-varlen cu_seqlens form (the kernels derive their chunk metadata on the device).

Functions:

mamba_chunk_scan_flashinfer(x, dt, A, B, C, D, dt_bias, cu_seqlens, out, ssm_state, state_indices, has_initial_states)

Scan the prefill tokens and write each sequence's final state.

x: (seqlen, nheads, headdim) bf16, may be a strided view dt: (seqlen, nheads) A, dt_bias: (nheads,) fp32 B, C: (seqlen, ngroups, dstate) bf16, may be strided views D: (nheads,) bf16 cu_seqlens: (batch + 1,) int32, starting at 0 out: (seqlen, nheads, headdim) bf16, contiguous; written in place ssm_state: (num_slots, nheads, headdim, dstate) state cache state_indices: (batch,) cache slot of each sequence has_initial_states: (batch,) bool, or None when no sequence has a state

Source code in vllm/model_executor/layers/mamba/ops/ssd_flashinfer.py 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 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
def mamba_chunk_scan_flashinfer(
    x: torch.Tensor,
    dt: torch.Tensor,
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    D: torch.Tensor,
    dt_bias: torch.Tensor,
    cu_seqlens: torch.Tensor,
    out: torch.Tensor,
    ssm_state: torch.Tensor,
    state_indices: torch.Tensor,
    has_initial_states: torch.Tensor | None,
) -> None:
    """Scan the prefill tokens and write each sequence's final state.

    x: (seqlen, nheads, headdim) bf16, may be a strided view
    dt: (seqlen, nheads)
    A, dt_bias: (nheads,) fp32
    B, C: (seqlen, ngroups, dstate) bf16, may be strided views
    D: (nheads,) bf16
    cu_seqlens: (batch + 1,) int32, starting at 0
    out: (seqlen, nheads, headdim) bf16, contiguous; written in place
    ssm_state: (num_slots, nheads, headdim, dstate) state cache
    state_indices: (batch,) cache slot of each sequence
    has_initial_states: (batch,) bool, or None when no sequence has a state
    """
    initial_states = (
        gather_initial_states(ssm_state, state_indices, has_initial_states)
        if has_initial_states is not None
        else None
    )
    runner = get_ssd_runner(
        x.device.index,
        torch.accelerator.current_stream(x.device).stream_id,
        x.size(1),
        B.size(1),
        ssm_state.dtype,
        initial_states is not None,
    )
    _, final_states = runner.run(
        x.unsqueeze(0),
        dt.unsqueeze(0),
        A,
        B.unsqueeze(0),
        C.unsqueeze(0),
        D=D,
        dt_bias=dt_bias,
        dt_softplus=True,
        initial_states=initial_states,
        cu_seqlens=cu_seqlens,
        out=out.unsqueeze(0),
        return_final_states=True,
    )
    scatter_states(ssm_state, final_states, state_indices)