From a901a5dcc26c4cc25df1f318195c36cac557c7a8 Mon Sep 17 00:00:00 2001
From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Date: Mon, 14 Sep 2026 22:13:20 -0700
Subject: [PATCH 1/2] [DSpark] Support pipeline-parallel targets in aggregated
 serving

Draft tokens under PP: the last stage runs the speculator and now
broadcasts the fresh drafts to earlier stages (whose next verification
step would otherwise embed stale buffer contents), with the double-post
on the pp_broadcast group gated out for the speculator-less diffusion
path. Non-last stages JIT-compile the deferred post-update kernel during
warmup so its first compile cannot deadlock the pipeline mid-serving.

Draft embedding under PP: the target's embedding table lives on the
first stage, so DeepSeek-V4/Kimi-K3 DSpark drafters load their own copy
from the checkpoint (loads_own_embed_under_pp) instead of aliasing.

Padded graph batch safety: the DFlash prepare-inputs kernel now clears
input_ids/positions and sets is_padding on CUDA-graph padding rows, and
the DSv4 top-k router zeroes padded-row selections instead of reading
uninitialized state.

Co-authored-by: Kimi Code <noreply@moonshot.cn>
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
---
 tests/kernels/moe/test_topk_softplus_sqrt.py  | 130 +++++++++++++++
 tests/models/kimi_k3/test_dspark_mla.py       |  23 +++
 tests/models/kimi_k3/test_eagle3.py           | 108 +++++++++++++
 .../spec_decode/test_dflash_prepare_inputs.py | 149 +++++++++++++++++
 tests/v1/worker/test_pp_utils.py              | 152 +++++++++++++++++-
 .../test_spec_decode_embed_sharing_pp.py      |  26 +++
 .../layers/fused_moe/router/dsv4_topk.py      |  18 +++
 .../router/fused_topk_bias_router.py          |   5 +
 vllm/models/deepseek_v4/nvidia/dspark.py      |  21 ++-
 vllm/models/kimi_k3/nvidia/dspark_mla.py      |  65 ++++++--
 vllm/v1/worker/gpu/input_batch.py             |  26 +++
 vllm/v1/worker/gpu/model_runner.py            |  68 +++++++-
 vllm/v1/worker/gpu/pp_utils.py                |  85 +++++++---
 .../gpu/spec_decode/dflash/speculator.py      |   7 +
 vllm/v1/worker/gpu/spec_decode/eagle/utils.py |  11 +-
 vllm/v1/worker/gpu/warmup.py                  |   5 +
 vllm/v1/worker/gpu_worker.py                  |  12 ++
 17 files changed, 862 insertions(+), 49 deletions(-)
 create mode 100644 tests/models/kimi_k3/test_dspark_mla.py

diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py
index bc438c66931b..bb8856f83a4d 100644
--- a/tests/kernels/moe/test_topk_softplus_sqrt.py
+++ b/tests/kernels/moe/test_topk_softplus_sqrt.py
@@ -339,6 +339,136 @@ 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.",
+)
+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/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 858622fa07da..4313aaa3f625 100644
--- a/tests/models/kimi_k3/test_eagle3.py
+++ b/tests/models/kimi_k3/test_eagle3.py
@@ -200,3 +200,111 @@ 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 under global
+    per-tap keys (EagleModelMixin.pack_local_aux_hidden_states) and the last
+    stage prepends the collected remote taps to its own.
+    """
+    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)],
+    )
+    # EagleModelMixin caches the PP aux layout in _set_aux_hidden_state_layers;
+    # the stubs set layers directly, so prime the caches by hand.
+    object.__setattr__(stage0, "_aux_slot_base_cached", 0)
+    object.__setattr__(stage1, "_aux_slot_base_cached", 1)
+    object.__setattr__(stage1, "_aux_upstream_total_cached", 1)
+
+    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 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)
+
+    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 captured on the stage -> no aux keys on the wire."""
+    stage0 = _make_stage(
+        start_layer=0,
+        taps=(2,),
+        layer_outputs=[(torch.ones(1, 2), None, torch.zeros(1, 2))],
+    )
+    object.__setattr__(stage0, "_aux_slot_base_cached", 0)
+
+    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 not any(key.startswith("aux_hidden_states_") for key in out.tensors)
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/tests/v1/worker/test_pp_utils.py b/tests/v1/worker/test_pp_utils.py
index cfc92479f57e..8dfdf2e6d480 100644
--- a/tests/v1/worker/test_pp_utils.py
+++ b/tests/v1/worker/test_pp_utils.py
@@ -5,16 +5,35 @@
 from unittest.mock import Mock
 
 import numpy as np
