vllm.v1.worker.mamba_utils

Classes:

• MambaBuffers –

  Single owner for all mamba-specific runner buffers.

• MambaSpecDecodeGPUContext –

  Context for GPU-side Mamba state copy operations during the

Functions:

• cleanup_mamba_state_idx –

  Pop stale mamba_state_idx entries for finished/preempted/resumed reqs.

• get_mamba_group_ids –

  Return the sorted group ids of all mamba groups, without duplicates.

• get_mamba_groups –

  Return a mapping from each distinct MambaSpec to its sorted group_ids.

• postprocess_mamba_fused_kernel –

  Fused GPU kernel for postprocess_mamba that computes decisions AND performs

• postprocess_mamba_gpu –

  Publish accepted-token results to model-wide Mamba state.

• precopy_mamba_align_fused_kernel –

  Pre-copy mamba "align" state across block boundaries.

• preprocess_mamba –

  Copy the mamba state of previous step to the last

• preprocess_mamba_align_fused_kernel –

  Fused align preprocess: emit the pre-copy src column/offset AND advance

• stage_postprocess_inputs_to_gpu –

  Stage the per-request decisions consumed after token acceptance.

MambaBuffers dataclass

Single owner for all mamba-specific runner buffers.

The two sub-objects have different gates: preprocess handles prefix-cache state migration; postprocess_align owns the fused lifecycle state used by hybrid speculative decoding and FlashInfer ReplaySSM, and is None otherwise.

Source code in vllm/v1/worker/mamba_utils.py

@dataclasses.dataclass
class MambaBuffers:
    """Single owner for all mamba-specific runner buffers.

    The two sub-objects have different gates:
    ``preprocess`` handles prefix-cache state migration; ``postprocess_align``
    owns the fused lifecycle state used by hybrid speculative decoding and
    FlashInfer ReplaySSM, and is ``None`` otherwise.
    """

    preprocess: MambaCopyBuffers
    postprocess_align: MambaSpecDecodeGPUContext | None

    @classmethod
    def create(
        cls,
        max_num_reqs: int,
        kv_cache_config: KVCacheConfig,
        copy_funcs: MambaStateCopyFuncsByType,
        forward_context: dict[str, Any],
        make_buffer: Callable[..., CpuGpuBuffer],
        device: torch.device,
        with_postprocess_align: bool,
    ) -> "MambaBuffers":
        return cls(
            preprocess=MambaCopyBuffers.create(
                max_num_reqs, kv_cache_config, copy_funcs, make_buffer
            ),
            postprocess_align=(
                MambaSpecDecodeGPUContext.create(
                    max_num_reqs=max_num_reqs,
                    kv_cache_config=kv_cache_config,
                    copy_funcs=copy_funcs,
                    forward_context=forward_context,
                    device=device,
                    make_buffer=make_buffer,
                )
                if with_postprocess_align
                else None
            ),
        )

MambaSpecDecodeGPUContext dataclass

Context for GPU-side Mamba state copy operations during the fused postprocess path.

Used for hybrid speculative state maintenance and for model-owned FlashInfer ReplaySSM prefix-cache lifecycle maintenance.

Precomputes memory layout metadata (base addresses, strides, element sizes) so the GPU kernel can perform state copies without CPU-GPU sync.

State types are distinguished by conv_width: >0 for conv states (sliding window with offset-based copies), 0 for temporal states (full block copies).

Methods:

• compute_aligned_state_indices –

  Compute every Mamba group's aligned physical state IDs in one launch.

• create –

  Create context with allocated buffers (metadata populated later).

• initialize_from_forward_context –

  Extract and cache memory layout metadata from Mamba state tensors.

• run_fused_postprocess –

  Run the fused postprocess_mamba kernel on GPU.

• run_fused_postprocess_align –

  V2 align postprocess: save the running state to the block-aligned

• run_fused_precopy –

  Pre-copy each request's previous running block into its new window

Source code in vllm/v1/worker/mamba_utils.py

