From c9e6442e0f8cbca5d53ad420eed088e00d8db9cc Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sun, 23 Aug 2026 19:39:56 -0700 Subject: [PATCH 01/13] feat(dspark): support PP prefill in disaggregated serving Materialize DSpark context KV on pipeline-parallel prefill workers without running the decode-only draft path, so the decode-side drafter starts warm in PD deployments. Connector-agnostic: works with connectors that register and align named per-layer regions (e.g. Mooncake); the NIXL connector raises a loud NotImplementedError until it describes packed KV regions per producer stage. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/config/test_dspark_prefill_only.py | 51 +++++ tests/models/test_deepseek_v4_mega_moe.py | 17 +- vllm/config/speculative.py | 38 +++- vllm/config/vllm.py | 2 + vllm/engine/arg_utils.py | 1 + vllm/models/deepseek_v4/amd/model.py | 6 +- vllm/models/deepseek_v4/nvidia/dspark.py | 208 +++++++++++++----- vllm/models/deepseek_v4/nvidia/model.py | 6 +- vllm/models/deepseek_v4/xpu/model.py | 6 +- vllm/v1/core/sched/scheduler.py | 5 +- vllm/v1/worker/gpu/model_runner.py | 141 +++++++----- .../gpu/spec_decode/dflash/speculator.py | 110 +++++---- .../gpu/spec_decode/dspark/speculator.py | 1 + .../v1/worker/gpu/spec_decode/dspark/utils.py | 87 +++++--- vllm/v1/worker/gpu/spec_decode/speculator.py | 24 +- 15 files changed, 506 insertions(+), 197 deletions(-) create mode 100644 tests/config/test_dspark_prefill_only.py diff --git a/tests/config/test_dspark_prefill_only.py b/tests/config/test_dspark_prefill_only.py new file mode 100644 index 000000000000..554e0d89fcc1 --- /dev/null +++ b/tests/config/test_dspark_prefill_only.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.config import KVTransferConfig, ParallelConfig, SpeculativeConfig +from vllm.config.kv_transfer import KVRole +from vllm.config.speculative import SpeculativeMethod + + +def _spec_config( + *, pp: int, role: KVRole, method: SpeculativeMethod = "dspark" +) -> SpeculativeConfig: + config = object.__new__(SpeculativeConfig) + config.method = method + config.target_parallel_config = ParallelConfig(pipeline_parallel_size=pp) + config.target_kv_transfer_config = KVTransferConfig( + kv_connector="NixlConnector", + kv_role=role, + ) + return config + + +@pytest.mark.parametrize( + ("pp", "role", "method", "expected"), + [ + (2, "kv_producer", "dspark", True), + (4, "kv_producer", "dspark", True), + (1, "kv_producer", "dspark", False), + (2, "kv_consumer", "dspark", False), + (2, "kv_both", "dspark", False), + (2, "kv_producer", "dflash", False), + ], +) +def test_dspark_prefill_only_role_detection(pp, role, method, expected): + assert _spec_config(pp=pp, role=role, method=method).is_dspark_prefill_only() is ( + expected + ) + + +def test_dspark_prefill_materializer_uses_pp1_draft_config(): + target = ParallelConfig(pipeline_parallel_size=4, tensor_parallel_size=1) + + draft = SpeculativeConfig.create_draft_parallel_config( + target, + speculative_draft_tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + assert draft.pipeline_parallel_size == 1 + assert draft.tensor_parallel_size == 1 diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index f3c70dc88a0d..e1870b78a688 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -503,12 +503,27 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp")) DeepSeekV4MTP.process_weights_after_loading(mtp) - dspark = SimpleNamespace(_finalize_moe=lambda: calls.append("dspark")) + dspark = SimpleNamespace( + model=SimpleNamespace(context_kv_only=False), + _finalize_moe=lambda: calls.append("dspark"), + ) DSparkDeepseekV4ForCausalLM.process_weights_after_loading(dspark) assert calls == ["mtp", "dspark"] +def test_dspark_context_materializer_skips_absent_confidence_head(): + """The context-only P model omits decode-only heads entirely.""" + dspark = object.__new__(DSparkDeepseekV4ForCausalLM) + dspark.model = SimpleNamespace() + + assert dspark._remap_dspark_name("mtp.2.confidence_head.weight") is None + + dspark.model.context_kv_only = True + dspark._finalize_moe = lambda: pytest.fail("context-only model has no MoE") + dspark.process_weights_after_loading() + + @pytest.mark.skipif( not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index d3b4af1cf060..f0a8d894ce79 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -12,6 +12,7 @@ from vllm.config import LoadConfig from vllm.config.cache import CacheDType from vllm.config.kernel import MoEBackend +from vllm.config.kv_transfer import KVTransferConfig from vllm.config.model import HfOverrides, ModelConfig from vllm.config.parallel import ParallelConfig from vllm.config.utils import config @@ -463,6 +464,8 @@ class SpeculativeConfig: """The configuration of the target model.""" target_parallel_config: SkipValidation[ParallelConfig] = None # type: ignore """The parallel configuration for the target model.""" + target_kv_transfer_config: SkipValidation[KVTransferConfig] = None # type: ignore + """The KV transfer configuration for the target model.""" # dynamic speculative decoding control num_speculative_tokens_per_batch_size: list[tuple[int, int, int]] | None = None @@ -1449,13 +1452,29 @@ def __post_init__(self): self.draft_parallel_config = ( SpeculativeConfig.create_draft_parallel_config( - self.target_parallel_config, self.draft_tensor_parallel_size + self.target_parallel_config, + self.draft_tensor_parallel_size, + pipeline_parallel_size=( + 1 if self.is_dspark_prefill_only() else None + ), ) ) if self.method != "dspark" and self.enable_adaptive_verification: raise ValueError("Adaptive verification only supported with DSpark") + if self.is_dspark_prefill_only(): + # is_dspark_prefill_only() implies target_kv_transfer_config is set. + assert self.target_kv_transfer_config is not None + connector = self.target_kv_transfer_config.kv_connector + if connector is not None and "Nixl" in connector: + raise NotImplementedError( + "DSpark prefill materialization with pipeline parallelism " + "requires a connector that transfers named per-layer KV " + "regions (e.g. MooncakeConnector). The NIXL connector does " + "not describe packed KV regions per producer stage yet." + ) + return self def _validate_suffix_decoding(self): @@ -1630,13 +1649,18 @@ def update_arch_(self): def create_draft_parallel_config( target_parallel_config: ParallelConfig, speculative_draft_tensor_parallel_size: int, + pipeline_parallel_size: int | None = None, ) -> ParallelConfig: """Create a parallel config for use by the draft worker. This is mostly a copy of the target parallel config, except the tp_size. """ draft_parallel_config = ParallelConfig( - pipeline_parallel_size=target_parallel_config.pipeline_parallel_size, + pipeline_parallel_size=( + target_parallel_config.pipeline_parallel_size + if pipeline_parallel_size is None + else pipeline_parallel_size + ), tensor_parallel_size=speculative_draft_tensor_parallel_size, distributed_executor_backend=target_parallel_config.distributed_executor_backend, max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers, @@ -1647,6 +1671,16 @@ def create_draft_parallel_config( return draft_parallel_config + def is_dspark_prefill_only(self) -> bool: + kv_transfer_config = self.target_kv_transfer_config + return ( + self.method == "dspark" + and self.target_parallel_config is not None + and self.target_parallel_config.pipeline_parallel_size > 1 + and kv_transfer_config is not None + and kv_transfer_config.kv_role == "kv_producer" + ) + @field_validator("attention_backend", mode="before") @classmethod def _parse_attention_backend(cls, value: Any) -> Any: diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 51b8b1019c63..69d8fd5e5a7a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -605,6 +605,8 @@ def num_speculative_tokens(self) -> int: self.speculative_config is not None and self.speculative_config.num_speculative_tokens is not None ): + if self.speculative_config.is_dspark_prefill_only(): + return 0 return self.speculative_config.num_speculative_tokens if ( self.diffusion_config is not None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index d117fedf37f7..e8874224be40 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1895,6 +1895,7 @@ def create_speculative_config( { "target_model_config": target_model_config, "target_parallel_config": target_parallel_config, + "target_kv_transfer_config": self.kv_transfer_config, } ) return SpeculativeConfig(**self.speculative_config) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 7a67c8f420e6..1e84bef0b076 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -640,8 +640,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.hc_head_op = HCHeadOp() spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() + needs_mtp_hidden_states = ( + spec_config is not None + and (spec_config.use_eagle() or spec_config.uses_draft_model()) + and not spec_config.is_dspark_prefill_only() ) if get_pp_group().is_last_rank and needs_mtp_hidden_states: self._mtp_hidden_buffer = torch.empty( diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 1a3c2c780d28..bd8daa1839d4 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -63,6 +63,34 @@ _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") +class _DSparkContextKVAttention(nn.Module): + """The subset of a DSV4 draft attention layer needed by prefill.""" + + def __init__(self, attn: nn.Module) -> None: + super().__init__() + self.fused_wqa_wkv = attn.fused_wqa_wkv + self.kv_norm = attn.kv_norm + self.rotary_emb = attn.rotary_emb + self.swa_cache_layer = attn.swa_cache_layer + self.q_lora_rank = attn.q_lora_rank + self.n_local_heads = attn.n_local_heads + self.head_dim = attn.head_dim + self.padded_heads = attn.padded_heads + self.eps = attn.eps + for name in ("_flashinfer_fp8_kv_scale", "_flashinfer_fp8_q_scale_inv"): + value = getattr(attn, name, None) + if value is not None: + self.register_buffer(name, value, persistent=False) + + +class _DSparkContextKVLayer(nn.Module): + """A draft layer stripped down to its context-KV projection.""" + + def __init__(self, layer: DeepseekV4DecoderLayer) -> None: + super().__init__() + self.attn = _DSparkContextKVAttention(layer.attn) + + class DSparkDeepseekV4Model(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -76,15 +104,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) self.use_sequence_parallel = _use_sequence_parallel(vllm_config) + self.context_kv_only = vllm_config.speculative_config.is_dspark_prefill_only() self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 - # Shared with the target (aliased by the speculator's loading utility). - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "embed_tokens"), - ) + if not self.context_kv_only: + # Shared with the target (aliased by the speculator's loading utility). + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) self.main_proj = ReplicatedLinear( config.hidden_size * len(self.target_layer_ids), @@ -96,53 +126,66 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.topk_indices_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - config.index_topk, - dtype=torch.int32, + self.topk_indices_buffer = ( + None + if self.context_kv_only + else torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + ) ) current_vllm_config = get_current_vllm_config() - self.layers = nn.ModuleList( - [ - DeepseekV4DecoderLayer( - current_vllm_config, - prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), - topk_indices_buffer=self.topk_indices_buffer, + layers: list[nn.Module] = [] + for i in range(self.num_dspark_layers): + layer_prefix = maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}") + layer = DeepseekV4DecoderLayer( + current_vllm_config, + prefix=layer_prefix, + topk_indices_buffer=self.topk_indices_buffer, + ) + if self.context_kv_only: + # The full attention object registers itself for metadata lookup, + # but only its SWA cache layer participates in materialization. + current_vllm_config.compilation_config.static_forward_context.pop( + f"{layer_prefix}.attn", None ) - for i in range(self.num_dspark_layers) - ] - ) - - # Heads: final norm + hc_head, and the Markov + confidence heads - # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - hc_dim = self.hc_mult * config.hidden_size - self.hc_head_fn = nn.Parameter( - torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_base = nn.Parameter( - torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False - ) - self.hc_head_scale = nn.Parameter( - torch.empty(1, dtype=torch.float32), requires_grad=False - ) - draft_vocab_size = ( - getattr(config, "draft_vocab_size", None) or config.vocab_size - ) - self.markov_head = DSparkMarkovHead( - config.vocab_size, - draft_vocab_size, - config.dspark_markov_rank, - prefix=maybe_prefix(prefix, "markov_head"), - ) - self.confidence_head: DSparkConfidenceHead | None = None - if getattr(config, "enable_confidence_head", True): - self.confidence_head = DSparkConfidenceHead( - config.hidden_size + config.dspark_markov_rank, - prefix=maybe_prefix(prefix, "confidence_head"), + layers.append(_DSparkContextKVLayer(layer)) + else: + layers.append(layer) + self.layers = nn.ModuleList(layers) + + if not self.context_kv_only: + # Heads: final norm + hc_head, and the Markov + confidence heads + # loaded from the final MTP layer in the target checkpoint. + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), + requires_grad=False, ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + draft_vocab_size = ( + getattr(config, "draft_vocab_size", None) or config.vocab_size + ) + self.markov_head = DSparkMarkovHead( + config.vocab_size, + draft_vocab_size, + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + self.confidence_head: DSparkConfidenceHead | None = None + if getattr(config, "enable_confidence_head", True): + self.confidence_head = DSparkConfidenceHead( + config.hidden_size + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "confidence_head"), + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -193,6 +236,8 @@ def forward( positions: torch.Tensor, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor: + if self.context_kv_only: + raise RuntimeError("The DSpark prefill materializer cannot draft tokens.") if inputs_embeds is None: inputs_embeds = self.embed_input_ids(input_ids) full_num_tokens = positions.shape[0] @@ -319,13 +364,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.model = DSparkDeepseekV4Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - # Shared with the target (aliased by the speculator's load utility). - self.lm_head = ParallelLMHead( - self.config.vocab_size, - self.config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) - self.logits_processor = LogitsProcessor(self.config.vocab_size) + if not self.model.context_kv_only: + # Shared with the target (aliased by the speculator's load utility). + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.config.vocab_size) # --- Hooks used by the speculator ------------------------------------- @@ -391,6 +437,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: Non-mtp weights (embed/head/main layers) belong to the target model and are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. """ + if self.model.context_kv_only: + return self._load_context_kv_weights(weights) + first_layer = self.model.layers[0] use_mega_moe = first_layer.ffn.use_mega_moe if use_mega_moe: @@ -510,12 +559,54 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params + def _load_context_kv_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> set[str]: + """Load only weights used to build DSpark prefix KV on a P worker.""" + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for checkpoint_name, loaded_weight in weights: + name = self._remap_dspark_name(checkpoint_name) + if name is None: + continue + is_context_param = name.startswith( + ("model.main_proj.", "model.main_norm.") + ) or ( + name.startswith("model.layers.") + and any( + part in name + for part in (".attn.wq_a.", ".attn.wkv.", ".attn.kv_norm.") + ) + ) + if not is_context_param: + continue + if name.endswith(".scale"): + name = name.removesuffix(".scale") + ".weight_scale_inv" + for weight_name, shard_id in (("attn.wq_a", 0), ("attn.wkv", 1)): + if weight_name not in name: + continue + name = name.replace(weight_name, "attn.fused_wqa_wkv") + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + break + else: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + logger.info_once( + "DSpark context-KV materializer loaded: %d params", len(loaded_params) + ) + return loaded_params + def _finalize_moe(self) -> None: for layer in self.model.layers: layer.ffn.finalize_mega_moe_weights() def process_weights_after_loading(self) -> None: - self._finalize_moe() + if not self.model.context_kv_only: + self._finalize_moe() def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. @@ -527,7 +618,10 @@ def _remap_dspark_name(self, name: str) -> str | None: return None stage = int(m.group(1)) rest = m.group(2) - if rest.startswith("confidence_head.") and self.model.confidence_head is None: + if ( + rest.startswith("confidence_head.") + and getattr(self.model, "confidence_head", None) is None + ): return None # Head-stack params live at model level (mtp.last), context combiner at # model level (mtp.0); everything else is a per-layer decoder block. diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index a4695263165f..5f29f8774709 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1322,8 +1322,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): requires_grad=False, ) spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() + needs_mtp_hidden_states = ( + spec_config is not None + and (spec_config.use_eagle() or spec_config.uses_draft_model()) + and not spec_config.is_dspark_prefill_only() ) if get_pp_group().is_last_rank and needs_mtp_hidden_states: self._mtp_hidden_buffer = torch.empty( diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index ddf8ae175b22..cbfb4ff3a351 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -1050,8 +1050,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.hc_head_op = HCHeadOp() spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() + needs_mtp_hidden_states = ( + spec_config is not None + and (spec_config.use_eagle() or spec_config.uses_draft_model()) + and not spec_config.is_dspark_prefill_only() ) if get_pp_group().is_last_rank and needs_mtp_hidden_states: self._mtp_hidden_buffer = torch.empty( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c373d5cf5c87..9fb11604a888 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -271,7 +271,10 @@ def __init__( vllm_max_batch_size=self.scheduler_config.max_num_seqs, vllm_num_speculative_tokens=self.num_spec_tokens, ) - self.use_eagle = speculative_config.use_eagle() + self.use_eagle = ( + speculative_config.use_eagle() + and not speculative_config.is_dspark_prefill_only() + ) if self.use_eagle: self.num_prefill_lookahead = ( self.num_spec_tokens diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index cde043d1cb31..e86281e3867e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -174,6 +174,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.parallel_config = vllm_config.parallel_config self.scheduler_config = vllm_config.scheduler_config self.speculative_config = vllm_config.speculative_config + self.dspark_prefill_only = bool( + self.speculative_config is not None + and self.speculative_config.is_dspark_prefill_only() + ) self._draft_workspace_lane = int( self.speculative_config is not None and self.speculative_config.use_dspark() ) @@ -261,7 +265,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ): # Drafting may require auxiliary hidden states from target model outputs self.use_aux_hidden_state_outputs = True - if self.use_pp: + if self.use_pp and not self.dspark_prefill_only: raise ValueError( f"{self.speculative_config.method} with pipeline parallel " "is not supported." @@ -388,9 +392,12 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: if isinstance(self.speculator, DraftModelSpeculator): with use_workspace_lane(self._draft_workspace_lane): self.speculator.load_model(self.model) - eplb_models_added = self.eplb.maybe_register_speculator( - self.speculator, self.speculative_config, load_dummy_weights - ) + if not self.dspark_prefill_only: + eplb_models_added = self.eplb.maybe_register_speculator( + self.speculator, + self.speculative_config, + load_dummy_weights, + ) time_after_load = time.perf_counter() self.model_memory_usage = m.consumed_memory @@ -443,7 +450,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: if custom: self.sampler, self.rejection_sampler = custom - elif self.speculative_config is not None: + elif self.speculative_config is not None and not self.dspark_prefill_only: self.rejection_sampler = RejectionSampler( self.sampler, self.speculative_config, @@ -577,8 +584,9 @@ def initialize_kv_cache( # The speculator clears the flag at load time when the checkpoint has # no confidence head, so it holds the effective value. self.adaptive_verification = maybe_create_adaptive_verification_manager( - enable_adaptive_verification=getattr( - self.speculator, "enable_adaptive_verification", False + enable_adaptive_verification=( + not self.dspark_prefill_only + and getattr(self.speculator, "enable_adaptive_verification", False) ), attn_groups=self.attn_groups, attn_cg_support=attn_cg_support, @@ -643,7 +651,7 @@ def initialize_kv_cache( self.input_buffers, self.attn_groups, ) - if self.speculator is not None: + if self.speculator is not None and not self.dspark_prefill_only: # After set_attn, so the speculator can size its cudagraph mode # to its own attention support. self.speculator.init_cudagraph_manager(cudagraph_mode) @@ -760,6 +768,7 @@ def _dummy_run( # dummy run the eagle speculator's propose to ensure DP/EP sync. if self.speculator is not None: assert self.sampler is not None + assert hidden_states is not None self.step_timing.drafter_start() mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None if self.speculator.supports_mm_inputs: @@ -779,30 +788,46 @@ def _dummy_run( spec_hidden_states = hidden_states if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + if pre_hc_hidden_states is not None: + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] with use_workspace_lane(self._draft_workspace_lane): - self.speculator.propose( - input_batch=input_batch, - attn_metadata=attn_metadata, - slot_mappings=slot_mappings_by_layer, - last_hidden_states=spec_hidden_states, - aux_hidden_states=aux_hidden_states, - num_sampled=torch.ones( - input_batch.num_reqs, dtype=torch.int32, device=self.device - ), - num_rejected=torch.zeros( - input_batch.num_reqs, dtype=torch.int32, device=self.device - ), - last_sampled=self.req_states.last_sampled_tokens, - next_prefill_tokens=self.req_states.next_prefill_tokens, - temperature=self.sampler.sampling_states.temperature.gpu, - seeds=self.sampler.sampling_states.seeds.gpu, - dp_sync=dp_sync, - dummy_run=True, - skip_attn_for_dummy_run=skip_attn, - mm_inputs=mm_inputs, - is_profile=is_profile, + num_sampled = torch.ones( + input_batch.num_reqs, dtype=torch.int32, device=self.device ) + num_rejected = torch.zeros_like(num_sampled) + if self.dspark_prefill_only: + self.speculator.materialize_context_kv( + input_batch=input_batch, + last_hidden_states=spec_hidden_states, + aux_hidden_states=aux_hidden_states, + num_sampled=num_sampled, + num_rejected=num_rejected, + last_sampled=self.req_states.last_sampled_tokens, + next_prefill_tokens=self.req_states.next_prefill_tokens, + temperature=self.sampler.sampling_states.temperature.gpu, + seeds=self.sampler.sampling_states.seeds.gpu, + dummy_run=True, + skip_attn_for_dummy_run=skip_attn, + ) + else: + self.speculator.propose( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings=slot_mappings_by_layer, + last_hidden_states=spec_hidden_states, + aux_hidden_states=aux_hidden_states, + num_sampled=num_sampled, + num_rejected=num_rejected, + last_sampled=self.req_states.last_sampled_tokens, + next_prefill_tokens=self.req_states.next_prefill_tokens, + temperature=self.sampler.sampling_states.temperature.gpu, + seeds=self.sampler.sampling_states.seeds.gpu, + dp_sync=dp_sync, + dummy_run=True, + skip_attn_for_dummy_run=skip_attn, + mm_inputs=mm_inputs, + is_profile=is_profile, + ) self.step_timing.drafter_end() assert hidden_states is not None # Last PP rank always has hidden_states @@ -924,7 +949,7 @@ def capture_model(self) -> int: use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, lora_capture_hook=create_lora_capture_hook(self.lora_config, self), ) - if self.speculator is not None: + if self.speculator is not None and not self.dspark_prefill_only: with use_workspace_lane(self._draft_workspace_lane): self.speculator.capture() if self.adaptive_verification is not None: @@ -1914,30 +1939,46 @@ def sample_tokens( spec_hidden_states = hidden_states if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + if pre_hc_hidden_states is not None: + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] with use_workspace_lane(self._draft_workspace_lane): - draft_tokens = self.speculator.propose( - input_batch, - attn_metadata, - slot_mappings_by_layer, - spec_hidden_states, - aux_hidden_states, - num_sampled, - num_rejected, - self.req_states.last_sampled_tokens, - self.req_states.next_prefill_tokens, - self.sampler.sampling_states.temperature.gpu, - self.sampler.sampling_states.seeds.gpu, - dp_sync=dp_sync, - mm_inputs=mm_inputs, - ) - self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - if self.adaptive_verification is not None: + if self.dspark_prefill_only: + self.speculator.materialize_context_kv( + input_batch, + spec_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + self.req_states.last_sampled_tokens, + self.req_states.next_prefill_tokens, + self.sampler.sampling_states.temperature.gpu, + self.sampler.sampling_states.seeds.gpu, + ) + draft_tokens = None + else: + draft_tokens = self.speculator.propose( + input_batch, + attn_metadata, + slot_mappings_by_layer, + spec_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + self.req_states.last_sampled_tokens, + self.req_states.next_prefill_tokens, + self.sampler.sampling_states.temperature.gpu, + self.sampler.sampling_states.seeds.gpu, + dp_sync=dp_sync, + mm_inputs=mm_inputs, + ) + if draft_tokens is not None: + self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens + if draft_tokens is not None and self.adaptive_verification is not None: self.adaptive_verification.record_confidences( self.speculator.draft_token_confidence_probs, input_batch ) - if self.num_speculative_steps > 0: + if self.num_speculative_steps > 0 and not self.dspark_prefill_only: # Spec-decode and diffusion LLMs both use draft tokens but the latter does # not have a speculator (i.e. self.speculator is None) self.draft_tokens_handler.set_draft_tokens( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 2ff7b22d0636..34133e4f9204 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -314,45 +314,27 @@ def _build_draft_attn_metadata( ) @torch.inference_mode() - def propose( + def materialize_context_kv( self, input_batch: InputBatch, - attn_metadata: dict[str, Any], - slot_mappings: dict[str, torch.Tensor], - # [num_tokens, hidden_size] last_hidden_states: torch.Tensor, - # num_layers x [num_tokens, hidden_size] aux_hidden_states: list[torch.Tensor] | None, - # [num_reqs] num_sampled: torch.Tensor, - # [num_reqs] num_rejected: torch.Tensor, - # [max_num_reqs] last_sampled: torch.Tensor, - # [max_num_reqs] next_prefill_tokens: torch.Tensor, - # [max_num_reqs] temperature: torch.Tensor, - # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, - mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, - is_profile: bool = False, - ) -> torch.Tensor: + ) -> None: num_reqs = input_batch.num_reqs num_target_tokens = input_batch.num_tokens - num_query_tokens = num_reqs * self.num_query_per_req max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() self.draft_max_seq_len = min( max_seq_len + self.num_query_per_req, self.max_model_len ) - # NOTE: To avoid CPU-GPU synchronization without CPU knowing the - # number of rejected tokens, we maintain the size of input_ids and - # hidden_states the same as the target model's. This means, we pad each - # request's query length to include any rejected positions. if aux_hidden_states: hidden_states = self.model.combine_hidden_states( torch.cat(aux_hidden_states, dim=-1) @@ -362,32 +344,13 @@ def propose( self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens]) if dummy_run and skip_attn_for_dummy_run: - # Memory profiling path: block_tables / kv_cache_config are not initialized. - # Since DFlash needs to build its own attention metadata, we must skip the - # preparation in this path and run a minimal forward pass. self.model.precompute_and_store_context_kv( self.hidden_states[:num_target_tokens], self.context_positions[:num_target_tokens], ) - # DFlash processes all speculative tokens in one forward pass, - # so the real token count is num_query_tokens. - self._prepare_eplb_forward(num_query_tokens) - self._generate_draft( - num_reqs, - num_query_tokens, - attn_metadata=None, - slot_mappings=None, - num_tokens_across_dp=( - dp_sync.num_tokens_across_dp if dp_sync is not None else None - ), - cudagraph_runtime_mode=CUDAGraphMode.NONE, - ) - return self.draft_tokens[:num_reqs] + return - # The query slot mapping is written into the shared BlockTables slot_mappings. - # That buffer's address is what the captured CUDA graph reads from at replay. assert self.draft_kv_cache_group_id >= 0 - # Support multiple draft KV cache groups by preparing inputs once for each for i, gid in enumerate(self.draft_kv_cache_group_ids): prepare_dflash_inputs( self.input_buffers, @@ -420,10 +383,6 @@ def propose( self.sample_from_anchor, ) - # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph - # because the context shape varies per step. During dummy runs the block tables - # are placeholders, so we skip the cache write to avoid clobbering real entries. - # Each layer uses the context slots of its own kv-cache group. if dummy_run: context_slots: torch.Tensor | list[torch.Tensor | None] | None = None elif self._layer_group_idx is not None: @@ -439,6 +398,69 @@ def propose( context_slots, ) + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + # [num_tokens, hidden_size] + last_hidden_states: torch.Tensor, + # num_layers x [num_tokens, hidden_size] + aux_hidden_states: list[torch.Tensor] | None, + # [num_reqs] + num_sampled: torch.Tensor, + # [num_reqs] + num_rejected: torch.Tensor, + # [max_num_reqs] + last_sampled: torch.Tensor, + # [max_num_reqs] + next_prefill_tokens: torch.Tensor, + # [max_num_reqs] + temperature: torch.Tensor, + # [max_num_reqs] + seeds: torch.Tensor, + dp_sync: DPSyncState | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + num_query_tokens = num_reqs * self.num_query_per_req + self.materialize_context_kv( + input_batch, + last_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + temperature, + seeds, + dummy_run=dummy_run, + skip_attn_for_dummy_run=skip_attn_for_dummy_run, + ) + + if dummy_run and skip_attn_for_dummy_run: + # Memory profiling path: block_tables / kv_cache_config are not initialized. + # Since DFlash needs to build its own attention metadata, we must skip the + # preparation in this path and run a minimal forward pass. + # DFlash processes all speculative tokens in one forward pass, + # so the real token count is num_query_tokens. + self._prepare_eplb_forward(num_query_tokens) + self._generate_draft( + num_reqs, + num_query_tokens, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=( + dp_sync.num_tokens_across_dp if dp_sync is not None else None + ), + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + return self.draft_tokens[:num_reqs] + # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs batch_desc, batch_sync = dispatch_cg_and_sync_dp( self.query_cudagraph_manager, diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 19fc45201846..3935f42644d2 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -84,6 +84,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ) self.enable_adaptive_verification = ( self.speculative_config.enable_adaptive_verification + and not self.speculative_config.is_dspark_prefill_only() ) def load_draft_model( diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 177060747ab8..0db7fc7526a3 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -36,6 +36,29 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + pp_group = get_pp_group() + if pp_group.world_size != 1 and not speculative_config.is_dspark_prefill_only(): + raise NotImplementedError("DSpark does not support pipeline parallelism.") + + target_language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + target_inner = target_language_model.model + if pp_group.world_size != 1: + start_layer = target_inner.start_layer + end_layer = target_inner.end_layer + aux_layers = tuple( + layer + 1 for layer in draft_model_config.hf_config.dspark_target_layer_ids + ) + if any(layer <= start_layer or layer > end_layer for layer in aux_layers): + raise ValueError( + "DSpark prefill materialization requires every auxiliary hidden " + "state on the last pipeline stage. " + f"Stage owns ({start_layer}, {end_layer}], requested {aux_layers}." + ) + from vllm.compilation.backends import set_model_tag from vllm.model_executor.model_loader import get_model from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal @@ -76,44 +99,38 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo vllm_config=draft_vllm_config, model_config=draft_model_config ) - if get_pp_group().world_size != 1: - raise NotImplementedError("DSpark does not support pipeline parallelism.") - - target_language_model = ( - target_model.get_language_model() - if hasattr(target_model, "get_language_model") - else target_model - ) - target_inner = target_language_model.model draft_inner = draft_model.model target_vocab_size = vllm_config.model_config.get_vocab_size() - target_embed = getattr(target_inner, "embed_tokens", None) - draft_embed = getattr(draft_inner, "embed_tokens", None) - if ( - target_embed is not None - and draft_model_config.get_vocab_size() <= target_vocab_size - and _should_share( - draft_model, "has_own_embed_tokens", draft_embed, target_embed + if not speculative_config.is_dspark_prefill_only(): + target_embed = getattr(target_inner, "embed_tokens", None) + draft_embed = getattr(draft_inner, "embed_tokens", None) + if ( + target_embed is not None + and draft_model_config.get_vocab_size() <= target_vocab_size + and _should_share( + draft_model, "has_own_embed_tokens", draft_embed, target_embed + ) + ): + if draft_embed is not None: + del draft_inner.embed_tokens + draft_inner.embed_tokens = target_embed + + target_lm_head = get_target_lm_head(target_model, target_language_model) + draft_lm_head = getattr(draft_model, "lm_head", None) + draft_output_vocab_size = ( + getattr(draft_model_config.hf_config, "draft_vocab_size", None) + or draft_model_config.get_vocab_size() ) - ): - if draft_embed is not None: - del draft_inner.embed_tokens - draft_inner.embed_tokens = target_embed - - target_lm_head = get_target_lm_head(target_model, target_language_model) - draft_lm_head = getattr(draft_model, "lm_head", None) - draft_output_vocab_size = ( - getattr(draft_model_config.hf_config, "draft_vocab_size", None) - or draft_model_config.get_vocab_size() - ) - if ( - target_lm_head is not None - and draft_output_vocab_size == target_vocab_size - and _should_share(draft_model, "has_own_lm_head", draft_lm_head, target_lm_head) - ): - if draft_lm_head is not None: - del draft_model.lm_head - draft_model.lm_head = target_lm_head + if ( + target_lm_head is not None + and draft_output_vocab_size == target_vocab_size + and _should_share( + draft_model, "has_own_lm_head", draft_lm_head, target_lm_head + ) + ): + if draft_lm_head is not None: + del draft_model.lm_head + draft_model.lm_head = target_lm_head return draft_model diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 303f16b79f3c..b163e2db87eb 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -69,6 +69,25 @@ def propose( ) -> torch.Tensor: pass + def materialize_context_kv( + self, + input_batch: InputBatch, + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + ) -> None: + """Materialize draft context KV without generating draft tokens.""" + raise NotImplementedError( + f"{type(self).__name__} does not support context-KV materialization." + ) + class DraftModelSpeculator(BaseSpeculator): def __init__(self, vllm_config: VllmConfig, device: torch.device): @@ -131,7 +150,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ) self.draft_logits: torch.Tensor | None = None - if self.speculative_config.draft_sample_method == "probabilistic": + if ( + self.speculative_config.draft_sample_method == "probabilistic" + and not self.speculative_config.is_dspark_prefill_only() + ): # Pre-temperature logits, cached from the previous decode step. dtype, fill = self.draft_logits_spec(vllm_config) self.draft_logits = torch.full( From 02fd2e70a4da900b4877fcce65d3fa8df3a562dc Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sat, 22 Aug 2026 22:41:43 +0800 Subject: [PATCH 02/13] fix(dspark): make padded graph batches safe Track and mask graph-padding rows through DFlash input preparation and DSV4 routing, preserve compact request-slot mappings, and allow an optional fixed DSpark graph batch size. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> (cherry picked from commit fc7061192e367d631cfe7642c84ea8ef5dc5ce44) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/kernels/moe/test_topk_softplus_sqrt.py | 73 +++++++++ .../spec_decode/test_dflash_prepare_inputs.py | 149 ++++++++++++++++++ .../layers/fused_moe/router/dsv4_topk.py | 18 +++ .../router/fused_topk_bias_router.py | 1 + .../gpu/spec_decode/dflash/speculator.py | 22 ++- .../gpu/spec_decode/dspark/speculator.py | 31 ++++ 6 files changed, 290 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index c5f00253e781..f30096d79f55 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -234,6 +234,79 @@ def test_dsv4_fast_topk( ) +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="The DeepSeek V4 fast path is CUDA-only.", +) +def test_dsv4_fast_topk_padding(monkeypatch: pytest.MonkeyPatch): + """Verify the DSV4 fast path removes graph-padding rows from routing.""" + torch.manual_seed(0) + num_tokens = 17 + num_experts = 256 + hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda") + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") + is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda") + is_padding[1::2] = True + gating_output[is_padding] = float("nan") + + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.router." + "fused_topk_bias_router._get_padding_mask", + lambda _: is_padding, + ) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=gating_output, + scoring_func="sqrtsoftplus", + e_score_correction_bias=correction_bias, + topk=6, + renormalize=True, + routed_scaling_factor=1.5, + ) + + assert torch.equal(topk_ids[is_padding], torch.full_like(topk_ids[is_padding], -1)) + assert torch.equal( + topk_weights[is_padding], torch.zeros_like(topk_weights[is_padding]) + ) + + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( + gating_output=gating_output[~is_padding], + topk=6, + renormalize=True, + routed_scaling_factor=1.5, + e_score_correction_bias=correction_bias, + ) + torch.testing.assert_close(topk_ids[~is_padding], topk_ids_ref, atol=0, rtol=0) + torch.testing.assert_close( + topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5 + ) + + # The mask buffer is persistent under CUDA graph replay, but its contents + # change with every batch. Verify that the kernel reads those contents at + # runtime rather than specializing on the mask captured above. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_weights, graph_ids = dsv4_topk( + gating_output, + correction_bias, + torch.int32, + 1.5, + is_padding=is_padding, + ) + + is_padding.logical_not_() + graph.replay() + assert torch.equal( + graph_ids[is_padding], torch.full_like(graph_ids[is_padding], -1) + ) + assert torch.equal( + graph_weights[is_padding], torch.zeros_like(graph_weights[is_padding]) + ) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="This test is skipped on non-CUDA platform.", diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 16d6d8e516df..f3f94e1a398f 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import contextmanager from types import SimpleNamespace import numpy as np @@ -8,7 +9,9 @@ import torch from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( + DFlashSpeculator, prepare_dflash_inputs, ) @@ -33,6 +36,7 @@ def _run_prepare( input_buffers = SimpleNamespace( input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device), positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device), + is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device), query_start_loc=torch.full( (max_num_reqs + 1,), -1, dtype=torch.int32, device=device ), @@ -140,6 +144,114 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix(): assert out.temperature[2].item() == 1.0 assert out.seeds[2].item() == 17 + assert not out.input_buffers.is_padding[:3].any() + assert out.input_buffers.is_padding[3:].all() + assert out.input_buffers.input_ids[3:].cpu().tolist() == [0] * 13 + assert out.input_buffers.positions[3:].cpu().tolist() == [0] * 13 + + +def test_prepare_dflash_inputs_compacts_noncontiguous_request_slots(): + device = torch.device("cuda") + max_num_reqs = 4 + max_num_tokens = 16 + num_speculative_steps = 3 + input_buffers = SimpleNamespace( + input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device), + positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device), + is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device), + query_start_loc=torch.full( + (max_num_reqs + 1,), -1, dtype=torch.int32, device=device + ), + seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device), + ) + input_batch = SimpleNamespace( + num_reqs=2, + num_scheduled_tokens=np.array([4, 4], dtype=np.int32), + positions=torch.tensor( + [10, 11, 12, 13, 20, 21, 22, 23], + dtype=torch.int64, + device=device, + ), + query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32, device=device), + # Active batch rows are compact, while request state remains in slots 3 and 1. + idx_mapping=torch.tensor([3, 1], dtype=torch.int32, device=device), + ) + query_slot_mapping = torch.full( + (max_num_tokens,), -2, dtype=torch.int64, device=device + ) + context_positions = torch.full( + (max_num_tokens,), -1, dtype=torch.int64, device=device + ) + context_slot_mapping = torch.full( + (max_num_tokens,), -2, dtype=torch.int64, device=device + ) + sample_indices = torch.full( + (max_num_reqs * num_speculative_steps,), + -1, + dtype=torch.int64, + device=device, + ) + sample_pos = torch.full_like(sample_indices, -1) + sample_idx_mapping = torch.full( + sample_indices.shape, -1, dtype=torch.int32, device=device + ) + temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device) + seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device) + input_temperature = torch.tensor( + [0.0, 0.5, 0.0, 1.0], dtype=torch.float32, device=device + ) + input_seeds = torch.tensor([0, 11, 0, 33], dtype=torch.int64, device=device) + last_sampled = torch.tensor([0, 77, 0, 99], dtype=torch.int64, device=device) + next_prefill_tokens = torch.zeros_like(last_sampled) + block_table = torch.tensor( + [[0, 0, 7, 8, 9, 10, 11, 12], [0, 0, 13, 14, 15, 16, 17, 18]], + dtype=torch.int32, + device=device, + ) + + prepare_dflash_inputs( + input_buffers, + query_slot_mapping, + context_positions, + context_slot_mapping, + sample_indices, + sample_pos, + sample_idx_mapping, + temperature, + seeds, + input_batch, + torch.tensor([1, 1], dtype=torch.int32, device=device), + torch.tensor([2, 1], dtype=torch.int32, device=device), + last_sampled, + next_prefill_tokens, + input_temperature, + input_seeds, + block_table, + 4, + 0, + 1, + 1, + 123, + num_speculative_steps, + num_speculative_steps, + max_num_reqs, + max_num_tokens, + 128, + sample_from_anchor=True, + ) + torch.accelerator.synchronize() + + # Query rows follow compact batch order, but every persistent state lookup + # follows idx_mapping instead of accidentally using the compact row index. + assert input_buffers.input_ids[:6].cpu().tolist() == [99, 123, 123, 77, 123, 123] + assert input_buffers.positions[:6].cpu().tolist() == [12, 13, 14, 23, 24, 25] + assert sample_indices[:6].cpu().tolist() == [0, 1, 2, 3, 4, 5] + assert sample_idx_mapping[:6].cpu().tolist() == [3, 3, 3, 1, 1, 1] + assert temperature.cpu().tolist() == [0.0, 0.5, 0.0, 1.0] + assert seeds.cpu().tolist() == [0, 11, 0, 33] + assert not input_buffers.is_padding[:6].any() + assert input_buffers.is_padding[6:].all() + def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp(): out = _run_prepare( @@ -174,3 +286,40 @@ def test_prepare_dflash_inputs_never_writes_the_null_block(): PAD_SLOT_ID, PAD_SLOT_ID, ] + + +def test_dflash_forward_context_receives_draft_padding_mask(monkeypatch): + device = torch.device("cuda") + input_buffers = SimpleNamespace( + input_ids=torch.tensor([11, 12, 0, 0], dtype=torch.int32, device=device), + positions=torch.tensor([7, 8, 0, 0], dtype=torch.int64, device=device), + is_padding=torch.tensor([False, False, True, True], device=device), + ) + observed = None + + @contextmanager + def fake_set_forward_context(*args, **kwargs): + nonlocal observed + observed = kwargs["is_padding"].clone() + yield + + monkeypatch.setattr( + dflash_speculator, "set_forward_context", fake_set_forward_context + ) + speculator = SimpleNamespace( + input_buffers=input_buffers, + vllm_config=SimpleNamespace(), + model=lambda **kwargs: kwargs["input_ids"], + ) + + result = DFlashSpeculator._run_model( + speculator, + num_tokens=4, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + ) + + assert result.tolist() == [11, 12, 0, 0] + assert observed is not None + assert observed.tolist() == [False, False, True, True] diff --git a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py index dd861d4ebcf7..c8332f021f1b 100644 --- a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py +++ b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py @@ -41,11 +41,13 @@ def can_use_dsv4_topk( def _dsv4_topk_kernel( gating_output_ptr, correction_bias_ptr, + is_padding_ptr, topk_weights_ptr, topk_ids_ptr, routed_scaling_factor, NUM_EXPERTS: tl.constexpr, BLOCK_N: tl.constexpr, + HAS_PADDING: tl.constexpr, launch_pdl: tl.constexpr, ): row = tl.program_id(0) @@ -89,6 +91,11 @@ def _dsv4_topk_kernel( output_mask = topk_offsets < 6 output_offsets = row * 6 + topk_offsets + if HAS_PADDING: + is_padding = tl.load(is_padding_ptr + row) + selected_weights = tl.where(is_padding, 0.0, selected_weights) + selected_ids = tl.where(is_padding, -1, selected_ids) + if launch_pdl: tl.extra.cuda.gdc_launch_dependents() @@ -101,8 +108,17 @@ def dsv4_topk( correction_bias: torch.Tensor, indices_dtype: torch.dtype, routed_scaling_factor: float, + is_padding: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens, num_experts = gating_output.shape + if is_padding is not None: + assert is_padding.dtype == torch.bool + assert is_padding.shape == (num_tokens,) + assert is_padding.device == gating_output.device + assert is_padding.is_contiguous() + assert indices_dtype in (torch.int32, torch.int64), ( + "Padding requires a signed indices dtype for the -1 sentinel." + ) shape = (num_tokens, _TOPK) topk_weights = gating_output.new_empty(shape, dtype=torch.float32) topk_ids = gating_output.new_empty(shape, dtype=indices_dtype) @@ -110,11 +126,13 @@ def dsv4_topk( _dsv4_topk_kernel[(num_tokens,)]( gating_output, correction_bias, + is_padding, topk_weights, topk_ids, routed_scaling_factor, NUM_EXPERTS=num_experts, BLOCK_N=triton.next_power_of_2(num_experts), + HAS_PADDING=is_padding is not None, num_warps=1, launch_pdl=current_platform.is_arch_support_pdl(), ) diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index f1c263615645..1b10d8cf1040 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -149,6 +149,7 @@ def fused_topk_bias( e_score_correction_bias, output_indices_dtype, routed_scaling_factor, + is_padding=_get_padding_mask(gating_output.shape[0]), ) M, _ = hidden_states.size() diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 34133e4f9204..c3e399b3c038 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -235,6 +235,7 @@ def _run_model( num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, + is_padding=self.input_buffers.is_padding[:num_tokens], ): last_hidden_states = self.model( input_ids=self.input_buffers.input_ids[:num_tokens], @@ -313,6 +314,11 @@ def _build_draft_attn_metadata( dcp_local_seq_lens=dcp_local_seq_lens, ) + def _get_graph_dispatch_shape( + self, num_reqs: int, num_query_tokens: int + ) -> tuple[int, int]: + return num_reqs, num_query_tokens + @torch.inference_mode() def materialize_context_kv( self, @@ -462,10 +468,13 @@ def propose( return self.draft_tokens[:num_reqs] # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs + dispatch_num_reqs, dispatch_num_tokens = self._get_graph_dispatch_shape( + num_reqs, num_query_tokens + ) batch_desc, batch_sync = dispatch_cg_and_sync_dp( self.query_cudagraph_manager, - num_reqs, - num_query_tokens, + dispatch_num_reqs, + dispatch_num_tokens, uniform_token_count=self.num_query_per_req, dp_size=self.dp_size, dp_rank=self.dp_rank, @@ -518,6 +527,7 @@ def _prepare_dflash_inputs_kernel( # Outputs out_input_ids_ptr, out_query_positions_ptr, + out_is_padding_ptr, out_query_start_loc_ptr, out_seq_lens_ptr, out_query_slot_mapping_ptr, @@ -646,10 +656,10 @@ def _prepare_dflash_inputs_kernel( local_q_slot, PAD_SLOT_ID, ) - tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query) + tl.store(out_is_padding_ptr + query_idx, False, mask=is_query) tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query) # --- Sample indices / positions / idx_mapping --- @@ -704,12 +714,15 @@ def _prepare_dflash_inputs_kernel( tl.store(out_sample_pos_ptr + block, 0, mask=mask) tl.store(out_sample_idx_mapping_ptr + block, -1, mask=mask) # Pad query slot mappings past num_query_tokens with PAD so the - # captured CG sees PAD slots (no K/V write) for replay sizes + # captured CG sees deterministic, inert rows for replay sizes # larger than the current request count. q_pad_start = num_reqs * num_query_per_req for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) mask = block < max_num_tokens + tl.store(out_input_ids_ptr + block, 0, mask=mask) + tl.store(out_query_positions_ptr + block, 0, mask=mask) + tl.store(out_is_padding_ptr + block, True, mask=mask) tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) @@ -761,6 +774,7 @@ def prepare_dflash_inputs( _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( input_buffers.input_ids, input_buffers.positions, + input_buffers.is_padding, input_buffers.query_start_loc, input_buffers.seq_lens, query_slot_mapping, diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 3935f42644d2..9314e2f0e8d8 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -23,6 +23,7 @@ backbone forward AND the sequential Markov sampling. """ +import os from typing import Any import torch @@ -87,6 +88,36 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): and not self.speculative_config.is_dspark_prefill_only() ) + fixed_graph_num_reqs = int(os.getenv("VLLM_DSPARK_FIXED_GRAPH_NUM_REQS", "0")) + if not 0 <= fixed_graph_num_reqs <= self.max_num_reqs: + raise ValueError( + "VLLM_DSPARK_FIXED_GRAPH_NUM_REQS must be between 0 and " + f"max_num_seqs ({self.max_num_reqs}), got {fixed_graph_num_reqs}." + ) + self.fixed_graph_num_reqs = fixed_graph_num_reqs + if fixed_graph_num_reqs: + logger.warning( + "Fixing DSpark draft CUDA graph batches at %d requests (%d tokens).", + fixed_graph_num_reqs, + fixed_graph_num_reqs * self.num_query_per_req, + ) + + def _get_graph_dispatch_shape( + self, num_reqs: int, num_query_tokens: int + ) -> tuple[int, int]: + if not self.fixed_graph_num_reqs: + return num_reqs, num_query_tokens + if num_reqs > self.fixed_graph_num_reqs: + raise RuntimeError( + "DSpark draft batch has " + f"{num_reqs} requests, exceeding the fixed CUDA graph batch " + f"of {self.fixed_graph_num_reqs}." + ) + return ( + self.fixed_graph_num_reqs, + self.fixed_graph_num_reqs * self.num_query_per_req, + ) + def load_draft_model( self, target_model: torch.nn.Module, From d5686d648eec9ba8ec2b6cfc93709db1c8d29863 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:59:40 -0700 Subject: [PATCH 03/13] feat(dspark): support pipeline-parallel targets in aggregated serving Extend the PD prefill-only DSpark-under-PP support to aggregated (IFB) serving: the drafter runs wholly on the last pipeline stage with a draft-local PP=1 config in both modes. - SpeculativeConfig.use_dspark_last_stage_drafter(): target-PP>1 forces draft PP=1 regardless of KV role (was kv_producer-only), fixing the draft model config's PP verification. - model_runner: allow the dspark/dflash aux-hidden-state spec path under PP when the drafter is last-stage (previously ValueError). - load_dspark_model: a PPMissingLayer target embed (PP first-stage layer) no longer aliases into the draft; the draft keeps its own embedding. - DSpark draft load_weights: under PP, load the shared embed.weight from the checkpoint (the checkpoint ships no mtp embed), since the target's table is unavailable on the last stage. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/config/test_dspark_prefill_only.py | 38 +++++++++++++++++++ vllm/config/speculative.py | 12 +++++- vllm/models/deepseek_v4/nvidia/dspark.py | 16 ++++++-- vllm/v1/worker/gpu/model_runner.py | 5 ++- .../v1/worker/gpu/spec_decode/dspark/utils.py | 15 +++++++- 5 files changed, 78 insertions(+), 8 deletions(-) diff --git a/tests/config/test_dspark_prefill_only.py b/tests/config/test_dspark_prefill_only.py index 554e0d89fcc1..01679f5a07c6 100644 --- a/tests/config/test_dspark_prefill_only.py +++ b/tests/config/test_dspark_prefill_only.py @@ -49,3 +49,41 @@ def test_dspark_prefill_materializer_uses_pp1_draft_config(): assert draft.pipeline_parallel_size == 1 assert draft.tensor_parallel_size == 1 + + +def _spec_config_no_kv( + *, pp: int, method: SpeculativeMethod = "dspark" +) -> SpeculativeConfig: + config = object.__new__(SpeculativeConfig) + config.method = method + config.target_parallel_config = ParallelConfig(pipeline_parallel_size=pp) + config.target_kv_transfer_config = None # type: ignore[assignment] + return config + + +@pytest.mark.parametrize( + ("pp", "method", "expected"), + [ + (2, "dspark", True), + (4, "dspark", True), + (1, "dspark", False), + (2, "dflash", False), + (2, "mtp", False), + ], +) +def test_dspark_last_stage_drafter_aggregated(pp, method, expected): + # Aggregated (IFB) serving: no KV transfer config at all. + assert _spec_config_no_kv(pp=pp, method=method).use_dspark_last_stage_drafter() is ( + expected + ) + + +def test_dspark_last_stage_drafter_covers_pd_roles(): + # kv_producer (prefill-only) and kv_consumer/aggregated all take the + # last-stage drafter path once the target is pipeline-parallel. + assert ( + _spec_config(pp=2, role="kv_producer").use_dspark_last_stage_drafter() is True + ) + assert ( + _spec_config(pp=2, role="kv_consumer").use_dspark_last_stage_drafter() is True + ) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 6720707de835..c042e0f65a12 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1495,7 +1495,7 @@ def __post_init__(self): self.target_parallel_config, self.draft_tensor_parallel_size, pipeline_parallel_size=( - 1 if self.is_dspark_prefill_only() else None + 1 if self.use_dspark_last_stage_drafter() else None ), ) ) @@ -1730,6 +1730,16 @@ def is_dspark_prefill_only(self) -> bool: and kv_transfer_config.kv_role == "kv_producer" ) + def use_dspark_last_stage_drafter(self) -> bool: + # A DSpark drafter under a pipeline-parallel target runs wholly on the + # last pipeline stage, so the draft model always uses PP=1. Holds for + # both the PD prefill-only producer and aggregated (IFB) serving. + return ( + self.method == "dspark" + and self.target_parallel_config is not None + and self.target_parallel_config.pipeline_parallel_size > 1 + ) + @field_validator("attention_backend", mode="before") @classmethod def _parse_attention_backend(cls, value: Any) -> Any: diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index bd8daa1839d4..a07f86f2b8ab 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -21,6 +21,7 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( @@ -478,11 +479,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: head_start = n_local_head * tp_rank head_end = n_local_head * (tp_rank + 1) + # Under pipeline parallelism the drafter only exists on the last stage, + # where the target's embedding table is a PPMissingLayer placeholder, so + # the draft loads its own copy of the shared embedding weight. + load_own_embed = get_pp_group().world_size > 1 for name, loaded_weight in weights: - mapped = self._remap_dspark_name(name) - if mapped is None: - continue - name = mapped + if load_own_embed and name == "embed.weight": + name = "model.embed_tokens.weight" + else: + mapped = self._remap_dspark_name(name) + if mapped is None: + continue + name = mapped if "confidence_head." in name: loaded_confidence_head = True diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index aa357642b0cb..247403319743 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -273,7 +273,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ): # Drafting may require auxiliary hidden states from target model outputs self.use_aux_hidden_state_outputs = True - if self.use_pp and not self.dspark_prefill_only: + if ( + self.use_pp + and not self.speculative_config.use_dspark_last_stage_drafter() + ): raise ValueError( f"{self.speculative_config.method} with pipeline parallel " "is not supported." diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 0db7fc7526a3..d97f37f6dc8e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -37,7 +37,10 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_model_config = speculative_config.draft_model_config pp_group = get_pp_group() - if pp_group.world_size != 1 and not speculative_config.is_dspark_prefill_only(): + if ( + pp_group.world_size != 1 + and not speculative_config.use_dspark_last_stage_drafter() + ): raise NotImplementedError("DSpark does not support pipeline parallelism.") target_language_model = ( @@ -62,7 +65,10 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo from vllm.compilation.backends import set_model_tag from vllm.model_executor.model_loader import get_model from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal - from vllm.model_executor.models.utils import get_draft_quant_config + from vllm.model_executor.models.utils import ( + PPMissingLayer, + get_draft_quant_config, + ) from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( _should_share, get_target_lm_head, @@ -104,6 +110,11 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo if not speculative_config.is_dspark_prefill_only(): target_embed = getattr(target_inner, "embed_tokens", None) + if isinstance(target_embed, PPMissingLayer): + # Under PP the target's embedding table only exists on the first + # pipeline stage; the last-stage drafter keeps and loads its own + # copy from the checkpoint instead. + target_embed = None draft_embed = getattr(draft_inner, "embed_tokens", None) if ( target_embed is not None From 153d74cd3513538713b54a3c6ed549baf7fb83b2 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Thu, 3 Sep 2026 06:15:42 -0700 Subject: [PATCH 04/13] fix(pp): make warmup deadlock-free under pipeline parallelism Warmup runs synthetic steps whose sampled outputs are discarded, but the PP sampled-token broadcast still enqueued side-stream NCCL ops that interlocked with the next step's activation p2p and deadlocked startup. Disabling the broadcast for the warmup window removes that interlock, but it also removes the only warmup coverage of the deferred post-update path (update_pp_decode_requests -> post_update). Its first triton compile then happened mid-serving on non-last ranks, where the in-flight broadcast recv keeps a NCCL kernel spinning on the device and blocks the CUDA module load: the last rank waits for activations the first rank never sends, and the pipeline deadlocks. - PPHandler: add a disabled flag gating receive/broadcast. - Worker: disable the broadcast for compile_or_warm_up_model and restore it before serving. - ModelRunner.warmup_pp_decode_update(): launch post_update with an all -1 idx_mapping (a no-op) during warmup_kernels so the kernel is compiled while no NCCL op is in flight. - PPHandler: pad the broadcast (2, N) combined buffer so its unbind views stay 16-byte aligned for any num_reqs. Triton specializes on pointer alignment, and the previously misaligned num_rejected view (num_reqs % 4 != 0) compiled a second post_update variant at serving time despite the warmup, recreating the same deadlock. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/v1/worker/test_pp_utils.py | 68 +++++++++++++++++++++++++++++- vllm/v1/worker/gpu/model_runner.py | 30 +++++++++++++ vllm/v1/worker/gpu/pp_utils.py | 35 +++++++++++++-- vllm/v1/worker/gpu/warmup.py | 5 +++ vllm/v1/worker/gpu_worker.py | 12 ++++++ 5 files changed, 146 insertions(+), 4 deletions(-) diff --git a/tests/v1/worker/test_pp_utils.py b/tests/v1/worker/test_pp_utils.py index cfc92479f57e..f76639a0557e 100644 --- a/tests/v1/worker/test_pp_utils.py +++ b/tests/v1/worker/test_pp_utils.py @@ -5,8 +5,9 @@ from unittest.mock import Mock import numpy as np +import torch -from vllm.v1.worker.gpu import pp_utils +from vllm.v1.worker.gpu import model_runner, pp_utils def _batch(num_computed, prefill_len, num_scheduled): @@ -81,3 +82,68 @@ def test_decode_row_ahead_of_a_prefill_chunk(): assert mask is not None assert mask.tolist() == [True, False] + + +def test_disabled_handler_skips_broadcast_and_receive(monkeypatch): + """While disabled (warmup), neither side enqueues a broadcast op.""" + sent = [] + monkeypatch.setattr( + pp_utils.torch.distributed, + "broadcast", + lambda *args, **kwargs: sent.append((args, kwargs)), + ) + + handler = object.__new__(pp_utils.PPHandler) + handler.set_disabled(True) + + handler.is_last_rank = False + assert handler.receive(Mock()) is False + + handler.is_last_rank = True + assert handler.broadcast(Mock(), Mock(), Mock(), Mock()) is None + + assert sent == [] + + handler.set_disabled(False) + assert handler.disabled is False + + +def test_alloc_combined_keeps_unbind_views_16_byte_aligned(): + """Triton specializes on pointer alignment: an unaligned `num_rejected` + would compile a second `_post_update_kernel` variant at serving time, + where the in-flight broadcast NCCL kernel can block the module load.""" + for num_reqs in range(1, 9): + combined = pp_utils._alloc_combined(num_reqs, torch.device("cpu")) + num_sampled, num_rejected = combined.unbind(dim=0) + assert num_sampled.data_ptr() % 16 == 0 + assert num_rejected.data_ptr() % 16 == 0 + assert combined.shape[1] >= num_reqs + + +def test_warmup_pp_decode_update_matches_serving_specialization(monkeypatch): + """The warmup launch must hit the same triton specialization as serving. + + A mismatch means the first real ``update_pp_decode_requests`` recompiles + mid-serving, where the in-flight broadcast NCCL kernel blocks the CUDA + module load and deadlocks the pipeline. + """ + calls = [] + monkeypatch.setattr(model_runner, "post_update", lambda *args: calls.append(args)) + + runner = object.__new__(model_runner.GPUModelRunner) + runner.device = torch.device("cpu") + runner.pp_handler = Mock(max_sample_len=3) + runner.req_states = Mock() + + runner.warmup_pp_decode_update() + + assert len(calls) == 1 + args = calls[0] + idx_mapping, _, _, output_bin_counts = args[:4] + sampled_tokens, num_sampled, num_rejected, query_start_loc = args[4:8] + assert idx_mapping.tolist() == [-1] and idx_mapping.dtype == torch.int64 + assert output_bin_counts is None + assert query_start_loc is None + assert sampled_tokens.shape == (1, 3) and sampled_tokens.dtype == torch.int64 + assert num_sampled.dtype == torch.int32 + assert num_rejected.dtype == torch.int32 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 247403319743..7be52562bef6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1050,6 +1050,36 @@ def update_pp_decode_requests(self): if outputs is not None: self.postprocess_sampled(**outputs) + def warmup_pp_decode_update(self) -> None: + """JIT-compile the kernel behind ``update_pp_decode_requests``. + + That path only runs on real steps, so the warmup steps never reach it + on non-last PP ranks. Its first triton compile must not happen + mid-serving: the in-flight sampled-token broadcast keeps a NCCL kernel + spinning on this device, which blocks the CUDA module load and + deadlocks the pipeline. An all -1 idx_mapping makes this a no-op. + The freshly allocated int32 tensors are 16-byte aligned, matching the + padded views `PPHandler` produces at serving time (triton specializes + on pointer alignment). + """ + assert self.pp_handler is not None + post_update( + torch.full((1,), -1, dtype=torch.int64, device=self.device), + self.req_states.num_computed_tokens.gpu, + self.req_states.last_sampled_tokens, + None, + torch.zeros( + (1, self.pp_handler.max_sample_len), + dtype=torch.int64, + device=self.device, + ), + torch.zeros(1, dtype=torch.int32, device=self.device), + torch.zeros(1, dtype=torch.int32, device=self.device), + None, + self.req_states.all_token_ids.gpu, + self.req_states.total_len.gpu, + ) + def add_requests(self, scheduler_output: SchedulerOutput) -> None: for new_req_data in scheduler_output.scheduled_new_reqs: assert new_req_data.prefill_token_ids is not None diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index cf52a6d3821e..790e6f24d957 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -45,6 +45,18 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: return produces_sample if produces_sample.any() else None +def _alloc_combined(num_reqs: int, device: torch.device) -> torch.Tensor: + """Allocate the (2, N) int32 buffer broadcast alongside sampled tokens. + + The inner dim is padded to a multiple of 4 so that both `unbind` views + stay 16-byte aligned for any `num_reqs`: triton specializes on pointer + alignment, and a misaligned `num_rejected` would JIT-compile a second + `_post_update_kernel` variant at serving time. The padding is broadcast + but never read. Sender and receiver must both use this allocation. + """ + return torch.empty(2, -(-num_reqs // 4) * 4, dtype=torch.int32, device=device) + + class PPHandler: """Runs the PP sampled-token broadcast/recv on a side stream so the default stream isn't gated by the matching peer call. Step T's recv is @@ -83,6 +95,14 @@ def __init__( group_desc="pp_broadcast" ) + # Warmup steps run the pipeline with synthetic batches whose outputs are + # discarded; the sampled-token broadcast is disabled there so its + # side-stream NCCL ops cannot overlap the next step's activation p2p. + self.disabled = False + + def set_disabled(self, disabled: bool) -> None: + self.disabled = disabled + def on_req_idx_freed(self, req_idx: int) -> None: self.req_idx_gen_np[req_idx] += 1 @@ -123,6 +143,8 @@ def receive(self, input_batch: InputBatch) -> bool: """Returns True iff sampled tokens need to be gathered from *all* requests in the batch.""" assert not self.is_last_rank + if self.disabled: + return False need_sampled_mask = compute_need_sampled_mask(input_batch) if need_sampled_mask is None: # Leave this step's reserved slot as None. @@ -138,7 +160,7 @@ def receive(self, input_batch: InputBatch) -> bool: sampled_tokens = torch.empty( num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device ) - combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device) + combined = _alloc_combined(num_reqs, self.device) torch.distributed.broadcast( sampled_tokens, src=self.last_rank, group=self.broadcast_group ) @@ -171,11 +193,16 @@ def broadcast( input_batch: InputBatch, ) -> None: assert self.is_last_rank - if compute_need_sampled_mask(input_batch) is None: + if self.disabled: + return + mask = compute_need_sampled_mask(input_batch) + if mask is None: # No request needs sampled outputs for a subsequent decode step. return assert sampled_token_ids.dtype == torch.int64 + assert num_sampled.dtype == torch.int32 + assert num_rejected.dtype == torch.int32 if current_platform.is_xpu(): self.main_stream.synchronize() @@ -187,7 +214,9 @@ def broadcast( src=self.last_rank, group=self.broadcast_group, ) - combined = torch.stack((num_sampled, num_rejected), dim=0) + combined = _alloc_combined(num_sampled.shape[0], self.device) + combined[0, : num_sampled.shape[0]] = num_sampled + combined[1, : num_sampled.shape[0]] = num_rejected torch.distributed.broadcast( combined, src=self.last_rank, group=self.broadcast_group ) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index f99c924f26d7..b56156f631c0 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -420,6 +420,11 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: for step_indices, step_spec_flags in decode_steps: _run_decode_step(step_indices, step_spec_flags) + # The deferred PP post-update path only runs on real steps, so the steps + # above never JIT-compile its kernel on non-last ranks. + if model_runner.pp_handler is not None and not model_runner.is_last_pp_rank: + model_runner.warmup_pp_decode_update() + # Clean up - process finish_req_ids. cleanup_output = SchedulerOutput.make_empty() cleanup_output.finished_req_ids = set(req_ids) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 9ee217936c4a..1447b3aa43d3 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -752,6 +752,15 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: @instrument(span_name="Warmup (GPU)") def compile_or_warm_up_model(self) -> CompilationTimes: + # All warmup phases below run synthetic steps whose sampled outputs are + # discarded. The PP sampled-token broadcast would carry no payload, and + # its side-stream NCCL ops can overlap the next step's activation p2p + # and deadlock the pipeline, so keep it disabled for the whole warmup + # window and restore it before serving. + pp_handler = getattr(self.model_runner, "pp_handler", None) + if pp_handler is not None: + pp_handler.set_disabled(True) + warmup_sizes: list[int] = [] if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE: @@ -926,6 +935,9 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # intra-op parallelism. set_torch_threads_for_runtime() + if pp_handler is not None: + pp_handler.set_disabled(False) + return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time, From 9aace31d6ecfefc217c1bdbd817048496d777db5 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:17:00 -0700 Subject: [PATCH 05/13] fix(pp): complete the sampled-token broadcast contract for spec decoding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aggregated (IFB) PP + last-stage-drafter serving stalled after the first prefill: broadcast() sent the sampler output at its natural width — [num_reqs, 1] from the plain sampler on steps without drafts — while receive() always expects [num_reqs, max_sample_len]. A NCCL broadcast with mismatched counts never completes on the receiver, wedging the first stage's broadcast stream at the final prefill chunk and cascading into the main-stream wait_event two steps later. And once the hang is gone, earlier stages would still embed token id 0 for every draft slot: nothing ever wrote req_states.draft_tokens off the last rank, since broadcast() runs before propose() and never carried drafts. Verification logits for those positions would be garbage. - broadcast(): pad the plain-sampler [N, 1] payload to the constant [N, max_sample_len] wire shape (consumers read at most num_sampled entries per row). - broadcast_drafts(): after propose(), send a clone of the fresh drafts (the speculator overwrites its persistent buffer next step), stream- ordered after the sampled-token sends. - receive(): enqueue the matching draft recv and carry it in PendingRecv; post_update() adopts the broadcast drafts into req_states.draft_tokens on non-last ranks. - warmup_pp_decode_update(): cover the draft-writing specialization. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/v1/worker/test_pp_utils.py | 84 +++++++++++++++++++++++++++++- vllm/v1/worker/gpu/input_batch.py | 26 +++++++++ vllm/v1/worker/gpu/model_runner.py | 18 +++++++ vllm/v1/worker/gpu/pp_utils.py | 62 ++++++++++++++++++++++ 4 files changed, 188 insertions(+), 2 deletions(-) diff --git a/tests/v1/worker/test_pp_utils.py b/tests/v1/worker/test_pp_utils.py index f76639a0557e..8dfdf2e6d480 100644 --- a/tests/v1/worker/test_pp_utils.py +++ b/tests/v1/worker/test_pp_utils.py @@ -5,17 +5,35 @@ from unittest.mock import Mock import numpy as np +import pytest import torch from vllm.v1.worker.gpu import model_runner, pp_utils -def _batch(num_computed, prefill_len, num_scheduled): +def _cuda_handler(max_sample_len=6): + handler = object.__new__(pp_utils.PPHandler) + handler.is_last_rank = True + handler.disabled = False + handler.max_sample_len = max_sample_len + handler.last_rank = 1 + handler.broadcast_group = Mock() + handler.device = torch.device("cuda") + handler.main_stream = torch.cuda.current_stream() + handler.broadcast_stream = torch.cuda.Stream() + return handler + + +def _batch(num_computed, prefill_len, num_scheduled, idx_mapping=None): + num_reqs = len(num_computed) + if idx_mapping is None: + idx_mapping = list(range(num_reqs)) return Mock( - num_reqs=len(num_computed), + num_reqs=num_reqs, num_computed_tokens_np=np.array(num_computed, dtype=np.int32), prefill_len_np=np.array(prefill_len, dtype=np.int32), num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32), + idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64), ) @@ -141,9 +159,71 @@ def test_warmup_pp_decode_update_matches_serving_specialization(monkeypatch): args = calls[0] idx_mapping, _, _, output_bin_counts = args[:4] sampled_tokens, num_sampled, num_rejected, query_start_loc = args[4:8] + broadcast_drafts, draft_tokens_out = args[10:12] assert idx_mapping.tolist() == [-1] and idx_mapping.dtype == torch.int64 assert output_bin_counts is None assert query_start_loc is None assert sampled_tokens.shape == (1, 3) and sampled_tokens.dtype == torch.int64 assert num_sampled.dtype == torch.int32 assert num_rejected.dtype == torch.int32 + # Spec-enabled PP handlers receive drafts over the broadcast; the warmup + # must compile that specialization (non-None draft pointers) too. + assert broadcast_drafts.shape == (1, 2) and broadcast_drafts.dtype == torch.int64 + assert draft_tokens_out is runner.req_states.draft_tokens + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream") +def test_broadcast_pads_plain_sampler_rows_to_max_sample_len(monkeypatch): + """The wire shape must not depend on whether the batch carried drafts: + the receiver always allocates [num_reqs, max_sample_len], and a NCCL + broadcast with mismatched counts hangs the receiver.""" + sent = [] + monkeypatch.setattr( + pp_utils.torch.distributed, + "broadcast", + lambda tensor, **kw: sent.append(tensor), + ) + handler = _cuda_handler() + batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1]) + + handler.broadcast( + torch.zeros(1, 1, dtype=torch.int64, device="cuda"), # plain sampler + torch.ones(1, dtype=torch.int32, device="cuda"), + torch.zeros(1, dtype=torch.int32, device="cuda"), + batch, + ) + + assert sent[0].shape == (1, 6) + assert sent[1].shape == (2, 4) + torch.accelerator.synchronize() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream") +def test_broadcast_drafts_gathers_fresh_rows_from_the_table(monkeypatch): + sent = [] + monkeypatch.setattr( + pp_utils.torch.distributed, + "broadcast", + lambda tensor, **kw: sent.append(tensor), + ) + handler = _cuda_handler() + # The batch's single row maps to request-state row 2. + batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1]) + batch.idx_mapping = torch.tensor([2], dtype=torch.int64, device="cuda") + table = torch.arange(20, dtype=torch.int64, device="cuda").view(4, 5) + + handler.broadcast_drafts(table, batch) + + assert sent[0].shape == (1, 5) + # The payload is a gather into a fresh tensor: propose() overwrites its + # persistent buffer on the next step, possibly before this send completes. + assert sent[0].data_ptr() != table.data_ptr() + assert sent[0].cpu().tolist() == [table[2].cpu().tolist()] + + # An all-prefill batch sends nothing (receive() enqueues nothing either). + sent.clear() + prefill_batch = _batch(num_computed=[0], prefill_len=[4096], num_scheduled=[448]) + prefill_batch.idx_mapping = prefill_batch.idx_mapping.cuda() + handler.broadcast_drafts(table, prefill_batch) + assert sent == [] + torch.accelerator.synchronize() diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 2a38c3d71061..9348dc1a2cb1 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -551,6 +551,11 @@ def _post_update_kernel( all_token_ids_ptr, all_token_ids_stride, total_len_ptr, + broadcast_drafts_ptr, + broadcast_drafts_stride, + draft_tokens_ptr, + draft_tokens_stride, + num_spec, ): req_id = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_id) @@ -558,6 +563,17 @@ def _post_update_kernel( # Filter rows with negative index entries. return + if broadcast_drafts_ptr is not None: + # PP path: adopt the draft tokens proposed by the last rank's + # speculator so the next verification step embeds the real drafts. + for i in range(num_spec): + token_id = tl.load( + broadcast_drafts_ptr + req_id * broadcast_drafts_stride + i + ) + tl.store( + draft_tokens_ptr + req_state_idx * draft_tokens_stride + i, token_id + ) + total_len = tl.load(total_len_ptr + req_state_idx) num_sampled = tl.load(num_sampled_ptr + req_id) if num_sampled > 0: @@ -618,6 +634,11 @@ def post_update( all_token_ids: torch.Tensor, # [max_num_reqs] total_len: torch.Tensor, + # [num_reqs, num_spec]; drafts broadcast from the last PP rank. Only + # passed on non-last PP ranks, which never run the speculator. + broadcast_drafts: torch.Tensor | None = None, + # [max_num_reqs, num_spec] + draft_tokens_out: torch.Tensor | None = None, ) -> None: num_reqs = idx_mapping.shape[0] _post_update_kernel[(num_reqs,)]( @@ -634,6 +655,11 @@ def post_update( all_token_ids, all_token_ids.stride(0), total_len, + broadcast_drafts, + broadcast_drafts.stride(0) if broadcast_drafts is not None else 0, + draft_tokens_out, + draft_tokens_out.stride(0) if draft_tokens_out is not None else 0, + broadcast_drafts.shape[1] if broadcast_drafts is not None else 0, num_warps=1, ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7be52562bef6..ba1cd68be707 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1063,6 +1063,12 @@ def warmup_pp_decode_update(self) -> None: on pointer alignment). """ assert self.pp_handler is not None + num_spec = self.pp_handler.max_sample_len - 1 + broadcast_drafts = ( + torch.zeros((1, num_spec), dtype=torch.int64, device=self.device) + if num_spec > 0 + else None + ) post_update( torch.full((1,), -1, dtype=torch.int64, device=self.device), self.req_states.num_computed_tokens.gpu, @@ -1078,6 +1084,8 @@ def warmup_pp_decode_update(self) -> None: None, self.req_states.all_token_ids.gpu, self.req_states.total_len.gpu, + broadcast_drafts, + self.req_states.draft_tokens if broadcast_drafts is not None else None, ) def add_requests(self, scheduler_output: SchedulerOutput) -> None: @@ -1554,6 +1562,7 @@ def postprocess_sampled( num_sampled: torch.Tensor, num_rejected: torch.Tensor, query_start_loc: torch.Tensor | None = None, + broadcast_drafts: torch.Tensor | None = None, ) -> None: # Update the number of computed tokens. if self.is_last_pp_rank: @@ -1572,6 +1581,8 @@ def postprocess_sampled( query_start_loc, self.req_states.all_token_ids.gpu, self.req_states.total_len.gpu, + broadcast_drafts, + self.req_states.draft_tokens if broadcast_drafts is not None else None, ) self.model_state.postprocess_state( @@ -2046,6 +2057,13 @@ def sample_tokens( ) if draft_tokens is not None: self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens + if self.pp_handler is not None: + # Earlier stages never run the speculator; ship the + # drafts so their next verification step embeds the real + # draft tokens instead of stale buffer contents. + self.pp_handler.broadcast_drafts( + self.req_states.draft_tokens, input_batch + ) if draft_tokens is not None and self.adaptive_verification is not None: self.adaptive_verification.record_confidences( self.speculator.draft_token_confidence_probs, input_batch diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 790e6f24d957..32f53b9416c3 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -30,6 +30,10 @@ class PendingRecv: # Snapshot of slot generation counters at receive time, used to # detect requests aborted since then. gen_at_receive_np: np.ndarray # [num_reqs] + # Draft tokens proposed by the last rank's speculator for this step, + # broadcast separately after propose() runs. None when spec decoding + # is disabled (max_sample_len == 1). + draft_tokens: torch.Tensor | None = None # [num_reqs, num_spec] def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: @@ -137,6 +141,7 @@ def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: num_sampled=slot.num_sampled, num_rejected=slot.num_rejected, idx_mapping=idx_mapping, + broadcast_drafts=slot.draft_tokens, ) def receive(self, input_batch: InputBatch) -> bool: @@ -167,6 +172,20 @@ def receive(self, input_batch: InputBatch) -> bool: torch.distributed.broadcast( combined, src=self.last_rank, group=self.broadcast_group ) + draft_tokens = None + if self.max_sample_len > 1: + # The sender enqueues this send in broadcast_drafts() after + # propose(); NCCL stream order keeps the recvs matched. + draft_tokens = torch.empty( + num_reqs, + self.max_sample_len - 1, + dtype=torch.int64, + device=self.device, + ) + torch.distributed.broadcast( + draft_tokens, src=self.last_rank, group=self.broadcast_group + ) + draft_tokens.record_stream(self.main_stream) event = self.broadcast_stream.record_event() num_sampled, num_rejected = combined.unbind(dim=0) # Must record_stream since these were allocated on broadcast stream but @@ -182,6 +201,7 @@ def receive(self, input_batch: InputBatch) -> bool: input_batch.idx_mapping_np, need_sampled_mask, gen_at_receive_np, + draft_tokens, ) return bool(need_sampled_mask.all()) @@ -209,6 +229,17 @@ def broadcast( with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) + if sampled_token_ids.shape[1] != self.max_sample_len: + # The plain sampler emits [num_reqs, 1] on steps without + # drafts while the rejection sampler emits the full + # [num_reqs, max_sample_len]; the receiver always expects the + # latter. Pad — consumers read at most num_sampled per row. + assert sampled_token_ids.shape[1] == 1 + padded = sampled_token_ids.new_zeros( + sampled_token_ids.shape[0], self.max_sample_len + ) + padded[:, :1] = sampled_token_ids + sampled_token_ids = padded torch.distributed.broadcast( sampled_token_ids.contiguous(), src=self.last_rank, @@ -222,3 +253,34 @@ def broadcast( ) for tensor in (sampled_token_ids, num_sampled, num_rejected): tensor.record_stream(self.broadcast_stream) + + def broadcast_drafts( + self, draft_token_table: torch.Tensor, input_batch: InputBatch + ) -> None: + """Broadcast the speculator's freshly proposed draft tokens. + + Runs after propose() on the last rank; the send is stream-ordered + after broadcast()'s sends, matching receive()'s enqueue order. The + payload is gathered from the runner's draft table (just updated from + propose()'s output) into a fresh compact tensor: the speculator + overwrites its own persistent buffer on the next step, possibly + before this async send completes. + """ + assert self.is_last_rank + if self.disabled or self.max_sample_len == 1: + return + if compute_need_sampled_mask(input_batch) is None: + return + assert draft_token_table.dtype == torch.int64 + assert draft_token_table.shape[1] == self.max_sample_len - 1 + + # Gather on the main stream so the payload is ordered after this + # step's table update and before any later one. + drafts = draft_token_table[input_batch.idx_mapping] + assert drafts.shape[0] == input_batch.num_reqs + with torch.cuda.stream(self.broadcast_stream): + self.broadcast_stream.wait_stream(self.main_stream) + torch.distributed.broadcast( + drafts, src=self.last_rank, group=self.broadcast_group + ) + drafts.record_stream(self.broadcast_stream) From 0f851e2c5dd33140db427e6511a570cde722036d Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:21:12 -0700 Subject: [PATCH 06/13] fix(dspark): address review findings on connector check, block drop, DP pin - SpeculativeConfig: detect NIXL via KVTransferConfig.has_connector() so a NixlConnector nested in a MultiConnector is also rejected for DSpark prefill-only materialization. - Scheduler: make use_eagle_block_drop follow the narrowed use_eagle, so a DSpark prefill-only producer no longer drops its trailing prefix-cache block for a read-ahead it never performs. - DFlashSpeculator.propose: when the fixed DSpark graph batch re-pins the dispatch shape, rebuild the reused DP sync's token count to match; dispatch_cg_and_sync_dp asserts on the mismatch for dp_size > 1. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- vllm/config/speculative.py | 3 +-- vllm/v1/core/sched/scheduler.py | 7 ++++++- vllm/v1/worker/gpu/spec_decode/dflash/speculator.py | 9 +++++++++ 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 385e1da37c08..f4d254c7f420 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1522,8 +1522,7 @@ def __post_init__(self): if self.is_dspark_prefill_only(): # is_dspark_prefill_only() implies target_kv_transfer_config is set. assert self.target_kv_transfer_config is not None - connector = self.target_kv_transfer_config.kv_connector - if connector is not None and "Nixl" in connector: + if self.target_kv_transfer_config.has_connector("NixlConnector"): raise NotImplementedError( "DSpark prefill materialization with pipeline parallelism " "requires a connector that transfers named per-layer KV " diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2737a19ac543..05ef045d1166 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -283,7 +283,12 @@ def __init__( if speculative_config.use_multi_module_mtp() else 1 ) - self.use_eagle_block_drop = speculative_config.use_eagle_block_drop() + # Follow the narrowed use_eagle: a DSpark prefill-only producer + # never does the EAGLE-style read-ahead, so it must not drop its + # trailing prefix-cache block either. + self.use_eagle_block_drop = ( + self.use_eagle and speculative_config.use_eagle_block_drop() + ) if self.use_eagle and not self.use_eagle_block_drop: logger.warning( "EAGLE trailing prefix-cache block dropping is disabled. " diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 4bb90245f092..d01d6458b219 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -466,6 +466,15 @@ def propose( dispatch_num_reqs, dispatch_num_tokens = self._get_graph_dispatch_shape( num_reqs, num_batch_tokens ) + if batch_sync is not None and dispatch_num_tokens != num_batch_tokens: + # The dispatch shape was re-pinned (fixed DSpark graph batch): keep + # the reused DP agreement consistent with the dispatched count. + batch_sync = replace( + batch_sync, + num_tokens_across_dp=torch.full_like( + batch_sync.num_tokens_across_dp, dispatch_num_tokens + ), + ) batch_desc, batch_sync = dispatch_cg_and_sync_dp( self.query_cudagraph_manager, dispatch_num_reqs, From d882cdf964424fcc93a1e1a4156228932cd53215 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Fri, 4 Sep 2026 09:01:53 -0700 Subject: [PATCH 07/13] feat(dspark): support Kimi-K3 targets with cross-stage aux hidden taps K3's DSpark speculator taps aux hidden states at layers spread across the whole model, so under pipeline parallelism earlier stages must forward their captured taps to the last stage where the drafter runs. At HEAD this crashes at startup: load_dspark_model reads dspark_target_layer_ids directly, which speculators-format draft configs (e.g. K3) do not define. - load_dspark_model: resolve aux taps via the shared get_eagle3_aux_layers_from_config fallback chain (value-identical for DeepSeek-V4's dspark_target_layer_ids) and let targets opt out of the last-stage-only tap restriction via supports_pp_aux_hidden_state_transport. - KimiLinearModel: pack aux taps captured on earlier stages into IntermediateTensors and unpack them on later stages; stage-entry capture stays first-stage-only so boundary taps are not captured twice. - tests: cover cross-stage transport ordering, dedup, and buffer width. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/models/kimi_k3/test_eagle3.py | 112 ++++++++++++++++++ vllm/models/kimi_k3/nvidia/model.py | 56 +++++++-- .../v1/worker/gpu/spec_decode/dspark/utils.py | 28 +++-- 3 files changed, 176 insertions(+), 20 deletions(-) diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 858622fa07da..fb62921cd4b2 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -200,3 +200,115 @@ def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch assert got_pending is layer_hidden_states assert got_residual is block_residual torch.testing.assert_close(aux_hidden_states[0], captured) + + +def _make_stage( + *, + start_layer: int, + taps: tuple[int, ...], + layer_outputs: list[tuple[torch.Tensor, None, torch.Tensor]], +) -> KimiLinearModel: + model = _make_kimi_linear_model() + end_layer = start_layer + len(layer_outputs) + object.__setattr__(model, "start_layer", start_layer) + object.__setattr__(model, "end_layer", end_layer) + # The real model keeps the global layer list and slices [start:end]. + layers = [Mock() for _ in range(end_layer)] + for i, out in enumerate(layer_outputs): + layers[start_layer + i] = Mock(return_value=out) + object.__setattr__(model, "layers", layers) + object.__setattr__(model, "aux_hidden_state_layers", taps) + object.__setattr__(model, "config", SimpleNamespace(hidden_size=2)) + return model + + +def test_kimi_linear_aux_hidden_states_flow_across_pp_stages(monkeypatch): + """A tap owned by an earlier PP stage must reach the last stage intact. + + The drafter's taps can reference layers outside the last stage (K3 taps + [24, 48, 72, 88, 92]); each stage packs the taps it owns into + IntermediateTensors and the last stage returns the full ordered set. + """ + stage0_hidden = torch.tensor([[1.0, 2.0]]) + stage0_residual = torch.tensor([[3.0, 4.0]]) + stage1_hidden = torch.tensor([[5.0, 6.0]]) + stage1_residual = torch.tensor([[7.0, 8.0]]) + + stage0 = _make_stage( + start_layer=0, + taps=(1, 2), + layer_outputs=[(stage0_hidden, None, stage0_residual)], + ) + stage1 = _make_stage( + start_layer=1, + taps=(1, 2), + layer_outputs=[(stage1_hidden, None, stage1_residual)], + ) + + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=False), + ) + stage0_out = stage0.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=torch.zeros(1, 2), + ) + + # Stage 0 owns the post-layer-1 tap; it is packed for the wire. + stage0_aux = stage0_hidden + stage0_residual + torch.testing.assert_close(stage0_out.tensors["aux_hidden_states"], stage0_aux) + + # The receiving buffer on stage 1 must be sized for exactly that one tap. + stage1_buffers = stage1.make_empty_intermediate_tensors( + batch_size=1, dtype=torch.bfloat16, device=torch.device("cpu") + ) + assert stage1_buffers.tensors["aux_hidden_states"].shape == (1, 2) + + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=False, is_last_rank=True), + ) + output, aux_hidden_states = stage1.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=stage0_out, + ) + + # The boundary tap (position 1 == stage1's start_layer) must not be + # duplicated by the stage-entry capture: two taps, in ascending order. + assert len(aux_hidden_states) == 2 + torch.testing.assert_close(aux_hidden_states[0], stage0_aux) + torch.testing.assert_close(aux_hidden_states[1], stage1_hidden + stage1_residual) + torch.testing.assert_close(output, stage1_hidden + stage1_residual) + + +def test_kimi_linear_first_stage_without_taps_sends_no_aux_buffer(monkeypatch): + """No taps at or below the stage boundary -> no aux key on the wire.""" + stage0 = _make_stage( + start_layer=0, + taps=(2,), + layer_outputs=[(torch.ones(1, 2), None, torch.zeros(1, 2))], + ) + assert ( + "aux_hidden_states" + not in stage0.make_empty_intermediate_tensors( + batch_size=1, dtype=torch.bfloat16, device=torch.device("cpu") + ).tensors + ) + + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=False), + ) + out = stage0.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=torch.zeros(1, 2), + ) + assert "aux_hidden_states" not in out.tensors diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 2fc20ecf2198..8d834aabb499 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1116,6 +1116,10 @@ def forward( class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): + # Aux hidden taps ride the pipeline packed inside IntermediateTensors, so a + # DSpark drafter on the last stage may tap target layers on any stage. + supports_pp_aux_hidden_state_transport = True + packed_modules_mapping = { "gate_up_proj": ["gate_proj", "up_proj"], "in_proj_qkvgfab": ["q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj"], @@ -1220,14 +1224,25 @@ def make_empty_intermediate_tensors( cdiv(self.start_layer, self.attn_res_block_size), self.config.hidden_size, ) - return IntermediateTensors( - { - "hidden_states": torch.zeros( - (batch_size, self.config.hidden_size), dtype=dtype, device=device - ), - "residual": torch.zeros(residual_shape, dtype=dtype, device=device), - } - ) + tensors: dict[str, torch.Tensor] = { + "hidden_states": torch.zeros( + (batch_size, self.config.hidden_size), dtype=dtype, device=device + ), + "residual": torch.zeros(residual_shape, dtype=dtype, device=device), + } + # Aux taps at or below this stage's first layer are captured by earlier + # stages and arrive packed along the feature dim. + num_incoming_aux = sum( + layer <= self.start_layer + for layer in getattr(self, "aux_hidden_state_layers", ()) + ) + if num_incoming_aux: + tensors["aux_hidden_states"] = torch.zeros( + (batch_size, num_incoming_aux * self.config.hidden_size), + dtype=dtype, + device=device, + ) + return IntermediateTensors(tensors) def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: super()._set_aux_hidden_state_layers(layers) @@ -1325,10 +1340,14 @@ def forward( else: hidden_states = self.embed_input_ids(input_ids) residual = None + packed_aux_hidden_states = None else: assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] + packed_aux_hidden_states = intermediate_tensors.tensors.get( + "aux_hidden_states" + ) assert hidden_states is not None full_num_tokens = positions.shape[0] @@ -1341,9 +1360,19 @@ def forward( hidden_states = sp_shard(hidden_states) assert residual is None, "Currently, SP is not supported with PP" - # sharded aux hidden states when sp is enabled aux_hidden_states: list[torch.Tensor] = [] - if self.start_layer in self.aux_hidden_state_layers: + if packed_aux_hidden_states is not None: + # Unpack the taps earlier stages captured (ascending layer order). + aux_hidden_states.extend( + packed_aux_hidden_states.split(self.config.hidden_size, dim=-1) + ) + # The stage-entry capture doubles the previous stage's own last-layer + # tap once aux states flow across stages, so it only makes sense on the + # first stage (e.g. a tap on the embedding output). + if ( + get_pp_group().is_first_rank + and self.start_layer in self.aux_hidden_state_layers + ): if self.use_attn_res or residual is None: aux_hidden_states.append(hidden_states) else: @@ -1393,9 +1422,10 @@ def forward( ) if prefix_sum is not None: hidden_states = hidden_states + prefix_sum - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) + tensors = {"hidden_states": hidden_states, "residual": residual} + if aux_hidden_states: + tensors["aux_hidden_states"] = torch.cat(aux_hidden_states, dim=-1) + return IntermediateTensors(tensors) if self.use_attn_res: assert prefix_sum is not None diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index d97f37f6dc8e..07fb181013b2 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -50,16 +50,30 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo ) target_inner = target_language_model.model if pp_group.world_size != 1: + from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, + ) + start_layer = target_inner.start_layer end_layer = target_inner.end_layer - aux_layers = tuple( - layer + 1 for layer in draft_model_config.hf_config.dspark_target_layer_ids - ) - if any(layer <= start_layer or layer > end_layer for layer in aux_layers): + # Same fallback chain the target model's aux taps are configured with: + # covers dspark_target_layer_ids (DeepSeek-V4) and the speculators-format + # fields (eagle_aux_hidden_state_layer_ids / target_layer_ids, e.g. K3). + aux_layers = get_eagle3_aux_layers_from_config(speculative_config) + if aux_layers is None: + raise ValueError( + "DSpark draft config declares no auxiliary target-layer taps." + ) + # Models with PP aux transport (e.g. Kimi-K3) forward taps captured on + # earlier stages to the last stage, so their taps may live anywhere. + if not getattr( + target_inner, "supports_pp_aux_hidden_state_transport", False + ) and any(layer <= start_layer or layer > end_layer for layer in aux_layers): raise ValueError( - "DSpark prefill materialization requires every auxiliary hidden " - "state on the last pipeline stage. " - f"Stage owns ({start_layer}, {end_layer}], requested {aux_layers}." + "DSpark prefill materialization requires every auxiliary " + "hidden state on the last pipeline stage. " + f"Stage owns ({start_layer}, {end_layer}], requested " + f"{aux_layers}." ) from vllm.compilation.backends import set_model_tag From 472875917c8b7c7f7764c74018e2cfbf198f4d64 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sun, 6 Sep 2026 10:10:07 -0700 Subject: [PATCH 08/13] fix(dspark): let PP drafters load their own embedding when the target's is stranded Under pipeline parallelism the target's embed_tokens lives on the first stage, so a drafter on a later stage cannot alias it. K3 and DSv4 DSpark checkpoints ship their own embed_tokens copy, but the weight mapper dropped it and maybe_share_target_embed raised once sharing was refused. - DSpark K3: build a VocabParallelEmbedding under PP, keep the embed weight in the mapper (drop_embed=False), and alias only outside PP. - Add a loads_own_embed_under_pp capability flag (DSv4 + K3 DSpark) so maybe_share_target_embed only raises when the drafter neither aliases nor loads its own table under PP. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/models/kimi_k3/test_dspark_mla.py | 23 +++++++ tests/models/kimi_k3/test_eagle3.py | 4 +- .../test_spec_decode_embed_sharing_pp.py | 26 ++++++++ vllm/models/deepseek_v4/nvidia/dspark.py | 5 +- vllm/models/kimi_k3/nvidia/dspark_mla.py | 65 +++++++++++++++---- vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 11 +++- 6 files changed, 114 insertions(+), 20 deletions(-) create mode 100644 tests/models/kimi_k3/test_dspark_mla.py diff --git a/tests/models/kimi_k3/test_dspark_mla.py b/tests/models/kimi_k3/test_dspark_mla.py new file mode 100644 index 000000000000..0791a5183ead --- /dev/null +++ b/tests/models/kimi_k3/test_dspark_mla.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K3 DSpark draft weight-mapping tests.""" + +from vllm.models.kimi_k3.nvidia.dspark_mla import _build_weights_mapper + + +def test_mapper_drops_embed_for_target_aliasing(): + mapper = _build_weights_mapper(drop_embed=True) + assert mapper.apply_list(["embed_tokens.weight"]) == [] + assert mapper.apply_list(["lm_head.weight"]) == [] + # Draft-owned weights still map into the model namespace. + assert mapper.apply_list(["layers.0.mlp.gate_proj.weight"]) == [ + "model.layers.0.mlp.gate_up_proj.weight" + ] + + +def test_mapper_keeps_embed_under_pp(): + # Under PP the drafter cannot alias the target's first-stage table, so the + # checkpoint's own embed_tokens.weight must flow through. + mapper = _build_weights_mapper(drop_embed=False) + assert mapper.apply_list(["embed_tokens.weight"]) == ["model.embed_tokens.weight"] + assert mapper.apply_list(["lm_head.weight"]) == [] diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index f293a502e7fa..4313aaa3f625 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -266,9 +266,7 @@ def test_kimi_linear_aux_hidden_states_flow_across_pp_stages(monkeypatch): # Stage 0 owns the post-layer-1 tap; it rides the wire under its global # slot key. stage0_aux = stage0_hidden + stage0_residual - torch.testing.assert_close( - stage0_out.tensors["aux_hidden_states_0"], stage0_aux - ) + torch.testing.assert_close(stage0_out.tensors["aux_hidden_states_0"], stage0_aux) monkeypatch.setattr( kimi_model, diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py index 1a1d96a6b4ef..586363f54b68 100644 --- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py +++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py @@ -63,6 +63,32 @@ def test_missing_target_embedding_raises_instead_of_running_on_garbage(monkeypat ) +def test_pp_drafter_loading_own_embedding_keeps_it(monkeypatch): + """DSv4/K3-style drafters load embed_tokens from their own checkpoint when + PP strands the target's table on the first stage; they must neither alias + nor raise.""" + monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) + draft_embed = _embed(fill=1.0) + draft_inner = _inner(draft_embed) + draft = SimpleNamespace(has_own_embed_tokens=False, loads_own_embed_under_pp=True) + + eagle_utils.maybe_share_target_embed(draft, draft_inner, _inner(PPMissingLayer())) + + assert draft_inner.embed_tokens is draft_embed + + +def test_pp_drafter_without_any_embedding_still_raises(monkeypatch): + monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) + draft_inner = _inner(None) + draft_inner.embed_tokens = None + draft = SimpleNamespace(has_own_embed_tokens=False, loads_own_embed_under_pp=True) + + with pytest.raises(RuntimeError, match="needs the target input embedding"): + eagle_utils.maybe_share_target_embed( + draft, draft_inner, _inner(PPMissingLayer()) + ) + + def test_drafter_with_distinct_weights_keeps_them(monkeypatch): monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) draft_embed = _embed(fill=1.0) diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index a07f86f2b8ab..451460fd4956 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -347,8 +347,11 @@ def _insert_context_kv( class DSparkDeepseekV4ForCausalLM(nn.Module): # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so - # load_dspark_model always aliases the target's. + # load_dspark_model aliases the target's — except under PP, where the + # target's table sits on the first stage and the drafter loads its own + # copy of the shared embed weight (see load_weights). has_own_embed_tokens = False + loads_own_embed_under_pp = True has_own_lm_head = False # Full-vocab draft: draft ids are target ids, no remapping needed. draft_id_to_target_id = None diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 97febda0c86e..2e243dc6c8b8 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -9,12 +9,19 @@ import vllm._custom_ops as ops from vllm.config import VllmConfig +from vllm.distributed.parallel_state import ( + get_pp_group, + model_parallel_is_initialized, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, ReplicatedLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -29,6 +36,12 @@ from vllm.v1.worker.workspace import current_workspace_manager +def _target_pp_world_size() -> int: + if not model_parallel_is_initialized(): + return 1 + return get_pp_group().world_size + + def _duplicate_context_kv_weights( weights: Iterable[tuple[str, torch.Tensor]], num_layers: int ) -> Iterable[tuple[str, torch.Tensor]]: @@ -137,8 +150,18 @@ def __init__( self.config = vllm_config.speculative_config.draft_model_config.hf_config self.quant_config = get_draft_quant_config(vllm_config) - # The frozen target embedding is aliased after the draft checkpoint loads. + # The frozen target embedding is aliased after the draft checkpoint + # loads. Under pipeline parallelism that table exists only on the + # first stage while the drafter runs on the last, so the draft builds + # its own table and loads embed_tokens.weight from its checkpoint + # (the K3 DSpark checkpoint always ships it). self.embed_tokens: nn.Module | None = None + if _target_pp_world_size() > 1: + self.embed_tokens = VocabParallelEmbedding( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) self.context_proj = ReplicatedLinear( self.config.target_hidden_size * self.config.num_target_layers, @@ -395,18 +418,19 @@ def forward( return hidden_states -class K3DSparkForCausalLM(nn.Module): - has_own_embed_tokens = False - has_own_lm_head = False - draft_id_to_target_id = None - hf_to_vllm_mapper = WeightsMapper( - # confidence_head is training-only. The frozen target embedding and LM - # head are shared after this draft-specific checkpoint is loaded. - orig_to_new_substr={ - "confidence_head": None, - "embed_tokens": None, - "lm_head": None, - }, +def _build_weights_mapper(*, drop_embed: bool) -> WeightsMapper: + # confidence_head is training-only. The frozen target LM head is shared + # after this draft-specific checkpoint is loaded; the embedding is shared + # too, except under pipeline parallelism where the drafter cannot reach + # the first-stage table and loads its own copy instead. + orig_to_new_substr = { + "confidence_head": None, + "lm_head": None, + } + if drop_embed: + orig_to_new_substr["embed_tokens"] = None + return WeightsMapper( + orig_to_new_substr=orig_to_new_substr, orig_to_new_prefix={"": "model."}, orig_to_new_stacked={ ".gate_proj": (".gate_up_proj", 0), @@ -416,6 +440,17 @@ class K3DSparkForCausalLM(nn.Module): }, ) + +class K3DSparkForCausalLM(nn.Module): + # The checkpoint ships embed_tokens.weight but no lm_head: the embedding + # is aliased from the target, except under PP where the drafter builds + # and loads its own table (the target's lives on the first stage). + has_own_embed_tokens = False + loads_own_embed_under_pp = True + has_own_lm_head = False + draft_id_to_target_id = None + hf_to_vllm_mapper = _build_weights_mapper(drop_embed=True) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() assert vllm_config.speculative_config is not None @@ -427,6 +462,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: start_layer_id=target_layer_num, prefix=maybe_prefix(prefix, "model"), ) + if _target_pp_world_size() > 1: + # The draft built its own embedding table; keep the checkpoint's + # embed_tokens.weight mapping instead of dropping it. + self.hf_to_vllm_mapper = _build_weights_mapper(drop_embed=False) # Assigned by load_dspark_model from the target. Keeping no placeholder # avoids a transient full-vocabulary allocation for this 163k-vocab model. diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 05f771ab8298..57e8f1a20b97 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -54,9 +54,14 @@ def maybe_share_target_embed( return if target_embed is None: - if hasattr(draft_inner, "embed_tokens") and not getattr( - draft_model, "has_own_embed_tokens", False - ): + # A drafter whose checkpoint ships the embedding can load its own copy + # on this stage (e.g. DeepSeek-V4/Kimi-K3 DSpark under PP, where the + # target's table lives on the first stage while the drafter runs on + # the last). Anything else would run on an uninitialized table. + loads_own = getattr(draft_model, "has_own_embed_tokens", False) or getattr( + draft_model, "loads_own_embed_under_pp", False + ) + if draft_embed is None or not loads_own: raise RuntimeError( f"{type(draft_model).__name__} needs the target input embedding, " "but it is unavailable on this PP stage" From 1bb49fbc933eca06f36ca73e2f74af201c4ca6f8 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sun, 6 Sep 2026 10:10:13 -0700 Subject: [PATCH 09/13] fix(dspark): address review findings on topk uint32 dispatch and context-KV construction - fused_topk_bias_router: compute the padding mask before dispatch and exclude the uint32+padding combination from the dsv4_topk fast path, whose -1 sentinel requires signed logits. - DSpark context-KV-only layers: construct the attention submodule directly via _select_dsv4_attn_cls instead of instantiating a whole DeepseekV4DecoderLayer, avoiding a transient MoE allocation that could OOM on memory-tight GPUs. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/kernels/moe/test_topk_softplus_sqrt.py | 54 +++++++++++++++++++ .../router/fused_topk_bias_router.py | 6 ++- vllm/models/deepseek_v4/nvidia/dspark.py | 29 ++++++---- 3 files changed, 79 insertions(+), 10 deletions(-) diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index d078d1ac530c..a1741596e6fc 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -280,6 +280,60 @@ def test_dsv4_fast_topk( ) +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="The DeepSeek V4 fast path is CUDA-only.", +) +def test_dsv4_fast_topk_padding_uint32_falls_back(monkeypatch: pytest.MonkeyPatch): + """Padded rows need the -1 sentinel, which uint32 cannot represent: the + router must skip the dsv4 fast path and still route the real rows.""" + torch.manual_seed(0) + num_tokens = 17 + num_experts = 256 + hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda") + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") + is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda") + is_padding[1::2] = True + gating_output[is_padding] = float("nan") + + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.router." + "fused_topk_bias_router._get_padding_mask", + lambda _: is_padding, + ) + # uint32 + padding would trip dsv4_topk's signed-indices assertion; the + # generic path must take over instead. + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=gating_output, + scoring_func="sqrtsoftplus", + e_score_correction_bias=correction_bias, + topk=6, + renormalize=True, + indices_type=torch.uint32, + routed_scaling_factor=1.5, + ) + + assert topk_ids.dtype == torch.uint32 + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( + gating_output=gating_output[~is_padding], + topk=6, + renormalize=True, + routed_scaling_factor=1.5, + e_score_correction_bias=correction_bias, + ) + # uint32 CUDA tensors do not support boolean-mask indexing; widen first. + torch.testing.assert_close( + topk_ids.to(torch.int64)[~is_padding], topk_ids_ref.to(torch.int64), atol=0, rtol=0 + ) + torch.testing.assert_close( + topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5 + ) + + @pytest.mark.skipif( not current_platform.is_cuda(), reason="The DeepSeek V4 fast path is CUDA-only.", diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index 8576caf47095..d464d73ba17d 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -148,9 +148,13 @@ def fused_topk_bias( ) output_indices_dtype = torch.int32 if indices_type is None else indices_type + padding_mask = _get_padding_mask(gating_output.shape[0]) if ( scoring_func == "sqrtsoftplus" and hash_indices_table is None + # dsv4_topk marks padded rows with a -1 sentinel, which unsigned + # indices cannot represent; keep the generic path for that pair. + and (padding_mask is None or output_indices_dtype != torch.uint32) and can_use_dsv4_topk( gating_output, e_score_correction_bias, @@ -165,7 +169,7 @@ def fused_topk_bias( e_score_correction_bias, output_indices_dtype, routed_scaling_factor, - is_padding=_get_padding_mask(gating_output.shape[0]), + is_padding=padding_mask, input_ids=input_tokens, bias_vl=bias_vl, image_sentinel_lo=image_sentinel_lo, diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 451460fd4956..dcf09b9706a1 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -53,6 +53,7 @@ from .model import ( DeepseekV4DecoderLayer, DeepseekV4Model, + _select_dsv4_attn_cls, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, ) @@ -87,9 +88,9 @@ def __init__(self, attn: nn.Module) -> None: class _DSparkContextKVLayer(nn.Module): """A draft layer stripped down to its context-KV projection.""" - def __init__(self, layer: DeepseekV4DecoderLayer) -> None: + def __init__(self, attn: nn.Module) -> None: super().__init__() - self.attn = _DSparkContextKVAttention(layer.attn) + self.attn = _DSparkContextKVAttention(attn) class DSparkDeepseekV4Model(nn.Module): @@ -141,20 +142,30 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: layers: list[nn.Module] = [] for i in range(self.num_dspark_layers): layer_prefix = maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}") - layer = DeepseekV4DecoderLayer( - current_vllm_config, - prefix=layer_prefix, - topk_indices_buffer=self.topk_indices_buffer, - ) if self.context_kv_only: + # Only the attention submodule participates in context-KV + # materialization; building the full decoder layer would + # transiently allocate its MoE experts, norms, and + # hyper-connection parameters on the device. + attn = _select_dsv4_attn_cls(current_vllm_config)( + current_vllm_config, + prefix=f"{layer_prefix}.attn", + topk_indices_buffer=self.topk_indices_buffer, + ) # The full attention object registers itself for metadata lookup, # but only its SWA cache layer participates in materialization. current_vllm_config.compilation_config.static_forward_context.pop( f"{layer_prefix}.attn", None ) - layers.append(_DSparkContextKVLayer(layer)) + layers.append(_DSparkContextKVLayer(attn)) else: - layers.append(layer) + layers.append( + DeepseekV4DecoderLayer( + current_vllm_config, + prefix=layer_prefix, + topk_indices_buffer=self.topk_indices_buffer, + ) + ) self.layers = nn.ModuleList(layers) if not self.context_kv_only: From 5b572da027bd363c7271ef326b1f943b1fe7a87c Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sun, 6 Sep 2026 15:46:10 -0700 Subject: [PATCH 10/13] fix(pp): do not double-post draft broadcasts when a speculator ran The merge of upstream #50514 kept two broadcast_drafts call sites on the last PP rank: one in the propose() path and one next to set_draft_tokens(). With a speculator present both fire on the same step, so the last rank posts two draft broadcasts while earlier stages post one recv. The extra send shifts the pp_broadcast FIFO: the next step's sampled-token recv pairs with a leftover draft broadcast of a different size, and the mismatched NCCL collective spins on-device, blocking any subsequent context-level CUDA call (module load, malloc) and deadlocking the pipeline. Only broadcast from the handler path when there is no speculator (e.g. diffusion-style drafts); when a speculator ran, propose() already shipped the fresh drafts. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- vllm/v1/worker/gpu/model_runner.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c5d0f33b54cf..1556c1ab2b9c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2134,7 +2134,11 @@ def sample_tokens( input_batch, self.req_states.draft_tokens[input_batch.idx_mapping], ) - if self.pp_handler is not None: + if self.pp_handler is not None and self.speculator is None: + # When a speculator ran, the propose() path above already + # broadcast the fresh drafts. Broadcasting here as well would + # double-post on the pp_broadcast group and misalign the + # recv FIFO on earlier stages, hanging the pipeline. self.pp_handler.broadcast_drafts( self.req_states.draft_tokens, input_batch ) From 0a2f85917fa0f074b9f0bb7bee28fd89c70259dc Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 7 Sep 2026 03:37:24 -0700 Subject: [PATCH 11/13] fix(k3): keep the DSpark draft's marker from flagging the target's KV group MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MLAAttentionSpec.merge ORs non_causal_multi_token_decode (since #55234), so when the K3-native DSpark draft's marked MLA layers merge with the target's identical-geometry MLA layers, the target's KV cache group reports non-causal multi-token capability as well. On TritonMLA that raises the group's reorder_batch_threshold to the spec block length: fresh short prefills and the target's causal verification blocks get misrouted into a decode path that expects one query row per request, producing NaN logits — output collapses to a repeated token from the very first step, and the drafter's NaN aux inputs drive acceptance to zero. Tag the draft layers' spec with a distinct model_version so the drafter keeps its own flagged group and the target group stays causal. Verified on Kimi-K3-pruned75 + Inferact/Kimi-K3-DSpark (SM120): tp8 eager and pp2tp4 with CUDA graphs both produce sane output (acceptance 0.85/0.64 per position); the graph-capture illegal memory access is gone as well. RedHatAI-format draft configuration unchanged and still green. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- vllm/models/kimi_k3/nvidia/mla.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 87bf59749c1d..4201d890354f 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -430,6 +430,15 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), # fp8_ds_mla: 656-byte custom layout; see flashmla_sparse.py. state_content_bytes=656 if self.kv_cache_dtype == "fp8_ds_mla" else None, + # Keep the non-causal DSpark draft out of the target's KV cache + # group: MLAAttentionSpec.merge ORs non_causal_multi_token_decode, + # so a merged group would flag the causal target too, raising its + # TritonMLA reorder threshold and misrouting its short prefills and + # causal verification blocks into a decode path that expects one + # query row per request. + model_version="kimi_k3_dspark" + if self.non_causal_multi_token_decode + else None, non_causal_multi_token_decode=self.non_causal_multi_token_decode, ) From 2962de54e06552dd0b21f8321d284b9fdc62d12b Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 7 Sep 2026 06:04:44 -0700 Subject: [PATCH 12/13] fix: address CI lint findings - pp_utils: drop the stale duplicate broadcast_drafts definition left by the merge (the later definition, with the disabled/short-circuit guards, is the live one). - test_topk_softplus_sqrt: wrap an over-length assert_close line. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/kernels/moe/test_topk_softplus_sqrt.py | 5 ++++- vllm/v1/worker/gpu/pp_utils.py | 15 --------------- 2 files changed, 4 insertions(+), 16 deletions(-) diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index a1741596e6fc..ae7d73780d83 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -327,7 +327,10 @@ def test_dsv4_fast_topk_padding_uint32_falls_back(monkeypatch: pytest.MonkeyPatc ) # uint32 CUDA tensors do not support boolean-mask indexing; widen first. torch.testing.assert_close( - topk_ids.to(torch.int64)[~is_padding], topk_ids_ref.to(torch.int64), atol=0, rtol=0 + topk_ids.to(torch.int64)[~is_padding], + topk_ids_ref.to(torch.int64), + atol=0, + rtol=0, ) torch.testing.assert_close( topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5 diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index f6b553a583a3..c4296be7fafd 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -181,21 +181,6 @@ def get_prev_sampled_outputs( broadcast_drafts=slot.draft_tokens, ) - def broadcast_drafts( - self, draft_tokens: torch.Tensor, input_batch: InputBatch - ) -> None: - """Broadcast draft proposals so non-last ranks can embed real token ids.""" - assert self.is_last_rank - if compute_need_sampled_mask(input_batch) is None: - return - with torch.cuda.stream(self.broadcast_stream): - self.broadcast_stream.wait_stream(self.main_stream) - send = draft_tokens[input_batch.idx_mapping].contiguous() - torch.distributed.broadcast( - send, src=self.last_rank, group=self.broadcast_group - ) - send.record_stream(self.broadcast_stream) - def receive(self, input_batch: InputBatch) -> bool: """Returns True iff sampled tokens need to be gathered from *all* requests in the batch.""" From be79b490dd01a5f884ec3add3957469b61abb893 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 14 Sep 2026 20:49:33 -0700 Subject: [PATCH 13/13] Fix test fakes broken by DSpark PP additions - warmup: short-circuit on is_last_pp_rank before reading pp_handler so SimpleNamespace runner stubs without the attribute keep working. - test_gpu_model_runner_v2: the QSA runner stub bypasses __init__, so set dspark_prefill_only explicitly. Co-authored-by: Kimi Code Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tests/v1/worker/test_gpu_model_runner_v2.py | 1 + vllm/v1/worker/gpu/warmup.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 86ff1a074538..52d500b6eb57 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -44,6 +44,7 @@ def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch): num_new_sampled_tokens_per_step=1, ) runner.speculator = None + runner.dspark_prefill_only = False runner.req_states = [] runner.input_buffers = SimpleNamespace(query_start_loc=None) runner.vocab_size = 1 diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 31056db2ac21..2b7f6b6d494e 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -447,7 +447,7 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: # The deferred PP post-update path only runs on real steps, so the steps # above never JIT-compile its kernel on non-last ranks. - if model_runner.pp_handler is not None and not model_runner.is_last_pp_rank: + if not model_runner.is_last_pp_rank and model_runner.pp_handler is not None: model_runner.warmup_pp_decode_update() # Clean up - process finish_req_ids.