+import pytest
+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):
+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),
     )
 
 
@@ -81,3 +100,130 @@ 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]
+    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/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/model_executor/layers/fused_moe/router/dsv4_topk.py b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py
index dbd29f3d08ab..1df95bc6119f 100644
--- a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py
+++ b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py
@@ -41,6 +41,7 @@ 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,
@@ -49,6 +50,7 @@ def _dsv4_topk_kernel(
         image_sentinel_lo,
         NUM_EXPERTS: tl.constexpr,
         BLOCK_N: tl.constexpr,
+        HAS_PADDING: tl.constexpr,
         HAS_VL: tl.constexpr,
         launch_pdl: tl.constexpr,
     ):
@@ -106,6 +108,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()
 
@@ -118,11 +125,20 @@ def dsv4_topk(
     correction_bias: torch.Tensor,
     indices_dtype: torch.dtype,
     routed_scaling_factor: float,
+    is_padding: torch.Tensor | None = None,
     input_ids: torch.Tensor | None = None,
     bias_vl: torch.Tensor | None = None,
     image_sentinel_lo: int = 0,
 ) -> 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."
+        )
     has_vl = bias_vl is not None and image_sentinel_lo > 0
     if bias_vl is not None:
         assert input_ids is not None, "bias_vl routing requires input_ids"
@@ -136,6 +152,7 @@ def dsv4_topk(
         _dsv4_topk_kernel[(num_tokens,)](
             gating_output,
             correction_bias,
+            is_padding,
             topk_weights,
             topk_ids,
             routed_scaling_factor,
@@ -144,6 +161,7 @@ def dsv4_topk(
             image_sentinel_lo,
             NUM_EXPERTS=num_experts,
             BLOCK_N=triton.next_power_of_2(num_experts),
+            HAS_PADDING=is_padding is not None,
             HAS_VL=has_vl,
             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 257ede27d5fa..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,6 +169,7 @@ def fused_topk_bias(
                 e_score_correction_bias,
                 output_indices_dtype,
                 routed_scaling_factor,
+                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 987909702cf4..c538a1201000 100644
--- a/vllm/models/deepseek_v4/nvidia/dspark.py
+++ b/vllm/models/deepseek_v4/nvidia/dspark.py
@@ -22,6 +22,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 (
@@ -346,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
@@ -475,11 +479,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         head_end = n_local_head * (tp_rank + 1)
 
         weights = _duplicate_context_wkv_weights(weights, len(self.model.layers))
+        # 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/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py
index 0baff884ccaa..39a9eeafd7bf 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 (
     DSparkConfidenceHead,
     DSparkMarkovHead,
@@ -36,6 +43,12 @@
 )
 
 
+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]]:
@@ -144,8 +157,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,
@@ -425,18 +448,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),
@@ -446,6 +470,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
@@ -463,6 +498,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/input_batch.py b/vllm/v1/worker/gpu/input_batch.py
index 32c083e45974..315549244489 100644
--- a/vllm/v1/worker/gpu/input_batch.py
+++ b/vllm/v1/worker/gpu/input_batch.py
@@ -564,6 +564,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)
@@ -571,6 +576,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:
@@ -631,6 +647,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,)](
@@ -647,6 +668,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 57f0180edc91..9d0fd47a9972 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -857,6 +857,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:
@@ -876,7 +877,8 @@ 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]]
             if isinstance(self.sampler, GPUWatermarkSampler):
                 self.speculator.prepare_watermarking(
                     self.sampler._get_contexts(input_batch.idx_mapping),
@@ -1110,6 +1112,44 @@ 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
+        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,
+            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,
+            broadcast_drafts,
+            self.req_states.draft_tokens if broadcast_drafts is not None else None,
+        )
+
     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