@dataclasses.dataclass
class MambaSpecDecodeGPUContext:
    """Context for GPU-side Mamba state copy operations during the
    fused postprocess path.

    Used for hybrid speculative state maintenance and for model-owned
    FlashInfer ReplaySSM prefix-cache lifecycle maintenance.

    Precomputes memory layout metadata (base addresses, strides, element sizes)
    so the GPU kernel can perform state copies without CPU-GPU sync.

    State types are distinguished by conv_width: >0 for conv states (sliding
    window with offset-based copies), 0 for temporal states (full block copies).
    """

    # Per-state metadata tensors (shape: [num_states]). Mamba groups may have
    # different numbers of state types (for example, GDN has two while PLE
    # short-conv has one).
    # These are populated from forward_context during the first forward pass
    state_base_addrs: torch.Tensor  # int64: base address of each state tensor
    state_block_strides: torch.Tensor  # int64: bytes per block
    state_elem_sizes: torch.Tensor  # int32: element size in bytes
    state_inner_sizes: torch.Tensor  # int64: elements in inner dimensions
    state_conv_widths: torch.Tensor  # int32: conv width (0 for temporal states)
    state_group_indices: torch.Tensor  # int32: maps state_idx to group index
    # DS conv row metadata. Zero keeps the single-region copy path.
    state_dim_row_count: torch.Tensor  # int32: per-block dim row count
    state_dim_row_stride: torch.Tensor  # int64: bytes between rows

    # Configuration
    block_size: int
    num_states: int
    mamba_group_ids: list[int]
    num_groups: int

    # Snapshot retained while the live accepted-token buffer is normalized.
    num_accepted_tokens_snapshot: torch.Tensor

    # Per-group block-table base addresses: int64[num_groups]. Populated in
    # initialize_from_forward_context from the persistent per-group block
    # table tensors (whose data_ptr is stable across steps). Rows are request
    # state slots for the state copies and batch rows for the aligned indices.
    block_table_ptrs: torch.Tensor
    aligned_index_block_table_ptrs: torch.Tensor
    block_table_stride_req: int = 0

    # persistent output for the once-per-step, all-group aligned-index launch.
    # shape: [num_groups, max_num_reqs, 1 + num_speculative_blocks].
    aligned_state_indices: torch.Tensor | None = None

    # Per-request staging buffers (CPU+GPU mirrors). The runner stages
    # values into the CPU view in ``_prepare_inputs`` and the fused kernel
    # reads the GPU side. These only exist when the postprocess kernel is
    # enabled (spec decode + hybrid, or FlashInfer ReplaySSM prefix caching).
    mamba_state_idx_buf: CpuGpuBuffer | None = None
    num_scheduled_tokens_buf: CpuGpuBuffer | None = None
    num_computed_tokens_buf: CpuGpuBuffer | None = None
    num_draft_tokens_buf: CpuGpuBuffer | None = None
    is_prefilling_buf: CpuGpuBuffer | None = None
    precopy_src_col_buf: CpuGpuBuffer | None = None
    precopy_token_bias_buf: CpuGpuBuffer | None = None

    # Flag to track if metadata has been populated
    is_initialized: bool = False
    # Persistent all-layer ReplaySSM descriptors, populated with the cache
    # addresses on first real forward. None for non-FlashInfer configurations.
    replayssm: ReplaySSMModelContext | None = None

    @classmethod
    def create(
        cls,
        max_num_reqs: int,
        kv_cache_config: KVCacheConfig,
        copy_funcs: MambaStateCopyFuncsByType,
        forward_context: dict[str, Any],
        device: torch.device,
        make_buffer: Callable[..., CpuGpuBuffer],
    ) -> "MambaSpecDecodeGPUContext":
        """Create context with allocated buffers (metadata populated later)."""
        mamba_groups = get_mamba_groups(kv_cache_config)
        mamba_group_ids = get_mamba_group_ids(mamba_groups)
        mamba_spec = next(iter(mamba_groups))
        assert all(
            spec.block_size == mamba_spec.block_size
            and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks
            and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode
            for spec in mamba_groups
        ), "all mamba groups must share cache scheduling parameters"
        # Count physical state tensors across all Mamba groups. Different
        # groups may expose different state specs.
        total_states = sum(
            len(
                _get_mamba_gpu_state_copy_funcs(
                    _get_mamba_spec_for_layer(
                        kv_cache_config.kv_cache_groups[gid], layer_name
                    ),
                    forward_context[layer_name],
                    copy_funcs,
                )
            )
            for gid in mamba_group_ids
            for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names
        )

        return cls(
            state_base_addrs=torch.zeros(
                total_states, dtype=torch.int64, device=device
            ),
            state_block_strides=torch.zeros(
                total_states, dtype=torch.int64, device=device
            ),
            state_elem_sizes=torch.zeros(
                total_states, dtype=torch.int32, device=device
            ),
            state_inner_sizes=torch.zeros(
                total_states, dtype=torch.int64, device=device
            ),
            state_conv_widths=torch.zeros(
                total_states, dtype=torch.int32, device=device
            ),
            state_group_indices=torch.zeros(
                total_states, dtype=torch.int32, device=device
            ),
            state_dim_row_count=torch.zeros(
                total_states, dtype=torch.int32, device=device
            ),
            state_dim_row_stride=torch.zeros(
                total_states, dtype=torch.int64, device=device
            ),
            block_size=mamba_spec.block_size,
            num_states=total_states,
            mamba_group_ids=mamba_group_ids,
            num_groups=len(mamba_group_ids),
            num_accepted_tokens_snapshot=torch.zeros(
                max_num_reqs, dtype=torch.int32, device=device
            ),
            block_table_ptrs=torch.zeros(
                len(mamba_group_ids), dtype=torch.int64, device=device
            ),
            aligned_index_block_table_ptrs=torch.zeros(
                len(mamba_group_ids), dtype=torch.int64, device=device
            ),
            aligned_state_indices=torch.empty(
                (
                    len(mamba_group_ids),
                    max_num_reqs,
                    1 + mamba_spec.num_speculative_blocks,
                ),
                dtype=torch.int32,
                device=device,
            ),
            mamba_state_idx_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            num_draft_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            is_prefilling_buf=make_buffer(max_num_reqs, dtype=torch.bool),
            precopy_src_col_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            precopy_token_bias_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            is_initialized=False,
            replayssm=None,
        )

    def initialize_from_forward_context(
        self,
        kv_cache_config: KVCacheConfig,
        forward_context: dict[str, Any],
        mamba_state_copy_funcs: MambaStateCopyFuncsByType,
        block_tables: list[torch.Tensor],
        aligned_index_block_tables: list[torch.Tensor] | None = None,
    ) -> None:
        """Extract and cache memory layout metadata from Mamba state tensors.

        This method populates the pre-allocated metadata tensors with information
        needed by `postprocess_mamba_fused_kernel` to perform state copies entirely
        on the GPU without CPU-GPU synchronization.

        For each Mamba layer and state type, the following metadata is extracted:
        - state_base_addrs: GPU memory address (data_ptr) of the state tensor
        - state_block_strides: Bytes between consecutive blocks (stride * elem_size)
        - state_elem_sizes: Element size in bytes (e.g., 2 for float16)
        - state_inner_sizes: For conv states, elements per conv position (stride(1)),
          used to compute offset when slicing state[block, offset:]. For temporal
          states, this field is unused (set to 1).
        - state_conv_widths: Conv dimension size for conv states, 0 for temporal states

        The conv vs temporal state type is detected by inspecting the copy function
        name: functions containing "conv" are treated as conv states.

        This method is idempotent - it only executes once (guarded by is_initialized
        flag) since the metadata is static after model loading.

        Args:
            kv_cache_config: Configuration containing KV cache group info and
                layer name mappings.
            forward_context: Dictionary mapping layer names to attention objects,
                populated after the model is loaded. Each attention object must
                have a `kv_cache` attribute containing the list of state tensors.
            mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to
                copy functions.
            block_tables: per-mamba-group persistent block-table tensors, in
                the same order as `mamba_group_ids`, with one row per request
                state slot. Their `data_ptr()` / `stride(0)` are captured once
                for the state-copy kernels to index into.
            aligned_index_block_tables: per-mamba-group persistent tables
                read by `compute_aligned_state_indices`, with one row per
                batch row. Defaults to `block_tables`.

        """
        if self.is_initialized:
            return
        # This only runs once per worker.
        with gpu_sync_allowed():
            self._populate_metadata(
                kv_cache_config,
                forward_context,
                mamba_state_copy_funcs,
                block_tables,
                aligned_index_block_tables or block_tables,
            )

    def _populate_metadata(
        self,
        kv_cache_config: KVCacheConfig,
        forward_context: dict[str, Any],
        mamba_state_copy_funcs: MambaStateCopyFuncsByType,
        block_tables: list[torch.Tensor],
        aligned_index_block_tables: list[torch.Tensor],
    ) -> None:
        idx = 0
        assert len(block_tables) == self.num_groups, (
            f"expected {self.num_groups} block tables, got {len(block_tables)}"
        )
        replayssm_groups: list[tuple[list[Any], MambaSpec, torch.Tensor]] = []
        has_baseline_layer = False
        for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids):
            kv_cache_group = kv_cache_config.kv_cache_groups[mamba_group_id]
            mixers: list[Any] = []
            layer_names = kv_cache_group.layer_names
            for layer_name in layer_names:
                mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name)
                attention = forward_context[layer_name]
                state_copy_funcs = _get_mamba_gpu_state_copy_funcs(
                    mamba_spec, attention, mamba_state_copy_funcs
                )
                kv_caches: list[torch.Tensor] = attention.kv_cache
                is_flashinfer_replayssm = (
                    isinstance(attention, MambaBase)
                    and attention.use_flashinfer_replayssm
                )
                has_baseline_layer |= not is_flashinfer_replayssm
                if is_flashinfer_replayssm:
                    if len(kv_caches) != len(state_copy_funcs):
                        raise ValueError(
                            "FlashInfer ReplaySSM requires five bound Mamba2 states"
                        )
                    mixers.append(attention)
                if len(kv_caches) < len(state_copy_funcs):
                    raise ValueError(
                        f"Expected at least {len(state_copy_funcs)} Mamba state "
                        f"tensors, got {len(kv_caches)}"
                    )
                for state_type_idx, copy_func in enumerate(state_copy_funcs):
                    state = kv_caches[state_type_idx]
                    # Base address
                    self.state_base_addrs[idx] = _reinterpret_u64_as_i64(
                        state.data_ptr()
                    )

                    # Block stride (bytes between consecutive blocks)
                    # state shape: [num_blocks, ...], stride(0) = elements per block
                    if state.dim() > 1:
                        block_stride_elems = state.stride(0)
                    else:
                        block_stride_elems = state.numel()
                    self.state_block_strides[idx] = (
                        block_stride_elems * state.element_size()
                    )

                    # Element size
                    self.state_elem_sizes[idx] = state.element_size()

                    assert (
                        copy_func is get_conv_copy_spec
                        or copy_func is get_temporal_copy_spec
                    ), f"unexpected copy func: {copy_func}"
                    if copy_func is get_conv_copy_spec:
                        if state.dim() != 3:
                            raise ValueError(
                                "Expected 3D conv state cache, got "
                                f"shape {tuple(state.shape)}"
                            )
                        if is_conv_state_dim_first():
                            # DS layout: state_len is the slide axis.
                            self.state_conv_widths[idx] = state.size(2)
                            self.state_inner_sizes[idx] = 1
                            self.state_dim_row_count[idx] = state.size(1)
                            self.state_dim_row_stride[idx] = (
                                state.stride(1) * state.element_size()
                            )
                        else:
                            # SD layout: dim is contiguous.
                            self.state_conv_widths[idx] = state.size(1)
                            self.state_inner_sizes[idx] = state.stride(1)
                    else:
                        # Temporal state: inner_size = natural elements per
                        # block (prod of inner dims).  The kernel uses this
                        # to compute copy_size = inner_size * elem_size,
                        # which gives the correct byte count even when the
                        # state tensor is as_strided with padded page strides
                        # (state_block_stride would be the page size, too big).
                        self.state_conv_widths[idx] = 0
                        self.state_inner_sizes[idx] = (
                            state[0].numel() if state.dim() > 1 else 1
                        )
                        # Temporal copies vectorize with uint64 loads/stores.
                        # The kernel's head/tail handles misalignment for
                        # correctness, but unaligned base/stride costs in
                        # throughput.
                        base_addr = state.data_ptr()
                        block_stride_bytes = block_stride_elems * state.element_size()
                        if base_addr % 8 != 0:
                            logger.warning_once(
                                "layer %s: state.data_ptr() = %#x is not "
                                "8B-aligned; _memcpy_u64_tiled uint64 "
                                "vectorization will pay misaligned load cost",
                                layer_name,
                                base_addr,
                            )
                        if block_stride_bytes % 8 != 0:
                            logger.warning_once(
                                "layer %s: block stride = %dB is not "
                                "8B-aligned; _memcpy_u64_tiled uint64 "
                                "vectorization will pay misaligned load cost",
                                layer_name,
                                block_stride_bytes,
                            )

                    self.state_group_indices[idx] = group_local_idx
                    idx += 1

            if mixers:
                spec = kv_cache_group.kv_cache_spec
                if not isinstance(spec, MambaSpec):
                    raise TypeError(
                        "FlashInfer ReplaySSM layers require a Mamba cache spec; "
                        f"got {type(spec).__name__}"
                    )
                replayssm_groups.append(
                    (mixers, spec, aligned_index_block_tables[group_local_idx])
                )

        assert idx == self.num_states
        if replayssm_groups and has_baseline_layer:
            raise ValueError(
                "mixed FlashInfer ReplaySSM and baseline Mamba layers are unsupported"
            )

        # Cache per-group block-table base addresses and per-request stride.
        # `block_tables[i]` is the persistent 2D int32 block-table tensor for
        # `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the
        # engine's lifetime, so we capture them once here.
        assert len(block_tables) == len(aligned_index_block_tables) == self.num_groups
        strides = {bt.stride(0) for bt in (*block_tables, *aligned_index_block_tables)}
        assert len(strides) == 1, (
            f"all mamba block tables must share stride(0), got {strides}"
        )
        self.block_table_stride_req = int(next(iter(strides)))
        for i, (bt, aligned_bt) in enumerate(
            zip(block_tables, aligned_index_block_tables)
        ):
            self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr())
            self.aligned_index_block_table_ptrs[i] = _reinterpret_u64_as_i64(
                aligned_bt.data_ptr()
            )

        if replayssm_groups:
            self.replayssm = ReplaySSMModelContext.create(
                replayssm_groups,
                self.num_accepted_tokens_snapshot.numel(),
            )

        self.is_initialized = True

    def compute_aligned_state_indices(
        self,
        seq_lens: torch.Tensor,
        num_reqs: int,
    ) -> torch.Tensor:
        """Compute every Mamba group's aligned physical state IDs in one launch."""
        assert self.is_initialized
        assert seq_lens.is_cuda or seq_lens.is_xpu
        assert 0 <= num_reqs <= seq_lens.shape[0]
        assert self.aligned_state_indices is not None
        assert num_reqs <= self.aligned_state_indices.shape[1]
        if num_reqs == 0:
            return self.aligned_state_indices[:, :0]

        num_state_slots = self.aligned_state_indices.shape[2]
        block_rows = 32
        grid = (triton.cdiv(num_reqs, block_rows),)
        get_aligned_state_indices_multi_group_kernel[grid](
            self.aligned_index_block_table_ptrs,
            seq_lens,
            self.aligned_state_indices,
            self.block_table_stride_req,
            seq_lens.stride(0),
            self.aligned_state_indices.stride(0),
            self.aligned_state_indices.stride(1),
            self.aligned_state_indices.stride(2),
            num_reqs,
            CACHE_BLOCK_SIZE=self.block_size,
            NUM_GROUPS=self.num_groups,
            BLOCK_GROUPS=triton.next_power_of_2(self.num_groups),
            NUM_STATE_SLOTS=num_state_slots,
            BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
            BLOCK_ROWS=block_rows,
            num_warps=1,
        )
        return self.aligned_state_indices[:, :num_reqs]

    def run_fused_postprocess(
        self,
        num_reqs: int,
        num_accepted_tokens_gpu: torch.Tensor,
        mamba_state_idx_gpu: torch.Tensor,
        num_scheduled_tokens_gpu: torch.Tensor,
        num_computed_tokens_gpu: torch.Tensor,
        num_draft_tokens_gpu: torch.Tensor,
    ) -> None:
        """Run the fused postprocess_mamba kernel on GPU.

        This computes decisions and performs mamba state copies entirely on GPU,
        eliminating the CPU-GPU sync that was previously needed.

        Args:
            num_reqs: Number of active requests
            num_accepted_tokens_gpu: [num_reqs] accepted token counts
            mamba_state_idx_gpu: [num_reqs] source block indices
            num_scheduled_tokens_gpu: [num_reqs] scheduled token counts
            num_computed_tokens_gpu: [num_reqs] computed token counts
            num_draft_tokens_gpu: [num_reqs] draft token counts

        """
        if num_reqs == 0 or not self.is_initialized:
            return

        self.num_accepted_tokens_snapshot[:num_reqs].copy_(
            num_accepted_tokens_gpu[:num_reqs]
        )

        total_states = self.num_states
        grid = (num_reqs, total_states, _TEMPORAL_TILES)

        postprocess_mamba_fused_kernel[grid](
            self.num_accepted_tokens_snapshot,
            mamba_state_idx_gpu,
            num_scheduled_tokens_gpu,
            num_computed_tokens_gpu,
            num_draft_tokens_gpu,
            self.block_table_ptrs,
            self.block_table_stride_req,
            self.state_base_addrs,
            self.state_block_strides,
            self.state_elem_sizes,
            self.state_inner_sizes,
            self.state_conv_widths,
            self.state_group_indices,
            self.state_dim_row_count,
            self.state_dim_row_stride,
            num_accepted_tokens_gpu,
            None,  # idx_mapping: V1 decision arrays are already in req order
            num_reqs,
            block_size=self.block_size,
            COPY_BLOCK_SIZE=1024,
            CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
            SKIP_TEMPORAL_STATE_COPY=self.replayssm is not None,
            TEMPORAL_TILES=_TEMPORAL_TILES,
        )

    def run_fused_precopy(
        self,
        num_reqs: int,
        state_idx_gpu: torch.Tensor,
        src_col_gpu: torch.Tensor,
        token_bias_gpu: torch.Tensor,
        idx_mapping: torch.Tensor | None,
    ) -> None:
        """Pre-copy each request's previous running block into its new window
        block before the forward pass (align boundary migration).

        Args:
            num_reqs: Number of active requests (batch order).
            state_idx_gpu: [max_reqs] post-advance dst block column per req slot.
            src_col_gpu: [max_reqs] pre-advance src block column (-1 = fresh).
            token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1).
            idx_mapping: optional [num_reqs] batch_idx -> req_state_idx.
                None means V1 batch order already equals request state order.

        """
        if num_reqs == 0 or not self.is_initialized:
            return
        total_states = self.num_states
        grid = (num_reqs, total_states, _TEMPORAL_TILES)
        precopy_mamba_align_fused_kernel[grid](
            state_idx_gpu,
            src_col_gpu,
            token_bias_gpu,
            self.block_table_ptrs,
            self.block_table_stride_req,
            self.state_base_addrs,
            self.state_block_strides,
            self.state_elem_sizes,
            self.state_inner_sizes,
            self.state_conv_widths,
            self.state_group_indices,
            self.state_dim_row_count,
            self.state_dim_row_stride,
            idx_mapping,
            num_reqs,
            COPY_BLOCK_SIZE=1024,
            CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
            HAS_IDX_MAPPING=idx_mapping is not None,
            ZERO_TEMPORAL_TOKEN_BIAS=self.replayssm is not None,
            TEMPORAL_TILES=_TEMPORAL_TILES,
        )

    def run_fused_postprocess_align(
        self,
        num_reqs: int,
        num_accepted_tokens_gpu: torch.Tensor,
        state_idx_gpu: torch.Tensor,
        new_num_computed_tokens_gpu: torch.Tensor,
        idx_mapping: torch.Tensor,
    ) -> None:
        """V2 align postprocess: save the running state to the block-aligned
        position after spec-decode acceptance leaves the sequence non-aligned.

        ``num_accepted_tokens_gpu`` is updated in place while the kernel reads
        from a snapshot to avoid cross-program races when the accepted position
        stays in the running block and the count is reset to 1.
        ``new_num_computed_tokens`` already holds the post-step computed count
        (PRECOMPUTED_NEW_COMPUTED).
        ``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING).
        """
        if num_reqs == 0 or not self.is_initialized:
            return

        # V2 reads non-contiguous idx_mapping positions, so snapshot the whole
        # decision buffer rather than only [:num_reqs].
        accepted_snapshot = self.num_accepted_tokens_snapshot
        accepted_snapshot.copy_(num_accepted_tokens_gpu)

        total_states = self.num_states
        grid = (num_reqs, total_states, _TEMPORAL_TILES)
        postprocess_mamba_fused_kernel[grid](
            accepted_snapshot,
            state_idx_gpu,
            None,  # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED
            new_num_computed_tokens_gpu,
            None,  # num_draft: unused under PRECOMPUTED_NEW_COMPUTED
            self.block_table_ptrs,
            self.block_table_stride_req,
            self.state_base_addrs,
            self.state_block_strides,
            self.state_elem_sizes,
            self.state_inner_sizes,
            self.state_conv_widths,
            self.state_group_indices,
            self.state_dim_row_count,
            self.state_dim_row_stride,
            num_accepted_tokens_gpu,
            idx_mapping,
            num_reqs,
            block_size=self.block_size,
            COPY_BLOCK_SIZE=1024,
            CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
            HAS_IDX_MAPPING=True,
            PRECOMPUTED_NEW_COMPUTED=True,
            SKIP_TEMPORAL_STATE_COPY=self.replayssm is not None,
            TEMPORAL_TILES=_TEMPORAL_TILES,
        )

