[Spec] Model-agnostic last-stage draft embedding under pipeline parallelism

[YAMY1234]

Motivation

Update after #30775 merged. The decode-side PP x spec path now carries its own inline checkpoint read of the draft embedding in eagle_worker_v2 (type scan for the draft's VocabParallelEmbedding, a candidate-name table, DefaultModelLoader._prepare_weights), gated on SGLANG_ENABLE_PP_SPEC. This PR replaces that inline code with a shared module so the env-gated decode path and the PD-prefill path (#39378's gate) use one loader, and extends it to multi_layer_eagle_worker_v2 / frozen_kv_mtp_worker_v2.

Part of #39634 (P0). Under pipeline parallelism the EAGLE/MTP draft runs only on the last stage, while the target's input embedding lives only on the first stage (embed_tokens is a PPMissingLayer elsewhere). EagleDraftWorker.init_lm_head calls target.get_embed_and_head(), which raises AttributeError on the last stage for every model whose getter is self.model.embed_tokens.weight, self.lm_head.weight (DeepSeek, GLM, Llama, ...).

Today each model works around this in its own file: Qwen3.5 (#37471) keeps a draft-owned embedding and loads it in qwen3_5_mtp.load_weights; #39378 materializes a second embed_tokens on the target's last stage for DeepSeek/GLM. That per-model work is why the PP+spec gates in #39378 / #39602 / #38698 are architecture allowlists.

Modifications

• speculative/pp_draft_embedding.py (new)
  • resolve_target_embed_and_head: PP-safe wrapper; a PPMissingLayer embedding yields embed=None and the head is read directly; any other AttributeError propagates.
  • load_draft_embedding_from_checkpoint: resolves the checkpoint through DefaultModelLoader._prepare_weights (so ModelScope, --download-dir, load_format and the .bin fallback match the target load; dummy skips), locates the input embedding via model.safetensors.index.json and reads only that shard (no index: keys are collected across all shards before choosing; layers.<n>. MTP embeddings are never picked), and loads it into the draft's own embedding through its weight_loader. The draft embedding is the draft's single VocabParallelEmbedding (not ParallelLMHead); a draft without one fails with an explicit error. Returns the parameter so the existing set_embed_and_head(embed, head) call is unchanged.
  • resolve_draft_embed_and_head takes no pipeline-rank arguments: a stage whose target embedding is a PPMissingLayer loads the draft's own embedding; pp_size == 1 never returns embed=None.
• eagle_worker_v2, multi_layer_eagle_worker_v2, frozen_kv_mtp_worker_v2: all three init_lm_head paths go through resolve_draft_embed_and_head and no longer read the PP group. The SGLANG_ENABLE_PP_SPEC early-return block and its helpers (_find_draft_input_embedding, _load_checkpoint_tensor, _EMBED_TENSOR_NAMES) are removed; their re-openable load-format guard and extra checkpoint spellings are folded into the module. No model file changes.
• arg_groups/validation_hook.py: the _PP_EAGLE_SUPPORTED_ARCHITECTURES allowlist and the model_architecture argument of check_pipeline_parallel_compat are removed. The list existed because only some model families could hand the draft an embedding on the last stage; with the checkpoint loader that is true for every model, and a draft without any embedding fails in the loader with an explicit error instead.
• test/registered/e2e/pp/test_pp_spec_embed_scan.py: imports retargeted to the module.

Behavior changes

• The removed SGLANG_ENABLE_PP_SPEC block carried assert hot_token_id is None and not is_eagle3(). EAGLE3 under PP is still rejected by validation_hook.py; --speculative-token-map now takes the generic head.clone()[hot_token_id] path instead of being rejected. set_lm_head_from_target now also runs on the env-gated path.
• PD-prefill PP + EAGLE/MTP is no longer limited to the DeepSeek / GLM-5.2 / Qwen3.5 architectures listed in Fix disagg PP MTP for GLM-5.2 #39378 (GLM-5.3-Flash, Glm5NextForConditionalGeneration, was rejected by that list on main).

Relationship to #30775: that PR carries its own checkpoint embedding read in init_lm_head, active only under SGLANG_ENABLE_PP_SPEC; this PR replaces it. The type-based draft lookup and the extra tensor spellings here are taken from #30775.

Known follow-ups once this lands: drop the model-specific embedding handling in qwen3_5_mtp.load_weights (until then Qwen3.5 loads its draft embedding twice, correct but one redundant shard read) and the target-side pp_stage_needs_embedding in #39378.

Accuracy / validation

• Unit tests (test/registered/unit/spec/test_pp_draft_embedding.py, CPU, 8 cases): PP-safe resolution incl. an unrelated AttributeError not being swallowed, index and no-index selection over a layers.61.embed_tokens decoy, .bin-only checkpoint, full-vocab tensor handed to weight_loader, load_format gating (dummy / streaming), loud failure for a draft without an embedding. The drafts use the real VocabParallelEmbedding under a published topology. test_server_args.py::TestPipelineParallelCompat updated for the removed allowlist.
• E2E, env-gated decode path: Pipeline parallelism x speculative decoding (EAGLE/MTP) compatibility #30775's test/registered/e2e/pp/test_pp_spec.py (4-GPU H100) exercises this module.
• E2E, PD-prefill path (GB300, this branch merged with main d229952): GLM-5.3-Flash, PD 1P(TP2xPP2 + MTP) : 1D(TP4 + MTP) over Mooncake, GSM8K-200 via sglang.test.run_eval: score 0.990, accept length 3.45 (TP4 + MTP reference 0.990 / 3.39). The last prefill stage logs Loaded draft embedding model.embed_tokens.weight from checkpoint key model.language_model.embed_tokens.weight (154880, 4096), a key the previous inline loader did not know.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed.

---

CI States

Latest PR Test (Base): ❌ Run #35758694318
Latest PR Test (Extra): ✅ Run #35758694131
Latest PR Test (AMD ROCm 10): ❌ Run #35758694379

[YAMY1234]

/tag-and-rerun-ci

[YAMY1234]

All NV CI Passed https://github.com/sgl-project/sglang/actions/runs/35559775182/job/106210438384?pr=39643