From 6fe67cbbf3e43da89bebf6ab0eeaca4ba6c75663 Mon Sep 17 00:00:00 2001 From: Chris Eastwood <106503529+eastwood-c@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:55:21 -0400 Subject: [PATCH] [Spec][V2] Support MTP speculative decoding under pipeline parallelism (#46994) Signed-off-by: Chris Eastwood Signed-off-by: Nick Hill Co-authored-by: Nick Hill Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../v1/attention/test_sparse_mla_backends.py | 45 +++++++++++++++++++ .../e2e/spec_decode/test_mtp_parallel_load.py | 9 +--- .../layers/attention/sparse_mla_attention.py | 44 ++++++++++++++---- vllm/model_executor/models/deepseek_mtp.py | 18 +++++++- vllm/model_executor/models/qwen3_5_mtp.py | 12 +++-- vllm/models/deepseek_v32/amd/mtp.py | 18 +++++++- vllm/models/deepseek_v32/nvidia/model.py | 5 ++- vllm/models/deepseek_v32/nvidia/mtp.py | 18 +++++++- .../mla/flashinfer_mla_sparse_sm120.py | 13 +++--- .../backends/mla/rocm_aiter_mla_sparse.py | 14 +++--- .../attention/backends/mla/xpu_mla_sparse.py | 12 +++-- vllm/v1/worker/gpu/model_runner.py | 8 +++- 12 files changed, 167 insertions(+), 49 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 002aa5a4b7..60f9fc83f9 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -24,6 +24,7 @@ from vllm.config import set_current_vllm_config from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, + SparseMLACommonImpl, _masked_mha_workspace_fits, _topk_mask_shape, ) @@ -44,6 +45,9 @@ import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as flashinfer_sparse from vllm.model_executor.layers.attention.mla_attention import ( _canonicalize_sparse_mla_kv_cache_dtype, ) +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, @@ -1864,3 +1868,44 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): assert out.is_contiguous() assert not out.isnan().any() assert not lse.isnan().any() + + +def test_sparse_impl_observes_repointed_indexer_buffer(): + """The MTP proposer repoints the draft's indexer at the target model's buffer + after the backend impl is built, so the impl must resolve the buffer per read. + Snapshotting it in __init__ leaves the layer reading indices nothing writes.""" + impl = object.__new__(FlashInferMLASparseImpl) + own = torch.zeros(4, 8, dtype=torch.int32) + target = torch.ones(4, 8, dtype=torch.int32) + indexer = SimpleNamespace(topk_indices_buffer=own) + impl.init_topk_indices_buffer(indexer, None) + + assert impl.topk_indices_buffer is own + + indexer.topk_indices_buffer = target + + assert impl.topk_indices_buffer is target + + +def test_explicit_topk_buffer_supersedes_indexer(): + """Backbone skip-topk layers have no indexer, and the proposer also assigns the + shared buffer directly onto draft submodules.""" + impl = object.__new__(FlashInferMLASparseImpl) + indexer = SimpleNamespace( + topk_indices_buffer=torch.zeros(2, 2, dtype=torch.int32), + ) + impl.init_topk_indices_buffer(indexer, None) + explicit = torch.ones(2, 2, dtype=torch.int32) + + impl.topk_indices_buffer = explicit + + assert impl.topk_indices_buffer is explicit + + +def test_sparse_mla_common_impl_resolves_buffer_lazily(): + """Guards the whole SparseMLACommonImpl family at once: a backend that + snapshots the buffer instead silently loses MTP buffer sharing.""" + assert issubclass(SparseMLACommonImpl, SharedTopkIndicesBuffer) + assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), ( + "topk_indices_buffer must stay a lazily-resolved property" + ) diff --git a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py index 2369972616..e7b2c36c4f 100644 --- a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py +++ b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py @@ -56,14 +56,7 @@ INLINE_CONFIGS = [ enable_expert_parallel=True, enable_eplb=True, ), - InlineConfig( - id="pp2", - pp_size=2, - skip_reason=( - "DeepSeek MTP pipeline-parallel support is in flight upstream; " - "see https://github.com/vllm-project/vllm/pull/44698" - ), - ), + InlineConfig(id="pp2", pp_size=2), ] diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 8c2bb72bf4..4de2c02d4f 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -501,7 +501,40 @@ def _build_topk_mask( return out[:batch_size, :max_q_len] -class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]): +class SharedTopkIndicesBuffer: + """Resolves the shared top-k index buffer for sparse MLA implementations. + + The indexer owns the buffer, but `LLMBaseProposer.load_model` repoints the + draft's indexers at the target's buffer after the impls are constructed, so + it must be resolved per read rather than snapshotted. Backbone skip-topk + layers have no indexer and pass the buffer explicitly. + """ + + _indexer: object | None = None + _topk_indices_buffer: torch.Tensor | None = None + + def init_topk_indices_buffer( + self, + indexer: object | None, + topk_indices_buffer: torch.Tensor | None, + ) -> None: + self._indexer = indexer + self._topk_indices_buffer = topk_indices_buffer + + @property + def topk_indices_buffer(self) -> torch.Tensor | None: + if self._indexer is not None: + return self._indexer.topk_indices_buffer # type: ignore[attr-defined] + return self._topk_indices_buffer + + @topk_indices_buffer.setter + def topk_indices_buffer(self, buffer: torch.Tensor | None) -> None: + # An explicit assignment supersedes the indexer. + self._indexer = None + self._topk_indices_buffer = buffer + + +class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic[T]): """Sparse MLA base with dense and masked-MHA prefill paths.""" is_sparse = True @@ -543,14 +576,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]): kv_b_proj, ) - # The indexer carries the shared buffer for normal layers and tests; - # the explicitly-passed buffer covers backbone skip layers, whose - # indexer is not constructed (see deepseek_v2.py). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer # type: ignore[attr-defined] - if indexer is not None - else topk_indices_buffer - ) + self.init_topk_indices_buffer(indexer, topk_indices_buffer) self._use_flashinfer_concat_mla_k = ( has_flashinfer() and which("ninja") is not None diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 8ac43e0bb2..09f4a847d0 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -39,9 +39,11 @@ from .deepseek_v2 import ( DeepseekV2MoE, _try_load_fp8_indexer_wk, ) +from .interfaces import SupportsPP from .utils import ( get_pp_missing_layer_names, get_spec_layer_idx_from_weight_name, + make_empty_intermediate_tensors_factory, maybe_prefix, ) @@ -230,7 +232,7 @@ class DeepSeekMultiTokenPredictor(nn.Module): @support_torch_compile -class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): +class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config @@ -240,6 +242,9 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ) # Set MoE hyperparameters self.set_moe_parameters() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], self.config.hidden_size + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -266,7 +271,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, @@ -328,6 +333,15 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: + # A tied top-level embed_tokens has no spec layer to rewrite + # from; the draft needs its own copy under PP. + param = params_dict.get(name) if "embed_tokens" in name else None + if param is not None: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) continue is_fusion_moe_shared_experts_layer = ( self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 4c064bead7..7fb699b183 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -38,6 +38,7 @@ from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig from .interfaces import ( MultiModalEmbeddings, SupportsMultiModal, + SupportsPP, _require_is_multimodal, ) from .utils import ( @@ -152,7 +153,9 @@ class Qwen3_5MultiTokenPredictor(nn.Module): inputs_embeds: torch.Tensor | None = None, spec_step_idx: int = 0, ) -> torch.Tensor: - if get_pp_group().is_first_rank: + # Branch on the inputs, not the rank: the drafter is built entirely on + # the last PP stage, where `is_first_rank` is False. + if intermediate_tensors is None: if inputs_embeds is None: inputs_embeds = self.embed_input_ids(input_ids) assert hidden_states.shape[-1] == inputs_embeds.shape[-1] @@ -162,7 +165,6 @@ class Qwen3_5MultiTokenPredictor(nn.Module): hidden_states = self.fc(hidden_states) residual = None else: - assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] @@ -212,7 +214,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module): "hidden_states": 0, } ) -class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): +class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP): packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -254,6 +256,10 @@ class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + def embed_input_ids( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index 9ea86ee26b..46c35dec98 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -32,8 +32,10 @@ from vllm.model_executor.models.deepseek_v2 import ( _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) +from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( get_pp_missing_layer_names, + make_empty_intermediate_tensors_factory, maybe_prefix, ) from vllm.models.deepseek_v32.common.kernels import fused_eh_norm @@ -160,7 +162,7 @@ class DeepseekV32MultiTokenPredictor(nn.Module): return self.logits_processor(mtp_layer.shared_head.head, hidden_states) -class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config @@ -169,6 +171,9 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.set_moe_parameters() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], self.config.hidden_size + ) self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( self.model.layers.values(), DeepseekV2MoE, @@ -192,7 +197,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, @@ -268,6 +273,15 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: + # A tied top-level embed_tokens has no spec layer to rewrite + # from; the draft needs its own copy under PP. + param = params_dict.get(name) if "embed_tokens" in name else None + if param is not None: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) continue is_fusion_moe_shared_experts_layer = ( self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 76aa6ca8e1..92b74e536f 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -182,7 +182,8 @@ class DeepseekV32Model(torch.nn.Module): # DSA is always sparse (has index_topk); allocate the shared top-k # buffer the indexer writes and the sparse MLA backend reads. self.is_v32 = True - topk_indices_buffer = torch.empty( + # On the model, not a local: the MTP proposer shares it with the draft. + self.topk_indices_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, config.index_topk, dtype=torch.int32, @@ -207,7 +208,7 @@ class DeepseekV32Model(torch.nn.Module): lambda prefix: DeepseekV32DecoderLayer( vllm_config=vllm_config, prefix=prefix, - topk_indices_buffer=topk_indices_buffer, + topk_indices_buffer=self.topk_indices_buffer, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 2ea01dbdcb..a5348ed220 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -33,8 +33,10 @@ from vllm.model_executor.models.deepseek_v2 import ( _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) +from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( get_pp_missing_layer_names, + make_empty_intermediate_tensors_factory, maybe_prefix, ) from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm @@ -260,7 +262,7 @@ class DeepseekV32MultiTokenPredictor(nn.Module): ) -class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config @@ -271,6 +273,9 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): if self.config.model_type == "glm_moe_dsa": enable_glm52_low_latency_gemm(self, vllm_config.model_config.dtype) self.set_moe_parameters() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], self.config.hidden_size + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -289,7 +294,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, @@ -368,6 +373,15 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: + # A tied top-level embed_tokens has no spec layer to rewrite + # from; the draft needs its own copy under PP. + param = params_dict.get(name) if "embed_tokens" in name else None + if param is not None: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) continue name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 4de2c2f53f..b7459203e0 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -6,6 +6,9 @@ from typing import TYPE_CHECKING, cast import torch +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.v1.attention.backend import ( AttentionLayer, AttentionType, @@ -29,7 +32,9 @@ def _kv_scale_format_for_model(model_type: str | None) -> str: return "pow2_fp32" -class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]): +class FlashInferMLASparseSM120Impl( + MLAAttentionImpl[FlashInferMLASparseMetadata], SharedTopkIndicesBuffer +): """SM120 FlashInfer sparse-MLA implementation.""" is_sparse = True @@ -86,11 +91,7 @@ class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata] # Skip-topk layers are built with indexer=None and get the shared # buffer via mla_args instead (cf. FLASHMLA_SPARSE). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer - if indexer is not None - else mla_args.get("topk_indices_buffer") - ) + self.init_topk_indices_buffer(indexer, mla_args.get("topk_indices_buffer")) from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 if not has_flashinfer_sparse_mla_sm120(): diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 5fba8f4237..b110f95d57 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -15,6 +15,9 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( get_mla_dims, ) +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -750,7 +753,9 @@ def reference_mla_sparse_prefill( return (result, lse) -class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]): +class ROCMAiterMLASparseImpl( + MLAAttentionImpl[ROCMAiterMLASparseMetadata], SharedTopkIndicesBuffer +): is_sparse = True supports_dense_mha_prefill = False supports_dcp = False @@ -795,12 +800,7 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]): self.kv_cache_dtype = kv_cache_dtype self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.softmax_scale = scale - # The indexer carries the shared buffer for normal layers and tests; - # the explicitly-passed buffer covers backbone skip layers, whose - # indexer is not constructed (see deepseek_v2.py). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer - ) + self.init_topk_indices_buffer(indexer, topk_indices_buffer) vllm_config = get_current_vllm_config() max_tokens = vllm_config.scheduler_config.max_num_batched_tokens diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index 8aaa4749d5..158fbfc415 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -13,6 +13,9 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( get_mla_dims, ) +from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SharedTopkIndicesBuffer, +) from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, @@ -183,7 +186,7 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata] return metadata -class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]): +class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata], SharedTopkIndicesBuffer): is_sparse = True def __init__( @@ -210,12 +213,7 @@ class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]): self.kv_cache_dtype = kv_cache_dtype self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.softmax_scale = scale - # The indexer carries the shared buffer for normal layers and tests; - # the explicitly-passed buffer covers backbone skip layers, whose - # indexer is not constructed (see deepseek_v2.py). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer - ) + self.init_topk_indices_buffer(indexer, topk_indices_buffer) def _forward_bf16_kv( self, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f41a5a8d3a..fcadf9e303 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2030,7 +2030,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None - if self.speculator is not None and self.speculator.supports_mm_inputs: + # The encoder runner exists only on the first PP rank, so later ranks + # have no cached embeddings to gather. + if ( + self.speculator is not None + and self.speculator.supports_mm_inputs + and self.model_state.supports_mm_inputs + ): # Get cached multimodal embeddings for draft forward. # NOTE: This is done here because postprocess updates # num_computed_prefill_tokens.