compute_aligned_state_indices(seq_lens, num_reqs)

Compute every Mamba group's aligned physical state IDs in one launch.

Source code in vllm/v1/worker/mamba_utils.py

def compute_aligned_state_indices(
    self,
    seq_lens: torch.Tensor,
    num_reqs: int,
) -> torch.Tensor:
    """Compute every Mamba group's aligned physical state IDs in one launch."""
    assert self.is_initialized
    assert seq_lens.is_cuda or seq_lens.is_xpu
    assert 0 <= num_reqs <= seq_lens.shape[0]
    assert self.aligned_state_indices is not None
    assert num_reqs <= self.aligned_state_indices.shape[1]
    if num_reqs == 0:
        return self.aligned_state_indices[:, :0]

    num_state_slots = self.aligned_state_indices.shape[2]
    block_rows = 32
    grid = (triton.cdiv(num_reqs, block_rows),)
    get_aligned_state_indices_multi_group_kernel[grid](
        self.aligned_index_block_table_ptrs,
        seq_lens,
        self.aligned_state_indices,
        self.block_table_stride_req,
        seq_lens.stride(0),
        self.aligned_state_indices.stride(0),
        self.aligned_state_indices.stride(1),
        self.aligned_state_indices.stride(2),
        num_reqs,
        CACHE_BLOCK_SIZE=self.block_size,
        NUM_GROUPS=self.num_groups,
        BLOCK_GROUPS=triton.next_power_of_2(self.num_groups),
        NUM_STATE_SLOTS=num_state_slots,
        BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
        BLOCK_ROWS=block_rows,
        num_warps=1,
    )
    return self.aligned_state_indices[:, :num_reqs]

create(max_num_reqs, kv_cache_config, copy_funcs, forward_context, device, make_buffer) classmethod

Create context with allocated buffers (metadata populated later).

Source code in vllm/v1/worker/mamba_utils.py

@classmethod
def create(
    cls,
    max_num_reqs: int,
    kv_cache_config: KVCacheConfig,
    copy_funcs: MambaStateCopyFuncsByType,
    forward_context: dict[str, Any],
    device: torch.device,
    make_buffer: Callable[..., CpuGpuBuffer],
) -> "MambaSpecDecodeGPUContext":
    """Create context with allocated buffers (metadata populated later)."""
    mamba_groups = get_mamba_groups(kv_cache_config)
    mamba_group_ids = get_mamba_group_ids(mamba_groups)
    mamba_spec = next(iter(mamba_groups))
    assert all(
        spec.block_size == mamba_spec.block_size
        and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks
        and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode
        for spec in mamba_groups
    ), "all mamba groups must share cache scheduling parameters"
    # Count physical state tensors across all Mamba groups. Different
    # groups may expose different state specs.
    total_states = sum(
        len(
            _get_mamba_gpu_state_copy_funcs(
                _get_mamba_spec_for_layer(
                    kv_cache_config.kv_cache_groups[gid], layer_name
                ),
                forward_context[layer_name],
                copy_funcs,
            )
        )
        for gid in mamba_group_ids
        for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names
    )

    return cls(
        state_base_addrs=torch.zeros(
            total_states, dtype=torch.int64, device=device
        ),
        state_block_strides=torch.zeros(
            total_states, dtype=torch.int64, device=device
        ),
        state_elem_sizes=torch.zeros(
            total_states, dtype=torch.int32, device=device
        ),
        state_inner_sizes=torch.zeros(
            total_states, dtype=torch.int64, device=device
        ),
        state_conv_widths=torch.zeros(
            total_states, dtype=torch.int32, device=device
        ),
        state_group_indices=torch.zeros(
            total_states, dtype=torch.int32, device=device
        ),
        state_dim_row_count=torch.zeros(
            total_states, dtype=torch.int32, device=device
        ),
        state_dim_row_stride=torch.zeros(
            total_states, dtype=torch.int64, device=device
        ),
        block_size=mamba_spec.block_size,
        num_states=total_states,
        mamba_group_ids=mamba_group_ids,
        num_groups=len(mamba_group_ids),
        num_accepted_tokens_snapshot=torch.zeros(
            max_num_reqs, dtype=torch.int32, device=device
        ),
        block_table_ptrs=torch.zeros(
            len(mamba_group_ids), dtype=torch.int64, device=device
        ),
        aligned_index_block_table_ptrs=torch.zeros(
            len(mamba_group_ids), dtype=torch.int64, device=device
        ),
        aligned_state_indices=torch.empty(
            (
                len(mamba_group_ids),
                max_num_reqs,
                1 + mamba_spec.num_speculative_blocks,
            ),
            dtype=torch.int32,
            device=device,
        ),
        mamba_state_idx_buf=make_buffer(max_num_reqs, dtype=torch.int32),
        num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
        num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
        num_draft_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
        is_prefilling_buf=make_buffer(max_num_reqs, dtype=torch.bool),
        precopy_src_col_buf=make_buffer(max_num_reqs, dtype=torch.int32),
        precopy_token_bias_buf=make_buffer(max_num_reqs, dtype=torch.int32),
        is_initialized=False,
        replayssm=None,
    )

initialize_from_forward_context(kv_cache_config, forward_context, mamba_state_copy_funcs, block_tables, aligned_index_block_tables=None)

Extract and cache memory layout metadata from Mamba state tensors.

This method populates the pre-allocated metadata tensors with information needed by postprocess_mamba_fused_kernel to perform state copies entirely on the GPU without CPU-GPU synchronization.

For each Mamba layer and state type, the following metadata is extracted: - state_base_addrs: GPU memory address (data_ptr) of the state tensor - state_block_strides: Bytes between consecutive blocks (stride * elem_size) - state_elem_sizes: Element size in bytes (e.g., 2 for float16) - state_inner_sizes: For conv states, elements per conv position (stride(1)), used to compute offset when slicing state[block, offset:]. For temporal states, this field is unused (set to 1). - state_conv_widths: Conv dimension size for conv states, 0 for temporal states

The conv vs temporal state type is detected by inspecting the copy function name: functions containing "conv" are treated as conv states.

This method is idempotent - it only executes once (guarded by is_initialized flag) since the metadata is static after model loading.

Parameters:

• kv_cache_config

  (KVCacheConfig) –

  Configuration containing KV cache group info and layer name mappings.

• forward_context

  (dict[str, Any]) –

  Dictionary mapping layer names to attention objects, populated after the model is loaded. Each attention object must have a kv_cache attribute containing the list of state tensors.

• mamba_state_copy_funcs

  (MambaStateCopyFuncsByType) –

  Mapping from MambaAttentionBackendEnum to copy functions.

• block_tables

  (list[Tensor]) –

  per-mamba-group persistent block-table tensors, in the same order as mamba_group_ids, with one row per request state slot. Their data_ptr() / stride(0) are captured once for the state-copy kernels to index into.

• aligned_index_block_tables

  (list[Tensor] | None, default: None ) –

  per-mamba-group persistent tables read by compute_aligned_state_indices, with one row per batch row. Defaults to block_tables.

Source code in vllm/v1/worker/mamba_utils.py

