[Core][MRV2] Support eagle3 spec decode with pipeline parallel

tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""EAGLE3 speculative decoding under pipeline parallelism."""
+
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.utils import multi_gpu_test
+from vllm import LLM, SamplingParams
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.v1.metrics.reader import Metric
+
+MODEL = "meta-llama/Llama-3.2-1B-Instruct"
+DRAFT = "nm-testing/Llama3_2_1B_speculator.eagle3"
+
+PROMPTS = [
+    "The capital of France is",
+    "2 + 2 equals",
+    "In one word, the color of the sky is",
+    "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:",
+]
+
+# Acceptance held within 2% across PP=1..4; a dropped tap cost ~11%.
+ACCEPTANCE_TOLERANCE = 0.95
+
+
+def _acceptance_length(metrics: list[Metric]) -> float:
+    """1 + accepted/drafts, the mean tokens emitted per target forward."""
+    by_name = {m.name: m for m in metrics}
+    drafts = by_name.get("vllm:spec_decode_num_drafts")
+    accepted = by_name.get("vllm:spec_decode_num_accepted_tokens")
+    assert drafts is not None and accepted is not None, (
+        "spec_decode metrics missing; check disable_log_stats=False"
+    )
+    assert int(drafts.value) > 0, "drafter never proposed anything"
+    return 1.0 + int(accepted.value) / int(drafts.value)
+
+
+def _run(pp_size: int, model: str, draft: str, cudagraph_mode: str | None) -> float:
+    kwargs = dict(
+        model=model,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=pp_size,
+        max_model_len=512,
+        gpu_memory_utilization=0.45,
+        disable_log_stats=False,
+        speculative_config={
+            "method": "eagle3",
+            "model": draft,
+            "num_speculative_tokens": 3,
+        },
+    )
+    if cudagraph_mode is not None:
+        kwargs["compilation_config"] = {"cudagraph_mode": cudagraph_mode}
+
+    llm = LLM(**kwargs)
+    try:
+        llm.generate(
+            PROMPTS,
+            SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
+        )
+        return _acceptance_length(llm.get_metrics())
+    finally:
+        del llm
+        torch.accelerator.empty_cache()
+        cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(scope="module")
+def draft_without_embed_tokens(tmp_path_factory) -> str:
+    """``DRAFT`` with its input embedding stripped out.
+
+    Most EAGLE3 checkpoints ship no ``embed_tokens`` and alias the target's --
+    yuhuili/EAGLE3-LLaMA3.1-Instruct-8B, which test_eagle_correctness.py runs, is
+    one. ``DRAFT`` carries its own, so on its own it never exercises that path.
+    Removing the key reproduces the shared-embedding class at 1B.
+    """
+    from safetensors.torch import load_file, save_file
+
+    from vllm.transformers_utils.repo_utils import hf_api
+
+    src = Path(
+        hf_api().snapshot_download(
+            repo_id=DRAFT,
+            repo_type="model",
+            allow_patterns=["*.json", "*.py", "*.safetensors"],
+        )
+    )
+    dst = tmp_path_factory.mktemp("eagle3_shared_embed")
+    for path in src.iterdir():
+        if path.is_file() and path.suffix in (".json", ".py"):
+            shutil.copy(path, dst / path.name)
+
+    tensors = load_file(src / "model.safetensors")
+    stripped = [name for name in tensors if "embed_tokens" in name]
+    assert stripped, f"{DRAFT} has no embed_tokens to strip"
+    for name in stripped:
+        del tensors[name]
+    save_file(tensors, str(dst / "model.safetensors"), metadata={"format": "pt"})
+    return str(dst)
+
+
+@multi_gpu_test(num_gpus=2)
+def test_eagle3_pipeline_parallel_shared_embedding(draft_without_embed_tokens: str):
+    """A drafter with no embedding of its own must still get the target's.
+
+    Sharing is what supplies it, and under PP the target's embedding lives on the
+    first stage while the drafter runs on the last. Get this wrong and the
+    embedding is never written at all: the load is skipped for a key the
+    checkpoint lacks, and the drafter proposes from uninitialized memory, which
+    costs acceptance without failing.
+    """
+    baseline = _run(1, MODEL, draft_without_embed_tokens, "FULL_AND_PIECEWISE")
+    parallel = _run(2, MODEL, draft_without_embed_tokens, "FULL_AND_PIECEWISE")
+
+    assert parallel >= baseline * ACCEPTANCE_TOLERANCE, (
+        f"acceptance length regressed under PP=2 for a drafter sharing the "
+        f"target's embedding: {parallel:.3f} < {baseline:.3f} * "
+        f"{ACCEPTANCE_TOLERANCE}"
+    )
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.parametrize("model,draft", [(MODEL, DRAFT)])
+@pytest.mark.parametrize("cudagraph_mode", [None, "FULL_AND_PIECEWISE"])
+def test_eagle3_pipeline_parallel_acceptance(
+    model: str,
+    draft: str,
+    cudagraph_mode: str | None,
+):
+    """Aux hidden states must survive the pipeline handoff.
+
+    Compares acceptance length at PP=2 against PP=1 on the same model. This
+    feature fails quietly: a stale, out-of-order or dropped tap still yields
+    well-formed proposals, so the engine boots and answers -- the proposals just
+    get rejected more often. Acceptance is what detects that.
+
+    Greedy text parity deliberately is not asserted. bf16 argmax ties break
+    differently once the batch shape changes, so even two runs with spec decode
+    off diverge from each other, which would make such an assertion flaky for
+    reasons unrelated to this feature.
+    """
+    baseline = _run(1, model, draft, cudagraph_mode)
+    parallel = _run(2, model, draft, cudagraph_mode)
+
+    assert parallel >= baseline * ACCEPTANCE_TOLERANCE, (
+        f"acceptance length regressed under PP=2 "
+        f"(cudagraph_mode={cudagraph_mode}): "
+        f"{parallel:.3f} < {baseline:.3f} * {ACCEPTANCE_TOLERANCE}"
+    )
+
+
+@multi_gpu_test(num_gpus=4)
+@pytest.mark.parametrize("model,draft", [(MODEL, DRAFT)])
+def test_eagle3_pipeline_parallel_far_stage_acceptance(model: str, draft: str):
+    """Cover the stages that do not hand off to the rank consuming their taps.
+
+    PP=2 exercises none of this: its only producer is the stage right before the
+    last one, whose taps ride the handoff it already sends. PP=4 is the smallest
+    size with two such producers, and it is where a tap first has to reach a rank
+    that is not its neighbour.
+
+    Full cudagraph is the mode to run this in. The layout is resolved at setup
+    precisely so the forward stays capturable, and 16 layers over 4 stages splits
+    evenly, so an uneven-split regression would not show up here.
+    """
+    baseline = _run(1, model, draft, "FULL_AND_PIECEWISE")
+    parallel = _run(4, model, draft, "FULL_AND_PIECEWISE")
+
+    assert parallel >= baseline * ACCEPTANCE_TOLERANCE, (
+        f"acceptance length regressed under PP=4: "
+        f"{parallel:.3f} < {baseline:.3f} * {ACCEPTANCE_TOLERANCE}"
+    )

tests/v1/worker/test_eagle3_aux_hidden_states_pp.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Accounting checks for EAGLE3 aux taps sent directly to the last PP rank.
+
+Each stage produces only its local taps; the last rank gathers them in
+pipeline order. These run on CPU with no distributed init.
+"""
+
+import pytest
+
+from vllm.distributed.utils import get_pp_indices
+from vllm.model_executor.models.interfaces import EagleModelMixin
+from vllm.model_executor.models.llama import LlamaModel
+from vllm.model_executor.models.qwen2 import Qwen2Model
+
+# Kimi-K3 / DSpark: 93 layers, target_layer_ids [2, 23, 47, 71, 89].
+# get_eagle3_aux_layers_from_config maps target_layer_ids -> +1.
+NUM_LAYERS = 93
+AUX_IDS = (3, 24, 48, 72, 90)
+
+
+def _taps_emitted(start_layer, end_layer, aux_ids, is_first_rank):
+    return list(
+        EagleModelMixin.local_aux_tap_ids(
+            start_layer, end_layer, aux_ids, is_first_rank
+        )
+    )
+
+
+def _simulate(num_layers, pp_size, aux_ids):
+    """Gather local taps from every stage in rank order (direct-to-last)."""
+    gathered: list[int] = []
+    for rank in range(pp_size):
+        start, end = get_pp_indices(num_layers, rank, pp_size)
+        gathered.extend(_taps_emitted(start, end, aux_ids, rank == 0))
+    return gathered
+
+
+@pytest.mark.parametrize("pp_size", [1, 2, 3, 4, 6, 8])
+def test_drafter_sees_all_taps_in_order(pp_size):
+    assert _simulate(NUM_LAYERS, pp_size, AUX_IDS) == list(AUX_IDS)
+
+
+def test_pp2_split_matches_expected_stages():
+    start0, end0 = get_pp_indices(NUM_LAYERS, 0, 2)
+    start1, end1 = get_pp_indices(NUM_LAYERS, 1, 2)
+    assert _taps_emitted(start0, end0, AUX_IDS, True) == [3, 24]
+    assert _taps_emitted(start1, end1, AUX_IDS, False) == [48, 72, 90]
+
+
+@pytest.mark.parametrize("pp_size", [1, 2, 4])
+def test_no_drafter_adds_no_payload_keys(pp_size):
+    assert _simulate(NUM_LAYERS, pp_size, ()) == []
+
+
+def test_tap_on_stage_boundary_is_not_double_counted():
+    """A boundary tap is emitted once by the upstream stage that computed it."""
+    _, end0 = get_pp_indices(NUM_LAYERS, 0, 2)
+    aux_ids = tuple(sorted(AUX_IDS + (end0,)))
+    assert _simulate(NUM_LAYERS, 2, aux_ids).count(end0) == 1
+
+
+def test_middle_stage_sends_only_local_taps():
+    """PP>2 middle ranks must not re-send upstream taps."""
+    start, end = get_pp_indices(NUM_LAYERS, 1, 4)
+    local = _taps_emitted(start, end, AUX_IDS, False)
+    upstream = [a for a in AUX_IDS if a <= start]
+    assert not set(local) & set(upstream)
+    assert local == [a for a in AUX_IDS if start < a <= end]
+
+
+@pytest.mark.parametrize("model_cls", [LlamaModel, Qwen2Model])
+def test_forward_does_not_name_update(model_cls):
+    """Packing taps must not go through dict.update.
+
+    TorchDynamoWrapper.bytecode_hook refuses any compiled forward whose
+    bytecode names `update`, so a model that packs its taps that way cannot
+    start under cudagraphs. The same holds for the other opted-in models.
+    """
+    assert "update" not in model_cls.forward.__code__.co_names