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 
55
56
57
58
59
60
61
62
63
64
65
66
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}")

override_replayssm_config(kernel, config)

Pin kernel's launch config for the duration of the context.

Source code in vllm/model_executor/layers/mamba/ops/replayssm_config.py 
29
30
31
32
33
34
35
36
37
38
39
40
@contextmanager
def override_replayssm_config(kernel: str, config: tuple):
    """Pin ``kernel``'s launch config for the duration of the context."""
    prev = _overrides.get(kernel)
    _overrides[kernel] = config
    try:
        yield
    finally:
        if prev is None:
            _overrides.pop(kernel, None)
        else:
            _overrides[kernel] = prev