def initialize_from_forward_context(
    self,
    kv_cache_config: KVCacheConfig,
    forward_context: dict[str, Any],
    mamba_state_copy_funcs: MambaStateCopyFuncsByType,
    block_tables: list[torch.Tensor],
    aligned_index_block_tables: list[torch.Tensor] | None = None,
) -> None:
    """Extract and cache memory layout metadata from Mamba state tensors.

    This method populates the pre-allocated metadata tensors with information
    needed by `postprocess_mamba_fused_kernel` to perform state copies entirely
    on the GPU without CPU-GPU synchronization.

    For each Mamba layer and state type, the following metadata is extracted:
    - state_base_addrs: GPU memory address (data_ptr) of the state tensor
    - state_block_strides: Bytes between consecutive blocks (stride * elem_size)
    - state_elem_sizes: Element size in bytes (e.g., 2 for float16)
    - state_inner_sizes: For conv states, elements per conv position (stride(1)),
      used to compute offset when slicing state[block, offset:]. For temporal
      states, this field is unused (set to 1).
    - state_conv_widths: Conv dimension size for conv states, 0 for temporal states

    The conv vs temporal state type is detected by inspecting the copy function
    name: functions containing "conv" are treated as conv states.

    This method is idempotent - it only executes once (guarded by is_initialized
    flag) since the metadata is static after model loading.

    Args:
        kv_cache_config: Configuration containing KV cache group info and
            layer name mappings.
        forward_context: Dictionary mapping layer names to attention objects,
            populated after the model is loaded. Each attention object must
            have a `kv_cache` attribute containing the list of state tensors.
        mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to
            copy functions.
        block_tables: per-mamba-group persistent block-table tensors, in
            the same order as `mamba_group_ids`, with one row per request
            state slot. Their `data_ptr()` / `stride(0)` are captured once
            for the state-copy kernels to index into.
        aligned_index_block_tables: per-mamba-group persistent tables
            read by `compute_aligned_state_indices`, with one row per
            batch row. Defaults to `block_tables`.

    """
    if self.is_initialized:
        return
    # This only runs once per worker.
    with gpu_sync_allowed():
        self._populate_metadata(
            kv_cache_config,
            forward_context,
            mamba_state_copy_funcs,
            block_tables,
            aligned_index_block_tables or block_tables,
        )

run_fused_postprocess(num_reqs, num_accepted_tokens_gpu, mamba_state_idx_gpu, num_scheduled_tokens_gpu, num_computed_tokens_gpu, num_draft_tokens_gpu)

Run the fused postprocess_mamba kernel on GPU.

This computes decisions and performs mamba state copies entirely on GPU, eliminating the CPU-GPU sync that was previously needed.

Parameters:

• num_reqs

  (int) –

  Number of active requests

• num_accepted_tokens_gpu

  (Tensor) –

  [num_reqs] accepted token counts

• mamba_state_idx_gpu

  (Tensor) –

  [num_reqs] source block indices

• num_scheduled_tokens_gpu

  (Tensor) –

  [num_reqs] scheduled token counts

• num_computed_tokens_gpu

  (Tensor) –

  [num_reqs] computed token counts

• num_draft_tokens_gpu

  (Tensor) –

  [num_reqs] draft token counts

Source code in vllm/v1/worker/mamba_utils.py

def run_fused_postprocess(
    self,
    num_reqs: int,
    num_accepted_tokens_gpu: torch.Tensor,
    mamba_state_idx_gpu: torch.Tensor,
    num_scheduled_tokens_gpu: torch.Tensor,
    num_computed_tokens_gpu: torch.Tensor,
    num_draft_tokens_gpu: torch.Tensor,
) -> None:
    """Run the fused postprocess_mamba kernel on GPU.

    This computes decisions and performs mamba state copies entirely on GPU,
    eliminating the CPU-GPU sync that was previously needed.

    Args:
        num_reqs: Number of active requests
        num_accepted_tokens_gpu: [num_reqs] accepted token counts
        mamba_state_idx_gpu: [num_reqs] source block indices
        num_scheduled_tokens_gpu: [num_reqs] scheduled token counts
        num_computed_tokens_gpu: [num_reqs] computed token counts
        num_draft_tokens_gpu: [num_reqs] draft token counts

    """
    if num_reqs == 0 or not self.is_initialized:
        return

    self.num_accepted_tokens_snapshot[:num_reqs].copy_(
        num_accepted_tokens_gpu[:num_reqs]
    )

    total_states = self.num_states
    grid = (num_reqs, total_states, _TEMPORAL_TILES)

    postprocess_mamba_fused_kernel[grid](
        self.num_accepted_tokens_snapshot,
        mamba_state_idx_gpu,
        num_scheduled_tokens_gpu,
        num_computed_tokens_gpu,
        num_draft_tokens_gpu,
        self.block_table_ptrs,
        self.block_table_stride_req,
        self.state_base_addrs,
        self.state_block_strides,
        self.state_elem_sizes,
        self.state_inner_sizes,
        self.state_conv_widths,
        self.state_group_indices,
        self.state_dim_row_count,
        self.state_dim_row_stride,
        num_accepted_tokens_gpu,
        None,  # idx_mapping: V1 decision arrays are already in req order
        num_reqs,
        block_size=self.block_size,
        COPY_BLOCK_SIZE=1024,
        CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
        SKIP_TEMPORAL_STATE_COPY=self.replayssm is not None,
        TEMPORAL_TILES=_TEMPORAL_TILES,
    )

run_fused_postprocess_align(num_reqs, num_accepted_tokens_gpu, state_idx_gpu, new_num_computed_tokens_gpu, idx_mapping)

V2 align postprocess: save the running state to the block-aligned position after spec-decode acceptance leaves the sequence non-aligned.

num_accepted_tokens_gpu is updated in place while the kernel reads from a snapshot to avoid cross-program races when the accepted position stays in the running block and the count is reset to 1. new_num_computed_tokens already holds the post-step computed count (PRECOMPUTED_NEW_COMPUTED). idx_mapping maps batch row -> req-state slot (HAS_IDX_MAPPING).

Source code in vllm/v1/worker/mamba_utils.py

def run_fused_postprocess_align(
    self,
    num_reqs: int,
    num_accepted_tokens_gpu: torch.Tensor,
    state_idx_gpu: torch.Tensor,
    new_num_computed_tokens_gpu: torch.Tensor,
    idx_mapping: torch.Tensor,
) -> None:
    """V2 align postprocess: save the running state to the block-aligned
    position after spec-decode acceptance leaves the sequence non-aligned.

    ``num_accepted_tokens_gpu`` is updated in place while the kernel reads
    from a snapshot to avoid cross-program races when the accepted position
    stays in the running block and the count is reset to 1.
    ``new_num_computed_tokens`` already holds the post-step computed count
    (PRECOMPUTED_NEW_COMPUTED).
    ``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING).
    """
    if num_reqs == 0 or not self.is_initialized:
        return

    # V2 reads non-contiguous idx_mapping positions, so snapshot the whole
    # decision buffer rather than only [:num_reqs].
    accepted_snapshot = self.num_accepted_tokens_snapshot
    accepted_snapshot.copy_(num_accepted_tokens_gpu)

    total_states = self.num_states
    grid = (num_reqs, total_states, _TEMPORAL_TILES)
    postprocess_mamba_fused_kernel[grid](
        accepted_snapshot,
        state_idx_gpu,
        None,  # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED
        new_num_computed_tokens_gpu,
        None,  # num_draft: unused under PRECOMPUTED_NEW_COMPUTED
        self.block_table_ptrs,
        self.block_table_stride_req,
        self.state_base_addrs,
        self.state_block_strides,
        self.state_elem_sizes,
        self.state_inner_sizes,
        self.state_conv_widths,
        self.state_group_indices,
        self.state_dim_row_count,
        self.state_dim_row_stride,
        num_accepted_tokens_gpu,
        idx_mapping,
        num_reqs,
        block_size=self.block_size,
        COPY_BLOCK_SIZE=1024,
        CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
        HAS_IDX_MAPPING=True,
        PRECOMPUTED_NEW_COMPUTED=True,
        SKIP_TEMPORAL_STATE_COPY=self.replayssm is not None,
        TEMPORAL_TILES=_TEMPORAL_TILES,
    )

run_fused_precopy(num_reqs, state_idx_gpu, src_col_gpu, token_bias_gpu, idx_mapping)

Pre-copy each request's previous running block into its new window block before the forward pass (align boundary migration).

Parameters:

• num_reqs

  (int) –

  Number of active requests (batch order).

• state_idx_gpu

  (Tensor) –

  [max_reqs] post-advance dst block column per req slot.

• src_col_gpu

  (Tensor) –

  [max_reqs] pre-advance src block column (-1 = fresh).

• token_bias_gpu

  (Tensor) –

  [max_reqs] accepted-token bias (num_accepted - 1).

• idx_mapping

  (Tensor | None) –

  optional [num_reqs] batch_idx -> req_state_idx. None means V1 batch order already equals request state order.

Source code in vllm/v1/worker/mamba_utils.py

def run_fused_precopy(
    self,
    num_reqs: int,
    state_idx_gpu: torch.Tensor,
    src_col_gpu: torch.Tensor,
    token_bias_gpu: torch.Tensor,
    idx_mapping: torch.Tensor | None,
) -> None:
    """Pre-copy each request's previous running block into its new window
    block before the forward pass (align boundary migration).

    Args:
        num_reqs: Number of active requests (batch order).
        state_idx_gpu: [max_reqs] post-advance dst block column per req slot.
        src_col_gpu: [max_reqs] pre-advance src block column (-1 = fresh).
        token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1).
        idx_mapping: optional [num_reqs] batch_idx -> req_state_idx.
            None means V1 batch order already equals request state order.

    """
    if num_reqs == 0 or not self.is_initialized:
        return
    total_states = self.num_states
    grid = (num_reqs, total_states, _TEMPORAL_TILES)
    precopy_mamba_align_fused_kernel[grid](
        state_idx_gpu,
        src_col_gpu,
        token_bias_gpu,
        self.block_table_ptrs,
        self.block_table_stride_req,
        self.state_base_addrs,
        self.state_block_strides,
        self.state_elem_sizes,
        self.state_inner_sizes,
        self.state_conv_widths,
        self.state_group_indices,
        self.state_dim_row_count,
        self.state_dim_row_stride,
        idx_mapping,
        num_reqs,
        COPY_BLOCK_SIZE=1024,
        CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
        HAS_IDX_MAPPING=idx_mapping is not None,
        ZERO_TEMPORAL_TOKEN_BIAS=self.replayssm is not None,
        TEMPORAL_TILES=_TEMPORAL_TILES,
    )

_FusedPrecopy

Bases: NamedTuple

Resolved fused align pre-copy resources (all non-None once resolved).

Source code in vllm/v1/worker/mamba_utils.py

class _FusedPrecopy(NamedTuple):
    """Resolved fused align pre-copy resources (all non-None once resolved)."""

    ctx: "MambaSpecDecodeGPUContext"
    state_idx: CpuGpuBuffer
    src_col: CpuGpuBuffer
    token_bias: CpuGpuBuffer

_copy_mamba_state_block(state_idx, bt_row_idx, src_col, dst_col, token_bias, block_table_ptrs_ptr, block_table_stride_req, state_base_addrs_ptr, state_block_strides_ptr, state_elem_sizes_ptr, state_inner_sizes_ptr, state_conv_widths_ptr, state_group_indices_ptr, state_dim_row_count_ptr, state_dim_row_stride_ptr, tile_idx, COPY_BLOCK_SIZE, CONV_STATE_DIM_FIRST, TEMPORAL_TILES)

Copy one (layer, state-type) mamba state block between block columns.

Shared copy body of postprocess_mamba_fused_kernel and precopy_mamba_align_fused_kernel, mirroring the V1 copy specs (get_conv_copy_spec / get_temporal_copy_spec): - conv state (conv_width > 0): shift the window by token_bias tokens, state[bt[src_col], token_bias:] -> state[bt[dst_col], :conv_width - token_bias] - temporal state: token_bias selects the accepted speculative column, state[bt[src_col + token_bias]] -> state[bt[dst_col]]

The caller owns the decision logic (which columns, whether to copy); this device function only performs the byte copy for the given metadata slot.