@@ -1596,6 +1636,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:
@@ -1614,6 +1655,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(
@@ -2140,7 +2183,10 @@ def sample_tokens(
             spec_hidden_states = draft_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[: draft_hidden_states.size(0)]
+                if pre_hc_hidden_states is not None:
+                    spec_hidden_states = pre_hc_hidden_states[
+                        : draft_hidden_states.size(0)
+                    ]
             if isinstance(self.sampler, GPUWatermarkSampler):
                 self.speculator.prepare_watermarking(
                     self.sampler._get_contexts(input_batch.idx_mapping),
@@ -2162,8 +2208,16 @@ def sample_tokens(
                     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 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
                 )
@@ -2175,7 +2229,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
                 )
diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py
index c2aaa1fba25a..862b86e5a107 100644
--- a/vllm/v1/worker/gpu/pp_utils.py
+++ b/vllm/v1/worker/gpu/pp_utils.py
@@ -46,6 +46,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
@@ -86,6 +98,14 @@ def __init__(
         )
         self.aux_hidden_state_relay_keys: tuple[str, ...] = ()
 
+        # 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
 
@@ -157,29 +177,15 @@ def get_prev_sampled_outputs(
             num_sampled=slot.num_sampled,
             num_rejected=slot.num_rejected,
             idx_mapping=idx_mapping,
+            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()
-            # Must record the idx_mapping tensor since it was allocated
-            # on the main stream.
-            input_batch.idx_mapping.record_stream(self.broadcast_stream)
-            torch.distributed.broadcast(
-                send, src=self.last_rank, group=self.broadcast_group
-            )
-
     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.
@@ -195,7 +201,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
             )
@@ -213,6 +219,7 @@ def receive(self, input_batch: InputBatch) -> bool:
                 torch.distributed.broadcast(
                     draft_tokens, src=self.last_rank, group=self.broadcast_group
                 )
+
             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
@@ -242,11 +249,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()
@@ -262,9 +274,42 @@ 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
             )
             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)
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
index a5710829e0ff..3f661bbda575 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
@@ -256,6 +256,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],
@@ -487,6 +488,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,
@@ -617,6 +619,7 @@ def _prepare_dflash_inputs_kernel(
     )
 
     tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
