# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses import itertools from collections.abc import Callable from typing import Any, NamedTuple import torch from vllm.config import CacheConfig from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFuncsByType, get_conv_copy_spec, get_temporal_copy_spec, is_conv_state_dim_first, ) from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, MambaSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch logger = init_logger(__name__) # 16 saturates HBM on H100/GB200 across the reqs=8..128 range in # microbenchmarks _TEMPORAL_TILES = 16 @triton.jit(do_not_specialize=["num_requests"]) def get_aligned_state_indices_multi_group_kernel( block_table_ptrs_ptr, seq_lens_ptr, state_indices_ptr, block_table_stride_req: tl.int64, seq_lens_stride: tl.constexpr, state_indices_stride_0: tl.constexpr, state_indices_stride_1: tl.constexpr, state_indices_stride_2: tl.constexpr, num_requests, CACHE_BLOCK_SIZE: tl.constexpr, NUM_GROUPS: tl.constexpr, BLOCK_GROUPS: tl.constexpr, NUM_STATE_SLOTS: tl.constexpr, BLOCK_STATE_SLOTS: tl.constexpr, BLOCK_ROWS: tl.constexpr, ): rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) valid_row = rows < num_requests seq_lens = tl.load( seq_lens_ptr + rows * seq_lens_stride, mask=valid_row, other=1, ) first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0) # load multiple block table for each group groups = tl.arange(0, BLOCK_GROUPS) valid_group = groups < NUM_GROUPS group_base_addrs = tl.load( block_table_ptrs_ptr + groups, mask=valid_group, other=0, ) block_tables = group_base_addrs.to(tl.pointer_type(tl.int32)) state_slots = tl.arange(0, BLOCK_STATE_SLOTS) valid_state_slot = state_slots < NUM_STATE_SLOTS state_indices = tl.load( block_tables[:, None, None] + rows[None, :, None] * block_table_stride_req + first_state_slot[None, :, None] + state_slots[None, None, :], mask=( valid_group[:, None, None] & valid_row[None, :, None] & valid_state_slot[None, None, :] ), ) tl.store( state_indices_ptr + groups[:, None, None] * state_indices_stride_0 + rows[None, :, None] * state_indices_stride_1 + state_slots[None, None, :] * state_indices_stride_2, state_indices, mask=( valid_group[:, None, None] & valid_row[None, :, None] & valid_state_slot[None, None, :] ), ) @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) def _reinterpret_u64_as_i64(value: int) -> int: """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" return value if value < (1 << 63) else value - (1 << 64) @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, ) @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 are in req-state-slot order; the block table # is in batch order, so HAS_IDX_MAPPING splits the two indexings. 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, # 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 aligned_new_computed = (new_num_computed // block_size) * block_size needs_copy = aligned_new_computed >= num_tokens_running_state if not needs_copy: return # Compute copy parameters accept_token_bias = aligned_new_computed - num_tokens_running_state dest_block_idx = aligned_new_computed // block_size - 1 # 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 bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx _copy_mamba_state_block( state_idx, bt_row_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, ) @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, ): """Fused align preprocess: emit the pre-copy src column/offset AND advance state_idx (with accepted-token reset) 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, and reset num_accepted to 1 when a block boundary is crossed (so the migrated state, now at the start of the new block, is read with the neutral bias). """ offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < num_reqs req_indices = tl.load(idx_mapping_ptr + offsets, mask=mask, other=0) 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) @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, # 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) _copy_mamba_state_block( state_idx, batch_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, ) @triton.jit def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): pid = tl.program_id(0) src_ptr = tl.load(src_ptrs + pid) dst_ptr = tl.load(dst_ptrs + pid) size = tl.load(sizes + pid) is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr offsets = tl.arange(0, BLOCK_SIZE) for i in range(0, size, BLOCK_SIZE): mask = (i + offsets) < size curr_src_ptr = (src_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) data = tl.load(curr_src_ptr, mask=mask) if is_left_overlap: # Preserve each lane's source before a lower-address lane stores # over it. The condition is uniform within the program. tl.debug_barrier() tl.store(curr_dst_ptr, data, mask=mask) def batch_memcpy(src_ptrs, dst_ptrs, sizes): batch = src_ptrs.shape[0] assert dst_ptrs.shape[0] == batch assert sizes.shape[0] == batch grid = (batch,) BLOCK_SIZE = 1024 batch_memcpy_kernel[grid](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=BLOCK_SIZE) def _get_mamba_spec_for_layer( kv_cache_group: KVCacheGroupSpec, layer_name: str ) -> MambaSpec: kv_cache_spec = kv_cache_group.kv_cache_spec if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name] assert isinstance(kv_cache_spec, MambaSpec) return kv_cache_spec 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} ) 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()} def validate_mamba_state_copy_funcs( mamba_groups: dict[MambaSpec, list[int]], copy_funcs: MambaStateCopyFuncsByType, ) -> None: for mamba_spec in mamba_groups: assert mamba_spec.mamba_type in copy_funcs, ( f"missing state copy funcs for {mamba_spec.mamba_type}" ) state_copy_funcs = copy_funcs[mamba_spec.mamba_type] assert 0 < len(state_copy_funcs) <= len(mamba_spec.shapes), ( f"{mamba_spec.mamba_type} declares {len(mamba_spec.shapes)} states, " f"but provides {len(state_copy_funcs)} state copy funcs; expected " "a non-empty copyable prefix" ) @dataclasses.dataclass class MambaCopyBuffers: src_ptrs: CpuGpuBuffer dst_ptrs: CpuGpuBuffer sizes: CpuGpuBuffer mamba_group_ids: list[int] mamba_spec: MambaSpec offset: int = 0 @classmethod def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, copy_funcs: MambaStateCopyFuncsByType, make_buffer: Callable[..., CpuGpuBuffer], ) -> "MambaCopyBuffers": mamba_groups = get_mamba_groups(kv_cache_config) 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" mamba_group_ids = get_mamba_group_ids(mamba_groups) entries_per_req = sum( len( copy_funcs[ _get_mamba_spec_for_layer( kv_cache_config.kv_cache_groups[gid], layer_name ).mamba_type ] ) for gid in mamba_group_ids for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names ) n = max_num_reqs * entries_per_req return cls( src_ptrs=make_buffer(n, dtype=torch.uint64), dst_ptrs=make_buffer(n, dtype=torch.uint64), sizes=make_buffer(n, dtype=torch.int32), mamba_group_ids=mamba_group_ids, mamba_spec=mamba_spec, ) @dataclasses.dataclass class MambaSpecDecodeGPUContext: """ Context for GPU-side Mamba state copy operations during the fused postprocess path. Only used when speculative decoding is enabled on a hybrid model (and the mamba_cache_config is in align mode). 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 # Output buffer for num_accepted_tokens updates num_accepted_tokens_out: 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). 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 + align mode). 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 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 @classmethod def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, copy_funcs: MambaStateCopyFuncsByType, 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" copy_funcs_by_spec = { spec: copy_funcs[spec.mamba_type] for spec in mamba_groups } # Count physical state tensors across all Mamba groups. Different # groups may expose different state specs. total_states = sum( len( copy_funcs_by_spec[ _get_mamba_spec_for_layer( kv_cache_config.kv_cache_groups[gid], layer_name ) ] ) 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_out=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_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), 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, ) 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], ) -> 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`. Their `data_ptr()` / `stride(0)` are captured once for the kernel to index into. """ 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, ) def _populate_metadata( self, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, block_tables: list[torch.Tensor], ) -> None: idx = 0 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] 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) state_copy_funcs = mamba_state_copy_funcs[mamba_spec.mamba_type] attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache 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 assert idx == self.num_states # 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) == self.num_groups, ( f"expected {self.num_groups} block tables, got {len(block_tables)}" ) strides = {bt.stride(0) for bt in 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 in enumerate(block_tables): self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr()) 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 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.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 # Initialize output to current values (unchanged unless src==dst) self.num_accepted_tokens_out[: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]( num_accepted_tokens_gpu, 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, self.num_accepted_tokens_out, 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(), 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, 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]. num_accepted_tokens_snapshot = self.num_accepted_tokens_out num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) postprocess_mamba_fused_kernel[grid]( num_accepted_tokens_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, TEMPORAL_TILES=_TEMPORAL_TILES, ) @dataclasses.dataclass class MambaBuffers: """Single owner for all mamba-specific runner buffers. The two sub-objects have different gates: ``preprocess`` is needed whenever ``mamba_cache_mode == "align"``; ``postprocess_align`` is needed only when align is combined with speculative decoding on a hybrid model, 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, 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, device=device, make_buffer=make_buffer, ) if with_postprocess_align else None ), ) def collect_mamba_copy_meta( copy_bufs: MambaCopyBuffers, kv_cache_config: KVCacheConfig, mamba_state_copy_funcs: MambaStateCopyFuncsByType, mamba_group_ids: list[int], src_block_idx: int, dest_block_idx: int, accept_token_bias: int, req_state: CachedRequestState, forward_context: dict[str, Any], ) -> None: if src_block_idx == dest_block_idx and accept_token_bias == 0: return src_ptrs_np = copy_bufs.src_ptrs.np dst_ptrs_np = copy_bufs.dst_ptrs.np sizes_np = copy_bufs.sizes.np offset = copy_bufs.offset for mamba_group_id in mamba_group_ids: block_ids = req_state.block_ids[mamba_group_id] dest_block_id = block_ids[dest_block_idx] kv_cache_group = kv_cache_config.kv_cache_groups[mamba_group_id] 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) state_copy_funcs = mamba_state_copy_funcs[mamba_spec.mamba_type] attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache for state, state_copy_func in zip(kv_caches, state_copy_funcs): copy_spec = state_copy_func( state, block_ids, src_block_idx, accept_token_bias + 1 ) src_ptrs_np[offset] = copy_spec.start_addr dst_ptrs_np[offset] = state[dest_block_id].data_ptr() sizes_np[offset] = copy_spec.num_elements * state.element_size() offset += 1 copy_bufs.offset = offset def do_mamba_copy_block(copy_bufs: MambaCopyBuffers): n = copy_bufs.offset if n == 0: return batch_memcpy( copy_bufs.src_ptrs.copy_to_gpu(n), copy_bufs.dst_ptrs.copy_to_gpu(n), copy_bufs.sizes.copy_to_gpu(n), ) 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) 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 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, ) def preprocess_mamba( scheduler_output: SchedulerOutput, kv_cache_config: KVCacheConfig, cache_config: CacheConfig, mamba_state_idx: dict[str, int], input_batch: GPUInputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, copy_bufs: MambaCopyBuffers, align_ctx: MambaSpecDecodeGPUContext | 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 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 prev_state_idx != -1 and prev_state_idx != curr_state_idx: 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) 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) def postprocess_mamba_all( scheduler_output: SchedulerOutput, kv_cache_config: KVCacheConfig, input_batch: GPUInputBatch, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], num_spec_tokens: int, num_reqs: int, ): """All-mode postprocess (only meaningful with num_spec_tokens > 0): record per-request the block index of the last token scheduled this step, so the next step can anchor its in-place writes when accepted drafts leave the sequence at a non-block-aligned position. """ if num_spec_tokens <= 0: return mamba_groups = get_mamba_groups(kv_cache_config) block_sizes = {mamba_spec.block_size for mamba_spec in mamba_groups} assert len(block_sizes) == 1, "all mamba groups must share block_size" block_size = next(iter(block_sizes)) full_decode_len = 1 + num_spec_tokens scheduled = scheduler_output.num_scheduled_tokens for req_id in input_batch.req_ids[:num_reqs]: num_query = scheduled.get(req_id, 0) if num_query == full_decode_len: req = requests[req_id] seq_len = req.num_computed_tokens + num_query mamba_state_idx[req_id] = max(0, (seq_len - 1) // block_size) else: mamba_state_idx.pop(req_id, None) def preprocess_mamba_all_specdec( scheduler_output: SchedulerOutput, input_batch: GPUInputBatch, mamba_state_idx: dict[str, int], num_reqs: int, prev_last_scheduled_idx_buf: CpuGpuBuffer, ) -> None: cleanup_mamba_state_idx(scheduler_output, mamba_state_idx) np_view = prev_last_scheduled_idx_buf.np for i, req_id in enumerate(input_batch.req_ids[:num_reqs]): np_view[i] = mamba_state_idx.get(req_id, -1) np_view[num_reqs:].fill(-1) prev_last_scheduled_idx_buf.copy_to_gpu() def postprocess_mamba_align_gpu( *, bufs: "MambaBuffers", num_reqs: int, num_accepted_tokens_gpu: torch.Tensor, num_accepted_tokens_cpu_tensor: torch.Tensor, input_batch: GPUInputBatch, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, ) -> None: """GPU-side mamba postprocess for spec decode + hybrid + align mode. Lazily binds the fused-kernel context to the persistent block tables and forward-context state pointers on the first call, runs the fused kernel, and async-copies the per-request accepted-token counts back to the input batch's CPU tensor for the next iteration's preprocess. """ ctx = bufs.postprocess_align # Caller is responsible for gating on spec decode + hybrid; this assert is # a tripwire if those gates ever drift apart. assert ctx is not None assert ctx.mamba_state_idx_buf 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 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 ], ) ctx.run_fused_postprocess( num_reqs=num_reqs, num_accepted_tokens_gpu=num_accepted_tokens_gpu, mamba_state_idx_gpu=ctx.mamba_state_idx_buf.gpu, 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, ) # ``num_accepted_tokens_out`` is pre-initialized from # ``num_accepted_tokens_gpu``; the kernel only overwrites entries to 1 # when src_block_idx == dest_block_idx (copy within the same block), so # the original count is preserved for everyone else. num_accepted_tokens_cpu_tensor[:num_reqs].copy_( ctx.num_accepted_tokens_out[:num_reqs], non_blocking=True ) def stage_postprocess_inputs_to_gpu( ctx: MambaSpecDecodeGPUContext, scheduler_output: SchedulerOutput, req_ids: list[str], num_reqs: int, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], ) -> None: """Stage all per-request inputs the fused mamba postprocess kernel reads. Walks ``req_ids[:num_reqs]`` once, writing each request's mamba block index and scheduled/computed/draft token counts into the matching pinned numpy views, then issues four non-blocking H→D copies. The fused kernel indexes the resulting GPU tensors by ``req_idx``. Buffers live on ``ctx`` and only exist when the postprocess kernel is enabled. Invariant: ``preprocess_mamba`` must have run first for the same batch so that every ``req_ids[i]`` has an entry in ``mamba_state_idx``. """ assert ctx.mamba_state_idx_buf 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 scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens num_scheduled = scheduler_output.num_scheduled_tokens state_idx_np = ctx.mamba_state_idx_buf.np 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] state_idx = mamba_state_idx.get(req_id) assert state_idx is not None, ( f"mamba_state_idx missing entry for {req_id!r}; " "preprocess_mamba must run before stage_postprocess_inputs_to_gpu" ) state_idx_np[i] = state_idx scheduled_np[i] = num_scheduled[req_id] computed_np[i] = requests[req_id].num_computed_tokens draft_np[i] = len(scheduled_spec_tokens.get(req_id, [])) ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) 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)