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 15, 03:30 UTC
[Spec Decode] Clean up spec decode example (#20240)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jun 30, 2025 at 08:28 UTC
2965c99c86b460ee819e4805764d769c7b7d3d8e
2 files changed
+21
-163
examples/offline_inference/eagle.py
deleted
-144
@@ -1,144 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import argparse
4
-import json
5
-import os
6
-
7
-from transformers import AutoTokenizer
8
-
9
-from vllm import LLM, SamplingParams
10
-from vllm.v1.metrics.reader import Counter, Vector
11
-
12
-
13
-def load_prompts(dataset_path, num_prompts):
14
- if os.path.exists(dataset_path):
15
- prompts = []
16
- try:
17
- with open(dataset_path) as f:
18
- for line in f:
19
- data = json.loads(line)
20
- prompts.append(data["turns"][0])
21
- except Exception as e:
22
- print(f"Error reading dataset: {e}")
23
- return []
24
- else:
25
- prompts = ["The future of AI is", "The president of the United States is"]
26
-
27
- return prompts[:num_prompts]
28
-
29
-
30
-def parse_args():
31
- parser = argparse.ArgumentParser()
32
- parser.add_argument(
33
- "--dataset",
34
- type=str,
35
- default="./examples/data/gsm8k.jsonl",
36
- help="downloaded from the eagle repo "
37
- "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/",
38
- )
39
- parser.add_argument(
40
- "--method", type=str, default="eagle", choices=["eagle", "eagle3"]
41
- )
42
- parser.add_argument("--max_num_seqs", type=int, default=8)
43
- parser.add_argument("--num_prompts", type=int, default=80)
44
- parser.add_argument("--num_spec_tokens", type=int, default=2)
45
- parser.add_argument("--tp", type=int, default=1)
46
- parser.add_argument("--draft_tp", type=int, default=1)
47
- parser.add_argument("--enforce_eager", action="store_true")
48
- parser.add_argument("--enable_chunked_prefill", action="store_true")
49
- parser.add_argument("--max_num_batched_tokens", type=int, default=2048)
50
- parser.add_argument("--temp", type=float, default=0)
51
- return parser.parse_args()
52
-
53
-
54
-def main():
55
- args = parse_args()
56
-
57
- model_dir = "meta-llama/Llama-3.1-8B-Instruct"
58
-
59
- if args.method == "eagle":
60
- eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
61
- elif args.method == "eagle3":
62
- eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
63
- else:
64
- raise ValueError(f"unknown method: {args.method}")
65
-
66
- max_model_len = 2048
67
-
68
- tokenizer = AutoTokenizer.from_pretrained(model_dir)
69
-
70
- prompts = load_prompts(args.dataset, args.num_prompts)
71
-
72
- prompt_ids = [
73
- tokenizer.apply_chat_template(
74
- [{"role": "user", "content": prompt}], add_generation_prompt=True
75
- )
76
- for prompt in prompts
77
- ]
78
-
79
- llm = LLM(
80
- model=model_dir,
81
- trust_remote_code=True,
82
- tensor_parallel_size=args.tp,
83
- enable_chunked_prefill=args.enable_chunked_prefill,
84
- max_num_batched_tokens=args.max_num_batched_tokens,
85
- enforce_eager=args.enforce_eager,
86
- max_model_len=max_model_len,
87
- max_num_seqs=args.max_num_seqs,
88
- gpu_memory_utilization=0.8,
89
- speculative_config={
90
- "method": args.method,
91
- "model": eagle_dir,
92
- "num_speculative_tokens": args.num_spec_tokens,
93
- "draft_tensor_parallel_size": args.draft_tp,
94
- "max_model_len": max_model_len,
95
- },
96
- disable_log_stats=False,
97
- )
98
-
99
- sampling_params = SamplingParams(temperature=args.temp, max_tokens=256)
100
-
101
- outputs = llm.generate(prompt_token_ids=prompt_ids, sampling_params=sampling_params)
102
-
103
- # print the generated text
104
- for output in outputs:
105
- print("-" * 50)
106
- print(f"prompt: {output.prompt}")
107
- print(f"generated text: {output.outputs[0].text}")
108
- print("-" * 50)
109
-
110
- try:
111
- metrics = llm.get_metrics()
112
- except AssertionError:
113
- print("Metrics are not supported in the V0 engine.")
114
- return
115
-
116
- num_drafts = num_accepted = 0
117
- acceptance_counts = [0] * args.num_spec_tokens
118
- for metric in metrics:
119
- if metric.name == "vllm:spec_decode_num_drafts":
120
- assert isinstance(metric, Counter)
121
- num_drafts += metric.value
122
- elif metric.name == "vllm:spec_decode_num_accepted_tokens":
123
- assert isinstance(metric, Counter)
124
- num_accepted += metric.value
125
- elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos":
126
- assert isinstance(metric, Vector)
127
- for pos in range(len(metric.values)):
128
- acceptance_counts[pos] += metric.values[pos]
129
-
130
- print("-" * 50)
131
- print(f"mean acceptance length: {1 + (num_accepted / num_drafts):.2f}")
132
- print("-" * 50)
133
-
134
- # print acceptance at each token position
135
- for i in range(len(acceptance_counts)):
136
- print(f"acceptance at token {i}:{acceptance_counts[i] / num_drafts:.2f}")
137
-
138
-
139
-if __name__ == "__main__":
140
- print(
141
- "[WARNING] Use examples/offline_inference/spec_decode.py"
142
- " instead of this script."
143
- )
144
- main()
examples/offline_inference/spec_decode.py
+21
-19
@@ -16,24 +16,17 @@ def parse_args():
16
parser = FlexibleArgumentParser()
17
add_dataset_parser(parser)
18
parser.add_argument(
19
- "--dataset",
19
+ "--method",
20
type=str,
21
- default="./examples/data/gsm8k.jsonl",
22
- help="downloaded from the eagle repo "
23
- "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/",
21
+ default="eagle",
22
+ choices=["ngram", "eagle", "eagle3", "mtp"],
23
)
25
- parser.add_argument(
26
- "--method", type=str, default="eagle", choices=["ngram", "eagle", "eagle3"]
27
- )
28
- parser.add_argument("--max-num-seqs", type=int, default=8)
24
parser.add_argument("--num-spec-tokens", type=int, default=2)
25
parser.add_argument("--prompt-lookup-max", type=int, default=5)
26
parser.add_argument("--prompt-lookup-min", type=int, default=2)
27
parser.add_argument("--tp", type=int, default=1)
33
- parser.add_argument("--draft-tp", type=int, default=1)
28
parser.add_argument("--enforce-eager", action="store_true")
29
parser.add_argument("--enable-chunked-prefill", action="store_true")
36
- parser.add_argument("--max-num-batched-tokens", type=int, default=2048)
30
parser.add_argument("--temp", type=float, default=0)
31
parser.add_argument("--top-p", type=float, default=1.0)
32
parser.add_argument("--top-k", type=int, default=-1)
@@ -41,7 +34,6 @@ def parse_args():
34
parser.add_argument("--output-len", type=int, default=256)
35
parser.add_argument("--model-dir", type=str, default=None)
36
parser.add_argument("--eagle-dir", type=str, default=None)
44
- parser.add_argument("--max-model-len", type=int, default=2048)
37
return parser.parse_args()
38
39
@@ -71,8 +63,6 @@ def main():
63
"method": args.method,
64
"model": eagle_dir,
65
"num_speculative_tokens": args.num_spec_tokens,
74
- "draft_tensor_parallel_size": args.draft_tp,
75
- "max_model_len": args.max_model_len,
66
}
67
elif args.method == "ngram":
68
speculative_config = {
@@ -80,7 +70,6 @@ def main():
70
"num_speculative_tokens": args.num_spec_tokens,
71
"prompt_lookup_max": args.prompt_lookup_max,
72
"prompt_lookup_min": args.prompt_lookup_min,
83
- "max_model_len": args.max_model_len,
73
}
74
else:
75
raise ValueError(f"unknown method: {args.method}")
@@ -92,7 +81,6 @@ def main():
81
enable_chunked_prefill=args.enable_chunked_prefill,
82
max_num_batched_tokens=args.max_num_batched_tokens,
83
enforce_eager=args.enforce_eager,
95
- max_model_len=args.max_model_len,
84
max_num_seqs=args.max_num_seqs,
85
gpu_memory_utilization=0.8,
86
speculative_config=speculative_config,
@@ -116,27 +104,41 @@ def main():
104
print("Metrics are not supported in the V0 engine.")
105
return
106
119
- num_drafts = num_accepted = 0
107
+ total_num_output_tokens = sum(
108
+ len(output.outputs[0].token_ids) for output in outputs
109
+ )
110
+ num_drafts = 0
111
+ num_draft_tokens = 0
112
+ num_accepted_tokens = 0
113
acceptance_counts = [0] * args.num_spec_tokens
114
for metric in metrics:
115
if metric.name == "vllm:spec_decode_num_drafts":
116
assert isinstance(metric, Counter)
117
num_drafts += metric.value
118
+ elif metric.name == "vllm:spec_decode_num_draft_tokens":
119
+ assert isinstance(metric, Counter)
120
+ num_draft_tokens += metric.value
121
elif metric.name == "vllm:spec_decode_num_accepted_tokens":
122
assert isinstance(metric, Counter)
127
- num_accepted += metric.value
123
+ num_accepted_tokens += metric.value
124
elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos":
125
assert isinstance(metric, Vector)
126
for pos in range(len(metric.values)):
127
acceptance_counts[pos] += metric.values[pos]
128
129
print("-" * 50)
134
- print(f"mean acceptance length: {1 + (num_accepted / num_drafts):.2f}")
130
+ print(f"total_num_output_tokens: {total_num_output_tokens}")
131
+ print(f"num_drafts: {num_drafts}")
132
+ print(f"num_draft_tokens: {num_draft_tokens}")
133
+ print(f"num_accepted_tokens: {num_accepted_tokens}")
134
+ acceptance_length = 1 + (num_accepted_tokens / num_drafts) if num_drafts > 0 else 1
135
+ print(f"mean acceptance length: {acceptance_length:.2f}")
136
print("-" * 50)
137
138
# print acceptance at each token position
139
for i in range(len(acceptance_counts)):
139
- print(f"acceptance at token {i}:{acceptance_counts[i] / num_drafts:.2f}")
140
+ acceptance_rate = acceptance_counts[i] / num_drafts if num_drafts > 0 else 0
141
+ print(f"acceptance at token {i}: {acceptance_rate:.2f}")
142
143
144
if __name__ == "__main__":