[WIP][Model Runner V2] Enable mamba align for spec decode

vllm/config/kernel.py

@@ -137,7 +137,7 @@ class KernelConfig:
     """Backend for MoE expert computation kernels. Available options:
 
     - "auto": Automatically select the best backend based on model and hardware
-    - "triton": Use Triton-based fused MoE kernels 
+    - "triton": Use Triton-based fused MoE kernels
     - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
     - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels
     - "cutlass": Use vLLM CUTLASS kernels

vllm/config/vllm.py

@@ -1873,6 +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"
             )
+            # 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":

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),

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

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(

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,13 +592,22 @@ 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,
             query_start_loc_d=query_start_loc_d,
             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(