[fix] Kimi-K3.1 DSpark under PP: share the target's embedding on the last stage

atom/config.py

@@ -3165,16 +3165,14 @@ def _check_pp_support(self) -> None:
         # draft layer's KV.
         if self.speculative_config is None or role == "kv_producer":
             return
-        from atom.models.registry import pp_speculative_methods
+        from atom.models.registry import pp_speculative_method, serves_pp_speculation
 
-        archs = getattr(self.hf_config, "architectures", None) or ()
-        spec = self.speculative_config
-        method = "dspark" if spec.use_dspark() else spec.method
-        if not any(method in pp_speculative_methods.get(a, ()) for a in archs):
+        if not serves_pp_speculation(self):
             raise ValueError(
                 "pipeline_parallel_size > 1 with speculative decoding supports "
                 "only the methods registry.pp_speculative_methods declares, got "
-                f"architectures={list(archs)}, method={method!r}."
+                f"architectures={getattr(self.hf_config, 'architectures', None)}, "
+                f"method={pp_speculative_method(self)!r}."
             )
         if self.dspark.confidence_schedule:
             raise ValueError(

atom/models/dspark_draft.py

@@ -685,8 +685,7 @@ def share_with_target(self, target_base: nn.Module, loaded: set) -> None:
                 f"vocab {target_vocab}. The draft shares the target's embedding "
                 "and LM head, so the two must agree."
             )
-        # A draft that loaded its own embedding (K3 DSpark on PP stage > 0)
-        # keeps it.
-        if self.embed_tokens is None:
-            self.embed_tokens = target_base.model.embed_tokens
+        # A stage without the target's embedding binds its PPMissingLayer,
+        # which the drafter's check_shared_embedding refuses.
+        self.embed_tokens = target_base.model.embed_tokens
         self.lm_head = target_base.lm_head

atom/models/kimi_k3/model.py

@@ -68,6 +68,7 @@
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
+    pp_holds_embedding,
 )
 from atom.quant_spec import effective_layer_quant
 from atom.utils import envs, mark_spliting_op
@@ -1868,7 +1869,7 @@ def __init__(self, atom_config: Config, prefix: str = ""):
         self.config = config
         self.vocab_size = config.vocab_size
 
-        if get_pp_group().is_first_rank:
+        if pp_holds_embedding(atom_config):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size, config.hidden_size
             )

atom/models/kimi_k3_dspark.py

@@ -49,8 +49,7 @@
   layers.{i}.mlp.down_proj.weight               [7168, 14336]
   markov_head.markov_w{1,2}.weight              [163840, 256]
   confidence_head.proj.{weight,bias}            SKIPPED (training-only)
-  embed_tokens.weight                           SKIPPED (target's is shared;
-                                                loaded on PP stages > 0)
+  embed_tokens.weight                           SKIPPED (target's is shared)
 
 Inferact/Kimi-K3.1-DSpark (67 tensors) is the same tree with 96 heads (q_b_proj
 [18432, 1536], kv_b_proj [24576, 512], o_proj [7168, 12288]), Markov rank 512
@@ -61,7 +60,6 @@
 
 import torch
 from aiter import QuantType, dtypes
-from aiter.dist.parallel_state import get_pp_group
 from aiter.rotary_embedding import get_rope
 from torch import nn
 
@@ -72,7 +70,6 @@
     qrep_tp_override,
 )
 from atom.model_ops.base_attention import Attention
