From 5aeb5213124a97970f55bcf80290304c5bd3de0d Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Mon, 30 Mar 2026 19:22:20 +0800 Subject: [PATCH 01/10] support qwen35 / mamba hybrid model for model runner v2 Signed-off-by: zhuhaoran --- vllm/config/vllm.py | 4 + vllm/v1/worker/gpu/attn_utils.py | 172 ++++++++++++++---- vllm/v1/worker/gpu/block_table.py | 8 +- vllm/v1/worker/gpu/cudagraph_utils.py | 29 ++- vllm/v1/worker/gpu/model_runner.py | 30 ++- vllm/v1/worker/gpu/model_states/default.py | 30 ++- vllm/v1/worker/gpu/model_states/interface.py | 2 + vllm/v1/worker/gpu/model_states/whisper.py | 2 + .../worker/gpu/spec_decode/eagle/cudagraph.py | 1 + vllm/v1/worker/gpu/states.py | 10 + 10 files changed, 239 insertions(+), 49 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 6229b44d52a8..5c60cd3f918a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1821,6 +1821,10 @@ def validate_block_size(self) -> None: "to schedule a multiple of block_size tokens even if they are " "in the middle of a mm input" ) + # TODO: support align mamba cache mode for model runner v2 + assert not envs.VLLM_USE_V2_MODEL_RUNNER, ( + "Model Runner V2 has not yet supported mamba_cache_mode='align'. " + ) @model_validator(mode="after") def validate_mamba_block_size(self) -> "VllmConfig": diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 34089a67b3be..fba7fe3a5f54 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -8,13 +8,16 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, KVCacheSpec, + MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup, bind_kv_cache @@ -112,8 +115,9 @@ def _reshape_kv_cache( kv_cache_raw_tensors: dict[str, torch.Tensor], attn_backends: dict[str, AttentionBackend], cache_dtype: str, -) -> dict[str, torch.Tensor]: - kv_caches: dict[str, torch.Tensor] = {} +) -> dict[str, Any]: + kv_caches: dict[str, Any] = {} + has_attn, has_mamba = False, False for kv_cache_group_spec in kv_cache_config.kv_cache_groups: for layer_name in kv_cache_group_spec.layer_names: kv_cache_spec = kv_cache_group_spec.kv_cache_spec @@ -125,35 +129,90 @@ def _reshape_kv_cache( assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes - attn_backend = attn_backends[layer_name] - kv_cache_shape = attn_backend.get_kv_cache_shape( - num_blocks, - kv_cache_spec.block_size, - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype, - ) + if isinstance(kv_cache_spec, AttentionSpec): + has_attn = True + attn_backend = attn_backends[layer_name] + kv_cache_shape = attn_backend.get_kv_cache_shape( + num_blocks, + kv_cache_spec.block_size, + kv_cache_spec.num_kv_heads, + kv_cache_spec.head_size, + ) + + # FIXME(woosuk): Add kv_cache_stride_order to all attn backends + try: + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() + assert len(kv_cache_stride_order) == len(kv_cache_shape) + except (AttributeError, NotImplementedError): + kv_cache_stride_order = tuple(range(len(kv_cache_shape))) + + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) + inv_order = [ + kv_cache_stride_order.index(i) + for i in range(len(kv_cache_stride_order)) + ] + + dtype = kv_cache_spec.dtype + reshaped = raw_tensor.view(dtype) + reshaped = reshaped.view(kv_cache_shape) + kv_caches[layer_name] = reshaped.permute(*inv_order) + + elif isinstance(kv_cache_spec, MambaSpec): + has_mamba = True + state_tensors = [] + storage_offset_bytes = 0 + for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes): + dtype_size = get_dtype_size(dtype) + num_element_per_page = kv_cache_spec.page_size_bytes // dtype_size + target_shape = (num_blocks, *shape) + stride = torch.empty(target_shape).stride() + target_stride = (num_element_per_page, *stride[1:]) + assert storage_offset_bytes % dtype_size == 0 + tensor = torch.as_strided( + raw_tensor.view(dtype), + size=target_shape, + stride=target_stride, + storage_offset=storage_offset_bytes // dtype_size, + ) + state_tensors.append(tensor) + storage_offset_bytes += stride[0] * dtype_size + kv_caches[layer_name] = state_tensors + else: + raise NotImplementedError( + f"Unsupported KV cache spec type: {type(kv_cache_spec)}" + ) + + if has_attn and has_mamba: + _update_hybrid_attention_layout(kv_caches, kv_cache_config) - # FIXME(woosuk): Add kv_cache_stride_order to all attention backends. - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() - assert len(kv_cache_stride_order) == len(kv_cache_shape) - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - inv_order = [ - kv_cache_stride_order.index(i) - for i in range(len(kv_cache_stride_order)) - ] - - dtype = kv_cache_spec.dtype - raw_tensor = raw_tensor.view(dtype) - raw_tensor = raw_tensor.view(kv_cache_shape) - kv_caches[layer_name] = raw_tensor.permute(*inv_order) return kv_caches +def _update_hybrid_attention_layout( + kv_caches: dict[str, Any], + kv_cache_config: KVCacheConfig, +) -> None: + for kv_cache_group_spec in kv_cache_config.kv_cache_groups: + kv_cache_spec = kv_cache_group_spec.kv_cache_spec + if not isinstance(kv_cache_spec, AttentionSpec): + continue + for layer_name in kv_cache_group_spec.layer_names: + kv_cache = kv_caches[layer_name] + if kv_cache.shape[0] == 2: + assert kv_cache.shape[1] != 2, ( + f"Cannot determine layout for tensor of shape {kv_cache.shape}" + ) + hidden_size = kv_cache.shape[2:].numel() + kv_cache.as_strided_( + size=kv_cache.shape, + stride=( + hidden_size, + 2 * hidden_size, + *kv_cache.stride()[2:], + ), + ) + + def init_kv_cache( runner_kv_caches: list[torch.Tensor], forward_context: dict[str, Any], @@ -161,7 +220,7 @@ def init_kv_cache( attn_backends: dict[str, AttentionBackend], device: torch.device, cache_dtype: str, -) -> dict[str, torch.Tensor]: +) -> dict[str, Any]: kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device) kv_caches = _reshape_kv_cache( kv_cache_config, kv_cache_raw_tensors, attn_backends, cache_dtype @@ -195,6 +254,9 @@ def build_attn_metadata( kv_cache_config: KVCacheConfig, dcp_local_seq_lens: torch.Tensor | None = None, encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] | None = None, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + for_cudagraph_capture: bool = False, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -224,11 +286,59 @@ def build_attn_metadata( common_attn_metadata.encoder_seq_lens = encoder_seq_lens_gpu common_attn_metadata.encoder_seq_lens_cpu = encoder_seq_lens_cpu + kv_cache_spec = kv_cache_config.kv_cache_groups[i].kv_cache_spec + is_mamba_group = isinstance(kv_cache_spec, MambaSpec) + for attn_group in attn_groups[i]: attn_metadata_builder = attn_group.get_metadata_builder(0) - metadata = attn_metadata_builder.build( - common_prefix_len=0, common_attn_metadata=common_attn_metadata - ) + if for_cudagraph_capture: + metadata = attn_metadata_builder.build_for_cudagraph_capture( + common_attn_metadata + ) + else: + extra_kwargs: dict[str, Any] = {} + if is_mamba_group: + extra_kwargs["num_accepted_tokens"] = num_accepted_tokens + extra_kwargs["num_decode_draft_tokens_cpu"] = ( + num_decode_draft_tokens_cpu + ) + metadata = attn_metadata_builder.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + **extra_kwargs, + ) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = metadata return attn_metadata + + +def prepare_mamba_hybrid_metadata( + req_states: RequestState, + idx_mapping: torch.Tensor, + num_reqs: int, + num_reqs_padded: int, + req_ids: list[str], + scheduled_spec_decode_tokens: dict[str, list[int]] | None, +) -> tuple[torch.Tensor, torch.Tensor]: + hybrid_accepted = torch.zeros( + num_reqs_padded, + dtype=req_states.num_accepted_tokens_gpu.dtype, + device=req_states.num_accepted_tokens_gpu.device, + ) + hybrid_accepted[:num_reqs] = req_states.num_accepted_tokens_gpu[idx_mapping] + + hybrid_draft_cpu = torch.full( + (num_reqs_padded,), -1, dtype=torch.int32, device="cpu" + ) + if scheduled_spec_decode_tokens: + for batch_idx, req_id in enumerate(req_ids): + draft_ids = scheduled_spec_decode_tokens.get(req_id) + if draft_ids is None: + continue + req_state_idx = req_states.req_id_to_index[req_id] + if ( + req_states.num_computed_prefill_tokens[req_state_idx] + >= req_states.prefill_len.np[req_state_idx] + ): + hybrid_draft_cpu[batch_idx] = len(draft_ids) + return hybrid_accepted, hybrid_draft_cpu diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index e79a7afbd81e..2628a8d30de9 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -21,6 +21,7 @@ def __init__( cp_size: int = 1, cp_rank: int = 0, cp_interleave: int = 1, + max_num_blocks_per_group: list[int] | None = None, ): self.block_sizes = block_sizes self.max_num_reqs = max_num_reqs @@ -36,11 +37,14 @@ def __init__( # num_kv_cache_groups x [max_num_reqs, max_num_blocks] self.block_tables: list[StagedWriteTensor] = [] for i in range(self.num_kv_cache_groups): - block_size = self.block_sizes[i] # When using DCP, each request's KV cache is sharded among different ranks. # As a result, one block on the current rank covers `block_size * cp_size` # tokens in the full, global (unsharded) sequence. - max_num_blocks = cdiv(self.max_model_len, block_size * self.cp_size) + if max_num_blocks_per_group is not None: + max_num_blocks = max_num_blocks_per_group[i] + else: + block_size = self.block_sizes[i] + max_num_blocks = cdiv(self.max_model_len, block_size * self.cp_size) block_table = StagedWriteTensor( (self.max_num_reqs, max_num_blocks), dtype=torch.int32, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d918131c68d4..d178d1c0c138 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -94,6 +94,7 @@ def __init__( self.decode_query_len = decode_query_len self.dp_size = vllm_config.parallel_config.data_parallel_size + self.tp_size = vllm_config.parallel_config.tensor_parallel_size self.is_first_pp_rank = get_pp_group().is_first_rank self.is_last_pp_rank = get_pp_group().is_last_rank @@ -103,6 +104,10 @@ def __init__( self._graphs_captured = False self._candidates: list[list[BatchExecutionDescriptor]] = [] self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} + # adjust the cudagraph sizes to be a multiple of the uniform decode query length + self.compilation_config.adjust_cudagraph_sizes_for_spec_decode( + self.decode_query_len, self.tp_size + ) self._init_candidates() def _init_candidates(self) -> None: @@ -324,6 +329,7 @@ def create_forward_fn( block_tables, attn_groups, kv_cache_config, + skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), ) def forward_fn(cg_mode: CUDAGraphMode) -> None: @@ -403,7 +409,8 @@ def prepare_inputs_to_capture( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, -) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: + skip_attn: bool = False, +) -> tuple[dict[str, Any] | None, dict[str, torch.Tensor]]: input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) input_block_tables = block_tables.get_dummy_block_tables(num_reqs) slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens) @@ -423,13 +430,15 @@ def prepare_inputs_to_capture( ) input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs] - attn_metadata = model_state.prepare_attn( - input_batch, - CUDAGraphMode.NONE, - input_block_tables, - slot_mappings, - attn_groups, - kv_cache_config, - for_capture=True, - ) + attn_metadata = None + if not skip_attn: + attn_metadata = model_state.prepare_attn( + input_batch, + CUDAGraphMode.NONE, + input_block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=True, + ) return attn_metadata, slot_mappings_by_layer diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a2f83c52e951..71122056c90f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -40,10 +40,11 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask +from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput @@ -178,6 +179,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) + # Mamba hybrid models. + self.is_mamba_hybrid = self.model_config.is_hybrid # Pooling models. self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None @@ -190,6 +193,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): num_speculative_steps=self.num_speculative_steps, vocab_size=self.vocab_size, device=self.device, + is_mamba_hybrid=self.is_mamba_hybrid, ) self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, @@ -337,10 +341,17 @@ def get_kv_cache_spec(self): def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: kv_cache_config = deepcopy(kv_cache_config) self.kv_cache_config = kv_cache_config - block_sizes = [ - kv_cache_group.kv_cache_spec.block_size - for kv_cache_group in kv_cache_config.kv_cache_groups - ] + block_sizes = [] + max_num_blocks_per_group = [] + for kv_cache_group in kv_cache_config.kv_cache_groups: + spec = kv_cache_group.kv_cache_spec + block_sizes.append(spec.block_size) + max_num_blocks = cdiv(self.max_model_len, spec.block_size * self.dcp_size) + if isinstance(spec, MambaSpec): + max_num_blocks = ( + max_num_blocks if self.cache_config.enable_prefix_caching else 1 + ) + spec.num_speculative_blocks + max_num_blocks_per_group.append(max_num_blocks) block_table_max_model_len = self.max_model_len if self.is_encoder_decoder: @@ -360,6 +371,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, + max_num_blocks_per_group=max_num_blocks_per_group, ) self.attn_backends, self.attn_groups = init_attn_backend( @@ -904,6 +916,12 @@ def postprocess( computed_prefill, self.req_states.prefill_len.np, out=computed_prefill ) + # Update accepted token counts on GPU for next step's GDN metadata. + if self.is_mamba_hybrid: + self.req_states.num_accepted_tokens_gpu[input_batch.idx_mapping] = ( + num_sampled + ) + @torch.inference_mode() def execute_model( self, @@ -1007,6 +1025,8 @@ def execute_model( slot_mappings, self.attn_groups, self.kv_cache_config, + self.req_states, + scheduler_output.scheduled_spec_decode_tokens, ) inputs_embeds = None diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 8e73867deb2e..26b8f03417a3 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -10,7 +10,10 @@ from vllm.tasks import GenerationTask from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.attn_utils import ( + build_attn_metadata, + prepare_mamba_hybrid_metadata, +) from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner @@ -40,6 +43,7 @@ def __init__( self.max_num_tokens = self.scheduler_config.max_num_batched_tokens self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() self.dtype = self.model_config.dtype + self.is_mamba_hybrid = self.model_config.is_hybrid if self.supports_mm_inputs: assert encoder_cache is not None @@ -161,6 +165,8 @@ def prepare_attn( slot_mappings: torch.Tensor, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, + req_states: RequestState | None = None, + scheduled_spec_decode_tokens: dict[str, list[int]] | None = None, for_capture: bool = False, ) -> dict[str, Any]: if cudagraph_mode == CUDAGraphMode.FULL: @@ -173,6 +179,25 @@ def prepare_attn( num_tokens = input_batch.num_tokens query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) max_query_len = input_batch.num_scheduled_tokens.max().item() + + num_accepted_tokens = None + num_decode_draft_tokens_cpu = None + if ( + self.is_mamba_hybrid + and scheduled_spec_decode_tokens + and req_states is not None + ): + num_accepted_tokens, num_decode_draft_tokens_cpu = ( + prepare_mamba_hybrid_metadata( + req_states, + input_batch.idx_mapping, + input_batch.num_reqs, + num_reqs, + input_batch.req_ids, + scheduled_spec_decode_tokens, + ) + ) + attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, @@ -186,5 +211,8 @@ def prepare_attn( slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + for_cudagraph_capture=for_capture, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index d83ab2fc515f..ae7d4efd6f14 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -66,6 +66,8 @@ def prepare_attn( slot_mappings: torch.Tensor, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, + req_states: RequestState | None = None, + scheduled_spec_decode_tokens: dict[str, list[int]] | None = None, for_capture: bool = False, ) -> dict[str, Any]: raise NotImplementedError diff --git a/vllm/v1/worker/gpu/model_states/whisper.py b/vllm/v1/worker/gpu/model_states/whisper.py index 1268fee88210..98eb1d53d6a8 100644 --- a/vllm/v1/worker/gpu/model_states/whisper.py +++ b/vllm/v1/worker/gpu/model_states/whisper.py @@ -103,6 +103,8 @@ def prepare_attn( slot_mappings: torch.Tensor, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, + req_states: RequestState | None = None, + scheduled_spec_decode_tokens: dict[str, list[int]] | None = None, for_capture: bool = False, ) -> dict[str, Any]: if cudagraph_mode == CUDAGraphMode.FULL: diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py index 1e75c48966b2..665b0f38963b 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py @@ -72,6 +72,7 @@ def create_forward_fn( block_tables, attn_groups, kv_cache_config, + skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), ) return lambda cg_mode: generate_fn( diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 24d225886106..6cbc3a8b5c78 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + import numpy as np import torch @@ -15,6 +17,7 @@ def __init__( num_speculative_steps: int, vocab_size: int, device: torch.device, + is_mamba_hybrid: bool = False, ): self.max_num_reqs = max_num_reqs self.max_model_len = max_model_len @@ -22,6 +25,7 @@ def __init__( self.num_speculative_steps = num_speculative_steps self.vocab_size = vocab_size self.device = device + self.is_mamba_hybrid = is_mamba_hybrid self.req_id_to_index: dict[str, int] = {} self.index_to_req_id: dict[int, str] = {} @@ -75,6 +79,12 @@ def __init__( self.max_num_reqs, dtype=torch.int32, device=device ) + # Mamba hybrid model (attention + mamba/GDN) state. + if self.is_mamba_hybrid: + self.num_accepted_tokens_gpu = torch.ones( + max_num_reqs, dtype=torch.int32, device=device + ) + @property def num_reqs(self) -> int: return len(self.req_id_to_index) From bf7b885da4f11845f3953e59eed6f4c4f4a1df40 Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Mon, 30 Mar 2026 23:20:27 +0800 Subject: [PATCH 02/10] use MambaHybridModelState to refactor code Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/attn_utils.py | 34 ------ vllm/v1/worker/gpu/model_runner.py | 32 +++-- vllm/v1/worker/gpu/model_states/__init__.py | 5 + vllm/v1/worker/gpu/model_states/default.py | 26 +--- vllm/v1/worker/gpu/model_states/interface.py | 7 ++ .../worker/gpu/model_states/mamba_hybrid.py | 112 ++++++++++++++++++ vllm/v1/worker/gpu/model_states/whisper.py | 1 + vllm/v1/worker/gpu/states.py | 8 -- 8 files changed, 140 insertions(+), 85 deletions(-) create mode 100644 vllm/v1/worker/gpu/model_states/mamba_hybrid.py diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index fba7fe3a5f54..21814e39352f 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -17,7 +17,6 @@ MambaSpec, UniformTypeKVCacheSpecs, ) -from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup, bind_kv_cache @@ -123,7 +122,6 @@ def _reshape_kv_cache( kv_cache_spec = kv_cache_group_spec.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, AttentionSpec) raw_tensor = kv_cache_raw_tensors[layer_name] assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 @@ -310,35 +308,3 @@ def build_attn_metadata( for layer_name in attn_group.layer_names: attn_metadata[layer_name] = metadata return attn_metadata - - -def prepare_mamba_hybrid_metadata( - req_states: RequestState, - idx_mapping: torch.Tensor, - num_reqs: int, - num_reqs_padded: int, - req_ids: list[str], - scheduled_spec_decode_tokens: dict[str, list[int]] | None, -) -> tuple[torch.Tensor, torch.Tensor]: - hybrid_accepted = torch.zeros( - num_reqs_padded, - dtype=req_states.num_accepted_tokens_gpu.dtype, - device=req_states.num_accepted_tokens_gpu.device, - ) - hybrid_accepted[:num_reqs] = req_states.num_accepted_tokens_gpu[idx_mapping] - - hybrid_draft_cpu = torch.full( - (num_reqs_padded,), -1, dtype=torch.int32, device="cpu" - ) - if scheduled_spec_decode_tokens: - for batch_idx, req_id in enumerate(req_ids): - draft_ids = scheduled_spec_decode_tokens.get(req_id) - if draft_ids is None: - continue - req_state_idx = req_states.req_id_to_index[req_id] - if ( - req_states.num_computed_prefill_tokens[req_state_idx] - >= req_states.prefill_len.np[req_state_idx] - ): - hybrid_draft_cpu[batch_idx] = len(draft_ids) - return hybrid_accepted, hybrid_draft_cpu diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 71122056c90f..2b1d50a25f5b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -179,8 +179,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) - # Mamba hybrid models. - self.is_mamba_hybrid = self.model_config.is_hybrid # Pooling models. self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None @@ -193,7 +191,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): num_speculative_steps=self.num_speculative_steps, vocab_size=self.vocab_size, device=self.device, - is_mamba_hybrid=self.is_mamba_hybrid, ) self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, @@ -341,27 +338,30 @@ def get_kv_cache_spec(self): def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: kv_cache_config = deepcopy(kv_cache_config) self.kv_cache_config = kv_cache_config + + block_table_max_model_len = self.max_model_len + if self.is_encoder_decoder: + # Cross-attention block tables need to index encoder tokens + # (e.g., Whisper), which can exceed decoder max_model_len. + block_table_max_model_len = max( + block_table_max_model_len, + getattr(self.model_config.hf_config, "max_source_positions", 0), + ) + block_sizes = [] max_num_blocks_per_group = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) - max_num_blocks = cdiv(self.max_model_len, spec.block_size * self.dcp_size) + max_num_blocks = cdiv( + block_table_max_model_len, spec.block_size * self.dcp_size + ) if isinstance(spec, MambaSpec): max_num_blocks = ( max_num_blocks if self.cache_config.enable_prefix_caching else 1 ) + spec.num_speculative_blocks max_num_blocks_per_group.append(max_num_blocks) - block_table_max_model_len = self.max_model_len - if self.is_encoder_decoder: - # Cross-attention block tables need to index encoder tokens - # (e.g., Whisper ~1500), which can exceed decoder max_model_len. - block_table_max_model_len = max( - block_table_max_model_len, - getattr(self.model_config.hf_config, "max_source_positions", 0), - ) - self.block_tables = BlockTables( block_sizes=block_sizes, max_num_reqs=self.max_num_reqs, @@ -916,11 +916,7 @@ def postprocess( computed_prefill, self.req_states.prefill_len.np, out=computed_prefill ) - # Update accepted token counts on GPU for next step's GDN metadata. - if self.is_mamba_hybrid: - self.req_states.num_accepted_tokens_gpu[input_batch.idx_mapping] = ( - num_sampled - ) + self.model_state.postprocess_state(input_batch, num_sampled) @torch.inference_mode() def execute_model( diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index 651452553332..06b5a92c3952 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -18,6 +18,11 @@ def init_model_state( return WhisperModelState(vllm_config, model, encoder_cache, device) + if vllm_config.model_config.is_hybrid: + from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState + + return MambaHybridModelState(vllm_config, model, encoder_cache, device) + from vllm.v1.worker.gpu.model_states.default import DefaultModelState return DefaultModelState(vllm_config, model, encoder_cache, device) diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 26b8f03417a3..60a78dc258de 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -10,10 +10,7 @@ from vllm.tasks import GenerationTask from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.worker.gpu.attn_utils import ( - build_attn_metadata, - prepare_mamba_hybrid_metadata, -) +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner @@ -43,7 +40,6 @@ def __init__( self.max_num_tokens = self.scheduler_config.max_num_batched_tokens self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() self.dtype = self.model_config.dtype - self.is_mamba_hybrid = self.model_config.is_hybrid if self.supports_mm_inputs: assert encoder_cache is not None @@ -180,24 +176,6 @@ def prepare_attn( query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) max_query_len = input_batch.num_scheduled_tokens.max().item() - num_accepted_tokens = None - num_decode_draft_tokens_cpu = None - if ( - self.is_mamba_hybrid - and scheduled_spec_decode_tokens - and req_states is not None - ): - num_accepted_tokens, num_decode_draft_tokens_cpu = ( - prepare_mamba_hybrid_metadata( - req_states, - input_batch.idx_mapping, - input_batch.num_reqs, - num_reqs, - input_batch.req_ids, - scheduled_spec_decode_tokens, - ) - ) - attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, @@ -211,8 +189,6 @@ def prepare_attn( slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, - num_accepted_tokens=num_accepted_tokens, - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, for_cudagraph_capture=for_capture, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index ae7d4efd6f14..57ec4f16e29a 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -38,6 +38,13 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: def apply_staged_writes(self) -> None: return None + def postprocess_state( + self, + input_batch: InputBatch, + num_sampled: torch.Tensor, + ) -> None: + return None + @abstractmethod def get_mm_embeddings( self, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py new file mode 100644 index 000000000000..bbe45772d0ab --- /dev/null +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +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 +from vllm.v1.worker.gpu.states import RequestState +from vllm.v1.worker.utils import AttentionGroup + + +class MambaHybridModelState(DefaultModelState): + """Model state for hybrid attention + Mamba / linear-attention models.""" + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ) -> None: + super().__init__(vllm_config, model, encoder_cache, device) + self.num_accepted_tokens_gpu = torch.ones( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + + def prepare_attn( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + req_states: RequestState | None = None, + scheduled_spec_decode_tokens: dict[str, list[int]] | None = None, + for_capture: bool = False, + ) -> dict[str, Any]: + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_tokens = input_batch.num_tokens + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + max_query_len = input_batch.num_scheduled_tokens.max().item() + + # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens + # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only + # compute them during actual (non-capture) forward execution. + num_decode_draft_tokens_cpu = None + num_accepted_tokens = None + if not for_capture: + assert req_states is not None + assert scheduled_spec_decode_tokens is not None + num_decode_draft_tokens_cpu = torch.full( + (input_batch.num_reqs_after_padding,), + -1, + dtype=torch.int32, + device="cpu", + ) + for batch_idx, req_id in enumerate(input_batch.req_ids): + draft_ids = scheduled_spec_decode_tokens.get(req_id) + if draft_ids is None: + continue + req_state_idx = req_states.req_id_to_index[req_id] + if ( + req_states.num_computed_prefill_tokens[req_state_idx] + >= req_states.prefill_len.np[req_state_idx] + ): + num_decode_draft_tokens_cpu[batch_idx] = len(draft_ids) + + num_accepted_tokens = torch.ones( + num_reqs, + dtype=self.num_accepted_tokens_gpu.dtype, + device=self.num_accepted_tokens_gpu.device, + ) + num_accepted_tokens[: input_batch.num_reqs] = self.num_accepted_tokens_gpu[ + input_batch.idx_mapping + ] + + return build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_tokens=num_tokens, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max_query_len, + seq_lens=input_batch.seq_lens, + max_seq_len=self.max_model_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + for_cudagraph_capture=for_capture, + ) + + def postprocess_state( + self, + input_batch: InputBatch, + num_sampled: torch.Tensor, + ) -> None: + self.num_accepted_tokens_gpu[input_batch.idx_mapping] = num_sampled diff --git a/vllm/v1/worker/gpu/model_states/whisper.py b/vllm/v1/worker/gpu/model_states/whisper.py index 98eb1d53d6a8..4f00c9c8aae4 100644 --- a/vllm/v1/worker/gpu/model_states/whisper.py +++ b/vllm/v1/worker/gpu/model_states/whisper.py @@ -133,6 +133,7 @@ def prepare_attn( kv_cache_config=kv_cache_config, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, encoder_seq_lens=encoder_seq_lens, + for_cudagraph_capture=for_capture, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 6cbc3a8b5c78..79c942df8298 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -17,7 +17,6 @@ def __init__( num_speculative_steps: int, vocab_size: int, device: torch.device, - is_mamba_hybrid: bool = False, ): self.max_num_reqs = max_num_reqs self.max_model_len = max_model_len @@ -25,7 +24,6 @@ def __init__( self.num_speculative_steps = num_speculative_steps self.vocab_size = vocab_size self.device = device - self.is_mamba_hybrid = is_mamba_hybrid self.req_id_to_index: dict[str, int] = {} self.index_to_req_id: dict[int, str] = {} @@ -79,12 +77,6 @@ def __init__( self.max_num_reqs, dtype=torch.int32, device=device ) - # Mamba hybrid model (attention + mamba/GDN) state. - if self.is_mamba_hybrid: - self.num_accepted_tokens_gpu = torch.ones( - max_num_reqs, dtype=torch.int32, device=device - ) - @property def num_reqs(self) -> int: return len(self.req_id_to_index) From 976dab90e133110f19cddf282013cce68929473f Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Wed, 8 Apr 2026 00:20:00 +0800 Subject: [PATCH 03/10] fix get_kv_cache_shape and dispatch attn build args for diff backend Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/attn_utils.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 21814e39352f..08bc21830208 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -10,6 +10,8 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, @@ -135,6 +137,7 @@ def _reshape_kv_cache( kv_cache_spec.block_size, kv_cache_spec.num_kv_heads, kv_cache_spec.head_size, + cache_dtype_str=cache_dtype, ) # FIXME(woosuk): Add kv_cache_stride_order to all attn backends @@ -284,26 +287,26 @@ def build_attn_metadata( common_attn_metadata.encoder_seq_lens = encoder_seq_lens_gpu common_attn_metadata.encoder_seq_lens_cpu = encoder_seq_lens_cpu - kv_cache_spec = kv_cache_config.kv_cache_groups[i].kv_cache_spec - is_mamba_group = isinstance(kv_cache_spec, MambaSpec) - for attn_group in attn_groups[i]: attn_metadata_builder = attn_group.get_metadata_builder(0) if for_cudagraph_capture: metadata = attn_metadata_builder.build_for_cudagraph_capture( common_attn_metadata ) + elif isinstance( + attn_metadata_builder, + (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder), + ): + metadata = attn_metadata_builder.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) else: - extra_kwargs: dict[str, Any] = {} - if is_mamba_group: - extra_kwargs["num_accepted_tokens"] = num_accepted_tokens - extra_kwargs["num_decode_draft_tokens_cpu"] = ( - num_decode_draft_tokens_cpu - ) metadata = attn_metadata_builder.build( common_prefix_len=0, common_attn_metadata=common_attn_metadata, - **extra_kwargs, ) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = metadata From 0ee4ea8ec2db9f8471845ba246568c6c23794c94 Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Wed, 8 Apr 2026 02:33:03 +0800 Subject: [PATCH 04/10] fix mrv2 qwen35 mtp multimodal registry Signed-off-by: zhuhaoran --- vllm/multimodal/registry.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/vllm/multimodal/registry.py b/vllm/multimodal/registry.py index fa414a5928d6..dcad685d6f06 100644 --- a/vllm/multimodal/registry.py +++ b/vllm/multimodal/registry.py @@ -111,7 +111,15 @@ def supports_multimodal_inputs(self, model_config: "ModelConfig") -> bool: return False mm_config = model_config.get_multimodal_config() - info = self._create_processing_info(model_config, tokenizer=None) + try: + info = self._create_processing_info(model_config, tokenizer=None) + except ValueError: + logger.warning_once( + "Model %s is treated as multimodal but has no registered " + "multimodal processor; running in text-only mode.", + model_config.model, + ) + return False # Check if all supported modalities have limit == 0 if all( @@ -170,7 +178,11 @@ def _get_model_cls(self, model_config: "ModelConfig") -> "SupportsMultiModal": from vllm.model_executor.model_loader import get_model_architecture model_cls, _ = get_model_architecture(model_config) - assert hasattr(model_cls, "_processor_factory") + if not hasattr(model_cls, "_processor_factory"): + raise ValueError( + f"Model class {model_cls.__name__} has no registered " + "multimodal processor" + ) return cast("SupportsMultiModal", model_cls) def _create_processing_ctx( From e852690706bb965bfd9da0002acd597850d736a4 Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Wed, 8 Apr 2026 10:53:53 +0800 Subject: [PATCH 05/10] remove unused future annotations import Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/states.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 79c942df8298..24d225886106 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from __future__ import annotations - import numpy as np import torch From 9fd6e04d8201737db4c51ca7bdf0170c0c773aea Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Wed, 8 Apr 2026 23:03:36 +0800 Subject: [PATCH 06/10] apply suggestions from @MengqingCao Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/attn_utils.py | 14 +++++++------- vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 08bc21830208..17c1a3a7c89e 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -125,9 +125,9 @@ def _reshape_kv_cache( if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name] - raw_tensor = kv_cache_raw_tensors[layer_name] - assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 - num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes + kv_raw_tensor = kv_cache_raw_tensors[layer_name] + assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 + num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes if isinstance(kv_cache_spec, AttentionSpec): has_attn = True @@ -154,9 +154,9 @@ def _reshape_kv_cache( ] dtype = kv_cache_spec.dtype - reshaped = raw_tensor.view(dtype) - reshaped = reshaped.view(kv_cache_shape) - kv_caches[layer_name] = reshaped.permute(*inv_order) + kv_tensor_reshaped = kv_raw_tensor.view(dtype) + kv_tensor_reshaped = kv_tensor_reshaped.view(kv_cache_shape) + kv_caches[layer_name] = kv_tensor_reshaped.permute(*inv_order) elif isinstance(kv_cache_spec, MambaSpec): has_mamba = True @@ -170,7 +170,7 @@ def _reshape_kv_cache( target_stride = (num_element_per_page, *stride[1:]) assert storage_offset_bytes % dtype_size == 0 tensor = torch.as_strided( - raw_tensor.view(dtype), + kv_raw_tensor.view(dtype), size=target_shape, stride=target_stride, storage_offset=storage_offset_bytes // dtype_size, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index bbe45772d0ab..8de60c7ec459 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -61,7 +61,7 @@ def prepare_attn( assert req_states is not None assert scheduled_spec_decode_tokens is not None num_decode_draft_tokens_cpu = torch.full( - (input_batch.num_reqs_after_padding,), + (num_reqs,), -1, dtype=torch.int32, device="cpu", From e1d3df566be6689b26ecabdecad7b48b7b82ceb6 Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Thu, 9 Apr 2026 00:51:40 +0800 Subject: [PATCH 07/10] fix num_accepted_tokens in chunked prefill and fix UniformTypeKVCacheSpecs in _update_hybrid_attention_layout Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/attn_utils.py | 8 +++++--- vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 6 +++++- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 17c1a3a7c89e..ed6cf13ab6a4 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -194,10 +194,12 @@ def _update_hybrid_attention_layout( kv_cache_config: KVCacheConfig, ) -> None: for kv_cache_group_spec in kv_cache_config.kv_cache_groups: - kv_cache_spec = kv_cache_group_spec.kv_cache_spec - if not isinstance(kv_cache_spec, AttentionSpec): - continue for layer_name in kv_cache_group_spec.layer_names: + kv_cache_spec = kv_cache_group_spec.kv_cache_spec + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name] + if not isinstance(kv_cache_spec, AttentionSpec): + continue kv_cache = kv_caches[layer_name] if kv_cache.shape[0] == 2: assert kv_cache.shape[1] != 2, ( diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 8de60c7ec459..3ef2426d6be4 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -109,4 +109,8 @@ def postprocess_state( input_batch: InputBatch, num_sampled: torch.Tensor, ) -> None: - self.num_accepted_tokens_gpu[input_batch.idx_mapping] = num_sampled + # Chunked prefill does not sample a token, so num_sampled can be 0. + # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. + self.num_accepted_tokens_gpu[input_batch.idx_mapping] = torch.clamp( + num_sampled, min=1 + ) From 778aa572ba720557403c513d797851b1ffd5529d Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Thu, 9 Apr 2026 11:26:32 +0800 Subject: [PATCH 08/10] feat: add is_prefill for MambaHybridModelState Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/attn_utils.py | 4 ++++ vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index ed6cf13ab6a4..7bb508dfbfee 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -257,6 +257,7 @@ def build_attn_metadata( kv_cache_config: KVCacheConfig, dcp_local_seq_lens: torch.Tensor | None = None, encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] | None = None, + is_prefilling: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, num_decode_draft_tokens_cpu: torch.Tensor | None = None, for_cudagraph_capture: bool = False, @@ -264,6 +265,8 @@ def build_attn_metadata( seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs] + if is_prefilling is not None: + is_prefilling = is_prefilling[:num_reqs] attn_metadata: dict[str, Any] = {} num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) @@ -283,6 +286,7 @@ def build_attn_metadata( slot_mapping=slot_mapping, causal=True, dcp_local_seq_lens=dcp_local_seq_lens, + is_prefilling=is_prefilling, ) if encoder_seq_lens and i in encoder_seq_lens: encoder_seq_lens_gpu, encoder_seq_lens_cpu = encoder_seq_lens[i] diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 3ef2426d6be4..fde2d2cfa7f6 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -55,11 +55,16 @@ def prepare_attn( # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only # compute them during actual (non-capture) forward execution. + is_prefilling = torch.zeros(num_reqs, dtype=torch.bool) num_decode_draft_tokens_cpu = None num_accepted_tokens = None if not for_capture: assert req_states is not None assert scheduled_spec_decode_tokens is not None + is_prefilling[: input_batch.num_reqs] = torch.from_numpy( + req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np] + < req_states.prefill_len.np[input_batch.idx_mapping_np] + ) num_decode_draft_tokens_cpu = torch.full( (num_reqs,), -1, @@ -99,6 +104,7 @@ def prepare_attn( slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + is_prefilling=is_prefilling, num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, for_cudagraph_capture=for_capture, From 2713f7d1baee855b655ef44afea7ce4961a2bed5 Mon Sep 17 00:00:00 2001 From: zhuhaoran Date: Mon, 27 Apr 2026 01:31:44 +0800 Subject: [PATCH 09/10] refactor: assert piecewise captures have no attention metadata Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/cudagraph_utils.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d178d1c0c138..9ef315928955 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -333,13 +333,12 @@ def create_forward_fn( ) def forward_fn(cg_mode: CUDAGraphMode) -> None: - batch_descriptor = ( - BatchDescriptor(num_tokens=num_tokens) - if cg_mode == CUDAGraphMode.PIECEWISE - else None - ) + batch_descriptor = None + if cg_mode == CUDAGraphMode.PIECEWISE: + assert attn_metadata is None + batch_descriptor = BatchDescriptor(num_tokens=num_tokens) with set_forward_context( - attn_metadata if cg_mode != CUDAGraphMode.PIECEWISE else None, + attn_metadata, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=cg_mode, From 895b7a81dae2319e8108f33eff58fb42de3d832b Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin Date: Mon, 27 Apr 2026 22:35:49 +0000 Subject: [PATCH 10/10] [Model Runner V2] enable spec decode + align mamba cache mode Signed-off-by: Giancarlo Delfin --- vllm/config/vllm.py | 7 +- .../layers/mamba/mamba_mixer.py | 11 ++- .../layers/mamba/mamba_mixer2.py | 17 ++-- vllm/v1/attention/backends/mamba_attn.py | 71 +++++++++++++++- .../worker/gpu/model_states/mamba_hybrid.py | 84 ++++++++++++++++++- 5 files changed, 170 insertions(+), 20 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 0f2369b39fee..9db73d38951c 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1873,10 +1873,9 @@ def validate_block_size(self) -> None: "to schedule a multiple of block_size tokens even if they are " "in the middle of a mm input" ) - # TODO: support align mamba cache mode for model runner v2 - assert not envs.VLLM_USE_V2_MODEL_RUNNER, ( - "Model Runner V2 has not yet supported mamba_cache_mode='align'. " - ) + # MRV2 supports align mode via dual-indexing: the SSM kernel + # reads from the committed block and writes per-token to staging + # blocks; postprocess copies the accepted state back. @model_validator(mode="after") def validate_mamba_block_size(self) -> "VllmConfig": diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 0e476755201e..5022dda2f18b 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -404,8 +404,17 @@ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): 1, block_idx_last_scheduled_token_d.unsqueeze(1) ).squeeze(1) else: - state_indices_tensor_d_input = state_indices_tensor_d + if attn_metadata.src_ssm_indices_tensor_d is None: + # Read and write in-place to the same blocks. + state_indices_tensor_d_input = state_indices_tensor_d + else: + # Read from separate set of blocks. Used for MRV2's "align" + # mamba cache mode. + state_indices_tensor_d_input = ( + attn_metadata.src_ssm_indices_tensor_d + ) state_indices_tensor_d_output = state_indices_tensor_d + # 2. Convolution sequence transformation conv_out_d = causal_conv1d_update( hidden_states_BC_d.transpose(0, 1), diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 2b4b1934f9b3..6c31c4c889e2 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -844,15 +844,16 @@ def conv_ssm_forward( state_indices_tensor_d_output = state_indices_tensor_d.gather( 1, block_idx_last_scheduled_token_d.unsqueeze(1) ).squeeze(1) - # for decode: - # block_idx_first_scheduled_token_d == - # block_idx_last_scheduled_token_d - # at block boundaries: - # block_idx_first_scheduled_token_d > - # block_idx_last_computed_token_d else: - # Without caching, read and write in-place to the same blocks: - state_indices_tensor_d_input = state_indices_tensor_d + if attn_metadata.src_ssm_indices_tensor_d is None: + # Read and write in-place to the same blocks. + state_indices_tensor_d_input = state_indices_tensor_d + else: + # Read from separate set of blocks. Used for MRV2's "align" + # mamba cache mode. + state_indices_tensor_d_input = ( + attn_metadata.src_ssm_indices_tensor_d + ) state_indices_tensor_d_output = state_indices_tensor_d # 2. Convolution sequence transformation diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index eec53032288d..d85ecc65c111 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -7,6 +7,7 @@ import torch +import vllm.envs as envs from vllm.config import VllmConfig from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import ( @@ -73,6 +74,11 @@ class BaseMambaAttentionMetadata: batch_ptr: torch.Tensor | None = None token_chunk_offset_ptr: torch.Tensor | None = None + # When set, the SSM kernel reads from src_ssm_indices_tensor_d and writes + # to state_indices_tensor_d. This overrides the default behavior of reading + # and writing in-place to state_indices_tensor_d. + src_ssm_indices_tensor_d: torch.Tensor | None = None + class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): metadata_cls: type[M] @@ -106,13 +112,18 @@ def __init__( self.compilation_config.max_cudagraph_capture_size, ) - if self.vllm_config.cache_config.mamba_cache_mode == "all": + mamba_cache_mode = self.vllm_config.cache_config.mamba_cache_mode + self.is_mrv2_spec_decode_align_mode: bool = ( + envs.VLLM_USE_V2_MODEL_RUNNER + and self.use_spec_decode + and mamba_cache_mode == "align" + ) + + if mamba_cache_mode == "all": max_num_blocks = cdiv( self.vllm_config.model_config.max_model_len, self.kv_cache_spec.block_size, ) - # Speculative decoding not supported with prefix caching, - # so keep shape consistent with prefill buffer # TODO: reduce this size as needed for decode-only cudagraph capture self.state_indices_tensor_d: torch.Tensor = torch.empty( ( @@ -148,6 +159,20 @@ def __init__( device=device, ) + self.src_ssm_indices_tensor_d: torch.Tensor | None = None + if self.is_mrv2_spec_decode_align_mode: + # Dual-indexing is used by MRV2 for spec decode + prefix caching. + # The SSM reads from a "committed everywhere" tensor and writes + # per-token to staging blocks; the conv kernel uses the narrow + # window in-place with num_accepted_tokens for offset. + # "all" mode auto-selects to "align" when spec decode is active, + # so dual-indexing is only needed for "align" mode. + self.src_ssm_indices_tensor_d = torch.empty( + (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens), + dtype=torch.int32, + device=device, + ) + self._init_reorder_batch_threshold(1, self.use_spec_decode) if self.use_spec_decode: self.supports_update_block_table = False @@ -416,6 +441,35 @@ def _compute_common_metadata( ] state_indices_tensor_p = state_indices_tensor_p[:, 0] + src_ssm_indices_tensor_d = None + # Construct separate read tensor for the SSM kernel for + # MRV2 + spec decode + align mode. + if ( + self.is_mrv2_spec_decode_align_mode + and num_decodes > 0 + and num_accepted_tokens is not None + ): + if num_computed_tokens is None: + num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() + # Get block containing last committed token. + committed_block_idx = torch.clamp( + (num_computed_tokens[:num_decodes] - 1) + // self.kv_cache_spec.block_size, + min=0, + ).to(torch.int64) + committed_phys = ( + common_attn_metadata.block_table_tensor[:num_decodes] + .gather(1, committed_block_idx.unsqueeze(1)) + .squeeze(1) + ) + # Every column for the SSM read is the committed block + # physical ID. + src_ssm_indices_tensor_d = ( + committed_phys.unsqueeze(1) + .expand_as(state_indices_tensor_d) + .contiguous() + ) + # Sometimes even with specdec enabled we get single-token prefill chunks that # should be treated as decodes but don't have num_accepted_tokens set. # These should be fine to process as non-spec decodes since there's only @@ -466,6 +520,7 @@ def _compute_common_metadata( has_initial_states_p=has_initial_states_p, state_indices_tensor_p=state_indices_tensor_p, state_indices_tensor_d=state_indices_tensor_d, + src_ssm_indices_tensor_d=src_ssm_indices_tensor_d, num_accepted_tokens=num_accepted_tokens, query_start_loc_d=query_start_loc_d, block_idx_last_scheduled_token=block_idx_last_scheduled_token, @@ -494,6 +549,7 @@ def _update_metadata_for_cudagraph_capture( num_accepted_tokens = metadata.num_accepted_tokens block_idx_last_scheduled_token = metadata.block_idx_last_scheduled_token block_idx_last_computed_token = metadata.block_idx_last_computed_token + src_ssm_indices_tensor_d = metadata.src_ssm_indices_tensor_d if ( metadata.num_prefills == 0 and metadata.num_decodes <= self.decode_cudagraph_max_bs @@ -536,6 +592,14 @@ def _update_metadata_for_cudagraph_capture( : metadata.num_decode_tokens ] + if src_ssm_indices_tensor_d is not None: + assert self.src_ssm_indices_tensor_d is not None + self.src_ssm_indices_tensor_d[: metadata.num_decodes].copy_( + src_ssm_indices_tensor_d, non_blocking=True + ) + src_ssm_indices_tensor_d = self.src_ssm_indices_tensor_d[:padded_bs] + src_ssm_indices_tensor_d[metadata.num_decodes :] = NULL_BLOCK_ID + return replace( metadata, state_indices_tensor_d=state_indices_tensor_d, @@ -543,6 +607,7 @@ def _update_metadata_for_cudagraph_capture( num_accepted_tokens=num_accepted_tokens, block_idx_last_scheduled_token=block_idx_last_scheduled_token, block_idx_last_computed_token=block_idx_last_computed_token, + src_ssm_indices_tensor_d=src_ssm_indices_tensor_d, ) def update_block_table( diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index fde2d2cfa7f6..00dc991e900c 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -7,7 +7,8 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -31,6 +32,15 @@ def __init__( self.max_num_reqs, dtype=torch.int32, device=self.device ) + self.is_spec_decode_align_mode = ( + vllm_config.num_speculative_tokens > 0 + and vllm_config.cache_config.mamba_cache_mode == "align" + ) + + # Used for align mode + spec decoding. + self.last_block_tables: tuple[torch.Tensor, ...] | None = None + self.last_kv_cache_config: KVCacheConfig | None = None + def prepare_attn( self, input_batch: InputBatch, @@ -91,6 +101,11 @@ def prepare_attn( input_batch.idx_mapping ] + if self.is_spec_decode_align_mode: + # Save state needed during postprocess_state. + self.last_block_tables = block_tables + self.last_kv_cache_config = kv_cache_config + return build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, @@ -115,8 +130,69 @@ def postprocess_state( input_batch: InputBatch, num_sampled: torch.Tensor, ) -> None: + num_accepted_tokens = torch.clamp(num_sampled, min=1) # Chunked prefill does not sample a token, so num_sampled can be 0. # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. - self.num_accepted_tokens_gpu[input_batch.idx_mapping] = torch.clamp( - num_sampled, min=1 - ) + self.num_accepted_tokens_gpu[input_batch.idx_mapping] = num_accepted_tokens + # The last accepted SSM state must be copied from the staging + # block to the running block to ensure that the next step's + # committed block read is correct. + self._copy_ssm_staging_to_committed(input_batch, num_accepted_tokens) + + def _copy_ssm_staging_to_committed( + self, + input_batch: InputBatch, + num_accepted_tokens: torch.Tensor, + ) -> None: + """Copy SSM state from the staging block that holds the last-accepted + token's state back to the running block (column 0 of the state indices + tensor). + """ + if not self.is_spec_decode_align_mode: + return + + assert self.last_kv_cache_config is not None + assert self.last_block_tables is not None + + needs_copy_mask = num_accepted_tokens > 1 + if not needs_copy_mask.any(): + # No draft tokens were accepted, and thus no draft staging + # block SSM states need to be copied over. + return + + fwd_ctx = self.vllm_config.compilation_config.static_forward_context + # Compute the narrow window for the Mamba block tables to get the + # staging physical block IDs. We iterate over Mamba kv-cache + # groups; they share the same MambaSpec. + for idx, group in enumerate(self.last_kv_cache_config.kv_cache_groups): + if not isinstance(group.kv_cache_spec, MambaSpec): + # Skip non-Mamba groups. + continue + + block_table = self.last_block_tables[idx] + cache_mode = self.vllm_config.cache_config.mamba_cache_mode + state_indices_tensor = mamba_get_block_table_tensor( + block_table, + input_batch.seq_lens, + group.kv_cache_spec, + cache_mode, + ) + num_spec_tokens = self.vllm_config.num_speculative_tokens + state_indices_tensor = state_indices_tensor[:, : 1 + num_spec_tokens] + + # Source is the staging block for the last accepted token. + src_block = ( + (num_accepted_tokens - 1) + .clamp(max=state_indices_tensor.size(1) - 1) + .to(torch.int64) + ) + src_phys = state_indices_tensor.gather(1, src_block.unsqueeze(1)).squeeze(1) + # Destination is the running block. + dst_phys = state_indices_tensor[:, 0] + # Copy for every layer in the group. + for layer_name in group.layer_names: + layer = fwd_ctx[layer_name] + ssm_state = layer.kv_cache[1] + src_idx = src_phys[needs_copy_mask].long() + dst_idx = dst_phys[needs_copy_mask].long() + ssm_state[dst_idx] = ssm_state[src_idx]