[gbyu-amd]

Motivation

ARCH_ALIASES maps KimiK31ForConditionalGeneration onto Kimi-K3, so the Kimi-K3 PP DSpark path (#2365) also accepts Kimi-K3.1. Under PP, the draft on the last stage loaded the embedding from its own checkpoint. Inferact/Kimi-K3.1-DSpark has no embed_tokens, and the loader does not fail on a missing weight, so the draft would run with an embedding that was never loaded.

Technical Details

• KimiLinearModel keeps a real embed_tokens on the last stage when speculative_config is set, so the target checkpoint's table loads there.
• KimiK3DSpark always shares the target's embedding, the same as without PP, and drops its per-stage own-embedding path.
• tap_is_resident(-1) is now true only on the first stage. The last stage's copy of the embedding never runs in the target forward, so the -1 tap must still come through the relay.

Test Plan

• tests/test_pp_aux_relay.py: the last-stage test now holds a real embedding copy and checks that the -1 tap still gets no buffer and no hook.
• Also run: test_pp.py, test_dspark.py, test_kimi_k31_config.py, models/test_kimi_k3_mono_plan.py, test_gdn_state_fork.py.

Test Result

All of the files above pass in the atom_1031 container (24/99/82/3/67/7 passed). Full GSM8K (1319), Kimi-K3 pp2 × tp4 DSpark with prefix caching and ptpc_fp8 (the Kimi-K3 DSpark PP2 CI flags), at 83e3223df (this commit before the rebase onto main): 0.9712 / 0.9719 flexible / strict (±0.0046), acceptance 53.60% (4.75 tok/fwd), against 0.9704 / 53.7% for #2365 at f34cb8a8e. The relay and ReplaySSM markers match (4 / 8 lines) and the server log has no missing-weight warning. Kimi-K3.1 is not run: no local checkpoint.

🤖 Generated with Claude Code

[valar-cop]

a few small things..

ci green.

[valar-cop]

this is the actual fix but nothing tests it? add a cpu test with/without speculative_config, last rank stub.

[gbyu-amd]

Moved the check to atom.models.utils.pp_holds_embedding so a CPU test can reach it; test_pp.py::test_pp_holds_embedding covers first/middle/last stage with and without a speculative config (7d3fa50).

[valar-cop]

this makes two comments stale now: _FakeTarget at line 39 and the pp_aux_relay docstring. please update both.

[gbyu-amd]

Updated both (63e2c61).

[valarLip]

The direction is right: under PP the target keeps a real embedding on the last stage and the draft always shares it, which removes #2365's own-embed build/load path and fixes K3.1-DSpark drafting from an unloaded embedding when the draft checkpoint has no embed_tokens. Deciding the -1 tap by stage rather than by module type is also correct, since the last stage's copy is never run by the target forward. I found nothing that breaks (106 + 85 tests pass, 1 skipped, on this head; ruff/black clean). The items below are about where the decisions live and what the tests cover.

Action items for the author

1. [F1] Keep tap_is_resident structural: decide -1 by layers[0] instead of reading global PP state (drafter.py:123).
2. [F2] Make pp_holds_embedding follow whether the drafter actually shares the embedding, not speculative_config is not None (utils.py:139).
3. [F3] One way to decide "this stage holds the embedding" across models, tied to pp_speculative_methods or checked in PP validation (utils.py:146).
4. [F4] Re-wrap the module docstring (pp_aux_relay.py:10).
5. [F5] Remove the now-inert with_embed knob from _FakeTarget (test_pp_aux_relay.py:40).
6. [F6] Test the regression actually fixed: on the last stage the draft's embed_tokens is the target's and is not a PPMissingLayer (test_pp.py:224).
7. [F7] Minor: confirm the skip_weight_prefixes tuple -> list change is intentional (kimi_k3_dspark.py:547).

[valarLip]

[F1] tap_is_resident used to be purely structural; for -1 it now reads global get_pp_group().is_first_rank. not isinstance(layers[0], PPMissingLayer) (the embedding runs on the stage that owns layer 0) gives exactly the same answer and keeps it structural. PPAuxRelay.__init__ already has pp and reads it again through this helper, and the tests need a new _pp_group ExitStack patching get_pp_group in both drafter and pp_aux_relay (test_pp_aux_relay.py:54-61); with layers[0] those patches go away.

[gbyu-amd]

Done: -1 is resident where layer 0 is, so the helper reads no PP state and the double get_pp_group patch is gone (2026075).

[valarLip]

[F2] speculative_config is not None stands in for "the drafter on this stage shares the target's embedding", so the target model is guessing the drafter's behaviour. EAGLE-style drafters keep their own embedding under PP (drafter.py:643 guards on world_size == 1); once pp_speculative_methods admits eagle/mtp, the last stage would silently build and load a whole vocab x hidden table for nothing (about 2.35 GB for K3, ~590 MB per GPU at TP4). The condition should follow whether the drafter shares the embedding (e.g. share_with_target, or a registry declaration), not just whether speculation is on.

[gbyu-amd]

Done: the last stage builds the embedding only when registry.serves_pp_speculation holds, i.e. the configured method is declared in pp_speculative_methods, whose comment now says each declared method shares the embedding. An undeclared eagle/mtp builds nothing (f65a847).

[valarLip]

[F3] This helper sits in the model-wide utils but only Kimi-K3 uses it; MiniMax-M3 (also a DSpark that shares the embedding) still writes is_first_rank directly, so the repo has two spellings of one concept. Any architecture later added to pp_speculative_methods must remember to switch its embedding construction to this helper, or its draft gets a PPMissingLayer on the last stage again, and that constraint is neither in the registry nor checked anywhere. Please tie it to pp_speculative_methods, or have PP validation check the draft received a real embedding.

[gbyu-amd]

Done in f65a847: MiniMax-M3 builds through pp_holds_embedding too, and _check_pp_support uses the same serves_pp_speculation. The drafter also checks after share_with_target: a draft that got a PPMissingLayer fails at load, naming the registry and the helper.

[valarLip]

[F4] The rewritten docstring was not re-wrapped: this line is 124 characters while the surrounding lines stay within 78. black/ruff do not check docstring width, so CI won't catch it; please re-wrap to match.

[gbyu-amd]

Re-wrapped (cb9f7f9).

[valarLip]

[F5] With -1 no longer decided from target.embed_tokens, with_embed (passed as with_embed=first in _relay) is a knob that changes nothing; any value gives the same result, and only test_last_stage_arms_only_its_own_taps uses it, to assert that holding an embedding does not add a hook. It suggests to readers that relay behaviour depends on the embed module. Please drop the parameter and let _FakeTarget always hold an embedding (or follow F1).

[gbyu-amd]

Dropped; _FakeTarget always holds an embedding now (2026075).

[valarLip]

[F6] This enumerates a three-line boolean helper with 7 parameter sets but does not cover the regression this PR fixes: KimiLinearModel building and loading the embedding on the last stage, and KimiK3DSpark sharing the target's embedding under PP instead of building its own. If someone reverts kimi_k3/model.py to get_pp_group().is_first_rank, or adds an own-embed branch back to the draft, this stays green while K3.1 drafts from an unloaded embedding again. Please test that on the last stage draft.embed_tokens is target.model.embed_tokens and is not a PPMissingLayer.

[gbyu-amd]

Added in 9f3b5be: the K3 draft binds the target's embedding object through share_with_target, and a draft left with a PPMissingLayer is refused at load. A revert of kimi_k3/model.py to is_first_rank now fails that load check at startup instead of drafting from an unloaded table; building a whole KimiLinearModel in a unit test needs the TP group and aiter, so I stopped there.

[valarLip]

[F7] Minor: this goes back from a tuple to a mutable ClassVar[list[str]] shared by all instances. The tuple only existed for the removed per-stage rebinding and other models use a list, so this is fine if the intent is to match the repo convention; just confirming it is deliberate rather than a mechanical revert.

[gbyu-amd]

Deliberate: the tuple only existed for the per-stage rebinding this PR removes, and a list matches the other drafts.

[valar-cop]

ci triage for 703592a:

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/43133845999

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/43133845934

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/43133846121

[Jasen2201]

PP is primarily used for prefill in our P/D deployments. The DSpark prefill producer needs to populate draft context KV, but does not need to generate candidate tokens. We should avoid adding an extra embedding on the last PP stage for this case, since it consumes GPU memory unnecessarily (~2.35 GB in BF16 for K3, or ~0.59 GB per GPU when sharded at TP4).

A simple fix, following #3038:

• Keep the target embedding on the first stage, and only create the last-stage copy when that stage actually needs to generate draft tokens.
• For a context-only P/D producer, validate the vocabulary using lm_head and allow draft.embed_tokens = None.
• Skip producer proposal and draft graph warmup, while retaining draft context-KV writes and required state/bookkeeping updates (including commit_speculative_state()).

Please add a regression test for the PP producer without a last-stage embedding, alongside ordinary PP speculation that still requires one.

[valarLip]

⚠️ Action required before merge (2)

1. Reply to Jasen2201's P/D-producer point. Main already builds a draft-owned embedding of the same size on the last stage under PP, so this PR adds no memory. Skipping it for a context-only kv_producer is a reasonable follow-up. Say whether you'll do it here or open an issue.
2. Get a real CI run. The only full heavy run at this head (42749859685, 11-10) failed in every model on infra (flydsl 0.4.2 in the image, >=0.5.0 needed). The 11-11 run was gate-only. Re-run ATOM Test, including the Kimi-K3 DSpark PP2 job, and confirm it is green.

---

_{Re-reviewed at 703592a. Already fixed: F1-F7 (layers[0] residency, registry-driven serves_pp_speculation, MiniMax shares the helper, load-time PPMissingLayer refusal, docstring, no with_embed, regression test).}

[gbyu-amd]

@Jasen2201 Agreed a context-only producer needs no embedding: the K3 draft only embeds in block_backbone (the proposal), not in write_context_kv. This PR keeps the size main already spends (#2365 loads a draft-owned copy on the last stage), so I will do the producer skip in a separate PR on top of #3038, which adds the producer's propose/warmup skip it relies on, with the two regression tests you asked for.

[valarLip]

⚠️ Action required before merge (2)

1. Make the load-time PPMissingLayer refusal reachable, or delete it. atom/models/kimi_k3_dspark.py:614 now reads target_base.model.embed_tokens.num_embeddings before it binds the embedding. PPMissingLayer is a bare nn.Identity (atom/models/utils.py:29) with no num_embeddings, so share_with_target raises AttributeError before check_shared_embedding (atom/spec_decode/drafter.py:694) can run. MiniMax's draft (minimax_m3/dspark.py:508) reads it the same way. The DeepSeek drafts bind embed, not embed_tokens. So no real draft can reach the guard, and the reply to F6 ("a revert … fails that load check") is not what happens. Fix: restore target_vocab = target_base.model.vocab_size (main's line; KimiLinearModel.vocab_size is set at kimi_k3/model.py:1810) so the bind happens and the guard fires with its message. Also add a test that calls the real share_with_target with a PPMissingLayer target embedding and expects the TypeError. Otherwise drop check_shared_embedding and test_a_draft_sharing_a_missing_embedding_is_refused (tests/test_pp_aux_relay.py:296).
2. Get the heavy CI at this head green. Run 43149686406 at b12cbdf is still queued for 9 jobs, Kimi-K3 DSpark PP2 among them, the job that exercises this change. Merge only after it passes. (valarLip's ask 2 from 11-12 is still open.)

Jasen2201's P/D-producer memory point: answered and deferred to a follow-up PR on top of #3038, which is fine (valarLip ask 1 is done).

---

_{Re-reviewed at b12cbdf. Already fixed: F1-F7 (layers[0] residency, registry-driven serves_pp_speculation, MiniMax shares the helper, docstring re-wrap, no with_embed, regression test, list prefixes), reply to the P/D-producer point.}

[valar-cop]

some changes needed..

ci green.

[valar-cop]

this refusal is unreachable: share_with_target reads target_base.model.embed_tokens.num_embeddings first, and a PPMissingLayer embed raises AttributeError there, so check_shared_embedding never fires. use target_base.model.vocab_size like main? and add a test with a real PPMissingLayer target embed..

[gbyu-amd]

Fixed in ac819d3: the vocab check reads target_base.model.vocab_size again, so the bind happens and check_shared_embedding refuses the PPMissingLayer. test_a_k3_draft_on_a_stage_without_the_embedding_is_refused goes through the real share_with_target with a PPMissingLayer target embedding and fails if the old embed_tokens.num_embeddings read comes back.

[gbyu-amd]

@valarLip 1. Fixed in ac819d3 (details in the inline reply): the vocab check reads target_base.model.vocab_size, and the new test calls the real share_with_target with a PPMissingLayer target embedding and expects the TypeError. 2. At b12cbdf every accuracy job passed except DeepSeek-V4-Pro DSpark, which died at startup (AssertionError: VllmBackend can only be called once, then a DP rank-1 SHUTDOWN) before any request; this PR does not touch its path (its draft binds embed, so the check skips it). The push re-runs the heavy CI and I will report the result here, Kimi-K3 DSpark PP2 included.

[Jasen2201]

A note from #3038 (MiniMax-M3.1 PP4 P/D with DSpark, in review), since both PRs edit the same share_with_target block in Drafter.load_model.

#3038 makes context-only P/D producers skip share_with_target (if not self.context_kv_only:). Whichever PR lands second needs two things:

• check_shared_embedding(self.model) stays inside that branch.
• pp_holds_embedding returns False on a P/D producer. Such a stage never shares the embedding, so building a full-vocab copy on its last stage only costs memory (e.g. a Kimi-K3 PP producer).

If #3023 lands first, we will make both changes when rebasing #3038.

[valarLip]

⚠️ Action required before merge (1)

1. Get the heavy CI at this head green before merging. ATOM Test run 43232417010 at 6ca42d1 has passed only its gate and smoke jobs. All 14 accuracy jobs are still queued, Kimi-K3 DSpark PP2 among them, and that is the only job that exercises this change end to end. Merge once it passes. If DeepSeek-V4-Pro DSpark fails again with VllmBackend can only be called once, link the same failure on main so it is clearly unrelated. (valarLip's CI ask is still open.)

The code itself is done: I found nothing else to change at this head. Note for #3038: whichever PR lands second must keep check_shared_embedding(self.model) inside the if not self.context_kv_only: branch of Drafter.load_model and make pp_holds_embedding return False on a P/D producer, as Jasen2201 noted.

---

_{Re-reviewed at 6ca42d1. Already fixed: the load-time PPMissingLayer refusal is now reachable (kimi_k3_dspark.py:616 reads target_base.model.vocab_size, and test_a_k3_draft_on_a_stage_without_the_embedding_is_refused goes through the real share_with_target). Also fixed: F1-F7 and the reply to the P/D-producer point.}