@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4a5299c93f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 23:00 UTC

feat: spec decode with draft models (#24322)

Signed-off-by: Tomas Ruiz <tomas.ruiz.te@gmail.com>

Tomas Ruiz committed Jan 19, 2026 at 15:05 UTC 4a5299c93ff97c26def537b92562df5ada530fea
21 files changed +897 -115
examples/offline_inference/spec_decode.py
+17 -2
@@ -54,7 +54,7 @@ def parse_args():
54 "--method",
55 type=str,
56 default="eagle",
57 - choices=["ngram", "eagle", "eagle3", "mtp"],
57 + choices=["ngram", "eagle", "eagle3", "mtp", "draft_model"],
58 )
59 parser.add_argument("--num-spec-tokens", type=int, default=2)
60 parser.add_argument("--prompt-lookup-max", type=int, default=5)
@@ -70,7 +70,11 @@ def parse_args():
70 parser.add_argument("--output-len", type=int, default=256)
71 parser.add_argument("--model-dir", type=str, default=None)
72 parser.add_argument("--eagle-dir", type=str, default=None)
73 + parser.add_argument("--draft-model", type=str, default=None)
74 parser.add_argument("--custom-mm-prompts", action="store_true")
75 + parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
76 + parser.add_argument("--disable-padded-drafter-batch", action="store_true")
77 + parser.add_argument("--max-num-seqs", type=int, default=None)
78 return parser.parse_args()
79
80
@@ -111,6 +115,7 @@ def main(args):
115 "method": args.method,
116 "model": eagle_dir,
117 "num_speculative_tokens": args.num_spec_tokens,
118 + "disable_padded_drafter_batch": args.disable_padded_drafter_batch,
119 }
120 elif args.method == "ngram":
121 speculative_config = {
@@ -119,6 +124,15 @@ def main(args):
124 "prompt_lookup_max": args.prompt_lookup_max,
125 "prompt_lookup_min": args.prompt_lookup_min,
126 }
127 + elif args.method == "draft_model":
128 + assert args.draft_model is not None and args.draft_model != ""
129 + speculative_config = {
130 + "method": args.method,
131 + "model": args.draft_model,
132 + "num_speculative_tokens": args.num_spec_tokens,
133 + "enforce_eager": args.enforce_eager,
134 + "max_model_len": args.max_model_len,
135 + }
136 elif args.method == "mtp":
137 speculative_config = {
138 "method": "mtp",
@@ -133,12 +147,13 @@ def main(args):
147 tensor_parallel_size=args.tp,
148 enable_chunked_prefill=args.enable_chunked_prefill,
149 enforce_eager=args.enforce_eager,
136 - gpu_memory_utilization=0.9,
150 + gpu_memory_utilization=args.gpu_memory_utilization,
151 speculative_config=speculative_config,
152 disable_log_stats=False,
153 max_model_len=args.max_model_len,
154 limit_mm_per_prompt={"image": 5},
155 disable_chunked_mm_input=True,
156 + max_num_seqs=args.max_num_seqs,
157 )
158
159 sampling_params = SamplingParams(temperature=args.temp, max_tokens=args.output_len)
examples/online_serving/disaggregated_serving/moriio_toy_proxy_server.py
+1 -1
@@ -4,13 +4,13 @@ import asyncio
4 import copy
5 import logging
6 import os
7 -import re
7 import socket
8 import threading
9 import uuid
10
11 import aiohttp
12 import msgpack
13 +import regex as re
14 import zmq
15 from quart import Quart, make_response, request
16
tests/v1/e2e/test_spec_decode.py
+306 -9
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import random
4 +from collections.abc import Iterable
5 +from dataclasses import dataclass
6 from typing import Any
7
8 import pytest
@@ -10,32 +12,45 @@ from tests.utils import get_attn_backend_list_based_on_platform, large_gpu_mark
12 from vllm import LLM, SamplingParams
13 from vllm.assets.base import VLLM_S3_BUCKET_URL
14 from vllm.assets.image import VLM_IMAGES_DIR
15 +from vllm.benchmarks.datasets import InstructCoderDataset
16 +from vllm.config.vllm import VllmConfig
17 from vllm.distributed import cleanup_dist_env_and_memory
18 +from vllm.engine.arg_utils import EngineArgs
19 from vllm.platforms import current_platform
20 +from vllm.v1.metrics.reader import Metric
21 +from vllm.v1.spec_decode.draft_model import (
22 + create_vllm_config_for_draft_model,
23 + merge_toks_kernel,
24 +)
25
26 MTP_SIMILARITY_RATE = 0.8
27
28
29 def _skip_if_insufficient_gpus_for_tp(tp_size: int):
30 """Skip test if available GPUs < tp_size on ROCm."""
21 - if current_platform.is_rocm():
22 - available_gpus = torch.cuda.device_count()
23 - if available_gpus < tp_size:
24 - pytest.skip(
25 - f"Test requires {tp_size} GPUs, but only {available_gpus} available"
26 - )
31 + available_gpus = torch.cuda.device_count()
32 + if available_gpus < tp_size:
33 + pytest.skip(
34 + f"Test requires {tp_size} GPUs, but only {available_gpus} available"
35 + )
36
37
29 -def get_test_prompts(mm_enabled: bool):
38 +Messages = list[dict[str, Any]]
39 +
40 +
41 +def get_test_prompts(
42 + mm_enabled: bool, quiet: bool = False, num_prompts: int = 100
43 +) -> list[Messages]:
44 prompt_types = ["repeat", "sentence"]
45 if mm_enabled:
46 prompt_types.append("mm")
33 - num_prompts = 100
47 prompts = []
48
49 random.seed(0)
50 random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
38 - print(f"Prompt types: {random_prompt_type_choices}")
51 +
52 + if not quiet:
53 + print(f"Prompt types: {random_prompt_type_choices}")
54
55 # Generate a mixed batch of prompts, some of which can be easily
56 # predicted by n-gram matching and some which likely cannot.
@@ -75,11 +90,27 @@ def get_test_prompts(mm_enabled: bool):
90 return prompts
91
92
93 +def get_instruct_coder_messages(n: int) -> list[Messages]:
94 + dataset = InstructCoderDataset(
95 + dataset_path="likaixin/InstructCoder", dataset_split="train"
96 + )
97 + prompts: Iterable[str] = dataset.sample_prompts(n=n)
98 + return [[{"role": "user", "content": prompt}] for prompt in prompts]
99 +
100 +
101 @pytest.fixture
102 def sampling_config():
103 + return greedy_sampling()
104 +
105 +
106 +def greedy_sampling() -> SamplingParams:
107 return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False)
108
109
110 +def stochastic_sampling() -> SamplingParams:
111 + return SamplingParams(temperature=1.0, max_tokens=10, ignore_eos=False)
112 +
113 +
114 @pytest.fixture
115 def model_name():
116 return "meta-llama/Llama-3.1-8B-Instruct"
@@ -583,3 +614,269 @@ def test_mtp_correctness(
614 del spec_llm
615 torch.cuda.empty_cache()
616 cleanup_dist_env_and_memory()
617 +
618 +
619 +@dataclass
620 +class ArgsTest:
621 + target_model: str
622 + draft_model: str
623 + sampling_config: SamplingParams
624 + num_speculative_tokens: int
625 + expected_acceptance_rate: float
626 + expected_acceptance_len: float
627 + # Defaults
628 + target_tensor_parallel_size: int = 1
629 + draft_tensor_parallel_size: int = 1
630 + max_model_len: int = 1024
631 + gpu_memory_utilization: float = 0.5
632 + dataset: str = "test_prompts"
633 + num_prompts: int = 100
634 +
635 +
636 +cases = [
637 + # Same model for draft and target, greedy sampling.
638 + ArgsTest(
639 + target_model="Qwen/Qwen3-0.6B",
640 + draft_model="Qwen/Qwen3-0.6B",
641 + sampling_config=greedy_sampling(),
642 + num_speculative_tokens=3, # K
643 + expected_acceptance_len=3 + 1, # K + 1
644 + expected_acceptance_rate=1.0,
645 + ),
646 + # Smaller draft model, stochastic sampling.
647 + ArgsTest(
648 + target_model="Qwen/Qwen3-1.7B",
649 + draft_model="Qwen/Qwen3-0.6B",
650 + sampling_config=stochastic_sampling(),
651 + num_speculative_tokens=3,
652 + expected_acceptance_len=2.8 + 1,
653 + expected_acceptance_rate=0.9,
654 + ),
655 +]
656 +
657 +
658 +@pytest.mark.parametrize("args", cases)
659 +@pytest.mark.parametrize("enforce_eager", [True, False])
660 +def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
661 + assert_draft_model_correctness(args, enforce_eager)
662 +
663 +
664 +def test_draft_model_realistic_example():
665 + args = ArgsTest(
666 + target_model="Qwen/Qwen3-1.7B",
667 + draft_model="Qwen/Qwen3-0.6B",
668 + dataset="likaixin/InstructCoder",
669 + num_speculative_tokens=3,
670 + sampling_config=greedy_sampling(),
671 + # values below are not derived, but just prevent a regression
672 + expected_acceptance_len=2.8,
673 + expected_acceptance_rate=0.55,
674 + )
675 + assert_draft_model_correctness(args, enforce_eager=False)
676 +
677 +
678 +@pytest.mark.parametrize(
679 + "models",
680 + [
681 + # target_model, draft_model
682 + ("Qwen/Qwen3-1.7B-FP8", "Qwen/Qwen3-0.6B"), # target quantized
683 + ("Qwen/Qwen3-1.7B", "Qwen/Qwen3-0.6B-FP8"), # draft quantized
684 + ],
685 + ids=["target_quantized", "draft_quantized"],
686 +)
687 +@pytest.mark.parametrize("enforce_eager", [True, False])
688 +def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
689 + tgt_model, draft_model = models
690 + sd_case = ArgsTest(
691 + target_model=tgt_model,
692 + draft_model=draft_model,
693 + **some_high_acceptance_metrics(),
694 + )
695 + assert_draft_model_correctness(sd_case, enforce_eager)
696 +
697 +
698 +def test_draft_model_tensor_parallelism():
699 + """Ensure spec decode works when running with TP > 1."""
700 + _skip_if_insufficient_gpus_for_tp(2)
701 + sd_case = ArgsTest(
702 + target_model="Qwen/Qwen3-1.7B",
703 + target_tensor_parallel_size=2,
704 + draft_model="Qwen/Qwen3-0.6B",
705 + draft_tensor_parallel_size=2,
706 + **some_high_acceptance_metrics(),
707 + )
708 + assert_draft_model_correctness(sd_case, enforce_eager=False)
709 +
710 +
711 +def test_draft_model_engine_args_tensor_parallelism():
712 + """Ensure the vllm_config for the draft model is created correctly,
713 + and independently of the target model (quantization, TP, etc.)"""
714 + _skip_if_insufficient_gpus_for_tp(2)
715 +
716 + engine_args = EngineArgs(
717 + model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized
718 + tensor_parallel_size=2,
719 + speculative_config={
720 + "model": "Qwen/Qwen3-0.6B", # <<< draft not quantized
721 + "method": "draft_model",
722 + "num_speculative_tokens": 3,
723 + "draft_tensor_parallel_size": 1, # <<< valid arg name
724 + },
725 + )
726 + tgt_vllm_config: VllmConfig = engine_args.create_engine_config()
727 + assert tgt_vllm_config.parallel_config.tensor_parallel_size == 2
728 + assert tgt_vllm_config.quant_config.get_name() == "fp8"
729 +
730 + draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model(tgt_vllm_config)
731 + assert draft_vllm_config.parallel_config.tensor_parallel_size == 1
732 + assert draft_vllm_config.quant_config is None
733 +
734 +
735 +def test_draft_model_engine_args_rejects_invalid_tp_argname():
736 + """The user should pass "draft_tensor_parallel_size" rather than
737 + "tensor_parallel_size". We enforce this with validation."""
738 +
739 + engine_args = EngineArgs(
740 + model="Qwen/Qwen3-1.7B",
741 + tensor_parallel_size=1,
742 + speculative_config={
743 + "model": "Qwen/Qwen3-0.6B",
744 + "method": "draft_model",
745 + "num_speculative_tokens": 3,
746 + "tensor_parallel_size": 1, # <<< invalid arg name
747 + },
748 + )
749 + with pytest.raises(ValueError):
750 + engine_args.create_engine_config()
751 +
752 +
753 +def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
754 + """Compare the outputs using and not using speculative decoding.
755 + In the greedy decoding case, the outputs must match EXACTLY."""
756 + test_prompts: list[Messages] = get_messages(
757 + dataset=args.dataset, n=args.num_prompts
758 + )
759 +
760 + spec_llm = LLM(
761 + model=args.target_model,
762 + speculative_config={
763 + "model": args.draft_model,
764 + "method": "draft_model",
765 + "num_speculative_tokens": args.num_speculative_tokens,
766 + "max_model_len": args.max_model_len,
767 + "enforce_eager": enforce_eager,
768 + "draft_tensor_parallel_size": args.draft_tensor_parallel_size,
769 + "max_num_seqs": 100, # limit cudagraph capture runtime
770 + },
771 + max_model_len=args.max_model_len,
772 + gpu_memory_utilization=args.gpu_memory_utilization,
773 + tensor_parallel_size=args.target_tensor_parallel_size,
774 + enforce_eager=enforce_eager,
775 + disable_log_stats=False, # enables get_metrics()
776 + )
777 + # we don't check the outputs, only check the metrics
778 + spec_llm.chat(test_prompts, args.sampling_config)
779 + metrics = spec_llm.get_metrics()
780 +
781 + acceptance_rate: float = compute_acceptance_rate(metrics)
782 + acceptance_len: float = compute_acceptance_len(metrics)
783 + del spec_llm # CLEANUP
784 + torch.cuda.empty_cache()
785 + cleanup_dist_env_and_memory()
786 +
787 + print(
788 + f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
789 + f"temperature={args.sampling_config.temperature:.2f}, "
790 + f"acceptance_rate={acceptance_rate:.2f}, "
791 + f"acceptance_len={acceptance_len:.2f}, "
792 + )
793 +
794 + assert acceptance_rate >= args.expected_acceptance_rate
795 + assert acceptance_len >= args.expected_acceptance_len
796 +
797 +
798 +def get_messages(dataset: str, n: int) -> list[Messages]:
799 + if dataset == "test_prompts":
800 + return get_test_prompts(mm_enabled=False, quiet=True, num_prompts=n)
801 + elif dataset == "likaixin/InstructCoder":
802 + return get_instruct_coder_messages(n=n)
803 + else:
804 + raise NotImplementedError(f"Dataset '{dataset}' not implemented")
805 +
806 +
807 +def some_high_acceptance_metrics() -> dict:
808 + return {
809 + "sampling_config": greedy_sampling(),
810 + "num_speculative_tokens": 3,
811 + "expected_acceptance_len": 2.90 + 1,
812 + "expected_acceptance_rate": 0.90,
813 + }
814 +
815 +
816 +def test_merge_toks_kernel():
817 + device = "cuda"
818 + merged_len = 5 + 2 # len(target_toks) = 5, batch_size = 2
819 + merged = torch.full((merged_len,), -100, device=device) # -100 is arbitrary
820 + is_rejected_tok = torch.full((merged_len,), True, device=device)
821 + grid = (2,)
822 + merge_toks_kernel[grid](
823 + target_toks_ptr=torch.tensor([0, 1, 2, 0, 1], device=device),
824 + next_toks_ptr=torch.tensor([3, 2], device=device),
825 + query_start_locs_ptr=torch.tensor([0, 3], device=device),
826 + query_end_locs_ptr=torch.tensor([2, 4], device=device),
827 + out_ptr_merged_toks=merged,
828 + out_ptr_is_rejected_tok=is_rejected_tok,
829 + target_toks_size=5,
830 + rejected_tok_fill=-1,
831 + )
832 + expected_merged = torch.tensor([0, 1, 2, 3, 0, 1, 2], device=device)
833 + assert torch.allclose(merged, expected_merged)
834 +
835 + expected_rejected_toks = torch.tensor([False] * merged_len, device=device)
836 + assert torch.allclose(is_rejected_tok, expected_rejected_toks)
837 +
838 +
839 +def test_merge_toks_kernel_with_rejected_tokens():
840 + device = "cuda"
841 + merged_size = 9 + 2 # len(target_toks) = 9, batch_size = 2
842 + merged = torch.full((merged_size,), -100, device=device)
843 + is_rejected_tok = torch.full((merged_size,), True, device=device)
844 + grid = (2,)
845 + merge_toks_kernel[grid](
846 + # rejected tokens
847 + # ↓ ↓ ↓ ↓
848 + target_toks_ptr=torch.tensor([0, 1, 2, 13, 14, 15, 0, 1, 22], device=device),
849 + next_toks_ptr=torch.tensor([3, 2], device=device),
850 + query_start_locs_ptr=torch.tensor([0, 6], device=device),
851 + query_end_locs_ptr=torch.tensor([2, 7], device=device),
852 + out_ptr_merged_toks=merged,
853 + out_ptr_is_rejected_tok=is_rejected_tok,
854 + target_toks_size=9,
855 + rejected_tok_fill=-1,
856 + )
857 + expected_merged = torch.tensor([0, 1, 2, 3, -1, -1, -1, 0, 1, 2, -1], device=device)
858 + assert torch.allclose(merged, expected_merged)
859 +
860 + expected_rejected_toks = torch.tensor(
861 + [False, False, False, False, True, True, True, False, False, False, True],
862 + device=device,
863 + )
864 + assert torch.allclose(is_rejected_tok, expected_rejected_toks)
865 +
866 +
867 +def compute_acceptance_rate(metrics: list[Metric]) -> float:
868 + name2metric = {metric.name: metric for metric in metrics}
869 + n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value # type: ignore
870 + if n_draft_toks == 0:
871 + return float("nan")
872 + n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value # type: ignore
873 + return n_accepted_toks / n_draft_toks
874 +
875 +
876 +def compute_acceptance_len(metrics: list[Metric]) -> float:
877 + name2metric = {metric.name: metric for metric in metrics}
878 + n_drafts = name2metric["vllm:spec_decode_num_drafts"].value # type: ignore
879 + n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value # type: ignore
880 + if n_drafts == 0:
881 + return 1
882 + return 1 + (n_accepted_toks / n_drafts)
tests/v1/worker/test_utils.py
+35
@@ -55,3 +55,38 @@ def test_bind_kv_cache_non_attention(default_vllm_config):
55
56 assert runner_kv_caches[0] is kv_cache["model.layers.20.attn"]
57 assert runner_kv_caches[1] is kv_cache["model.layers.28.attn"]
58 +
59 +
60 +def test_bind_kv_cache_draft_model(default_vllm_config):
61 + from vllm.attention.layer import Attention
62 +
63 + layer_names = [
64 + "model.layers.0.attn",
65 + "model.layers.1.attn",
66 + "draft_model.layers.0.attn",
67 + "draft_model.layers.1.attn",
68 + ]
69 + ctx = {
70 + layer_name: Attention(32, 128, 0.1, prefix=layer_name)
71 + for layer_name in layer_names
72 + }
73 + kv_cache = {layer_name: torch.zeros((1,)) for layer_name in layer_names}
74 + runner_kv_caches: list[torch.Tensor] = []
75 + bind_kv_cache(kv_cache, ctx, runner_kv_caches)
76 +
77 + assert ctx["model.layers.0.attn"].kv_cache[0] is kv_cache["model.layers.0.attn"]
78 + assert ctx["model.layers.1.attn"].kv_cache[0] is kv_cache["model.layers.1.attn"]
79 + assert (
80 + ctx["draft_model.layers.0.attn"].kv_cache[0]
81 + is kv_cache["draft_model.layers.0.attn"]
82 + )
83 + assert (
84 + ctx["draft_model.layers.1.attn"].kv_cache[0]
85 + is kv_cache["draft_model.layers.1.attn"]
86 + )
87 +
88 + # caches are ordered by layer_index, interleaving target and draft model
89 + assert runner_kv_caches[0] is kv_cache["model.layers.0.attn"]
90 + assert runner_kv_caches[1] is kv_cache["draft_model.layers.0.attn"]
91 + assert runner_kv_caches[2] is kv_cache["model.layers.1.attn"]
92 + assert runner_kv_caches[3] is kv_cache["draft_model.layers.1.attn"]
vllm/benchmarks/datasets.py
+10 -9
@@ -2593,17 +2593,10 @@ class InstructCoderDataset(HuggingFaceDataset):
2593 request_id_prefix: str = "",
2594 no_oversample: bool = False,
2595 **kwargs,
2596 - ) -> list:
2596 + ) -> list[SampleRequest]:
2597 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2598 sampled_requests = []
2599 - for i, item in enumerate(self.data):
2600 - if len(sampled_requests) >= num_requests:
2601 - break
2602 - prompt = (
2603 - f"{item['input']}\n\n{item['instruction']} Just output "
2604 - "the code, do not include any explanation."
2605 - )
2606 -
2599 + for i, prompt in enumerate(self.sample_prompts(n=num_requests)):
2600 # apply template
2601 if not skip_chat_template:
2602 prompt = tokenizer.apply_chat_template(
@@ -2626,6 +2619,14 @@ class InstructCoderDataset(HuggingFaceDataset):
2619 )
2620 return sampled_requests
2621
2622 + def sample_prompts(self, n: int) -> Iterator[str]:
2623 + for item in self.data.take(n):
2624 + prompt = (
2625 + f"{item['input']}\n\n{item['instruction']} Just output "
2626 + "the code, do not include any explanation."
2627 + )
2628 + yield prompt
2629 +
2630
2631 # -----------------------------------------------------------------------------
2632 # MT-Bench Dataset Implementation
vllm/benchmarks/lib/ready_checker.py
+6
@@ -8,8 +8,12 @@ import time
8 import aiohttp
9 from tqdm.asyncio import tqdm
10
11 +from vllm.logger import init_logger
12 +
13 from .endpoint_request_func import RequestFunc, RequestFuncInput, RequestFuncOutput
14
15 +logger = init_logger(__name__)
16 +
17
18 async def wait_for_endpoint(
19 request_func: RequestFunc,
@@ -61,6 +65,8 @@ async def wait_for_endpoint(
65 if output.success:
66 pbar.close()
67 return output
68 + else:
69 + logger.warning("Endpoint is not ready. Error='%s'", output.error)
70 except aiohttp.ClientConnectorError:
71 pass
72
vllm/config/parallel.py
+4
@@ -3,6 +3,7 @@
3
4 import os
5 from collections.abc import Callable
6 +from dataclasses import replace
7 from typing import TYPE_CHECKING, Any, Literal
8
9 import torch
@@ -709,3 +710,6 @@ class ParallelConfig:
710 )
711
712 return self
713 +
714 + def replace(self, **kwargs) -> Self:
715 + return replace(self, **kwargs)
vllm/config/speculative.py
+33 -6
@@ -77,6 +77,9 @@ class SpeculativeConfig:
77 draft_tensor_parallel_size: int | None = Field(default=None, ge=1)
78 """The degree of the tensor parallelism for the draft model. Can only be 1
79 or the same as the target model's tensor parallel size."""
80 + tensor_parallel_size: int | None = None
81 + """Users should pass "draft_tensor_parallel_size". This parameter's purpose is to
82 + warn users when they mistakenly provide the wrong argument."""
83
84 # Draft model configuration
85 quantization: me_quant.QuantizationMethods | None = None
@@ -397,13 +400,11 @@ class SpeculativeConfig:
400 "one layer. Might need some code changes "
401 "to support multiple layers."
402 )
403 + elif self.method == "draft_model":
404 + pass
405 else:
401 - self.method = "draft_model"
406 raise NotImplementedError(
403 - "Speculative decoding with draft model is not "
404 - "supported yet. Please consider using other "
405 - "speculative decoding methods such as ngram, medusa, "
406 - "eagle, or mtp."
407 + f"Unsupported speculative method: '{self.method}'"
408 )
409
410 # Replace hf_config for EAGLE draft_model
@@ -631,6 +632,12 @@ class SpeculativeConfig:
632
633 @model_validator(mode="after")
634 def _verify_args(self) -> Self:
635 + if self.tensor_parallel_size is not None:
636 + raise ValueError(
637 + "'tensor_parallel_size' is not a valid argument in the "
638 + "speculative_config. Please pass 'draft_tensor_parallel_size' instead."
639 + )
640 +
641 if self.num_speculative_tokens is None:
642 raise ValueError(
643 "num_speculative_tokens must be provided with "
@@ -669,12 +676,32 @@ class SpeculativeConfig:
676 f"Eagle3 is only supported for {eagle3_target_supported} models. " # noqa: E501
677 f"Got {self.target_model_config.hf_text_config.model_type=}"
678 )
672 -
679 + self.verify_equal_vocab_size_if_draft_model()
680 return self
681
682 + def verify_equal_vocab_size_if_draft_model(self):
683 + if (
684 + self.method == "draft_model"
685 + and self.target_model_config is not None
686 + and self.draft_model_config is not None
687 + ):
688 + target_vocab_size = self.target_model_config.get_vocab_size()
689 + draft_vocab_size = self.draft_model_config.get_vocab_size()
690 + if target_vocab_size != draft_vocab_size:
691 + raise ValueError(
692 + f"Target and draft model should have the same vocabulary size. "
693 + f"Target model vocab_size={target_vocab_size}. "
694 + f"Draft model vocab_size={draft_vocab_size}. "
695 + f"Using models with different tokenizers can cause out-of-bounds "
696 + f"errors during speculative decoding."
697 + )
698 +
699 def use_eagle(self) -> bool:
700 return self.method in ("eagle", "eagle3", "mtp")
701
702 + def uses_draft_model(self) -> bool:
703 + return self.method == "draft_model"
704 +
705 def __repr__(self) -> str:
706 method = self.method
707 model = None if method in ("ngram", "suffix") else self.draft_model_config.model
vllm/config/vllm.py
+23 -13
@@ -1214,10 +1214,19 @@ class VllmConfig:
1214 compilation_config = self.compilation_config
1215 computed_compile_ranges_split_points = []
1216
1217 - # The upper bound of the compile ranges is the max_num_batched_tokens
1218 - max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
1219 - if max_num_batched_tokens is not None:
1220 - computed_compile_ranges_split_points.append(max_num_batched_tokens)
1217 + # The upper bound of the compile ranges is the max_num_batched_tokens.
1218 + # For speculative decoding with draft model, the compile range must be extended
1219 + # by 1 for each sequence.
1220 + compile_range_end = self.scheduler_config.max_num_batched_tokens
1221 + if compile_range_end is not None:
1222 + do_extend: bool = (
1223 + self.speculative_config is not None
1224 + and self.speculative_config.uses_draft_model()
1225 + )
1226 + if do_extend:
1227 + compile_range_end += self.scheduler_config.max_num_seqs
1228 +
1229 + computed_compile_ranges_split_points.append(compile_range_end)
1230
1231 # Add the compile ranges for flashinfer
1232 if compilation_config.pass_config.fuse_allreduce_rms:
@@ -1228,10 +1237,7 @@ class VllmConfig:
1237 self.model_config.get_hidden_size()
1238 * self.model_config.dtype.itemsize
1239 )
1231 - if (
1232 - max_num_batched_tokens is not None
1233 - and max_token_num < max_num_batched_tokens
1234 - ):
1240 + if compile_range_end is not None and max_token_num < compile_range_end:
1241 computed_compile_ranges_split_points.append(max_token_num)
1242 else:
1243 logger.debug(
@@ -1243,11 +1249,7 @@ class VllmConfig:
1249 for x in compilation_config.compile_ranges_split_points:
1250 assert isinstance(x, int)
1251 assert x > 0, f"Invalid compile range split point: {x}"
1246 - if (
1247 - max_num_batched_tokens is not None
1248 - and x < max_num_batched_tokens
1249 - and x > 1
1250 - ):
1252 + if compile_range_end is not None and x < compile_range_end and x > 1:
1253 computed_compile_ranges_split_points.append(x)
1254 compilation_config.compile_ranges_split_points = sorted(
1255 computed_compile_ranges_split_points
@@ -1316,6 +1318,14 @@ class VllmConfig:
1318 path = self.compilation_config.debug_dump_path / append_path
1319 return path
1320
1321 + def replace(self, **kwargs):
1322 + """
1323 + Replace attributes of the config, and 'recompute' the config.
1324 + dataclass.replace() calls __init__() and __post_init__(), source:
1325 + https://docs.python.org/3/library/dataclasses.html#dataclasses.replace
1326 + """
1327 + return replace(self, **kwargs)
1328 +
1329 def __str__(self):
1330 return (
1331 f"model={self.model_config.model!r}, "
vllm/engine/arg_utils.py
-15
@@ -1776,21 +1776,6 @@ class EngineArgs:
1776 ):
1777 _raise_unsupported_error(feature_name="Concurrent Partial Prefill")
1778
1779 - # N-gram, Medusa, and Eagle are supported for speculative decoding.
1780 - if self.speculative_config is not None:
1781 - # speculative_config could still be a dict at this point
1782 - if isinstance(self.speculative_config, dict):
1783 - method = self.speculative_config.get("method", None)
1784 - else:
1785 - method = self.speculative_config.method
1786 -
1787 - if method == "draft_model":
1788 - raise NotImplementedError(
1789 - "Draft model speculative decoding is not supported yet. "
1790 - "Please consider using other speculative decoding methods "
1791 - "such as ngram, medusa, eagle, or mtp."
1792 - )
1793 -
1779 if self.pipeline_parallel_size > 1:
1780 supports_pp = getattr(
1781 self.distributed_executor_backend, "supports_pp", False
vllm/model_executor/model_loader/__init__.py
+7 -2
@@ -124,12 +124,17 @@ def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
124
125
126 def get_model(
127 - *, vllm_config: VllmConfig, model_config: ModelConfig | None = None
127 + *,
128 + vllm_config: VllmConfig,
129 + model_config: ModelConfig | None = None,
130 + prefix: str = "",
131 ) -> nn.Module:
132 loader = get_model_loader(vllm_config.load_config)
133 if model_config is None:
134 model_config = vllm_config.model_config
132 - return loader.load_model(vllm_config=vllm_config, model_config=model_config)
135 + return loader.load_model(
136 + vllm_config=vllm_config, model_config=model_config, prefix=prefix
137 + )
138
139
140 __all__ = [
vllm/model_executor/model_loader/base_loader.py
+2 -2
@@ -36,7 +36,7 @@ class BaseModelLoader(ABC):
36 raise NotImplementedError
37
38 def load_model(
39 - self, vllm_config: VllmConfig, model_config: ModelConfig
39 + self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
40 ) -> nn.Module:
41 """Load a model with the given configurations."""
42 device_config = vllm_config.device_config
@@ -48,7 +48,7 @@ class BaseModelLoader(ABC):
48 with set_default_torch_dtype(model_config.dtype):
49 with target_device:
50 model = initialize_model(
51 - vllm_config=vllm_config, model_config=model_config
51 + vllm_config=vllm_config, model_config=model_config, prefix=prefix
52 )
53
54 log_model_inspection(model)
vllm/model_executor/model_loader/gguf_loader.py
+2 -2
@@ -335,7 +335,7 @@ class GGUFModelLoader(BaseModelLoader):
335 )
336
337 def load_model(
338 - self, vllm_config: VllmConfig, model_config: ModelConfig
338 + self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
339 ) -> nn.Module:
340 device_config = vllm_config.device_config
341 local_model_path = self._prepare_weights(model_config)
@@ -364,7 +364,7 @@ class GGUFModelLoader(BaseModelLoader):
364 target_device = torch.device(device_config.device)
365 with set_default_torch_dtype(model_config.dtype):
366 with target_device:
367 - model = initialize_model(vllm_config=vllm_config)
367 + model = initialize_model(vllm_config=vllm_config, prefix=prefix)
368 self.load_weights(model, model_config)
369
370 process_weights_after_loading(model, model_config, target_device)
vllm/model_executor/model_loader/tensorizer_loader.py
+4 -3
@@ -68,6 +68,7 @@ class TensorizerLoader(BaseModelLoader):
68 def _load_model_serialized_cpu(
69 self,
70 vllm_config: VllmConfig,
71 + prefix: str = "",
72 ) -> nn.Module:
73 """Load a serialized model with tensorizer to the CPU.
74
@@ -80,7 +81,7 @@ class TensorizerLoader(BaseModelLoader):
81 model_config = vllm_config.model_config
82 with set_default_torch_dtype(model_config.dtype):
83 with torch.device(device_config.device):
83 - model = initialize_model(vllm_config=vllm_config)
84 + model = initialize_model(vllm_config=vllm_config, prefix=prefix)
85
86 model.load_weights(self._get_weights_iterator())
87 return model.eval()
@@ -112,7 +113,7 @@ class TensorizerLoader(BaseModelLoader):
113 model.load_weights(self._get_weights_iterator())
114
115 def load_model(
115 - self, vllm_config: VllmConfig, model_config: ModelConfig
116 + self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
117 ) -> nn.Module:
118 parallel_config = vllm_config.parallel_config
119 self._verify_config(model_config, parallel_config)
@@ -134,7 +135,7 @@ class TensorizerLoader(BaseModelLoader):
135 )
136 self.load_weights(model, model_config)
137 return model
137 - return self._load_model_serialized_cpu(vllm_config=vllm_config)
138 + return self._load_model_serialized_cpu(vllm_config=vllm_config, prefix=prefix)
139
140 @staticmethod
141 def save_model(
vllm/v1/attention/backend.py
+11 -1
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 -from dataclasses import dataclass
5 +from dataclasses import dataclass, replace
6 from enum import Enum
7 from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar, get_args
8
@@ -329,6 +329,16 @@ class CommonAttentionMetadata:
329
330 _num_computed_tokens_cache: torch.Tensor | None = None
331
332 + def batch_size(self) -> int:
333 + return self.seq_lens.shape[0]
334 +
335 + def naive_query_lens(self) -> torch.Tensor:
336 + """Naive because it assumes that query ends where the next query starts."""
337 + return self.query_start_loc[1:] - self.query_start_loc[:-1]
338 +
339 + def replace(self, **kwargs) -> "CommonAttentionMetadata":
340 + return replace(self, **kwargs)
341 +
342 @property
343 @deprecated(
344 """
vllm/v1/attention/backends/utils.py
+32
@@ -818,3 +818,35 @@ def get_dcp_local_seq_lens(
818 )
819 dcp_local_seq_lens = base + remainder
820 return dcp_local_seq_lens.squeeze(1)
821 +
822 +
823 +def extend_all_queries_by_1(
824 + common_attn_metadata: CommonAttentionMetadata,
825 + arange: torch.Tensor,
826 + new_slot_mapping: torch.Tensor,
827 +) -> CommonAttentionMetadata:
828 + """
829 + Creates a new CommonAttentionMetadata with all query lengths increased by 1.
830 + Also all seq lens are increased by 1.
831 + This is useful e.g. in speculative decoding with draft models, where we
832 + extend each sequence by 1 token.
833 + The slot mapping is computed externally, as it requires more information.
834 + """
835 + cad = common_attn_metadata
836 + # query start loc must be increased by [+0, +1, +2, ..., +batch_size]
837 + new_query_start_loc = cad.query_start_loc + arange[: len(cad.query_start_loc)]
838 + new_query_start_loc_cpu = cad.query_start_loc_cpu + torch.arange(
839 + len(cad.query_start_loc_cpu), dtype=torch.int32
840 + )
841 + new_cad = cad.replace(
842 + query_start_loc=new_query_start_loc,
843 + query_start_loc_cpu=new_query_start_loc_cpu,
844 + seq_lens=cad.seq_lens + 1,
845 + # each request is extended by 1 token -> batch_size tokens are added
846 + num_actual_tokens=cad.num_actual_tokens + cad.batch_size(),
847 + # All query lens increase by 1, so max query len increases by 1
848 + max_query_len=cad.max_query_len + 1,
849 + max_seq_len=cad.max_seq_len + 1,
850 + slot_mapping=new_slot_mapping,
851 + )
852 + return new_cad
vllm/v1/core/sched/scheduler.py
+2
@@ -208,6 +208,8 @@ class Scheduler(SchedulerInterface):
208 if speculative_config.use_eagle():
209 self.use_eagle = True
210 self.num_lookahead_tokens = self.num_spec_tokens
211 + if speculative_config.uses_draft_model():
212 + self.num_lookahead_tokens = self.num_spec_tokens
213
214 # Create the KV cache manager.
215 self.kv_cache_manager = KVCacheManager(
vllm/v1/spec_decode/draft_model.py new
+271
@@ -0,0 +1,271 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +import torch
6 +
7 +from vllm.attention.layer import Attention
8 +from vllm.config import VllmConfig, get_layers_from_vllm_config
9 +from vllm.config.speculative import SpeculativeConfig
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.model_loader import get_model
12 +from vllm.triton_utils import tl, triton
13 +from vllm.v1.attention.backends.utils import (
14 + CommonAttentionMetadata,
15 + extend_all_queries_by_1,
16 +)
17 +from vllm.v1.spec_decode.eagle import PADDING_SLOT_ID, SpecDecodeBaseProposer
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +class DraftModelProposer(SpecDecodeBaseProposer):
23 + def __init__(
24 + self,
25 + vllm_config: VllmConfig,
26 + device: torch.device,
27 + runner=None,
28 + ):
29 + super().__init__(
30 + vllm_config=vllm_config,
31 + device=device,
32 + pass_hidden_states_to_model=False,
33 + runner=runner,
34 + )
35 + self._raise_if_multimodal()
36 + self._raise_if_mrope()
37 + self._raise_if_padded_drafter_batch_disabled()
38 + self._raise_if_vocab_size_mismatch()
39 + self._raise_if_draft_tp_mismatch()
40 +
41 + def _block_size(self) -> int:
42 + builder = self._get_attention_metadata_builder()
43 + return builder.kv_cache_spec.block_size
44 +
45 + def _raise_if_multimodal(self):
46 + if self.supports_mm_inputs:
47 + raise NotImplementedError(
48 + "Speculative Decoding with draft models "
49 + "does not support multimodal models yet"
50 + )
51 +
52 + def _raise_if_mrope(self):
53 + if self.draft_model_config.uses_mrope:
54 + raise NotImplementedError(
55 + "Speculative Decoding with draft models does not support M-RoPE yet"
56 + )
57 +
58 + def _raise_if_padded_drafter_batch_disabled(self):
59 + if self.vllm_config.speculative_config.disable_padded_drafter_batch:
60 + raise NotImplementedError(
61 + "Speculative Decoding with draft models only supports "
62 + "padded drafter batch. Please don't pass --disable-padded-drafter-batch"
63 + " in the speculative_config."
64 + )
65 +
66 + def _raise_if_vocab_size_mismatch(self):
67 + self.vllm_config.speculative_config.verify_equal_vocab_size_if_draft_model()
68 +
69 + def _raise_if_draft_tp_mismatch(self):
70 + # Note(Tomas Ruiz) If we run the target model with TP > 1 and
71 + # the draft model with TP = 1, then the different TP ranks collide.
72 + # Specifically when all ranks compile the draft model on rank 0
73 + # (because TP=1), then the torch compile cache is overwritten and corrupted.
74 + # We need a mechanism like this: https://github.com/vllm-project/vllm/pull/5414
75 + # To prevent this error, we assert that both TP sizes must be the same.
76 + spec_cfg: SpeculativeConfig = self.vllm_config.speculative_config
77 + tgt_tp = spec_cfg.target_parallel_config.tensor_parallel_size
78 + draft_tp = spec_cfg.draft_parallel_config.tensor_parallel_size
79 + if draft_tp != tgt_tp:
80 + raise ValueError(
81 + f"Currently, 'draft_tensor_parallel_size' and 'tensor_parallel_size' "
82 + f"must be the same. Got {draft_tp} and {tgt_tp}. "
83 + "Please pass 'draft_tensor_parallel_size' in the speculative_config."
84 + )
85 +
86 + def set_inputs_first_pass(
87 + self,
88 + target_token_ids: torch.Tensor,
89 + next_token_ids: torch.Tensor,
90 + target_positions: torch.Tensor,
91 + last_token_indices: torch.Tensor | None,
92 + cad: CommonAttentionMetadata,
93 + num_rejected_tokens_gpu: torch.Tensor | None,
94 + ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
95 + batch_size = cad.batch_size()
96 + grid = (batch_size,)
97 + start_locs = cad.query_start_loc[:-1]
98 + end_locs = cad.query_start_loc[1:] - 1
99 + if num_rejected_tokens_gpu is not None:
100 + end_locs -= num_rejected_tokens_gpu
101 +
102 + num_tokens = target_token_ids.shape[0] + batch_size
103 + is_rejected_tok = torch.empty(
104 + (num_tokens,), device=self.input_ids.device, dtype=torch.bool
105 + )
106 + merge_toks_kernel[grid](
107 + target_toks_ptr=target_token_ids,
108 + next_toks_ptr=next_token_ids,
109 + query_start_locs_ptr=start_locs,
110 + query_end_locs_ptr=end_locs,
111 + out_ptr_merged_toks=self.input_ids,
112 + out_ptr_is_rejected_tok=is_rejected_tok,
113 + target_toks_size=target_token_ids.shape[0],
114 + # passing a negative rejected_tok_fill value will raise an error
115 + # when the value is used to index into embeddings.
116 + # Therefore, we pass a valid integer, e.g. 0.
117 + rejected_tok_fill=0,
118 + )
119 + merge_toks_kernel[grid](
120 + target_toks_ptr=target_positions,
121 + next_toks_ptr=target_positions[end_locs] + 1,
122 + query_start_locs_ptr=start_locs,
123 + query_end_locs_ptr=end_locs,
124 + out_ptr_merged_toks=self.positions,
125 + out_ptr_is_rejected_tok=is_rejected_tok,
126 + target_toks_size=target_positions.shape[0],
127 + rejected_tok_fill=0,
128 + )
129 +
130 + # recompute slot mapping
131 + new_slot_mapping = compute_new_slot_mapping(
132 + cad=cad,
133 + new_positions=self.positions[:num_tokens],
134 + is_rejected_token_mask=is_rejected_tok,
135 + block_size=self._block_size(),
136 + max_model_len=self.max_model_len,
137 + )
138 + # update common_attn_metadata
139 + new_cad: CommonAttentionMetadata = extend_all_queries_by_1(
140 + cad,
141 + arange=self.arange,
142 + new_slot_mapping=new_slot_mapping,
143 + )
144 +
145 + new_last_token_indices = new_cad.query_start_loc[1:] - 1
146 + if num_rejected_tokens_gpu is not None:
147 + new_last_token_indices -= num_rejected_tokens_gpu
148 +
149 + return num_tokens, new_last_token_indices, new_cad
150 +
151 + def load_model(self, target_model: Any) -> None:
152 + """Takes target_model to satisfy the type checker."""
153 +
154 + # This must be computed before loading the draft model
155 + # because that mutates the forward_context of the vllm_config
156 + target_attn_layer_names = set(
157 + get_layers_from_vllm_config(self.vllm_config, Attention).keys()
158 + )
159 +
160 + from vllm.compilation.backends import set_model_tag
161 +
162 + draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model(
163 + target_model_vllm_config=self.vllm_config
164 + )
165 + logger.info(
166 + "Starting to load draft model %s. TP=%d, rank=%d",
167 + draft_vllm_config.model_config.model,
168 + draft_vllm_config.parallel_config.tensor_parallel_size,
169 + draft_vllm_config.parallel_config.rank,
170 + )
171 + with set_model_tag("draft_model"):
172 + self.model = get_model(vllm_config=draft_vllm_config, prefix="draft_model")
173 +
174 + # This must be computed after loading the draft model
175 + # because that mutates the forward_context of the vllm_config
176 + draft_attn_layer_names = (
177 + get_layers_from_vllm_config(self.vllm_config, Attention).keys()
178 + - target_attn_layer_names
179 + )
180 + self.attn_layer_names = list(draft_attn_layer_names)
181 +
182 +
183 +def create_vllm_config_for_draft_model(
184 + target_model_vllm_config: VllmConfig,
185 +) -> VllmConfig:
186 + """The vllm_config is configured for the target model, e.g.
187 + its quant_config and parallel_config. But the draft model is potentially
188 + quantized differently, and has potentially different tensor_parallel_size.
189 + This function creates a new vllm_config configured for the draft model.
190 + The vllm_config is useful when loading the draft model with get_model().
191 + """
192 + old = target_model_vllm_config
193 + new_parallel_config = old.speculative_config.draft_parallel_config.replace(
194 + rank=old.parallel_config.rank
195 + )
196 + new: VllmConfig = old.replace(
197 + quant_config=None, # quant_config is recomputed in __init__()
198 + model_config=old.speculative_config.draft_model_config,
199 + parallel_config=new_parallel_config,
200 + )
201 + return new
202 +
203 +
204 +def compute_new_slot_mapping(
205 + cad: CommonAttentionMetadata,
206 + new_positions: torch.Tensor,
207 + is_rejected_token_mask: torch.Tensor,
208 + block_size: int,
209 + max_model_len: int,
210 +):
211 + batch_size, n_blocks_per_req = cad.block_table_tensor.shape
212 + req_indices = torch.arange(batch_size, device=cad.query_start_loc.device)
213 + req_indices = torch.repeat_interleave(
214 + req_indices, cad.naive_query_lens() + 1, output_size=len(new_positions)
215 + )
216 + # Clamp the positions to prevent an out-of-bounds error when indexing
217 + # into block_table_tensor.
218 + clamped_positions = torch.clamp(new_positions, max=max_model_len - 1)
219 + block_table_indices = (
220 + req_indices * n_blocks_per_req + clamped_positions // block_size
221 + )
222 + block_nums = cad.block_table_tensor.view(-1)[block_table_indices]
223 + block_offsets = clamped_positions % block_size
224 + new_slot_mapping = block_nums * block_size + block_offsets
225 + # Mask out the position ids that exceed the max model length.
226 + exceeds_max_model_len = new_positions >= max_model_len
227 + new_slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID)
228 + # Mask out rejected tokens to prevent saves to the KV cache.
229 + new_slot_mapping.masked_fill_(is_rejected_token_mask, PADDING_SLOT_ID)
230 + return new_slot_mapping
231 +
232 +
233 +@triton.jit
234 +def merge_toks_kernel(
235 + target_toks_ptr,
236 + next_toks_ptr,
237 + query_start_locs_ptr,
238 + query_end_locs_ptr,
239 + out_ptr_merged_toks,
240 + out_ptr_is_rejected_tok,
241 + target_toks_size,
242 + rejected_tok_fill,
243 +):
244 + """
245 + Merges the `target_toks_ptr` and the `next_toks_ptr` into a new tensor
246 + called `out_ptr_merged_toks`. Rejected tokens are those after the
247 + `query_end_locs_ptr` and before the next `query_start_locs_ptr`. Fills the
248 + rejected tokens positions with the value `rejected_tok_fill`. Also fills a mask
249 + of the rejected tokens in `out_ptr_is_rejected_tok`.
250 + """
251 + pid = tl.program_id(0)
252 + start_loc = tl.load(query_start_locs_ptr + pid)
253 + is_last_program = pid == tl.num_programs(0) - 1
254 + if is_last_program:
255 + next_start_loc = target_toks_size.to(tl.int32)
256 + else:
257 + next_start_loc = tl.load(query_start_locs_ptr + pid + 1).to(tl.int32)
258 +
259 + end_loc = tl.load(query_end_locs_ptr + pid)
260 + new_val = tl.load(next_toks_ptr + pid)
261 + for i in range(start_loc, next_start_loc + 1):
262 + if i <= end_loc: # copy existing tokens
263 + old_val = tl.load(target_toks_ptr + i)
264 + tl.store(out_ptr_merged_toks + pid + i, old_val)
265 + tl.store(out_ptr_is_rejected_tok + pid + i, False)
266 + elif i == end_loc + 1: # copy bonus token
267 + tl.store(out_ptr_merged_toks + pid + i, new_val)
268 + tl.store(out_ptr_is_rejected_tok + pid + i, False)
269 + else: # fill rejected tokens
270 + tl.store(out_ptr_merged_toks + pid + i, rejected_tok_fill)
271 + tl.store(out_ptr_is_rejected_tok + pid + i, True)
vllm/v1/spec_decode/eagle.py
+98 -38
@@ -53,11 +53,12 @@ logger = init_logger(__name__)
53 PADDING_SLOT_ID = -1
54
55
56 -class EagleProposer:
56 +class SpecDecodeBaseProposer:
57 def __init__(
58 self,
59 vllm_config: VllmConfig,
60 device: torch.device,
61 + pass_hidden_states_to_model: bool,
62 runner=None,
63 ):
64 self.vllm_config = vllm_config
@@ -65,6 +66,7 @@ class EagleProposer:
66 assert self.speculative_config is not None
67 self.draft_model_config = self.speculative_config.draft_model_config
68 self.method = self.speculative_config.method
69 + self.pass_hidden_states_to_model = pass_hidden_states_to_model
70
71 self.runner = runner
72 self.device = device
@@ -72,7 +74,11 @@ class EagleProposer:
74 self.max_model_len = vllm_config.model_config.max_model_len
75 self.dp_rank = vllm_config.parallel_config.data_parallel_rank
76 self.num_speculative_tokens = self.speculative_config.num_speculative_tokens
75 - self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
77 + # The drafter can get longer sequences than the target model.
78 + max_batch_size = vllm_config.scheduler_config.max_num_seqs
79 + self.max_num_tokens = (
80 + vllm_config.scheduler_config.max_num_batched_tokens + max_batch_size
81 + )
82 self.token_arange_np = np.arange(self.max_num_tokens)
83 # We need to get the hidden size from the draft model config because
84 # the draft model's hidden size can be different from the target model's
@@ -143,7 +149,6 @@ class EagleProposer:
149
150 # We need +1 here because the arange is used to set query_start_loc,
151 # which has one more element than batch_size.
146 - max_batch_size = vllm_config.scheduler_config.max_num_seqs
152 max_num_slots_for_arange = max(max_batch_size + 1, self.max_num_tokens)
153 self.arange = torch.arange(
154 max_num_slots_for_arange, device=device, dtype=torch.int32
@@ -245,11 +250,7 @@ class EagleProposer:
250 mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
251 num_rejected_tokens_gpu: torch.Tensor | None = None,
252 ) -> torch.Tensor:
248 - num_tokens = target_token_ids.shape[0]
249 - batch_size = next_token_ids.shape[0]
250 -
251 - if last_token_indices is None:
252 - last_token_indices = common_attn_metadata.query_start_loc[1:] - 1
253 + batch_size = common_attn_metadata.batch_size()
254
255 if self.method == "eagle3":
256 assert isinstance(self.model, Eagle3LlamaForCausalLM)
@@ -257,12 +258,17 @@ class EagleProposer:
258 target_hidden_states
259 )
260 assert target_hidden_states.shape[-1] == self.hidden_size
260 - # Shift the input ids by one token.
261 - # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
262 - self.input_ids[: num_tokens - 1] = target_token_ids[1:]
263 - # Replace the last token with the next token.
264 - # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
265 - self.input_ids[last_token_indices] = next_token_ids
261 +
262 + num_tokens, last_token_indices, common_attn_metadata = (
263 + self.set_inputs_first_pass(
264 + target_token_ids=target_token_ids,
265 + next_token_ids=next_token_ids,
266 + target_positions=target_positions,
267 + last_token_indices=last_token_indices,
268 + cad=common_attn_metadata,
269 + num_rejected_tokens_gpu=num_rejected_tokens_gpu,
270 + )
271 + )
272
273 assert self.runner is not None
274
@@ -311,9 +317,10 @@ class EagleProposer:
317 if num_tokens_across_dp is not None:
318 num_tokens_across_dp[self.dp_rank] = num_input_tokens
319
314 - # copy inputs to buffer for cudagraph
315 - self._set_positions(num_tokens, target_positions)
316 - self.hidden_states[:num_tokens] = target_hidden_states
320 + if self.pass_hidden_states_to_model:
321 + # target_hidden_states and self.hidden_states can have different
322 + # hidden dims. E.g. large target model and small draft model.
323 + self.hidden_states[:num_tokens] = target_hidden_states
324
325 if self.supports_mm_inputs:
326 mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)
@@ -330,6 +337,14 @@ class EagleProposer:
337 input_ids = self.input_ids[:num_input_tokens]
338 inputs_embeds = None
339
340 + model_kwargs = {
341 + "input_ids": input_ids,
342 + "positions": self._get_positions(num_input_tokens),
343 + "inputs_embeds": inputs_embeds,
344 + }
345 + if self.pass_hidden_states_to_model:
346 + model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]
347 +
348 with set_forward_context(
349 per_layer_attn_metadata,
350 self.vllm_config,
@@ -337,17 +352,13 @@ class EagleProposer:
352 num_tokens_across_dp=num_tokens_across_dp,
353 cudagraph_runtime_mode=cudagraph_runtime_mode,
354 ):
340 - ret_hidden_states = self.model(
341 - input_ids=input_ids,
342 - positions=self._get_positions(num_input_tokens),
343 - hidden_states=self.hidden_states[:num_input_tokens],
344 - inputs_embeds=inputs_embeds,
345 - )
346 - if self.method == "mtp":
355 + ret_hidden_states = self.model(**model_kwargs)
356 + if not self.model_returns_tuple():
357 last_hidden_states = ret_hidden_states
358 hidden_states = last_hidden_states
359 else:
360 last_hidden_states, hidden_states = ret_hidden_states
361 +
362 sample_hidden_states = last_hidden_states[last_token_indices]
363 logits = self.model.compute_logits(sample_hidden_states)
364
@@ -357,9 +368,9 @@ class EagleProposer:
368 return draft_token_ids.view(-1, 1)
369
370 if self.uses_mrope:
360 - positions = target_positions[:, last_token_indices]
371 + positions = self.positions[:, last_token_indices]
372 else:
362 - positions = target_positions[last_token_indices]
373 + positions = self.positions[last_token_indices]
374 if self.method in (
375 "deepseek_mtp",
376 "ernie_mtp",
@@ -527,6 +538,14 @@ class EagleProposer:
538 inputs_embeds = None
539
540 # Run the model.
541 + model_kwargs = {
542 + "input_ids": input_ids,
543 + "positions": self._get_positions(input_batch_size),
544 + "inputs_embeds": inputs_embeds,
545 + }
546 + if self.pass_hidden_states_to_model:
547 + model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]
548 +
549 with set_forward_context(
550 per_layer_attn_metadata,
551 self.vllm_config,
@@ -534,17 +553,13 @@ class EagleProposer:
553 num_tokens_across_dp=batch_size_across_dp,
554 cudagraph_runtime_mode=cudagraph_runtime_mode,
555 ):
537 - ret_hidden_states = self.model(
538 - input_ids=input_ids,
539 - positions=self._get_positions(input_batch_size),
540 - hidden_states=self.hidden_states[:input_batch_size],
541 - inputs_embeds=inputs_embeds,
542 - )
543 - if self.method == "mtp":
556 + ret_hidden_states = self.model(**model_kwargs)
557 + if not self.model_returns_tuple():
558 last_hidden_states = ret_hidden_states
559 hidden_states = ret_hidden_states
560 else:
561 last_hidden_states, hidden_states = ret_hidden_states
562 +
563 hidden_states = hidden_states[:batch_size]
564 logits = self.model.compute_logits(last_hidden_states[:batch_size])
565 draft_token_ids = logits.argmax(dim=-1)
@@ -554,6 +569,34 @@ class EagleProposer:
569 draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
570 return draft_token_ids
571
572 + def set_inputs_first_pass(
573 + self,
574 + target_token_ids: torch.Tensor,
575 + next_token_ids: torch.Tensor,
576 + target_positions: torch.Tensor,
577 + last_token_indices: torch.Tensor | None,
578 + cad: CommonAttentionMetadata,
579 + num_rejected_tokens_gpu: torch.Tensor | None,
580 + ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
581 + if last_token_indices is None:
582 + last_token_indices = cad.query_start_loc[1:] - 1
583 +
584 + num_tokens = target_token_ids.shape[0]
585 + # Shift the input ids by one token.
586 + # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
587 + self.input_ids[: num_tokens - 1] = target_token_ids[1:]
588 + # Replace the last token with the next token.
589 + # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
590 + self.input_ids[last_token_indices] = next_token_ids
591 +
592 + # copy inputs to buffer for cudagraph
593 + self._set_positions(num_tokens, target_positions)
594 +
595 + return num_tokens, last_token_indices, cad
596 +
597 + def model_returns_tuple(self) -> bool:
598 + return self.method not in ("mtp", "draft_model")
599 +
600 def prepare_next_token_ids_cpu(
601 self,
602 sampled_token_ids: list[list[int]],
@@ -1214,12 +1257,14 @@ class EagleProposer:
1257 input_ids = self.input_ids[:num_input_tokens]
1258 inputs_embeds = None
1259
1217 - self.model(
1260 + kwargs = dict(
1261 input_ids=input_ids,
1262 positions=self._get_positions(num_input_tokens),
1220 - hidden_states=self.hidden_states[:num_input_tokens],
1263 inputs_embeds=inputs_embeds,
1264 )
1265 + if self.pass_hidden_states_to_model:
1266 + kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]
1267 + self.model(**kwargs)
1268
1269 def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:
1270 """Find and return the attention metadata builders for EAGLE layers.
@@ -1264,8 +1309,8 @@ class EagleProposer:
1309
1310 def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
1311 """
1267 - Validate that all eagle layers belong to the same KVCacheGroup.
1268 - Need this assumption to ensure all eagle layers can use the
1312 + Validate that all drafting layers belong to the same KVCacheGroup.
1313 + Need this assumption to ensure all drafting layers can use the
1314 same AttentionMetadata.
1315 May extend to multiple AttentionMetadata in the future.
1316 """
@@ -1283,7 +1328,7 @@ class EagleProposer:
1328 )
1329 )
1330 == 1
1286 - ), "All eagle layers should belong to the same kv cache group"
1331 + ), "All drafting layers should belong to the same kv cache group"
1332
1333 def _pad_batch_across_dp(
1334 self,
@@ -1308,6 +1353,21 @@ class EagleProposer:
1353 return num_tokens_dp_padded, num_toks_across_dp
1354
1355
1356 +class EagleProposer(SpecDecodeBaseProposer):
1357 + def __init__(
1358 + self,
1359 + vllm_config: VllmConfig,
1360 + device: torch.device,
1361 + runner=None,
1362 + ):
1363 + super().__init__(
1364 + vllm_config,
1365 + device,
1366 + pass_hidden_states_to_model=True,
1367 + runner=runner,
1368 + )
1369 +
1370 +
1371 # NOTE(woosuk): Currently, the below code is not used and we always use argmax
1372 # to sample the draft tokens. We will use this after we find a way to manage
1373 # the draft prob tensor.
vllm/v1/worker/gpu_model_runner.py
+31 -10
@@ -145,6 +145,7 @@ from vllm.v1.sample.logits_processor.interface import LogitsProcessor
145 from vllm.v1.sample.metadata import SamplingMetadata
146 from vllm.v1.sample.rejection_sampler import RejectionSampler
147 from vllm.v1.sample.sampler import Sampler
148 +from vllm.v1.spec_decode.draft_model import DraftModelProposer
149 from vllm.v1.spec_decode.eagle import EagleProposer
150 from vllm.v1.spec_decode.medusa import MedusaProposer
151 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
@@ -432,10 +433,20 @@ class GPUModelRunner(
433 # layers in the draft model.
434 if self.speculative_config and get_pp_group().is_last_rank:
435 self.drafter: (
435 - NgramProposer | SuffixDecodingProposer | EagleProposer | MedusaProposer
436 + NgramProposer
437 + | SuffixDecodingProposer
438 + | EagleProposer
439 + | DraftModelProposer
440 + | MedusaProposer
441 )
442 if self.speculative_config.method == "ngram":
443 self.drafter = NgramProposer(self.vllm_config)
444 + elif self.speculative_config.uses_draft_model():
445 + self.drafter = DraftModelProposer(
446 + vllm_config=self.vllm_config,
447 + device=self.device,
448 + runner=self,
449 + )
450 elif self.speculative_config.method == "suffix":
451 self.drafter = SuffixDecodingProposer(self.vllm_config)
452 elif self.speculative_config.use_eagle():
@@ -3443,10 +3454,13 @@ class GPUModelRunner(
3454 spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
3455 <= self.effective_drafter_max_model_len
3456 )
3446 - if spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch:
3447 - # EAGLE speculative decoding can use the GPU sampled tokens
3457 + use_gpu_toks = (
3458 + spec_config.use_eagle() or spec_config.uses_draft_model()
3459 + ) and not spec_config.disable_padded_drafter_batch
3460 + if use_gpu_toks:
3461 + # EAGLE/DraftModel speculative decoding can use the GPU sampled tokens
3462 # as inputs, and does not need to wait for bookkeeping to finish.
3449 - assert isinstance(self.drafter, EagleProposer)
3463 + assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
3464 sampled_token_ids = sampler_output.sampled_token_ids
3465 if input_fits_in_drafter:
3466 propose_draft_token_ids(sampled_token_ids)
@@ -3679,8 +3693,8 @@ class GPUModelRunner(
3693 target_hidden_states=hidden_states,
3694 sampling_metadata=sampling_metadata,
3695 )
3682 - elif spec_config.use_eagle():
3683 - assert isinstance(self.drafter, EagleProposer)
3696 + elif spec_config.use_eagle() or spec_config.uses_draft_model():
3697 + assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
3698
3699 if spec_config.disable_padded_drafter_batch:
3700 # When padded-batch is disabled, the sampled_token_ids should be
@@ -4475,8 +4489,12 @@ class GPUModelRunner(
4489 else:
4490 hidden_states = outputs
4491
4478 - if self.speculative_config and self.speculative_config.use_eagle():
4479 - assert isinstance(self.drafter, EagleProposer)
4492 + if self.speculative_config and (
4493 + self.speculative_config.use_eagle()
4494 + or self.speculative_config.uses_draft_model()
4495 + ):
4496 + assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
4497 + assert self.speculative_config is not None
4498 # Eagle currently only supports PIECEWISE cudagraphs.
4499 # Therefore only use cudagraphs if the main model uses PIECEWISE
4500 # NOTE(lucas): this is a hack, need to clean up.
@@ -5652,8 +5670,11 @@ class GPUModelRunner(
5670 kv_cache_config, kernel_block_sizes
5671 )
5672
5655 - if self.speculative_config and self.speculative_config.use_eagle():
5656 - assert isinstance(self.drafter, EagleProposer)
5673 + if self.speculative_config and (
5674 + self.speculative_config.use_eagle()
5675 + or self.speculative_config.uses_draft_model()
5676 + ):
5677 + assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
5678 # validate all draft model layers belong to the same kv cache
5679 # group
5680 self.drafter.validate_same_kv_cache_group(kv_cache_config)
vllm/v1/worker/utils.py
+2 -2
@@ -352,8 +352,8 @@ def bind_kv_cache(
352 pass
353 else:
354 raise NotImplementedError
355 - layer_name = layer_names[0]
356 - runner_kv_caches.append(kv_caches[layer_name])
355 + for layer_name in layer_names:
356 + runner_kv_caches.append(kv_caches[layer_name])
357
358 # Bind kv_caches to forward context
359 for layer_name, kv_cache in kv_caches.items():