tile_idx in [0, TEMPORAL_TILES) partitions the temporal state's u64 range into TEMPORAL_TILES contiguous, COPY_BLOCK_SIZE-aligned slices, giving more CTAs to fill the SMs at small batch (multi-MiB temporal copies otherwise leave the GPU under-filled). Conv states are small; only tile_idx == 0 copies them. TEMPORAL_TILES == 1 and tile_idx == 0 reproduces the untiled behavior.

Source code in vllm/v1/worker/mamba_utils.py

@triton.jit
def _copy_mamba_state_block(
    state_idx,
    bt_row_idx,
    src_col,
    dst_col,
    token_bias,
    block_table_ptrs_ptr,
    block_table_stride_req,
    state_base_addrs_ptr,
    state_block_strides_ptr,
    state_elem_sizes_ptr,
    state_inner_sizes_ptr,
    state_conv_widths_ptr,
    state_group_indices_ptr,
    # DS conv row metadata. Zero keeps the single-region copy path.
    state_dim_row_count_ptr,
    state_dim_row_stride_ptr,
    tile_idx,
    COPY_BLOCK_SIZE: tl.constexpr,
    CONV_STATE_DIM_FIRST: tl.constexpr,
    TEMPORAL_TILES: tl.constexpr,
):
    """Copy one (layer, state-type) mamba state block between block columns.

    Shared copy body of ``postprocess_mamba_fused_kernel`` and
    ``precopy_mamba_align_fused_kernel``, mirroring the V1 copy specs
    (``get_conv_copy_spec`` / ``get_temporal_copy_spec``):
    - conv state (conv_width > 0): shift the window by ``token_bias`` tokens,
      ``state[bt[src_col], token_bias:] ->
      state[bt[dst_col], :conv_width - token_bias]``
    - temporal state: ``token_bias`` selects the accepted speculative column,
      ``state[bt[src_col + token_bias]] -> state[bt[dst_col]]``

    The caller owns the decision logic (which columns, whether to copy); this
    device function only performs the byte copy for the given metadata slot.

    ``tile_idx`` in ``[0, TEMPORAL_TILES)`` partitions the temporal state's
    u64 range into ``TEMPORAL_TILES`` contiguous, COPY_BLOCK_SIZE-aligned
    slices, giving more CTAs to fill the SMs at small batch (multi-MiB
    temporal copies otherwise leave the GPU under-filled). Conv states are
    small; only ``tile_idx == 0`` copies them. ``TEMPORAL_TILES == 1`` and
    ``tile_idx == 0`` reproduces the untiled behavior.
    """
    state_base_addr = tl.load(state_base_addrs_ptr + state_idx)
    state_block_stride = tl.load(state_block_strides_ptr + state_idx)
    state_elem_size = tl.load(state_elem_sizes_ptr + state_idx)
    state_inner_size = tl.load(state_inner_sizes_ptr + state_idx)
    conv_width = tl.load(state_conv_widths_ptr + state_idx)

    # Load the group index for this state, then index into the correct
    # group's block table. Each mamba group has independently allocated
    # physical blocks. Reinterpret as int32* since block ids are int32.
    group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64)
    group_base_addr = tl.load(block_table_ptrs_ptr + group_idx)
    block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32))
    block_table_base = block_table_typed + bt_row_idx * block_table_stride_req

    # Widen block ids to int64 before they reach `block_id * state_block_stride`
    # below: state_block_stride can exceed 2**31 bytes for large mamba caches,
    # and Triton would otherwise do the multiply in int32 and wrap.
    dest_block_id = tl.load(block_table_base + dst_col).to(tl.int64)
    dst_addr = state_base_addr + dest_block_id * state_block_stride

    is_conv_state = conv_width > 0

    if CONV_STATE_DIM_FIRST and is_conv_state:
        # Conv states are small; only tile 0 does the copy. Higher tiles
        # early-return so they contribute nothing beyond a bounds check.
        if tile_idx > 0:
            return
        # DS conv layout: state_len is the slide axis; copy per dim row.
        src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
        dim_rows = tl.load(state_dim_row_count_ptr + state_idx)
        row_stride = tl.load(state_dim_row_stride_ptr + state_idx)
        src_block_addr = state_base_addr + src_block_id * state_block_stride
        offsets = tl.arange(0, COPY_BLOCK_SIZE)

        # Stable row-to-lane ownership makes left shifts memmove-safe while
        # exposing the dimension rows in parallel. All addresses retain
        # state_elem_size alignment: tensor strides and token offsets are
        # measured in whole elements before conversion to bytes.
        num_dst_tokens = conv_width - token_bias
        for token_idx in range(0, num_dst_tokens):
            for row_base in range(0, dim_rows, COPY_BLOCK_SIZE):
                rows = row_base + offsets
                mask = rows < dim_rows
                src_byte_addr = (
                    src_block_addr
                    + rows * row_stride
                    + (token_idx + token_bias) * state_elem_size
                )
                dst_byte_addr = (
                    dst_addr + rows * row_stride + token_idx * state_elem_size
                )
                if state_elem_size == 2:
                    src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16))
                    dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16))
                    data_u16 = tl.load(src_u16, mask=mask)
                    tl.store(dst_u16, data_u16, mask=mask)
                elif state_elem_size == 4:
                    src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32))
                    dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32))
                    data_u32 = tl.load(src_u32, mask=mask)
                    tl.store(dst_u32, data_u32, mask=mask)
                else:
                    for byte_idx in range(0, state_elem_size):
                        src_u8 = (src_byte_addr + byte_idx).to(
                            tl.pointer_type(tl.uint8)
                        )
                        dst_u8 = (dst_byte_addr + byte_idx).to(
                            tl.pointer_type(tl.uint8)
                        )
                        data_u8 = tl.load(src_u8, mask=mask)
                        tl.store(dst_u8, data_u8, mask=mask)
        return

    if is_conv_state:
        if tile_idx > 0:
            return
        # SD conv: copy
        #   state[bt[src_col], token_bias:] ->
        #   state[bt[dst_col], :conv_width - token_bias]
        src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
        src_block_addr = state_base_addr + src_block_id * state_block_stride
        token_bytes = state_inner_size * state_elem_size
        num_dst_tokens = conv_width - token_bias

        # Distinct blocks and exact self-copies cannot have a destructive
        # overlap, so retain the u64-vectorized single-CTA copy.
        if src_block_id != dest_block_id or token_bias == 0:
            src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes
            copy_size = num_dst_tokens.to(tl.int64) * token_bytes
            _memcpy_u64_tiled(
                src_addr,
                dst_addr,
                copy_size,
                tile_idx,
                COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
                NUM_TILES=1,
            )
            return

        # Copy tokens from low to high. Each token-sized source and destination
        # region is disjoint, so same-block left shifts are memmove-safe
        # without a barrier.
        for token_idx in range(0, num_dst_tokens):
            src_token = src_block_addr + (token_idx + token_bias) * token_bytes
            dst_token = dst_addr + token_idx * token_bytes
            _memcpy_u64_tiled(
                src_token,
                dst_token,
                token_bytes,
                tile_idx,
                COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
                NUM_TILES=1,
            )
        return

    # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]]
    # Body u64 range is partitioned across TEMPORAL_TILES CTAs to keep the
    # SMs filled at small batch.
    actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(tl.int64)
    src_addr = state_base_addr + actual_src_block_id * state_block_stride
    # Use natural block data size (inner_size * elem_size), NOT
    # state_block_stride which is the page stride and can exceed the
    # actual data when the state tensor uses as_strided page padding.
    copy_size = state_inner_size * state_elem_size
    _memcpy_u64_tiled(
        src_addr,
        dst_addr,
        copy_size,
        tile_idx,
        COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
        NUM_TILES=TEMPORAL_TILES,
    )

_memcpy_u64_tiled(src_addr, dst_addr, copy_size, tile_idx, COPY_BLOCK_SIZE, NUM_TILES)

Head/body/tail memcpy with the u64 body split across NUM_TILES CTAs.

Fast path (src and dst share sub-8B alignment): tile 0 owns the byte head that lifts dst to 8B and the 0-7 byte tail; body tiles vectorize as u64 over the aligned interior. NUM_TILES=1 collapses to a single- CTA memcpy. Production callers derive both addresses from the same state_base_addr + block_id * stride and always take this path.

Slow path (mismatched sub-8B alignment): byte-wide tiled copy. Some NVIDIA parts (e.g. GB200) reject misaligned u64 loads with cudaErrorMisalignedAddress instead of accepting the misaligned-sector cost, so we can't just let the fast path run with a misaligned src.

Source code in vllm/v1/worker/mamba_utils.py

@triton.jit
def _memcpy_u64_tiled(
    src_addr,
    dst_addr,
    copy_size,
    tile_idx,
    COPY_BLOCK_SIZE: tl.constexpr,
    NUM_TILES: tl.constexpr,
):
    """Head/body/tail memcpy with the u64 body split across ``NUM_TILES`` CTAs.

    Fast path (``src`` and ``dst`` share sub-8B alignment): tile 0 owns the
    byte head that lifts dst to 8B and the 0-7 byte tail; body tiles vectorize
    as u64 over the aligned interior. ``NUM_TILES=1`` collapses to a single-
    CTA memcpy. Production callers derive both addresses from the same
    ``state_base_addr + block_id * stride`` and always take this path.

    Slow path (mismatched sub-8B alignment): byte-wide tiled copy. Some
    NVIDIA parts (e.g. GB200) reject misaligned u64 loads with
    ``cudaErrorMisalignedAddress`` instead of accepting the misaligned-sector
    cost, so we can't just let the fast path run with a misaligned src.
    """
    src_addr_i = src_addr.to(tl.int64)
    dst_addr_i = dst_addr.to(tl.int64)

    if ((src_addr_i ^ dst_addr_i) & 7) == 0:
        head_bytes = tl.minimum(((-dst_addr_i) & 7).to(tl.int64), copy_size)
        if tile_idx == 0:
            head_off = tl.arange(0, 8)
            head_mask = head_off < head_bytes
            head_src = src_addr.to(tl.pointer_type(tl.uint8))
            head_dst = dst_addr.to(tl.pointer_type(tl.uint8))
            tl.store(
                head_dst + head_off,
                tl.load(head_src + head_off, mask=head_mask),
                mask=head_mask,
            )

        # Body: u64 tiled. Rounding per_tile up to COPY_BLOCK_SIZE keeps every
        # inner-loop iteration full-width vectorized; only the last non-empty
        # tile can be masked. Late tiles fall off the end and iterate zero
        # times.
        body_bytes = copy_size - head_bytes
        body_u64 = body_bytes // 8
        per_tile_u64_raw = tl.cdiv(body_u64, NUM_TILES)
        per_tile_u64 = tl.cdiv(per_tile_u64_raw, COPY_BLOCK_SIZE) * COPY_BLOCK_SIZE
        tile_start = tile_idx.to(tl.int64) * per_tile_u64
        tile_end = tl.minimum(tile_start + per_tile_u64, body_u64)

        src_body_u64 = (src_addr + head_bytes).to(tl.pointer_type(tl.uint64))
        dst_body_u64 = (dst_addr + head_bytes).to(tl.pointer_type(tl.uint64))
        offsets = tl.arange(0, COPY_BLOCK_SIZE)
        for i in range(tile_start, tile_end, COPY_BLOCK_SIZE):
            mask = (i + offsets) < tile_end
            data = tl.load(src_body_u64 + i + offsets, mask=mask)
            tl.store(dst_body_u64 + i + offsets, data, mask=mask)

        if tile_idx == 0:
            tail_start = head_bytes + body_u64 * 8
            tail_bytes = copy_size - tail_start
            tail_off = tl.arange(0, 8)
            tail_src = (src_addr + tail_start).to(tl.pointer_type(tl.uint8))
            tail_dst = (dst_addr + tail_start).to(tl.pointer_type(tl.uint8))
            tail_mask = tail_off < tail_bytes
            tl.store(
                tail_dst + tail_off,
                tl.load(tail_src + tail_off, mask=tail_mask),
                mask=tail_mask,
            )
    else:
        src_u8 = src_addr.to(tl.pointer_type(tl.uint8))
        dst_u8 = dst_addr.to(tl.pointer_type(tl.uint8))
        per_tile_bytes_raw = tl.cdiv(copy_size, NUM_TILES)
        per_tile_bytes = tl.cdiv(per_tile_bytes_raw, COPY_BLOCK_SIZE) * COPY_BLOCK_SIZE
        tile_start = tile_idx.to(tl.int64) * per_tile_bytes
        tile_end = tl.minimum(tile_start + per_tile_bytes, copy_size)
        offsets = tl.arange(0, COPY_BLOCK_SIZE)
        for i in range(tile_start, tile_end, COPY_BLOCK_SIZE):
            mask = (i + offsets) < tile_end
            data = tl.load(src_u8 + i + offsets, mask=mask)
            tl.store(dst_u8 + i + offsets, data, mask=mask)

