Skip to content

vllm.model_executor.layers.mamba.ops.replayssm_config

Launch-config selection for the ReplaySSM Mamba2 output_only decode kernel.

Mirrors mamba_ssm.py: a hard-coded heuristic per kernel, plus an override context manager for benchmarks/tests/config sweeps. Hardware is auto-detected (Blackwell vs not) so call sites need not thread it through.

Functions:

get_replayssm_config(kernel, **shape)

Return the launch config for kernel (override > tuned default).

kernel: "mamba2_output_only". shape carries the keying dims (dstate; L for the buffer length, default 16); hardware is auto-detected.

Source code in vllm/model_executor/layers/mamba/ops/replayssm_config.py 
40
41
42
43
44
45
46
47
48
49
50
51
def get_replayssm_config(kernel: str, **shape) -> tuple:
    """Return the launch config for ``kernel`` (override > tuned default).

    kernel: "mamba2_output_only". ``shape`` carries the keying dims (dstate;
    ``L`` for the buffer length, default 16); hardware is auto-detected.
    """
    if kernel in _overrides:
        return _overrides[kernel]
    bw = _is_blackwell()
    if kernel == "mamba2_output_only":
        return _mamba2_output_only(shape["dstate"], shape.get("L", 16), bw)
    raise ValueError(f"unknown ReplaySSM kernel config key: {kernel}")