From 8a26869bcad1ae4ab3069a7d2aece86c17c1c695 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Wed, 7 Oct 2026 10:32:46 +0800 Subject: [PATCH] [Bugfix] Index Mamba align state copies by request slot in MRV2 (#60210) Signed-off-by: zjy0516 Co-authored-by: Karl0007 <44745209+Karl0007@users.noreply.github.com> Co-authored-by: Yongqin Wang --- .../kernels/mamba/test_precopy_mamba_align.py | 80 ++++++++++++++++++- tests/v1/worker/test_mamba_utils.py | 46 +++++++++++ vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu/model_states/interface.py | 6 ++ .../worker/gpu/model_states/mamba_hybrid.py | 77 +++++++++--------- vllm/v1/worker/mamba_utils.py | 42 ++++++---- 6 files changed, 195 insertions(+), 57 deletions(-) diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py index 22b73c936e..8a91f8bdac 100644 --- a/tests/kernels/mamba/test_precopy_mamba_align.py +++ b/tests/kernels/mamba/test_precopy_mamba_align.py @@ -31,7 +31,11 @@ import torch from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils from vllm.platforms import current_platform from vllm.v1.worker import mamba_utils as worker_mamba_utils -from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, precopy_mamba_align_fused_kernel +from vllm.v1.worker.mamba_utils import ( + _TEMPORAL_TILES, + postprocess_mamba_fused_kernel, + precopy_mamba_align_fused_kernel, +) _parametrize: Callable[..., Callable[[Any], Any]] @@ -195,7 +199,8 @@ def test_precopy_matches_v1_copy_specs( convs, ssms, device, conv_state_dim_first ) bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + # Non-identity mapping: block-table rows are request slots, not batch rows. + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device).flip(0) grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles) precopy_mamba_align_fused_kernel[grid]( dst_col, @@ -225,6 +230,77 @@ def test_precopy_matches_v1_copy_specs( torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) +@_parametrize("conv_state_dim_first", [False, True]) +@_cuda_required +def test_v2_postprocess_indexes_block_tables_by_req_slot(conv_state_dim_first): + """Under PP the V2 postprocess runs after later steps have re-gathered the + batch-ordered tables, so it must index the request-slot tables by slot.""" + device = torch.device("cuda") + torch.manual_seed(0) + num_reqs = 4 + block_size = 8 + num_blocks = num_reqs * MAX_COLS + 1 + bt = torch.arange(1, num_blocks, dtype=torch.int32, device=device).view( + num_reqs, MAX_COLS + ) + # slot 0: copy col 1 -> 0; slot 1: src == dst; slot 2: not aligned; + # slot 3: copy col 2 -> 0 with one accepted draft token. + num_accepted = torch.tensor([1, 1, 1, 2], dtype=torch.int32, device=device) + state_idx = torch.tensor([1, 0, 1, 2], dtype=torch.int32, device=device) + new_num_computed = torch.tensor([8, 8, 5, 8], dtype=torch.int32, device=device) + ref_src = torch.tensor([1, 0, -1, 2], dtype=torch.int32) + ref_dst = torch.zeros(num_reqs, dtype=torch.int32) + ref_bias = torch.tensor([0, 0, 0, 1], dtype=torch.int32) + + convs, ssms = _build_state(num_blocks, device, conv_state_dim_first) + conv_ref, ssm_ref = _reference( + convs, + ssms, + bt.cpu(), + ref_src, + ref_dst, + ref_bias, + num_reqs, + conv_state_dim_first, + ) + base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( + convs, ssms, device, conv_state_dim_first + ) + bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) + idx_mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int32, device=device) + postprocess_mamba_fused_kernel[(num_reqs, NUM_LAYERS * 2, 1)]( + num_accepted.clone(), + state_idx, + None, + new_num_computed, + None, + bt_ptrs, + bt.stride(0), + base, + blk_stride, + elem, + inner, + width, + group, + drc, + drs, + num_accepted, + idx_mapping, + num_reqs, + block_size=block_size, + COPY_BLOCK_SIZE=1024, + CONV_STATE_DIM_FIRST=conv_state_dim_first, + HAS_IDX_MAPPING=True, + PRECOMPUTED_NEW_COMPUTED=True, + TEMPORAL_TILES=1, + ) + torch.accelerator.synchronize() + + for layer in range(NUM_LAYERS): + torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0) + torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) + + def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch): monkeypatch.setattr( layer_mamba_utils, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 50faf0fb8b..4407d6fc54 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -22,6 +22,7 @@ from vllm.v1.kv_cache_interface import ( MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, @@ -638,6 +639,51 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group(): assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_v2_align_ctx_binds_request_slot_tables_for_state_copies(): + """V2 state copies may run steps after their batch (PP), so they must read + the per-request-slot tables; aligned state indices are computed for the + current batch and keep reading the gathered tables.""" + cfg = _TestConfig(num_layers=1) + device = torch.device("cuda") + kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) + forward_context = { + "layer_0": _make_mock_attention( + torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), + torch.empty(cfg.num_blocks, cfg.temporal_state_dim), + ) + } + block_tables = BlockTables( + block_sizes=[cfg.block_size], + max_num_reqs=cfg.max_num_reqs, + max_num_batched_tokens=64, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[cfg.block_size], + ) + + state = object.__new__(MambaHybridModelState) + state._align_mode = True + state._mamba_group_ids = [] + state._mamba_spec = None + state._mamba_state_copy_funcs = _COPY_FUNCS + state.max_num_reqs = cfg.max_num_reqs + state.device = device + state.vllm_config = MagicMock() + state.vllm_config.compilation_config.static_forward_context = forward_context + + state.initialize_kv_cache(kv_cache_config, block_tables) + + ctx = state._mamba_ctx + assert ctx is not None and ctx.is_initialized + assert ctx.block_table_ptrs.tolist() == [ + block_tables.block_tables[0].gpu.data_ptr() + ] + assert ctx.aligned_index_block_table_ptrs.tolist() == [ + block_tables.input_block_tables[0].data_ptr() + ] + + # ----------------------------------------------------------------------------- # stage_postprocess_inputs_to_gpu: single-pass staging into pinned views # ----------------------------------------------------------------------------- diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c3ac543ff2..d0c08707f4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -778,6 +778,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.kv_caches = [ cache for cache in kv_caches_dict.values() if cache.device == self.device ] + self.model_state.initialize_kv_cache(self.kv_cache_config, self.block_tables) if is_profiling: self.kv_connector = NO_OP_KV_CONNECTOR else: diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index f03caf4dba..3db3fc9e7a 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -127,6 +127,12 @@ class ModelState(ABC): def apply_staged_writes(self) -> None: return None + def initialize_kv_cache( + self, kv_cache_config: KVCacheConfig, block_tables: BlockTables + ) -> None: + """Hook run after the KV cache tensors are allocated and bound.""" + return None + def capture_inner_cudagraphs( self, input_buffers: InputBuffers, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index f3f3c666e0..1f8af07898 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -21,6 +21,7 @@ from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState @@ -134,47 +135,42 @@ class MambaHybridModelState(DefaultModelState): self._mamba_spec = mamba_spec return self._mamba_group_ids, self._mamba_spec - def _ensure_align_ctx( - self, - kv_cache_config: KVCacheConfig, - mamba_group_ids: list[int], - block_tables: tuple[torch.Tensor, ...], - ) -> MambaSpecDecodeGPUContext: + def initialize_kv_cache( + self, kv_cache_config: KVCacheConfig, block_tables: BlockTables + ) -> None: + if not self._align_mode: + return + mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) if self._mamba_state_copy_funcs is None: mamba_groups = get_mamba_groups(kv_cache_config) mamba_types = {spec.mamba_type for spec in mamba_groups} copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types) validate_mamba_state_copy_funcs(mamba_groups, copy_funcs) self._mamba_state_copy_funcs = copy_funcs - copy_funcs = self._mamba_state_copy_funcs - if self._mamba_ctx is None: - # Both SD and DS conv layouts support a >0 spec-decode shift: the - # fused pre-copy kernel (``_copy_mamba_state_block``) applies the - # ``token_bias = num_accepted - 1`` window shift per conv layout - # (SD: contiguous slice; DS: per-dim-row strided slice), matching - # the V1 ``get_conv_copy_spec`` semantics. - self._mamba_ctx = MambaSpecDecodeGPUContext.create( - max_num_reqs=self.max_num_reqs, - kv_cache_config=kv_cache_config, - copy_funcs=copy_funcs, - device=self.device, - make_buffer=lambda n, dtype: CpuGpuBuffer( - n, dtype=dtype, device=self.device - ), - ) - ctx = self._mamba_ctx - if not ctx.is_initialized: - forward_context = self.vllm_config.compilation_config.static_forward_context - # block_tables are batch-order slices of the persistent - # input_block_tables (stable data_ptr), so the metadata is captured - # once here and reused across steps. - ctx.initialize_from_forward_context( - kv_cache_config, - forward_context, - copy_funcs, - [block_tables[gid] for gid in mamba_group_ids], - ) - return ctx + # Both SD and DS conv layouts support a >0 spec-decode shift: the + # fused pre-copy kernel (``_copy_mamba_state_block``) applies the + # ``token_bias = num_accepted - 1`` window shift per conv layout + # (SD: contiguous slice; DS: per-dim-row strided slice), matching + # the V1 ``get_conv_copy_spec`` semantics. + self._mamba_ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=self.max_num_reqs, + kv_cache_config=kv_cache_config, + copy_funcs=self._mamba_state_copy_funcs, + device=self.device, + make_buffer=lambda n, dtype: CpuGpuBuffer( + n, dtype=dtype, device=self.device + ), + ) + # Under PP a step's postprocess runs after later steps re-gathered the + # batch-ordered input tables, so state copies read the per-request-slot + # tables. Aligned state indices are for the current batch only. + self._mamba_ctx.initialize_from_forward_context( + kv_cache_config, + self.vllm_config.compilation_config.static_forward_context, + self._mamba_state_copy_funcs, + [block_tables.block_tables[gid].gpu for gid in mamba_group_ids], + [block_tables.input_block_tables[gid] for gid in mamba_group_ids], + ) def preprocess_state( self, @@ -194,8 +190,9 @@ class MambaHybridModelState(DefaultModelState): num_reqs = input_batch.num_reqs if num_reqs == 0: return - mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) - ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + _, mamba_spec = self._get_mamba_group_info(kv_cache_config) + ctx = self._mamba_ctx + assert ctx is not None # The state-advance + pre-copy kernels run every step; they fast-exit per # request when src_col < 0 or src_col == dst_col, so no copy happens on @@ -301,10 +298,8 @@ class MambaHybridModelState(DefaultModelState): if hasattr(builder, "mamba_aligned_state_indices"): aligned_index_builders.append((group_idx, builder)) if aligned_index_builders: - ctx = self._ensure_align_ctx( - kv_cache_config, mamba_group_ids, block_tables - ) - all_group_indices = ctx.compute_aligned_state_indices( + assert self._mamba_ctx is not None + all_group_indices = self._mamba_ctx.compute_aligned_state_indices( input_batch.seq_lens, num_reqs ) for group_idx, builder in aligned_index_builders: diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 33a2b8b480..a3b4eab340 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -392,8 +392,8 @@ def postprocess_mamba_fused_kernel( # 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. + # 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, @@ -476,10 +476,9 @@ def postprocess_mamba_fused_kernel( 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, + req_idx, src_block_idx, dest_block_idx, accept_token_bias, @@ -610,7 +609,7 @@ def precopy_mamba_align_fused_kernel( token_bias = tl.load(token_bias_ptr + req_idx) _copy_mamba_state_block( state_idx, - batch_idx, + req_idx, src_col, dst_col, token_bias, @@ -806,8 +805,10 @@ class MambaSpecDecodeGPUContext: # 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). + # 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. @@ -900,6 +901,9 @@ class MambaSpecDecodeGPUContext: 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), @@ -924,6 +928,7 @@ class MambaSpecDecodeGPUContext: 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. @@ -955,8 +960,12 @@ class MambaSpecDecodeGPUContext: 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. + 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: @@ -968,6 +977,7 @@ class MambaSpecDecodeGPUContext: forward_context, mamba_state_copy_funcs, block_tables, + aligned_index_block_tables or block_tables, ) def _populate_metadata( @@ -976,6 +986,7 @@ class MambaSpecDecodeGPUContext: 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 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): @@ -1076,16 +1087,19 @@ class MambaSpecDecodeGPUContext: # `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(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 in enumerate(block_tables): + 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() + ) self.is_initialized = True @@ -1107,7 +1121,7 @@ class MambaSpecDecodeGPUContext: block_rows = 32 grid = (triton.cdiv(num_reqs, block_rows),) get_aligned_state_indices_multi_group_kernel[grid]( - self.block_table_ptrs, + self.aligned_index_block_table_ptrs, seq_lens, self.aligned_state_indices, self.block_table_stride_req,