-from atom.model_ops.embed_head import VocabParallelEmbedding
 from atom.model_ops.layernorm import RMSNorm
 from atom.model_ops.linear import (
     ColumnParallelLinear,
@@ -535,14 +532,11 @@ class KimiK3DSpark(DSparkDraftModel):
     #   embed_tokens: a 2.35GB copy of the target's table. The target's is
     #     shared instead (share_with_target), so loading it would just burn
     #     memory and risk drifting from the target's.
-    #
-    # Under PP, __init__ drops embed_tokens from this list on non-first stages.
-    # A tuple, not a ClassVar: the per-stage override rebinds, never mutates.
-    skip_weight_prefixes: tuple[str, ...] = (
+    skip_weight_prefixes: ClassVar[list[str]] = [
         "confidence_head.",
         "embed_tokens.",
         "lm_head.",
-    )
+    ]
 
     def __init__(self, atom_config: "Config", layer_offset: int = 0) -> None:
         super().__init__()
@@ -582,17 +576,7 @@ def __init__(self, atom_config: "Config", layer_offset: int = 0) -> None:
         self.vocab_size = int(config.vocab_size)
 
         # Bound by share_with_target(); both are skipped at load.
-        # Exception: under PP the drafter lives on the last stage, whose target
-        # embedding is a PPMissingLayer, so load the checkpoint's own copy.
-        if not get_pp_group().is_first_rank:
-            self.skip_weight_prefixes = tuple(
-                p for p in self.skip_weight_prefixes if p != "embed_tokens."
-            )
-            self.embed_tokens = VocabParallelEmbedding(
-                config.vocab_size, config.hidden_size, prefix="embed_tokens"
-            )
-        else:
-            self.embed_tokens = None
+        self.embed_tokens = None
         self.lm_head = None
         # vLLM 0.28 probes this attribute before enabling adaptive verification.
         # This checkpoint deliberately skips the training-only head.

atom/models/minimax_m3/model.py

@@ -54,6 +54,7 @@
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
+    pp_holds_embedding,
 )
 from atom.quant_spec import effective_layer_quant
 from atom.utils import envs
@@ -762,7 +763,7 @@ def __init__(
         index_cache_config = atom_config.index_cache_dtype
         quant_config = atom_config.quant_config
 
-        if get_pp_group().is_first_rank:
+        if pp_holds_embedding(atom_config):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 config.hidden_size,

atom/models/registry.py

@@ -67,6 +67,25 @@
 }
 # Speculative methods each architecture serves under pipeline_parallel_size > 1
 # (its drafts verified per stage, its drafter's aux taps relayed across stages).
+# Each one's drafter shares the target's embedding, which the last stage then
+# builds (models.utils.pp_holds_embedding).
 pp_speculative_methods = {
     "KimiK3ForConditionalGeneration": ("dspark",),
 }
+
+
+def pp_speculative_method(atom_config) -> str | None:
+    """The configured speculative method, or None when there is none."""
+    spec = atom_config.speculative_config
+    if spec is None:
+        return None
+    return "dspark" if spec.use_dspark() else spec.method
+
+
+def serves_pp_speculation(atom_config) -> bool:
+    """Whether the architecture declares the configured method under PP."""
+    method = pp_speculative_method(atom_config)
+    archs = getattr(atom_config.hf_config, "architectures", None) or ()
+    return method is not None and any(
+        method in pp_speculative_methods.get(a, ()) for a in archs
+    )

atom/models/utils.py