_resolve_fused_precopy(align_ctx)

Bundle the fused-path buffers, or None for the scalar path.

Returning one non-None bundle lets callers narrow all four members with a single is not None check instead of re-asserting each buffer per use.

Source code in vllm/v1/worker/mamba_utils.py

def _resolve_fused_precopy(
    align_ctx: "MambaSpecDecodeGPUContext | None",
) -> _FusedPrecopy | None:
    """Bundle the fused-path buffers, or None for the scalar path.

    Returning one non-None bundle lets callers narrow all four members with a
    single ``is not None`` check instead of re-asserting each buffer per use.
    """
    if align_ctx is None:
        return None
    assert align_ctx.mamba_state_idx_buf is not None
    assert align_ctx.precopy_src_col_buf is not None
    assert align_ctx.precopy_token_bias_buf is not None
    return _FusedPrecopy(
        align_ctx,
        align_ctx.mamba_state_idx_buf,
        align_ctx.precopy_src_col_buf,
        align_ctx.precopy_token_bias_buf,
    )

cleanup_mamba_state_idx(scheduler_output, mamba_state_idx)

Pop stale mamba_state_idx entries for finished/preempted/resumed reqs.

Force-preempted requests (e.g., during reset_prefix_cache / KV cache flush) appear in resumed_req_ids without a corresponding entry in preempted_req_ids, leaving stale entries that can point to block indices beyond the new (smaller) block allocation.

Source code in vllm/v1/worker/mamba_utils.py

def cleanup_mamba_state_idx(
    scheduler_output: SchedulerOutput,
    mamba_state_idx: dict[str, int],
) -> None:
    """Pop stale `mamba_state_idx` entries for finished/preempted/resumed reqs.

    Force-preempted requests (e.g., during reset_prefix_cache / KV cache
    flush) appear in resumed_req_ids without a corresponding entry in
    preempted_req_ids, leaving stale entries that can point to block
    indices beyond the new (smaller) block allocation.
    """
    finished_req_ids = scheduler_output.finished_req_ids
    preempted_req_ids = scheduler_output.preempted_req_ids or set()
    resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
    for req_id in itertools.chain(finished_req_ids, preempted_req_ids, resumed_req_ids):
        mamba_state_idx.pop(req_id, None)

get_mamba_group_ids(mamba_groups)

Return the sorted group ids of all mamba groups, without duplicates.

A group whose layers use different MambaSpecs is listed under one key per spec, so the same group id can appear more than once.

Source code in vllm/v1/worker/mamba_utils.py

def get_mamba_group_ids(mamba_groups: dict[MambaSpec, list[int]]) -> list[int]:
    """Return the sorted group ids of all mamba groups, without duplicates.

    A group whose layers use different ``MambaSpec``s is listed under one key
    per spec, so the same group id can appear more than once.
    """
    return sorted(
        {group_id for group_ids in mamba_groups.values() for group_id in group_ids}
    )

get_mamba_groups(kv_cache_config)

Return a mapping from each distinct MambaSpec to its sorted group_ids.

Source code in vllm/v1/worker/mamba_utils.py

def get_mamba_groups(kv_cache_config: KVCacheConfig) -> dict[MambaSpec, list[int]]:
    """Return a mapping from each distinct MambaSpec to its sorted group_ids."""
    mamba_groups: dict[MambaSpec, set[int]] = {}
    for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
        kv_cache_spec = kv_cache_group.kv_cache_spec
        if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
            kv_cache_spec = kv_cache_spec.first_spec
        if not isinstance(kv_cache_spec, MambaSpec):
            continue
        for layer_name in kv_cache_group.layer_names:
            mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name)
            mamba_groups.setdefault(mamba_spec, set()).add(i)
    assert len(mamba_groups) > 0, "no mamba layers in the model"
    return {spec: sorted(ids) for spec, ids in mamba_groups.items()}

postprocess_mamba_fused_kernel(num_accepted_tokens_ptr, mamba_state_idx_ptr, num_scheduled_tokens_ptr, num_computed_tokens_ptr, num_draft_tokens_ptr, block_table_ptrs_ptr, block_table_stride_req, state_base_addrs_ptr, state_block_strides_ptr, state_elem_sizes_ptr, state_inner_sizes_ptr, state_conv_widths_ptr, state_group_indices_ptr, state_dim_row_count_ptr, state_dim_row_stride_ptr, num_accepted_tokens_out_ptr, idx_mapping_ptr, num_reqs, block_size, COPY_BLOCK_SIZE, CONV_STATE_DIM_FIRST, HAS_IDX_MAPPING=False, PRECOMPUTED_NEW_COMPUTED=False, SKIP_TEMPORAL_STATE_COPY=False, TEMPORAL_TILES=1)

Fused GPU kernel for postprocess_mamba that computes decisions AND performs mamba state copies without any CPU-GPU synchronization.

Grid: (num_reqs, num_states [, TEMPORAL_TILES]) - program_id(0) = request/batch index - program_id(1) = state_idx (flattened index into layer/state_type metadata) - program_id(2) = temporal-copy tile index (0 when TEMPORAL_TILES == 1)

The kernel indexes directly into pre-flattened metadata arrays using program_id(1). The grid dimensions encode the total state count.

Source code in vllm/v1/worker/mamba_utils.py

@triton.jit(do_not_specialize=["num_reqs"])
def postprocess_mamba_fused_kernel(
    # Decision inputs (per-request)
    num_accepted_tokens_ptr,
    mamba_state_idx_ptr,
    num_scheduled_tokens_ptr,
    num_computed_tokens_ptr,
    num_draft_tokens_ptr,
    # Per-group block table base addresses: int64[num_groups]. Each entry is
    # the data_ptr of that group's persistent [max_reqs, max_blocks] int32
    # block table.
    block_table_ptrs_ptr,
    block_table_stride_req: tl.int64,  # stride between requests (in elements)
    # Mamba state metadata, flattened across all layer/state tensors. Different
    # Mamba groups may contribute different numbers of state tensors.
    state_base_addrs_ptr,  # base address of each state tensor
    state_block_strides_ptr,  # bytes per block for each state
    state_elem_sizes_ptr,  # element size for each state
    state_inner_sizes_ptr,  # number of elements in inner dimensions
    state_conv_widths_ptr,  # conv width for conv states (0 for temporal)
    state_group_indices_ptr,  # maps state_idx to group index in block table
    # DS conv row metadata. Zero keeps the single-region copy path.
    state_dim_row_count_ptr,  # int32: per-block dim row count for DS conv
    state_dim_row_stride_ptr,  # int64: bytes between rows for DS conv
    # Output: num_accepted_tokens update (for src==dst case)
    num_accepted_tokens_out_ptr,
    # Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The
    # per-request decision arrays and the block tables are both in
    # req-state-slot order.
    idx_mapping_ptr,
    # Runtime parameter (varies per batch - NOT constexpr to avoid recompilation)
    num_reqs,
    # Compile-time constants (fixed after model initialization)
    # block_size: determined by model config, constant for all invocations
    block_size: tl.constexpr,
    # COPY_BLOCK_SIZE: fixed tuning parameter for memory copy loop
    COPY_BLOCK_SIZE: tl.constexpr,
    CONV_STATE_DIM_FIRST: tl.constexpr,
    # HAS_IDX_MAPPING: when True, program_id(0) is a batch index resolved to a
    # req-state slot via idx_mapping_ptr (V2). When False, it is the req index.
    HAS_IDX_MAPPING: tl.constexpr = False,
    # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds
    # the post-step new_num_computed value (V2 supplies the advanced count).
    PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False,
    # FlashInfer ReplaySSM migrates temporal state through its ring materializer.
    # Skip the generic temporal copy in that case so it cannot overwrite the
    # ReplaySSM-owned transition; canonical convolution state still needs the
    # generic snapshot at the same block boundary.
    SKIP_TEMPORAL_STATE_COPY: tl.constexpr = False,
    # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across
    # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a
    # 3D grid (num_reqs, total_states, TEMPORAL_TILES). Default 1 preserves
    # the existing 2D-grid contract.
    TEMPORAL_TILES: tl.constexpr = 1,
):
    """Fused GPU kernel for postprocess_mamba that computes decisions AND performs
    mamba state copies without any CPU-GPU synchronization.

    Grid: (num_reqs, num_states [, TEMPORAL_TILES])
    - program_id(0) = request/batch index
    - program_id(1) = state_idx (flattened index into layer/state_type metadata)
    - program_id(2) = temporal-copy tile index (0 when TEMPORAL_TILES == 1)

    The kernel indexes directly into pre-flattened metadata arrays using
    program_id(1). The grid dimensions encode the total state count.
    """
    batch_idx = tl.program_id(0)
    state_idx = tl.program_id(1)
    tile_idx = tl.program_id(2)

    # Bounds check
    if batch_idx >= num_reqs:
        return

    if HAS_IDX_MAPPING:
        req_idx = tl.load(idx_mapping_ptr + batch_idx)
        if req_idx < 0:
            return
    else:
        req_idx = batch_idx

    # Compute decision logic (mirrors postprocess_mamba Python reference)
    num_accepted = tl.load(num_accepted_tokens_ptr + req_idx)
    src_block_idx = tl.load(mamba_state_idx_ptr + req_idx)

    if PRECOMPUTED_NEW_COMPUTED:
        new_num_computed = tl.load(num_computed_tokens_ptr + req_idx)
        num_tokens_running_state = new_num_computed - num_accepted + 1
    else:
        num_scheduled = tl.load(num_scheduled_tokens_ptr + req_idx)
        num_computed = tl.load(num_computed_tokens_ptr + req_idx)
        num_draft = tl.load(num_draft_tokens_ptr + req_idx)
        num_tokens_running_state = num_computed + num_scheduled - num_draft
        new_num_computed = num_tokens_running_state + num_accepted - 1

    needs_copy, accept_token_bias, dest_block_idx = mamba_state_copy_boundary(
        num_tokens_running_state,
        new_num_computed,
        block_size,
    )

    if not needs_copy:
        return

    # Update accepted-token count before early exits (per-request, so only
    # state_idx == 0 writes). Also guard on tile_idx == 0 so tiles > 0
    # (when TEMPORAL_TILES > 1) do not duplicate the store.
    if src_block_idx == dest_block_idx and state_idx == 0 and tile_idx == 0:
        tl.store(num_accepted_tokens_out_ptr + req_idx, 1)

    # Skip no-op self-copy.
    if src_block_idx == dest_block_idx and accept_token_bias == 0:
        return

    if SKIP_TEMPORAL_STATE_COPY:
        # The flattened metadata interleaves convolution and temporal states.
        # state_conv_widths doubles as the state-kind tag: a positive width is
        # canonical convolution state, while zero marks ReplaySSM-owned temporal
        # state and must bypass _copy_mamba_state_block.
        conv_width = tl.load(state_conv_widths_ptr + state_idx)
        if conv_width == 0:
            return

    _copy_mamba_state_block(
        state_idx,
        req_idx,
        src_block_idx,
        dest_block_idx,
        accept_token_bias,
        block_table_ptrs_ptr,
        block_table_stride_req,
        state_base_addrs_ptr,
        state_block_strides_ptr,
        state_elem_sizes_ptr,
        state_inner_sizes_ptr,
        state_conv_widths_ptr,
        state_group_indices_ptr,
        state_dim_row_count_ptr,
        state_dim_row_stride_ptr,
        tile_idx,
        COPY_BLOCK_SIZE,
        CONV_STATE_DIM_FIRST,
        TEMPORAL_TILES,
    )