+    tl.store(out_is_padding_ptr + query_idx, False, 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_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
@@ -679,6 +682,9 @@ def _prepare_dflash_inputs_kernel(
             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)
 
 
@@ -730,6 +736,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/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
index 06bed9391e03..0cc85c6ffa03 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
@@ -55,9 +55,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"
diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py
index 56f7d438fc3f..2b7f6b6d494e 100644
--- a/vllm/v1/worker/gpu/warmup.py
+++ b/vllm/v1/worker/gpu/warmup.py
@@ -445,6 +445,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 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.
     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 3d4c5f3d5cc5..58af8211278e 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -805,6 +805,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 (
@@ -983,6 +992,9 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
         # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
         enable_gpu_sync_check()
 
+        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 b6cc5a2ca9c7058467da5a8fdf4a38e589e56d3e Mon Sep 17 00:00:00 2001
From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Date: Tue, 15 Sep 2026 04:09:58 -0700
Subject: [PATCH 2/2] Drop the draft self-embedding path in favor of #50514

On current main the target model retains its embedding table on the last
PP rank when speculative decoding is active (spec_decode_needs_target_embed
covers dspark), so the last-stage drafter aliases it via
maybe_share_target_embed. loads_own_embed_under_pp is redundant.

Co-authored-by: Kimi Code <noreply@moonshot.cn>
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           | 108 ------------------
 .../test_spec_decode_embed_sharing_pp.py      |  26 -----
 vllm/models/deepseek_v4/nvidia/dspark.py      |  21 +---
 vllm/models/kimi_k3/nvidia/dspark_mla.py      |  65 +++--------
 vllm/v1/worker/gpu/spec_decode/eagle/utils.py |  11 +-
 6 files changed, 21 insertions(+), 233 deletions(-)
 delete 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
deleted file mode 100644
index 0791a5183ead..000000000000
--- a/tests/models/kimi_k3/test_dspark_mla.py
+++ /dev/null
@@ -1,23 +0,0 @@
-# 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 4313aaa3f625..858622fa07da 100644
--- a/tests/models/kimi_k3/test_eagle3.py
+++ b/tests/models/kimi_k3/test_eagle3.py
@@ -200,111 +200,3 @@ 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 under global
-    per-tap keys (EagleModelMixin.pack_local_aux_hidden_states) and the last
-    stage prepends the collected remote taps to its own.
-    """
-    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)],
-    )
-    # EagleModelMixin caches the PP aux layout in _set_aux_hidden_state_layers;
-    # the stubs set layers directly, so prime the caches by hand.
-    object.__setattr__(stage0, "_aux_slot_base_cached", 0)
-    object.__setattr__(stage1, "_aux_slot_base_cached", 1)
-    object.__setattr__(stage1, "_aux_upstream_total_cached", 1)
-
-    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 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)
-
-    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 captured on the stage -> no aux keys on the wire."""
-    stage0 = _make_stage(
-        start_layer=0,
-        taps=(2,),
-        layer_outputs=[(torch.ones(1, 2), None, torch.zeros(1, 2))],
-    )
-    object.__setattr__(stage0, "_aux_slot_base_cached", 0)
-
-    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 not any(key.startswith("aux_hidden_states_") for key in out.tensors)
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 586363f54b68..1a1d96a6b4ef 100644
--- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py
+++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py
@@ -63,32 +63,6 @@ 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 c538a1201000..987909702cf4 100644
--- a/vllm/models/deepseek_v4/nvidia/dspark.py
+++ b/vllm/models/deepseek_v4/nvidia/dspark.py
@@ -22,7 +22,6 @@
     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 (
@@ -347,11 +346,8 @@ def _insert_context_kv(
 
 class DSparkDeepseekV4ForCausalLM(nn.Module):
     # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
-    # 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).
+    # load_dspark_model always aliases the target's.
     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
@@ -479,18 +475,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         head_end = n_local_head * (tp_rank + 1)
 
         weights = _duplicate_context_wkv_weights(weights, len(self.model.layers))
-        # 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:
-            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
+            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/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py
index 39a9eeafd7bf..0baff884ccaa 100644
--- a/vllm/models/kimi_k3/nvidia/dspark_mla.py
+++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py
@@ -9,19 +9,12 @@
 
 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 (
     DSparkConfidenceHead,
     DSparkMarkovHead,
@@ -43,12 +36,6 @@
 )
 
 
-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]]:
@@ -157,18 +144,8 @@ 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. 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).
+        # The frozen target embedding is aliased after the draft checkpoint loads.
         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,
@@ -448,19 +425,18 @@ def forward(
         return hidden_states
 
 
-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,
+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,
+        },
         orig_to_new_prefix={"": "model."},
         orig_to_new_stacked={
             ".gate_proj": (".gate_up_proj", 0),
@@ -470,17 +446,6 @@ def _build_weights_mapper(*, drop_embed: bool) -> WeightsMapper:
         },
     )
 
-
-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
@@ -498,10 +463,6 @@ 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 0cc85c6ffa03..06bed9391e03 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
@@ -55,14 +55,9 @@ def maybe_share_target_embed(
         return
 
     if target_embed is None:
-        # 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:
+        if hasattr(draft_inner, "embed_tokens") and not getattr(
+            draft_model, "has_own_embed_tokens", False
+        ):
             raise RuntimeError(
                 f"{type(draft_model).__name__} needs the target input embedding, "
                 "but it is unavailable on this PP stage"