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, 16:00 UTC
[Misc] improve example mlpspeculator and llm_engine_example (#16175)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
Reid committed
Apr 7, 2025 at 19:53 UTC
dc3529dbf65786fe25cce8144c76260266061c9d
2 files changed
+21
-7
examples/offline_inference/llm_engine_example.py
+6
-1
@@ -1,5 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-
2
+"""
3
+This file demonstrates using the `LLMEngine`
4
+for processing prompts with various sampling parameters.
5
+"""
6
import argparse
7
8
from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
@@ -26,6 +29,7 @@ def process_requests(engine: LLMEngine,
29
"""Continuously process a list of prompts and handle the outputs."""
30
request_id = 0
31
32
+ print('-' * 50)
33
while test_prompts or engine.has_unfinished_requests():
34
if test_prompts:
35
prompt, sampling_params = test_prompts.pop(0)
@@ -37,6 +41,7 @@ def process_requests(engine: LLMEngine,
41
for request_output in request_outputs:
42
if request_output.finished:
43
print(request_output)
44
+ print('-' * 50)
45
46
47
def initialize_engine(args: argparse.Namespace) -> LLMEngine:
examples/offline_inference/mlpspeculator.py
+15
-6
@@ -1,4 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+This file demonstrates the usage of text generation with an LLM model,
4
+comparing the performance with and without speculative decoding.
5
+
6
+Note that still not support `v1`:
7
+VLLM_USE_V1=0 python examples/offline_inference/mlpspeculator.py
8
+"""
9
10
import gc
11
import time
@@ -7,7 +14,7 @@ from vllm import LLM, SamplingParams
14
15
16
def time_generation(llm: LLM, prompts: list[str],
10
- sampling_params: SamplingParams):
17
+ sampling_params: SamplingParams, title: str):
18
# Generate texts from the prompts. The output is a list of RequestOutput
19
# objects that contain the prompt, generated text, and other information.
20
# Warmup first
@@ -16,11 +23,15 @@ def time_generation(llm: LLM, prompts: list[str],
23
start = time.time()
24
outputs = llm.generate(prompts, sampling_params)
25
end = time.time()
19
- print((end - start) / sum([len(o.outputs[0].token_ids) for o in outputs]))
26
+ print("-" * 50)
27
+ print(title)
28
+ print("time: ",
29
+ (end - start) / sum(len(o.outputs[0].token_ids) for o in outputs))
30
# Print the outputs.
31
for output in outputs:
32
generated_text = output.outputs[0].text
33
print(f"text: {generated_text!r}")
34
+ print("-" * 50)
35
36
37
if __name__ == "__main__":
@@ -41,8 +52,7 @@ if __name__ == "__main__":
52
# Create an LLM without spec decoding
53
llm = LLM(model="meta-llama/Llama-2-13b-chat-hf")
54
44
- print("Without speculation")
45
- time_generation(llm, prompts, sampling_params)
55
+ time_generation(llm, prompts, sampling_params, "Without speculation")
56
57
del llm
58
gc.collect()
@@ -55,5 +65,4 @@ if __name__ == "__main__":
65
},
66
)
67
58
- print("With speculation")
59
- time_generation(llm, prompts, sampling_params)
68
+ time_generation(llm, prompts, sampling_params, "With speculation")