postprocess_mamba_gpu(*, bufs, num_reqs, num_accepted_tokens_gpu, num_accepted_tokens_cpu_tensor, input_batch, kv_cache_config, forward_context, mamba_state_copy_funcs, run_prefix_state_migration)

Publish accepted-token results to model-wide Mamba state.

The first call binds persistent block tables and state pointers. When prefix migration is enabled, the fused kernel snapshots the original accepted counts, copies canonical state at crossed boundaries, and normalizes the live counts for the next step. ReplaySSM then commits its group-shared trackers from the original counts and materializes any planned prefix snapshots; mode none skips the copy phase and commits directly. Finally, the normalized live counts are copied back only when the next iteration has a CPU consumer.

Source code in vllm/v1/worker/mamba_utils.py

def postprocess_mamba_gpu(
    *,
    bufs: "MambaBuffers",
    num_reqs: int,
    num_accepted_tokens_gpu: torch.Tensor,
    num_accepted_tokens_cpu_tensor: torch.Tensor | None,
    input_batch: InputBatch,
    kv_cache_config: KVCacheConfig,
    forward_context: dict[str, Any],
    mamba_state_copy_funcs: MambaStateCopyFuncsByType,
    run_prefix_state_migration: bool,
) -> None:
    """Publish accepted-token results to model-wide Mamba state.

    The first call binds persistent block tables and state pointers. When prefix
    migration is enabled, the fused kernel snapshots the original accepted
    counts, copies canonical state at crossed boundaries, and normalizes the live
    counts for the next step. ReplaySSM then commits its group-shared trackers
    from the original counts and materializes any planned prefix snapshots; mode
    ``none`` skips the copy phase and commits directly. Finally, the normalized
    live counts are copied back only when the next iteration has a CPU consumer.
    """
    ctx = bufs.postprocess_align
    # The caller enables this context for spec-decode hybrid state copies or
    # for model-owned FlashInfer ReplaySSM lifecycle maintenance under STP.
    assert ctx is not None
    assert ctx.num_scheduled_tokens_buf is not None
    assert ctx.num_computed_tokens_buf is not None
    assert ctx.num_draft_tokens_buf is not None
    assert ctx.is_prefilling_buf is not None

    if not ctx.is_initialized:
        ctx.initialize_from_forward_context(
            kv_cache_config,
            forward_context,
            mamba_state_copy_funcs,
            [
                input_batch.block_table[gid].get_device_tensor(num_reqs)
                for gid in ctx.mamba_group_ids
            ],
        )

    accepted_tokens_for_postprocess = num_accepted_tokens_gpu
    live_cols = None
    if run_prefix_state_migration:
        assert ctx.mamba_state_idx_buf is not None
        live_cols = ctx.mamba_state_idx_buf.gpu
        ctx.run_fused_postprocess(
            num_reqs=num_reqs,
            num_accepted_tokens_gpu=num_accepted_tokens_gpu,
            mamba_state_idx_gpu=live_cols,
            num_scheduled_tokens_gpu=ctx.num_scheduled_tokens_buf.gpu,
            num_computed_tokens_gpu=ctx.num_computed_tokens_buf.gpu,
            num_draft_tokens_gpu=ctx.num_draft_tokens_buf.gpu,
        )
        accepted_tokens_for_postprocess = ctx.num_accepted_tokens_snapshot
    if ctx.replayssm is not None:
        # Keep this separate from run_fused_postprocess: that kernel is a
        # per-(request, canonical-state, tile) copy and is skipped when prefix
        # migration is disabled. ReplaySSM must instead publish its shared
        # per-(request, cache-group) ring trackers in every cache mode, then
        # optionally compact/materialize one plan per group. Folding the two
        # would either duplicate tracker writes for every state tensor or make
        # the generic Mamba kernel depend on FlashInfer-only group descriptors.
        ctx.replayssm.postprocess(
            idx_mapping=None,
            query_metadata=ctx.num_scheduled_tokens_buf.gpu,
            query_metadata_is_cumulative=False,
            num_computed_tokens=ctx.num_computed_tokens_buf.gpu,
            num_computed_is_post_step=False,
            num_accepted_tokens=accepted_tokens_for_postprocess,
            is_prefilling=ctx.is_prefilling_buf.gpu,
            live_cols=live_cols,
            num_reqs=num_reqs,
        )

    if num_accepted_tokens_cpu_tensor is not None and ctx.replayssm is None:
        # CPU consumers need the normalized live counts for the next step, not
        # the original snapshot retained for ReplaySSM.
        num_accepted_tokens_cpu_tensor[:num_reqs].copy_(
            num_accepted_tokens_gpu[:num_reqs], non_blocking=True
        )

precopy_mamba_align_fused_kernel(mamba_state_idx_ptr, src_col_ptr, token_bias_ptr, block_table_ptrs_ptr, block_table_stride_req, state_base_addrs_ptr, state_block_strides_ptr, state_elem_sizes_ptr, state_inner_sizes_ptr, state_conv_widths_ptr, state_group_indices_ptr, state_dim_row_count_ptr, state_dim_row_stride_ptr, idx_mapping_ptr, num_reqs, COPY_BLOCK_SIZE, CONV_STATE_DIM_FIRST, HAS_IDX_MAPPING=True, ZERO_TEMPORAL_TOKEN_BIAS=False, TEMPORAL_TILES=1)

Pre-copy mamba "align" state across block boundaries.

Before the forward pass, copy each request's last SSM/conv state from its previous block column into the new window block column, so the kernels read the initial state from the write-side block as usual (V1 align semantics). Same per-(layer, state) copy semantics as postprocess_mamba_fused_kernel (shared _copy_mamba_state_block body, i.e. the V1 preprocess_mamba copy specs), but driven by the GPU-resident src columns so it needs no CPU-GPU sync (async-scheduling safe).

Grid: (num_reqs, num_states [, TEMPORAL_TILES]). V2 passes a batch-to-state idx_mapping; V1 already stores the staged arrays in batch order and uses HAS_IDX_MAPPING=False.

Source code in vllm/v1/worker/mamba_utils.py

@triton.jit(do_not_specialize=["num_reqs"])
def precopy_mamba_align_fused_kernel(
    # Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by
    # the V2 fused align preprocess kernel for the current step:
    mamba_state_idx_ptr,  # post-advance dst block column
    src_col_ptr,  # pre-advance src block column (-1 = fresh)
    token_bias_ptr,  # accepted-token bias = num_accepted - 1 (pre-reset)
    # Same flattened state-layout metadata as postprocess_mamba_fused_kernel
    block_table_ptrs_ptr,
    block_table_stride_req: tl.int64,
    state_base_addrs_ptr,
    state_block_strides_ptr,
    state_elem_sizes_ptr,
    state_inner_sizes_ptr,
    state_conv_widths_ptr,
    state_group_indices_ptr,
    state_dim_row_count_ptr,
    state_dim_row_stride_ptr,
    idx_mapping_ptr,  # [num_reqs] batch_idx -> req_state_idx (-1 to skip)
    num_reqs,
    COPY_BLOCK_SIZE: tl.constexpr,
    CONV_STATE_DIM_FIRST: tl.constexpr,
    HAS_IDX_MAPPING: tl.constexpr = True,
    ZERO_TEMPORAL_TOKEN_BIAS: tl.constexpr = False,
    # TEMPORAL_TILES: see postprocess_mamba_fused_kernel. Default 1 preserves
    # the 2D-grid contract; > 1 requires a 3D grid.
    TEMPORAL_TILES: tl.constexpr = 1,
):
    """Pre-copy mamba "align" state across block boundaries.

    Before the forward pass, copy each request's last SSM/conv state from its
    previous block column into the new window block column, so the kernels read
    the initial state from the write-side block as usual (V1 align semantics).
    Same per-(layer, state) copy semantics as ``postprocess_mamba_fused_kernel``
    (shared ``_copy_mamba_state_block`` body, i.e. the V1 ``preprocess_mamba``
    copy specs), but driven by the GPU-resident src columns so it needs no
    CPU-GPU sync (async-scheduling safe).

    Grid: (num_reqs, num_states [, TEMPORAL_TILES]). V2 passes a
    batch-to-state idx_mapping; V1 already stores the staged arrays in batch
    order and uses HAS_IDX_MAPPING=False.
    """
    batch_idx = tl.program_id(0)
    state_idx = tl.program_id(1)
    tile_idx = tl.program_id(2)
    if batch_idx >= num_reqs:
        return
    if HAS_IDX_MAPPING:
        req_idx = tl.load(idx_mapping_ptr + batch_idx)
        if req_idx < 0:
            return
    else:
        req_idx = batch_idx

    src_col = tl.load(src_col_ptr + req_idx)
    dst_col = tl.load(mamba_state_idx_ptr + req_idx)
    # Fresh state, or still writing the same block: kernels locate the initial
    # state in-block via num_accepted (preserved when no boundary is crossed),
    # so there is nothing to copy.
    if src_col < 0 or src_col == dst_col:
        return

    token_bias = tl.load(token_bias_ptr + req_idx)
    if ZERO_TEMPORAL_TOKEN_BIAS:
        conv_width = tl.load(state_conv_widths_ptr + state_idx)
        token_bias = tl.where(conv_width > 0, token_bias, 0)
    _copy_mamba_state_block(
        state_idx,
        req_idx,
        src_col,
        dst_col,
        token_bias,
        block_table_ptrs_ptr,
        block_table_stride_req,
        state_base_addrs_ptr,
        state_block_strides_ptr,
        state_elem_sizes_ptr,
        state_inner_sizes_ptr,
        state_conv_widths_ptr,
        state_group_indices_ptr,
        state_dim_row_count_ptr,
        state_dim_row_stride_ptr,
        tile_idx,
        COPY_BLOCK_SIZE,
        CONV_STATE_DIM_FIRST,
        TEMPORAL_TILES,
    )

