[hmellor]

Overview

Support EAGLE3 aux hidden states with pipeline parallelism, and aux hidden state layer 0, in the Transformers modeling backend, matching native Llama. Part of serving Llama with the backend (#66547). The layer 0 changes moved here from #66413.

Claims

• Backend models set supports_aux_hidden_states_over_pp and pass aux hidden states between PP ranks with EagleModelMixin's helpers, as native Llama does.
• The last PP rank keeps the target input embedding when spec decode needs it (spec_decode_needs_target_embed), so drafters without embed_tokens can share it.
• Fixes _target_class becoming PPMissingLayer on rank 0, which made the aux recorders hook the PP placeholders instead of the decoder layers.
• Multimodal backend models still reject EAGLE3 with PP with the existing error. They return no aux hidden states even at PP=1, because HF installs the aux recorders from the VLM's _can_record_outputs and never reaches the decoder layers inside the language_model sub-model; that's a separate fix.
• Aux hidden state layer 0 is the input embeddings on the backend; previously _layer_names[-1] raised. It's recorded with an OutputRecorder on the input embedding, so it stays inside the compiled graph like the other layers (computing it in eager code hands the piecewise CUDA graphs a new tensor each step and crashes with an illegal memory access on replay). When the caller passes inputs_embeds, those are used directly.
• Native models also return the input embeddings for layer 0; previously the first layer's in-place fused add+RMSNorm overwrote the captured tensor.
• extract_hidden_states works with model_impl="transformers"; the draft config no longer inherits the target's model_impl.

Validation

B200, Llama-3.1-8B-Instruct + yuhuili/EAGLE3-LLaMA3.1-Instruct-8B, 3 speculative tokens, 6 prompts × 128 greedy tokens:

	PP=1	PP=2
Native	1.460	1.473
Backend	1.432	1.434

Mean acceptance length. On main the backend at PP=2 fails at startup (TransformersForCausalLM does not support eagle3 with pipeline parallelism).

pytest tests/v1/worker/test_eagle3_aux_hidden_states_pp.py
pytest tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py -k test_extract_layer_zero_is_input_embeddings
pytest "tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py::test_eagle3_pipeline_parallel_acceptance"

test_eagle3_aux_hidden_states_pp.py: 20 passed. test_extract_layer_zero_is_input_embeddings fails on main ([vllm]: max abs diff 158 vs the embeddings; [transformers]: ExtractHiddenStatesModel not found) and passes here; it runs with compilation, which also covers the CUDA graph crash above. test_eagle3_pp.py (4 GPUs, PP=1/2/4): both model_impls pass. With its original 4 prompts × 32 tokens the backend failed at PP=2 (1.628 < 2.048 × 0.95): compiled numerics changed one prompt's greedy path, which moved the mean that much. Across 64 prompts × 64 tokens on Llama-3.2-1B, acceptance is 1.450 to 1.471 for every impl, PP size (1/2/4) and CUDA graph mode, so the test now uses 16 prompts × 64 tokens (PP=1/2/4 acceptance: native 1.443/1.487/1.460, backend 1.468/1.475/1.467).

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hmellor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[hmellor]

/ci run

[github-actions]

✅ Triggered Buildkite CI #101185 for commit df5b5f6074d3.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as df5b5f6074d3 before starting CI. Pull the branch before pushing more changes.

[hmellor]

/ci run

[github-actions]

✅ Triggered Buildkite CI #101187 for commit 21d399ce3f14.