vllm.model_executor.layers.mamba.ops.replayssm_cutedsl

SM103 ReplaySSM with an FP16 recurrence checkpoint and quantized reads.

The native kernels support one token, 16 replay entries, 64 head dimensions, 128 state dimensions, and 8 or 16 heads per group. Slot zero is reserved for padding. Physical cache pages may have padding between slots.

Functions:

• initialize_replayssm_state –

  Initialize both checkpoints independently from FP32 prefill states.

• selective_state_update_replayssm –

  Decode one token and advance device-resident positions.

initialize_replayssm_state(cache, states, state_indices, *, state_dtype, seed=0)

Initialize both checkpoints independently from FP32 prefill states.

This operation also clears replay rings and event positions for reused slots. Nonpositive and out-of-range slot indices are ignored.

Source code in vllm/model_executor/layers/mamba/ops/replayssm_cutedsl.py

def initialize_replayssm_state(
    cache: Sequence[torch.Tensor],
    states: torch.Tensor,
    state_indices: torch.Tensor,
    *,
    state_dtype: str,
    seed: int = 0,
) -> None:
    """Initialize both checkpoints independently from FP32 prefill states.

    This operation also clears replay rings and event positions for reused slots.
    Nonpositive and out-of-range slot indices are ignored.
    """
    nvfp4 = _nvfp4(state_dtype)
    slots, heads, groups = _check_cache(cache, nvfp4)
    if states.dtype != torch.float32 or states.shape[1:] != (heads, 64, 128):
        raise ValueError("Prefill states must be FP32 [batch, heads, 64, 128]")
    if states.device != cache[1].device or state_indices.device != states.device:
        raise ValueError("Prefill states and slot indices must use the cache device")
    if state_indices.dtype not in (torch.int32, torch.int64):
        raise ValueError("Slot indices must be int32 or int64")
    ids = state_indices.reshape(-1).contiguous()
    if ids.numel() != states.shape[0]:
        raise ValueError("Each prefill state requires one slot index")
    if ids.numel() == 0:
        return
    if states.stride()[1:] != (64 * 128, 128, 1):
        states = states.contiguous()
    if nvfp4:
        # Warmup can broadcast one zero state over all request slots.
        source = states[:1] if states.stride(0) == 0 else states
        group_max = (
            source.abs().reshape(-1, groups, (heads // groups) * 64 * 128).amax(dim=-1)
        )
        global_scale = group_max.amax(dim=-1) / (6.0 * 448.0)
        global_scale = torch.where(global_scale > 0, global_scale, 1.0)
        max_stride = 0 if states.stride(0) == 0 else group_max.stride(0)
        global_stride = 0 if states.stride(0) == 0 else global_scale.stride(0)
        meta, parity = cache[7:9]
    else:
        group_max = global_scale = meta = parity = cache[6]
        max_stride = global_stride = 0
    _initialize_rows[(ids.numel(), heads * 64)](
        states,
        ids,
        cache[1].view(torch.int32),
        cache[5],
        cache[6],
        group_max,
        global_scale,
        meta,
        parity,
        states.stride(0),
        cache[1].stride(0) // 2,
        cache[5].stride(0),
        cache[6].stride(0),
        meta.stride(0),
        parity.stride(0),
        max_stride,
        global_stride,
        slots,
        heads,
        groups,
        nvfp4,
        seed,
        cache[-1].shape == (slots, 2),
        enable_fp_fusion=False,
    )
    counts = [math.prod(t.shape[1:]) for t in cache[2:5]]
    _clear_slots[(ids.numel(), triton.cdiv(max(counts), 256))](
        ids,
        cache[2],
        cache[3],
        cache[4],
        cache[-1],
        cache[2].stride(0),
        cache[3].stride(0),
        cache[4].stride(0),
        cache[-1].stride(0),
        *counts,
        2 if cache[-1].shape == (slots, 2) else 1,
        slots,
        256,
    )

selective_state_update_replayssm(cache, x, dt, A, B, C, D, dt_bias, state_indices, out, *, state_dtype, seed=0, num_accepted_tokens=None, cu_seqlens=None, max_seqlen=1)

Decode one token and advance device-resident positions.

An eager call compiles both flush paths before CUDA graph capture. Calls apply softplus to dt plus dt_bias. Active slot indices must be unique.

Source code in vllm/model_executor/layers/mamba/ops/replayssm_cutedsl.py

def selective_state_update_replayssm(
    cache: Sequence[torch.Tensor],
    x: torch.Tensor,
    dt: torch.Tensor,
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    D: torch.Tensor,
    dt_bias: torch.Tensor,
    state_indices: torch.Tensor,
    out: torch.Tensor,
    *,
    state_dtype: str,
    seed: int = 0,
    num_accepted_tokens: torch.Tensor | None = None,
    cu_seqlens: torch.Tensor | None = None,
    max_seqlen: int = 1,
) -> torch.Tensor:
    """Decode one token and advance device-resident positions.

    An eager call compiles both flush paths before CUDA graph capture. Calls
    apply softplus to dt plus dt_bias. Active slot indices must be unique.
    """
    nvfp4 = _nvfp4(state_dtype)
    slots, heads, groups = _check_cache(cache, nvfp4)
    if cache[-1].shape == (slots, 2):
        return _pending_state_update(
            cache,
            x,
            dt,
            A,
            B,
            C,
            D,
            dt_bias,
            state_indices,
            out,
            nvfp4=nvfp4,
            seed=seed,
            num_accepted_tokens=num_accepted_tokens,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
        )
    if max_seqlen != 1:
        raise ValueError("SM103 separate ReplaySSM supports one token")
    batch = state_indices.numel()
    if batch == 0:
        return out
    if (
        x.dtype != torch.bfloat16
        or B.dtype != torch.bfloat16
        or C.dtype != torch.bfloat16
    ):
        raise ValueError("Separate ReplaySSM requires BF16 x, B, and C")
    if state_indices.dtype not in (torch.int32, torch.int64):
        raise ValueError("Slot indices must be int32 or int64")
    if (
        out.dtype != torch.bfloat16
        or out.numel() != batch * heads * 64
        or not out.is_contiguous()
    ):
        raise ValueError("Output must be contiguous BF16 [batch, heads, 64]")
    tensors = (x, dt, A, B, C, D, dt_bias, state_indices, out)
    if any(t.device != cache[1].device for t in tensors):
        raise ValueError("Decode tensors must use the cache device")
    x = x.reshape(batch, heads, 64).contiguous()
    B = B.reshape(batch, groups, 128).contiguous()
    C = C.reshape(batch, groups, 128).contiguous()
    # The serving layer may supply tied head dimensions as stride-zero views.
    dt = dt.reshape(batch, heads, -1)[:, :, 0].float().contiguous()
    A = A.reshape(heads, -1)[:, 0].float().contiguous()
    dt_bias = dt_bias.reshape(heads, -1)[:, 0].float().contiguous()
    D = D.reshape(heads, -1).expand(heads, 64).float().contiguous()
    ids = state_indices.reshape(-1).contiguous()
    native_ids = torch.empty(batch, device=x.device, dtype=torch.int32)
    tags = torch.empty(batch, device=x.device, dtype=torch.int64)
    positions = torch.empty(batch, device=x.device, dtype=torch.int32)
    flush = torch.empty(batch, device=x.device, dtype=torch.int8)
    native_out = out.view(batch, heads, 64)
    _prepare_decode[(batch,)](
        ids,
        cache[-1],
        native_ids,
        tags,
        positions,
        flush,
        native_out,
        slots,
        cache[-1].stride(0),
        heads * 64,
        256,
    )
    common = (x, dt, A, B, C, D, dt_bias, cache[2], cache[3], cache[4], native_ids)
    if nvfp4:
        from vllm.third_party.replayssm.mamba2.nvfp4_flush import mamba2_nvfp4_flush
        from vllm.third_party.replayssm.mamba2.nvfp4_nonflush import (
            mamba2_nvfp4_nonflush,
        )

        mamba2_nvfp4_nonflush(
            *common,
            positions,
            flush,
            native_out,
            state_lo=cache[5],
            sf_lo=cache[6],
            sg_meta=cache[7],
        )
        mamba2_nvfp4_flush(
            cache[1],
            *common,
            tags,
            flush,
            native_out,
            state_lo=cache[5],
            sf_lo=cache[6],
            sg_meta=cache[7],
            sg_parity=cache[8],
            stochastic_rounding_seed=seed,
        )
    else:
        from vllm.third_party.replayssm.mamba2.fp8_flush import mamba2_fp8_flush
        from vllm.third_party.replayssm.mamba2.fp8_nonflush import mamba2_fp8_nonflush

        mamba2_fp8_nonflush(
            *common,
            positions,
            flush,
            native_out,
            state_lo=cache[5],
            state_lo_scale=cache[6],
        )
        mamba2_fp8_flush(
            cache[1],
            *common,
            tags,
            flush,
            native_out,
            state_lo=cache[5],
            state_lo_scale=cache[6],
            stochastic_rounding_seed=seed,
        )
    return out