@@ -136,6 +136,17 @@ def make_layers(
     return start_layer, end_layer, modules
 
 
+def pp_holds_embedding(atom_config) -> bool:
+    """Whether this PP stage builds the token embedding: the first stage embeds
+    the prompt, and the last holds a copy for a drafter that shares it."""
+    from aiter.dist.parallel_state import get_pp_group
+
+    from atom.models.registry import serves_pp_speculation
+
+    pp = get_pp_group()
+    return pp.is_first_rank or (pp.is_last_rank and serves_pp_speculation(atom_config))
+
+
 # NOTE: don't use lru_cache here because it can prevent garbage collection
 _model_to_pp_missing_layer_names: dict[int, list[str]] = {}
 

atom/spec_decode/drafter.py

@@ -117,16 +117,22 @@ def _pick(module):
     return embed
 
 
-def tap_is_resident(target_model: nn.Module, layers: nn.Module, lid: int) -> bool:
-    """Whether aux tap ``lid`` runs on this PP stage: a tap owned by another
-    stage sits on a ``PPMissingLayer``. Out-of-range ids are not resident."""
-    if lid == -1:
-        module = _resolve_embedding(target_model)
-    elif 0 <= lid < len(layers):
-        module = layers[lid]
-    else:
-        return False
-    return not isinstance(module, PPMissingLayer)
+def tap_is_resident(layers: nn.Module, lid: int) -> bool:
+    """Whether aux tap ``lid`` runs on this PP stage: a layer owned by another
+    stage is a ``PPMissingLayer``, and the embedding (``-1``) runs on the stage
+    that owns layer 0. Out-of-range ids are not resident."""
+    lid = 0 if lid == -1 else lid
+    return 0 <= lid < len(layers) and not isinstance(layers[lid], PPMissingLayer)
+
+
+def check_shared_embedding(draft: nn.Module) -> None:
+    """Refuse a draft that shared a ``PPMissingLayer`` as its embedding."""
+    if isinstance(getattr(draft, "embed_tokens", None), PPMissingLayer):
+        raise TypeError(
+            "The draft shares the target's embedding, which this PP stage does "
+            "not build: declare the method in registry.pp_speculative_methods "
+            "and build the embedding with models.utils.pp_holds_embedding."
+        )
 
 
 def _identity_extract(output: Any, _module: nn.Module) -> torch.Tensor:
@@ -536,9 +542,7 @@ def arm_aux_capture(self, target_model: nn.Module) -> None:
         layers = _resolve_decoder_layers(target_model)
         # A tap on an earlier PP stage reaches `take_relayed_aux` instead, so it
         # gets no rows and no hook. install_aux_taps still rejects bad ids.
-        own = tuple(
-            lid for lid in spec.layer_ids if tap_is_resident(target_model, layers, lid)
-        )
+        own = tuple(lid for lid in spec.layer_ids if tap_is_resident(layers, lid))
         self._aux_buffers = [
             torch.zeros(
                 self.max_num_tokens if lid in own else 0,
@@ -692,6 +696,7 @@ def load_model(self, target_model: nn.Module) -> None:
         # their own setattr-rebinding and short-circuit the default path.
         if hasattr(self.model, "share_with_target"):
             self.model.share_with_target(target_base, loaded)
+            check_shared_embedding(self.model)
             if self.is_block_drafter and hasattr(self.model, "reset_kv_cache"):
                 # Allocate DSpark's private rolling target-KV window now that the
                 # device/dtype and max concurrency are known.

atom/spec_decode/pp_aux_relay.py

@@ -4,12 +4,14 @@
 """Relay a block drafter's target aux hidden states across PP stages.
 
 The drafter lives only on the last PP stage, so hooks on taps owned by earlier
-stages land on ``PPMissingLayer`` and never fire, leaving zero aux buffers and
-garbage drafts. Earlier stages capture their own taps and forward them in the
-intermediate-tensor payload, concatenated along the hidden axis in the
-spec's ``layer_ids`` order, which ``install_aux_taps`` requires strictly
-increasing (-1, the embedding, first). Stages own contiguous ascending layer
-ranges, so every stage derives the same columns from ``layer_ids`` alone.
+stages would never fire: their layers are ``PPMissingLayer``, and the last
+stage's copy of the embedding (held for the drafter) is not run by the target
+forward. That would leave zero aux buffers and garbage drafts. Earlier stages
+capture their own taps and forward them in the intermediate-tensor payload,
+concatenated along the hidden axis in the spec's ``layer_ids`` order, which
+``install_aux_taps`` requires strictly increasing (-1, the embedding, first).
+Stages own contiguous ascending layer ranges, so every stage derives the same
+columns from ``layer_ids`` alone.
 """
 
 import logging
@@ -49,15 +51,13 @@ def __init__(
         taps = spec.layer_ids
 
         self.spec = spec
-        self.own_ids = tuple(
-            lid for lid in taps if tap_is_resident(target_model, layers, lid)
-        )
+        self.own_ids = tuple(lid for lid in taps if tap_is_resident(layers, lid))
         # Every tap below this stage's first layer was captured upstream, which
         # is exactly what the previous stage sent, even if it owned none itself.
         self.incoming_ids = tuple(
             lid
             for lid in taps
-            if not tap_is_resident(target_model, layers, lid) and lid < stage_start
+            if not tap_is_resident(layers, lid) and lid < stage_start
         )
         self.outgoing_ids = () if pp.is_last_rank else self.incoming_ids + self.own_ids
 

tests/test_dspark_draft.py

@@ -129,27 +129,16 @@ def test_default_share_with_target_refuses_a_vocab_mismatch():
     assert draft.embed_tokens is None and draft.lm_head is None
 
 
-def test_default_share_with_target_keeps_a_drafts_own_embedding():
-    """A draft that loaded its own embedding (K3 DSpark past PP stage 0)
-    keeps it and still scores with the target's head."""
-    draft = _stub_draft(163840)
-    own = draft.embed_tokens = _target_base(163840).model.embed_tokens
-    target = _target_base(163840)
-    draft.share_with_target(target, set())
-    assert draft.embed_tokens is own
-    assert draft.lm_head is target.lm_head
-
-
 def test_default_share_with_target_reads_the_vocab_off_the_head_under_pp():
     """Under PP the drafter's stage holds a PPMissingLayer for the target's
     embedding; the vocab check reads the head the draft scores with."""
     from atom.models.utils import PPMissingLayer
 
     draft = _stub_draft(163840)
-    draft.embed_tokens = _target_base(163840).model.embed_tokens
     target = _target_base(163840)
     target.model.embed_tokens = PPMissingLayer()
     draft.share_with_target(target, set())
+    assert draft.embed_tokens is target.model.embed_tokens
     assert draft.lm_head is target.lm_head
     with pytest.raises(ValueError, match="target vocab 200064"):
         _stub_draft(163840).share_with_target(_target_base(200064), set())

tests/test_pp.py

@@ -66,6 +66,7 @@
     PPMissingLayer,
     get_pp_indices,
     make_layers,
+    pp_holds_embedding,
 )
 from tests.conftest import MockConfig
 
@@ -207,6 +208,36 @@ def test_send_packs_only_proxy_keys_to_group_relative_next(monkeypatch):
     assert torch.equal(sent["tensors"]["hidden_states"], torch.ones(3, 4))
 
 
+@pytest.mark.parametrize(
+    "pp_rank,pp_size,arch,method,holds",
+    [
+        (0, 2, "KimiK3ForConditionalGeneration", None, True),
+        (1, 2, "KimiK3ForConditionalGeneration", None, False),
+        # A declared PP method's drafter shares the target's embedding.
+        (1, 2, "KimiK3ForConditionalGeneration", "dspark", True),
+        (1, 4, "KimiK3ForConditionalGeneration", "dspark", False),
+        (3, 4, "KimiK3ForConditionalGeneration", "dspark", True),
+        # Undeclared: the drafter keeps its own, so build none for it.
+        (1, 2, "KimiK3ForConditionalGeneration", "eagle", False),
+        (1, 2, "MiniMaxM3SparseForCausalLM", "dspark", False),
+        (0, 1, "KimiK3ForConditionalGeneration", "dspark", True),
+    ],
+)
+def test_pp_holds_embedding(monkeypatch, pp_rank, pp_size, arch, method, holds):
+    grp = MagicMock(is_first_rank=pp_rank == 0, is_last_rank=pp_rank == pp_size - 1)
+    fake_ps = MagicMock()
+    fake_ps.get_pp_group.return_value = grp
+    monkeypatch.setitem(sys.modules, "aiter.dist.parallel_state", fake_ps)
+    spec = method and types.SimpleNamespace(
+        method=method, use_dspark=lambda: method == "dspark"
+    )
+    config = types.SimpleNamespace(
+        hf_config=types.SimpleNamespace(architectures=[arch]),
+        speculative_config=spec,
+    )
+    assert pp_holds_embedding(config) is holds
+
+
 def test_send_tolerates_missing_residual(monkeypatch):
     _, sent = _fake_pp_group(monkeypatch)
     pp_comm.send_intermediate_tensors(