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 14, 05:30 UTC
[Spec Decode][CI] Add e2e test for `examples/spec_decode.py` and prevent breaking Acceptance Length (#24531)
Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Ekagra Ranjan committed
Sep 23, 2025 at 13:46 UTC
867ecdd1c87ba17cecf8a97b262e1a8117795e94
2 files changed
+40
-4
.buildkite/test-pipeline.yaml
+2
@@ -329,6 +329,8 @@ steps:
329
- python3 offline_inference/basic/classify.py
330
- python3 offline_inference/basic/embed.py
331
- python3 offline_inference/basic/score.py
332
+ - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
333
+ - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
334
335
- label: Platform Tests (CUDA) # 4min
336
timeout_in_minutes: 15
examples/offline_inference/spec_decode.py
+38
-4
@@ -49,6 +49,7 @@ def get_custom_mm_prompts(num_prompts):
49
def parse_args():
50
parser = FlexibleArgumentParser()
51
add_dataset_parser(parser)
52
+ parser.add_argument("--test", action="store_true")
53
parser.add_argument(
54
"--method",
55
type=str,
@@ -60,6 +61,7 @@ def parse_args():
61
parser.add_argument("--tp", type=int, default=1)
62
parser.add_argument("--enforce-eager", action="store_true")
63
parser.add_argument("--enable-chunked-prefill", action="store_true")
64
+ parser.add_argument("--max-model-len", type=int, default=16384)
65
parser.add_argument("--temp", type=float, default=0)
66
parser.add_argument("--top-p", type=float, default=1.0)
67
parser.add_argument("--top-k", type=int, default=-1)
@@ -71,8 +73,7 @@ def parse_args():
73
return parser.parse_args()
74
75
74
-def main():
75
- args = parse_args()
76
+def main(args):
77
args.endpoint_type = "openai-chat"
78
79
model_dir = args.model_dir
@@ -134,7 +135,7 @@ def main():
135
gpu_memory_utilization=0.8,
136
speculative_config=speculative_config,
137
disable_log_stats=False,
137
- max_model_len=16384,
138
+ max_model_len=args.max_model_len,
139
limit_mm_per_prompt={"image": 5},
140
disable_chunked_mm_input=True,
141
)
@@ -198,6 +199,39 @@ def main():
199
acceptance_rate = acceptance_counts[i] / num_drafts if num_drafts > 0 else 0
200
print(f"acceptance at token {i}: {acceptance_rate:.2f}")
201
202
+ return acceptance_length
203
+
204
205
if __name__ == "__main__":
203
- main()
206
+ args = parse_args()
207
+ acceptance_length = main(args)
208
+
209
+ if args.test:
210
+ # takes ~30s to run on 1xH100
211
+ assert args.method in ["eagle", "eagle3"]
212
+ assert args.tp == 1
213
+ assert args.num_spec_tokens == 3
214
+ assert args.dataset_name == "hf"
215
+ assert args.dataset_path == "philschmid/mt-bench"
216
+ assert args.num_prompts == 80
217
+ assert args.temp == 0
218
+ assert args.top_p == 1.0
219
+ assert args.top_k == -1
220
+ assert args.enable_chunked_prefill
221
+
222
+ # check acceptance length is within 2% of expected value
223
+ rtol = 0.02
224
+ expected_acceptance_length = 2.296 if args.method == "eagle" else 2.811
225
+
226
+ assert (
227
+ acceptance_length <= (1 + rtol) * expected_acceptance_length
228
+ and acceptance_length >= (1 - rtol) * expected_acceptance_length
229
+ ), (
230
+ f"acceptance_length {acceptance_length} is not "
231
+ f"within {rtol * 100}% of {expected_acceptance_length}"
232
+ )
233
+
234
+ print(
235
+ f"Test passed! Expected AL: "
236
+ f"{expected_acceptance_length}, got {acceptance_length}"
237
+ )