preprocess_mamba(scheduler_output, kv_cache_config, cache_config, mamba_state_idx, input_batch, requests, forward_context, mamba_state_copy_funcs, copy_bufs, align_ctx=None, num_computed_tokens_gpu=None, query_start_loc_gpu=None, num_accepted_tokens_gpu=None)

Copy the mamba state of previous step to the last (1 + num_speculative_blocks) block.

Source code in vllm/v1/worker/mamba_utils.py

def preprocess_mamba(
    scheduler_output: SchedulerOutput,
    kv_cache_config: KVCacheConfig,
    cache_config: CacheConfig,
    mamba_state_idx: dict[str, int],
    input_batch: InputBatch,
    requests: dict[str, CachedRequestState],
    forward_context: dict[str, Any],
    mamba_state_copy_funcs: MambaStateCopyFuncsByType,
    copy_bufs: MambaCopyBuffers,
    align_ctx: MambaSpecDecodeGPUContext | None = None,
    num_computed_tokens_gpu: torch.Tensor | None = None,
    query_start_loc_gpu: torch.Tensor | None = None,
    num_accepted_tokens_gpu: torch.Tensor | None = None,
):
    """Copy the mamba state of previous step to the last
    (1 + num_speculative_blocks) block.
    """
    fused = _resolve_fused_precopy(align_ctx)
    mamba_group_ids = copy_bufs.mamba_group_ids
    mamba_spec = copy_bufs.mamba_spec
    num_speculative_blocks = mamba_spec.num_speculative_blocks
    # TODO(Chen): we need to optimize this function a lot
    assert cache_config.enable_prefix_caching
    block_size = mamba_spec.block_size
    cleanup_mamba_state_idx(scheduler_output, mamba_state_idx)

    copy_bufs.offset = 0
    num_reqs = len(input_batch.req_ids)
    if fused is not None:
        if num_reqs == 0:
            return
        if not fused.ctx.is_initialized:
            fused.ctx.initialize_from_forward_context(
                kv_cache_config,
                forward_context,
                mamba_state_copy_funcs,
                [
                    input_batch.block_table[gid].get_device_tensor(num_reqs)
                    for gid in fused.ctx.mamba_group_ids
                ],
            )

        fused.src_col.np[:num_reqs] = -1
        if fused.ctx.replayssm is None:
            fused.token_bias.np[:num_reqs] = 0

    for i, req_id in enumerate(input_batch.req_ids):
        req_state = requests[req_id]
        prev_state_idx = mamba_state_idx.get(req_id)
        if prev_state_idx is None:
            # New / resumed request; num_computed_tokens == 0 gives -1.
            prev_state_idx = (req_state.num_computed_tokens - 1) // block_size

        num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
        num_blocks = (
            cdiv(req_state.num_computed_tokens + num_scheduled_tokens, block_size)
            + num_speculative_blocks
        )
        # We always save the current running state at the last
        # (1 + num_speculative_blocks) block.
        # A corner case worth mention here: assume we have block_size = 4 and
        # num_speculative_tokens = 2. The request is [A, B, C] and contains 2 draft
        # tokens [draft 1, draft 2]. Then we will have:
        # Block 0: [A, B, C, draft 1]
        # Block 1: [draft 2, TOFILL, TOFILL, TOFILL]
        # Block 2: speculative block
        # Block 3: speculative block
        # And use block 1 to save the running state.
        curr_state_idx = num_blocks - 1 - num_speculative_blocks
        mamba_state_idx[req_id] = curr_state_idx
        if fused is not None:
            fused.state_idx.np[i] = curr_state_idx
            if fused.ctx.replayssm is not None:
                # The GPU decision uses the actual computed count and the
                # restored accepted count in current batch order.
                fused.state_idx.np[i] = prev_state_idx

        # Generic Mamba retains its CPU accepted-token staging. ReplaySSM
        # computes the convolution offset from the restored GPU count below.
        if (
            prev_state_idx != -1
            and prev_state_idx != curr_state_idx
            and (fused is None or fused.ctx.replayssm is None)
        ):
            accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1
            if fused is not None:
                assert accept_token_bias >= 0
                fused.src_col.np[i] = prev_state_idx
                fused.token_bias.np[i] = accept_token_bias
            else:
                collect_mamba_copy_meta(
                    copy_bufs,
                    kv_cache_config,
                    mamba_state_copy_funcs,
                    mamba_group_ids,
                    prev_state_idx,
                    curr_state_idx,
                    accept_token_bias,
                    req_state,
                    forward_context,
                )
            input_batch.num_accepted_tokens_cpu[i] = 1

    if fused is not None:
        fused.state_idx.copy_to_gpu(num_reqs)
        if fused.ctx.replayssm is not None:
            assert num_computed_tokens_gpu is not None
            assert query_start_loc_gpu is not None
            assert num_accepted_tokens_gpu is not None
            preprocess_mamba_align_fused_kernel[(triton.cdiv(num_reqs, 256),)](
                None,
                fused.state_idx.gpu,
                num_computed_tokens_gpu,
                query_start_loc_gpu,
                num_accepted_tokens_gpu,
                fused.src_col.gpu,
                fused.token_bias.gpu,
                num_reqs,
                BLOCK_SIZE=256,
                MAMBA_BLOCK_SIZE=block_size,
                HAS_IDX_MAPPING=False,
            )
            fused.ctx.replayssm.reset_new_slots(
                idx_mapping=None,
                src_cols=fused.src_col.gpu,
                dst_cols=fused.state_idx.gpu,
                num_reqs=num_reqs,
            )
        else:
            fused.src_col.copy_to_gpu(num_reqs)
            fused.token_bias.copy_to_gpu(num_reqs)
        fused.ctx.run_fused_precopy(
            num_reqs=num_reqs,
            state_idx_gpu=fused.state_idx.gpu,
            src_col_gpu=fused.src_col.gpu,
            token_bias_gpu=fused.token_bias.gpu,
            idx_mapping=None,
        )
    else:
        do_mamba_copy_block(copy_bufs)

preprocess_mamba_align_fused_kernel(idx_mapping_ptr, state_idx_ptr, num_computed_tokens_ptr, query_start_loc_ptr, num_accepted_tokens_ptr, src_col_ptr, src_off_ptr, num_reqs, BLOCK_SIZE, MAMBA_BLOCK_SIZE, HAS_IDX_MAPPING=True)

Fused align preprocess: emit the pre-copy src column/offset AND advance state_idx in a single launch (V2 align).

Per batch_idx (0..num_reqs-1), resolving req slot via idx_mapping: 1. Read pre-advance state_idx and num_accepted (last step's values). 2. Store the pre-copy src columns for precopy_mamba_align_fused_kernel: - src_col = state_idx (the previous running block column) - src_off = max(num_accepted - 1, 0) (the accepted-token bias) 3. Advance state_idx to the new running block. Reset num_accepted to 1 when the block changes: the migrated convolution state now sits at the start of the new block and must be read with the neutral accepted-token bias. ReplaySSM ring cursors are migrated separately.

Source code in vllm/v1/worker/mamba_utils.py

@triton.jit(do_not_specialize=["num_reqs"])
def preprocess_mamba_align_fused_kernel(
    idx_mapping_ptr,
    state_idx_ptr,
    num_computed_tokens_ptr,
    query_start_loc_ptr,
    num_accepted_tokens_ptr,
    src_col_ptr,
    src_off_ptr,
    num_reqs,
    BLOCK_SIZE: tl.constexpr,
    MAMBA_BLOCK_SIZE: tl.constexpr,
    HAS_IDX_MAPPING: tl.constexpr = True,
):
    """Fused align preprocess: emit the pre-copy src column/offset AND advance
    state_idx in a single launch (V2 align).

    Per batch_idx (0..num_reqs-1), resolving req slot via idx_mapping:
      1. Read pre-advance state_idx and num_accepted (last step's values).
      2. Store the pre-copy src columns for ``precopy_mamba_align_fused_kernel``:
         - src_col = state_idx (the previous running block column)
         - src_off = max(num_accepted - 1, 0) (the accepted-token bias)
      3. Advance state_idx to the new running block. Reset num_accepted to 1
         when the block changes: the migrated convolution state
         now sits at the start of the new block and must be read with the neutral
         accepted-token bias. ReplaySSM ring cursors are migrated separately.
    """
    offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_reqs
    if HAS_IDX_MAPPING:
        req_indices = tl.load(idx_mapping_ptr + offsets, mask=mask, other=-1)
        mask = mask & (req_indices >= 0)
    else:
        req_indices = offsets

    state_idx = tl.load(state_idx_ptr + req_indices, mask=mask, other=-1)
    num_accepted = tl.load(num_accepted_tokens_ptr + req_indices, mask=mask, other=1)

    src_off = tl.maximum(num_accepted - 1, 0)
    tl.store(src_col_ptr + req_indices, state_idx, mask=mask)
    tl.store(src_off_ptr + req_indices, src_off, mask=mask)

    num_computed = tl.load(num_computed_tokens_ptr + req_indices, mask=mask, other=0)
    query_start = tl.load(query_start_loc_ptr + offsets, mask=mask, other=0)
    query_end = tl.load(query_start_loc_ptr + offsets + 1, mask=mask, other=0)
    computed_after = num_computed + query_end - query_start
    new_state_idx = (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1
    tl.store(state_idx_ptr + req_indices, new_state_idx, mask=mask)
    should_reset = (state_idx >= 0) & (state_idx != new_state_idx)
    tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset)

stage_postprocess_inputs_to_gpu(ctx, scheduler_output, req_ids, num_reqs, requests, *, num_decodes=None)

Stage the per-request decisions consumed after token acceptance.

Generic Mamba uses scheduled, computed, and draft counts. ReplaySSM also uses the final Mamba attention split to reset trackers after prefill.

Source code in vllm/v1/worker/mamba_utils.py

def stage_postprocess_inputs_to_gpu(
    ctx: MambaSpecDecodeGPUContext,
    scheduler_output: SchedulerOutput,
    req_ids: list[str],
    num_reqs: int,
    requests: dict[str, CachedRequestState],
    *,
    num_decodes: int | None = None,
) -> None:
    """Stage the per-request decisions consumed after token acceptance.

    Generic Mamba uses scheduled, computed, and draft counts. ReplaySSM also
    uses the final Mamba attention split to reset trackers after prefill.
    """
    assert ctx.num_scheduled_tokens_buf is not None
    assert ctx.num_computed_tokens_buf is not None
    assert ctx.num_draft_tokens_buf is not None
    assert ctx.is_prefilling_buf is not None

    scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
    num_scheduled = scheduler_output.num_scheduled_tokens
    scheduled_np = ctx.num_scheduled_tokens_buf.np
    computed_np = ctx.num_computed_tokens_buf.np
    draft_np = ctx.num_draft_tokens_buf.np
    for i in range(num_reqs):
        req_id = req_ids[i]
        scheduled = num_scheduled[req_id]
        req_state = requests[req_id]
        computed = req_state.num_computed_tokens
        num_draft = len(scheduled_spec_tokens.get(req_id, []))
        scheduled_np[i] = scheduled
        computed_np[i] = computed
        draft_np[i] = num_draft
    ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs)
    ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs)
    ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs)
    if num_decodes is not None:
        # Full decode graphs can include padded rows after the real requests.
        num_decodes = min(num_decodes, num_reqs)
        ctx.is_prefilling_buf.np[:num_decodes] = False
        ctx.is_prefilling_buf.np[num_decodes:num_reqs] = True
        ctx.is_prefilling_buf.copy_to_gpu(num_reqs)