@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 27bebcd897
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 15, 08:30 UTC

Convert `examples` to `ruff-format` (#18400)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed May 26, 2025 at 17:57 UTC 27bebcd89792d5c4b08af7a65095759526f2f9e1
83 files changed +2535 -2411
.pre-commit-config.yaml
+1 -1
@@ -17,7 +17,7 @@ repos:
17 - id: ruff
18 args: [--output-format, github, --fix]
19 - id: ruff-format
20 - files: ^(.buildkite|benchmarks)/.*
20 + files: ^(.buildkite|benchmarks|examples)/.*
21 - repo: https://github.com/codespell-project/codespell
22 rev: v2.4.1
23 hooks:
examples/offline_inference/audio_language.py
+80 -70
@@ -1,11 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 -This example shows how to use vLLM for running offline inference
3 +This example shows how to use vLLM for running offline inference
4 with the correct prompt format on audio language models.
5
6 For most models, the prompt format should follow corresponding examples
7 on HuggingFace model repository.
8 """
9 +
10 import os
11 from dataclasses import asdict
12 from typing import NamedTuple, Optional
@@ -22,7 +23,7 @@ audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
23 question_per_audio_count = {
24 0: "What is 1+1?",
25 1: "What is recited in the audio?",
25 - 2: "What sport and what nursery rhyme are referenced?"
26 + 2: "What sport and what nursery rhyme are referenced?",
27 }
28
29
@@ -72,8 +73,7 @@ def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
73 # MiniCPM-O
74 def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
75 model_name = "openbmb/MiniCPM-o-2_6"
75 - tokenizer = AutoTokenizer.from_pretrained(model_name,
76 - trust_remote_code=True)
76 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
77 engine_args = EngineArgs(
78 model=model_name,
79 trust_remote_code=True,
@@ -82,19 +82,18 @@ def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
82 limit_mm_per_prompt={"audio": audio_count},
83 )
84
85 - stop_tokens = ['<|im_end|>', '<|endoftext|>']
85 + stop_tokens = ["<|im_end|>", "<|endoftext|>"]
86 stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
87
88 audio_placeholder = "(<audio>./</audio>)" * audio_count
89 audio_chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>' }}{% endif %}" # noqa: E501
90 - messages = [{
91 - 'role': 'user',
92 - 'content': f'{audio_placeholder}\n{question}'
93 - }]
94 - prompt = tokenizer.apply_chat_template(messages,
95 - tokenize=False,
96 - add_generation_prompt=True,
97 - chat_template=audio_chat_template)
90 + messages = [{"role": "user", "content": f"{audio_placeholder}\n{question}"}]
91 + prompt = tokenizer.apply_chat_template(
92 + messages,
93 + tokenize=False,
94 + add_generation_prompt=True,
95 + chat_template=audio_chat_template,
96 + )
97
98 return ModelRequestData(
99 engine_args=engine_args,
@@ -113,7 +112,7 @@ def run_phi4mm(question: str, audio_count: int) -> ModelRequestData:
112 # Since the vision-lora and speech-lora co-exist with the base model,
113 # we have to manually specify the path of the lora weights.
114 speech_lora_path = os.path.join(model_path, "speech-lora")
116 - placeholders = "".join([f"<|audio_{i+1}|>" for i in range(audio_count)])
115 + placeholders = "".join([f"<|audio_{i + 1}|>" for i in range(audio_count)])
116
117 prompts = f"<|user|>{placeholders}{question}<|end|><|assistant|>"
118
@@ -145,15 +144,19 @@ def run_qwen2_audio(question: str, audio_count: int) -> ModelRequestData:
144 limit_mm_per_prompt={"audio": audio_count},
145 )
146
148 - audio_in_prompt = "".join([
149 - f"Audio {idx+1}: "
150 - f"<|audio_bos|><|AUDIO|><|audio_eos|>\n" for idx in range(audio_count)
151 - ])
147 + audio_in_prompt = "".join(
148 + [
149 + f"Audio {idx + 1}: <|audio_bos|><|AUDIO|><|audio_eos|>\n"
150 + for idx in range(audio_count)
151 + ]
152 + )
153
153 - prompt = ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
154 - "<|im_start|>user\n"
155 - f"{audio_in_prompt}{question}<|im_end|>\n"
156 - "<|im_start|>assistant\n")
154 + prompt = (
155 + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
156 + "<|im_start|>user\n"
157 + f"{audio_in_prompt}{question}<|im_end|>\n"
158 + "<|im_start|>assistant\n"
159 + )
160
161 return ModelRequestData(
162 engine_args=engine_args,
@@ -172,19 +175,22 @@ def run_qwen2_5_omni(question: str, audio_count: int):
175 limit_mm_per_prompt={"audio": audio_count},
176 )
177
175 - audio_in_prompt = "".join([
176 - "<|audio_bos|><|AUDIO|><|audio_eos|>\n" for idx in range(audio_count)
177 - ])
178 + audio_in_prompt = "".join(
179 + ["<|audio_bos|><|AUDIO|><|audio_eos|>\n" for idx in range(audio_count)]
180 + )
181
182 default_system = (
183 "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
184 "Group, capable of perceiving auditory and visual inputs, as well as "
182 - "generating text and speech.")
185 + "generating text and speech."
186 + )
187
184 - prompt = (f"<|im_start|>system\n{default_system}<|im_end|>\n"
185 - "<|im_start|>user\n"
186 - f"{audio_in_prompt}{question}<|im_end|>\n"
187 - "<|im_start|>assistant\n")
188 + prompt = (
189 + f"<|im_start|>system\n{default_system}<|im_end|>\n"
190 + "<|im_start|>user\n"
191 + f"{audio_in_prompt}{question}<|im_end|>\n"
192 + "<|im_start|>assistant\n"
193 + )
194 return ModelRequestData(
195 engine_args=engine_args,
196 prompt=prompt,
@@ -196,13 +202,10 @@ def run_ultravox(question: str, audio_count: int) -> ModelRequestData:
202 model_name = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
203
204 tokenizer = AutoTokenizer.from_pretrained(model_name)
199 - messages = [{
200 - 'role': 'user',
201 - 'content': "<|audio|>\n" * audio_count + question
202 - }]
203 - prompt = tokenizer.apply_chat_template(messages,
204 - tokenize=False,
205 - add_generation_prompt=True)
205 + messages = [{"role": "user", "content": "<|audio|>\n" * audio_count + question}]
206 + prompt = tokenizer.apply_chat_template(
207 + messages, tokenize=False, add_generation_prompt=True
208 + )
209
210 engine_args = EngineArgs(
211 model=model_name,
@@ -220,8 +223,7 @@ def run_ultravox(question: str, audio_count: int) -> ModelRequestData:
223
224 # Whisper
225 def run_whisper(question: str, audio_count: int) -> ModelRequestData:
223 - assert audio_count == 1, (
224 - "Whisper only support single audio input per prompt")
226 + assert audio_count == 1, "Whisper only support single audio input per prompt"
227 model_name = "openai/whisper-large-v3-turbo"
228
229 prompt = "<|startoftranscript|>"
@@ -252,27 +254,33 @@ model_example_map = {
254
255 def parse_args():
256 parser = FlexibleArgumentParser(
255 - description='Demo on using vLLM for offline inference with '
256 - 'audio language models')
257 - parser.add_argument('--model-type',
258 - '-m',
259 - type=str,
260 - default="ultravox",
261 - choices=model_example_map.keys(),
262 - help='Huggingface "model_type".')
263 - parser.add_argument('--num-prompts',
264 - type=int,
265 - default=1,
266 - help='Number of prompts to run.')
267 - parser.add_argument("--num-audios",
268 - type=int,
269 - default=1,
270 - choices=[0, 1, 2],
271 - help="Number of audio items per prompt.")
272 - parser.add_argument("--seed",
273 - type=int,
274 - default=None,
275 - help="Set the seed when initializing `vllm.LLM`.")
257 + description="Demo on using vLLM for offline inference with "
258 + "audio language models"
259 + )
260 + parser.add_argument(
261 + "--model-type",
262 + "-m",
263 + type=str,
264 + default="ultravox",
265 + choices=model_example_map.keys(),
266 + help='Huggingface "model_type".',
267 + )
268 + parser.add_argument(
269 + "--num-prompts", type=int, default=1, help="Number of prompts to run."
270 + )
271 + parser.add_argument(
272 + "--num-audios",
273 + type=int,
274 + default=1,
275 + choices=[0, 1, 2],
276 + help="Number of audio items per prompt.",
277 + )
278 + parser.add_argument(
279 + "--seed",
280 + type=int,
281 + default=None,
282 + help="Set the seed when initializing `vllm.LLM`.",
283 + )
284
285 return parser.parse_args()
286
@@ -283,29 +291,30 @@ def main(args):
291 raise ValueError(f"Model type {model} is not supported.")
292
293 audio_count = args.num_audios
286 - req_data = model_example_map[model](question_per_audio_count[audio_count],
287 - audio_count)
294 + req_data = model_example_map[model](
295 + question_per_audio_count[audio_count], audio_count
296 + )
297
298 # Disable other modalities to save memory
299 default_limits = {"image": 0, "video": 0, "audio": 0}
300 req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
292 - req_data.engine_args.limit_mm_per_prompt or {})
301 + req_data.engine_args.limit_mm_per_prompt or {}
302 + )
303
304 engine_args = asdict(req_data.engine_args) | {"seed": args.seed}
305 llm = LLM(**engine_args)
306
307 # We set temperature to 0.2 so that outputs can be different
308 # even when all prompts are identical when running batch inference.
299 - sampling_params = SamplingParams(temperature=0.2,
300 - max_tokens=64,
301 - stop_token_ids=req_data.stop_token_ids)
309 + sampling_params = SamplingParams(
310 + temperature=0.2, max_tokens=64, stop_token_ids=req_data.stop_token_ids
311 + )
312
313 mm_data = {}
314 if audio_count > 0:
315 mm_data = {
316 "audio": [
307 - asset.audio_and_sample_rate
308 - for asset in audio_assets[:audio_count]
317 + asset.audio_and_sample_rate for asset in audio_assets[:audio_count]
318 ]
319 }
320
@@ -315,8 +324,9 @@ def main(args):
324 # Batch inference
325 inputs = [inputs] * args.num_prompts
326 # Add LoRA request if applicable
318 - lora_request = (req_data.lora_requests *
319 - args.num_prompts if req_data.lora_requests else None)
327 + lora_request = (
328 + req_data.lora_requests * args.num_prompts if req_data.lora_requests else None
329 + )
330
331 outputs = llm.generate(
332 inputs,
examples/offline_inference/automatic_prefix_caching.py
+10 -6
@@ -16,13 +16,16 @@ but ask different questions.
16 Run:
17 python examples/offline_inference/automatic_prefix_caching.py
18 """
19 +
20 import time
21
22 from vllm import LLM, SamplingParams
23
24 # ruff: noqa: E501
25 # A prompt containing a large markdown table. The table is randomly generated by GPT-4.
25 -LONG_PROMPT = "You are a helpful assistant in recognizes the content of tables in markdown format. Here is a table as follows.\n# Table\n" + """
26 +LONG_PROMPT = (
27 + "You are a helpful assistant in recognizes the content of tables in markdown format. Here is a table as follows.\n# Table\n"
28 + + """
29 | ID | Name | Age | Occupation | Country | Email | Phone Number | Address |
30 |-----|---------------|-----|---------------|---------------|------------------------|----------------|------------------------------|
31 | 1 | John Doe | 29 | Engineer | USA | john.doe@example.com | 555-1234 | 123 Elm St, Springfield, IL |
@@ -56,6 +59,7 @@ LONG_PROMPT = "You are a helpful assistant in recognizes the content of tables i
59 | 29 | Amy White | 33 | Musician | New Zealand | amy.w@example.com | 555-5658 | 159 Maple St, Wellington, NZ |
60 | 30 | Ben Black | 38 | Chef | Ireland | ben.b@example.com | 555-7870 | 246 Fir St, Waterford, IE |
61 """
62 +)
63
64
65 def get_generation_time(llm, sampling_params, prompts):
@@ -72,7 +76,7 @@ def get_generation_time(llm, sampling_params, prompts):
76
77 def main():
78 # set enable_prefix_caching=True to enable APC
75 - llm = LLM(model='lmsys/longchat-13b-16k', enable_prefix_caching=True)
79 + llm = LLM(model="lmsys/longchat-13b-16k", enable_prefix_caching=True)
80
81 sampling_params = SamplingParams(temperature=0, max_tokens=100)
82
@@ -80,8 +84,8 @@ def main():
84 get_generation_time(
85 llm,
86 sampling_params,
83 - LONG_PROMPT +
84 - "Question: what is the age of John Doe? Your answer: The age of John Doe is ",
87 + LONG_PROMPT
88 + + "Question: what is the age of John Doe? Your answer: The age of John Doe is ",
89 )
90
91 # Querying the age of Zack Blue
@@ -89,8 +93,8 @@ def main():
93 get_generation_time(
94 llm,
95 sampling_params,
92 - LONG_PROMPT +
93 - "Question: what is the age of Zack Blue? Your answer: The age of Zack Blue is ",
96 + LONG_PROMPT
97 + + "Question: what is the age of Zack Blue? Your answer: The age of Zack Blue is ",
98 )
99
100
examples/offline_inference/basic/chat.py
+4 -14
@@ -56,22 +56,12 @@ def main(args: dict):
56
57 # In this script, we demonstrate how to pass input to the chat method:
58 conversation = [
59 - {
60 - "role": "system",
61 - "content": "You are a helpful assistant"
62 - },
63 - {
64 - "role": "user",
65 - "content": "Hello"
66 - },
67 - {
68 - "role": "assistant",
69 - "content": "Hello! How can I assist you today?"
70 - },
59 + {"role": "system", "content": "You are a helpful assistant"},
60 + {"role": "user", "content": "Hello"},
61 + {"role": "assistant", "content": "Hello! How can I assist you today?"},
62 {
63 "role": "user",
73 - "content":
74 - "Write an essay about the importance of higher education.",
64 + "content": "Write an essay about the importance of higher education.",
65 },
66 ]
67 outputs = llm.chat(conversation, sampling_params, use_tqdm=False)
examples/offline_inference/basic/classify.py
+8 -7
@@ -10,9 +10,9 @@ def parse_args():
10 parser = FlexibleArgumentParser()
11 parser = EngineArgs.add_cli_args(parser)
12 # Set example specific arguments
13 - parser.set_defaults(model="jason9693/Qwen2.5-1.5B-apeach",
14 - task="classify",
15 - enforce_eager=True)
13 + parser.set_defaults(
14 + model="jason9693/Qwen2.5-1.5B-apeach", task="classify", enforce_eager=True
15 + )
16 return parser.parse_args()
17
18
@@ -36,10 +36,11 @@ def main(args: Namespace):
36 print("\nGenerated Outputs:\n" + "-" * 60)
37 for prompt, output in zip(prompts, outputs):
38 probs = output.outputs.probs
39 - probs_trimmed = ((str(probs[:16])[:-1] +
40 - ", ...]") if len(probs) > 16 else probs)
41 - print(f"Prompt: {prompt!r} \n"
42 - f"Class Probabilities: {probs_trimmed} (size={len(probs)})")
39 + probs_trimmed = (str(probs[:16])[:-1] + ", ...]") if len(probs) > 16 else probs
40 + print(
41 + f"Prompt: {prompt!r} \n"
42 + f"Class Probabilities: {probs_trimmed} (size={len(probs)})"
43 + )
44 print("-" * 60)
45
46
examples/offline_inference/basic/embed.py
+7 -7
@@ -10,9 +10,9 @@ def parse_args():
10 parser = FlexibleArgumentParser()
11 parser = EngineArgs.add_cli_args(parser)
12 # Set example specific arguments
13 - parser.set_defaults(model="intfloat/e5-mistral-7b-instruct",
14 - task="embed",
15 - enforce_eager=True)
13 + parser.set_defaults(
14 + model="intfloat/e5-mistral-7b-instruct", task="embed", enforce_eager=True
15 + )
16 return parser.parse_args()
17
18
@@ -36,10 +36,10 @@ def main(args: Namespace):
36 print("\nGenerated Outputs:\n" + "-" * 60)
37 for prompt, output in zip(prompts, outputs):
38 embeds = output.outputs.embedding
39 - embeds_trimmed = ((str(embeds[:16])[:-1] +
40 - ", ...]") if len(embeds) > 16 else embeds)
41 - print(f"Prompt: {prompt!r} \n"
42 - f"Embeddings: {embeds_trimmed} (size={len(embeds)})")
39 + embeds_trimmed = (
40 + (str(embeds[:16])[:-1] + ", ...]") if len(embeds) > 16 else embeds
41 + )
42 + print(f"Prompt: {prompt!r} \nEmbeddings: {embeds_trimmed} (size={len(embeds)})")
43 print("-" * 60)
44
45
examples/offline_inference/basic/score.py
+3 -3
@@ -10,9 +10,9 @@ def parse_args():
10 parser = FlexibleArgumentParser()
11 parser = EngineArgs.add_cli_args(parser)
12 # Set example specific arguments
13 - parser.set_defaults(model="BAAI/bge-reranker-v2-m3",
14 - task="score",
15 - enforce_eager=True)
13 + parser.set_defaults(
14 + model="BAAI/bge-reranker-v2-m3", task="score", enforce_eager=True
15 + )
16 return parser.parse_args()
17
18
examples/offline_inference/batch_llm_inference.py
+11 -11
@@ -17,12 +17,14 @@ Ray Data provides functionality for:
17 Learn more about Ray Data's LLM integration:
18 https://docs.ray.io/en/latest/data/working-with-llms.html
19 """
20 +
21 import ray
22 from packaging.version import Version
23 from ray.data.llm import build_llm_processor, vLLMEngineProcessorConfig
24
24 -assert Version(ray.__version__) >= Version(
25 - "2.44.1"), "Ray version must be at least 2.44.1"
25 +assert Version(ray.__version__) >= Version("2.44.1"), (
26 + "Ray version must be at least 2.44.1"
27 +)
28
29 # Uncomment to reduce clutter in stdout
30 # ray.init(log_to_driver=False)
@@ -53,20 +55,18 @@ config = vLLMEngineProcessorConfig(
55 vllm_processor = build_llm_processor(
56 config,
57 preprocess=lambda row: dict(
56 - messages=[{
57 - "role": "system",
58 - "content": "You are a bot that responds with haikus."
59 - }, {
60 - "role": "user",
61 - "content": row["text"]
62 - }],
58 + messages=[
59 + {"role": "system", "content": "You are a bot that responds with haikus."},
60 + {"role": "user", "content": row["text"]},
61 + ],
62 sampling_params=dict(
63 temperature=0.3,
64 max_tokens=250,
66 - )),
65 + ),
66 + ),
67 postprocess=lambda row: dict(
68 answer=row["generated_text"],
69 - **row # This will return all the original columns in the dataset.
69 + **row, # This will return all the original columns in the dataset.
70 ),
71 )
72
examples/offline_inference/chat_with_tools.py
+58 -52
@@ -50,87 +50,93 @@ model_name = "mistralai/Mistral-7B-Instruct-v0.3"
50 # or any other mistral model with function calling ability
51
52 sampling_params = SamplingParams(max_tokens=8192, temperature=0.0)
53 -llm = LLM(model=model_name,
54 - tokenizer_mode="mistral",
55 - config_format="mistral",
56 - load_format="mistral")
53 +llm = LLM(
54 + model=model_name,
55 + tokenizer_mode="mistral",
56 + config_format="mistral",
57 + load_format="mistral",
58 +)
59
60
61 def generate_random_id(length=9):
62 characters = string.ascii_letters + string.digits
61 - random_id = ''.join(random.choice(characters) for _ in range(length))
63 + random_id = "".join(random.choice(characters) for _ in range(length))
64 return random_id
65
66
67 # simulate an API that can be called
66 -def get_current_weather(city: str, state: str, unit: 'str'):
67 - return (f"The weather in {city}, {state} is 85 degrees {unit}. It is "
68 - "partly cloudly, with highs in the 90's.")
68 +def get_current_weather(city: str, state: str, unit: "str"):
69 + return (
70 + f"The weather in {city}, {state} is 85 degrees {unit}. It is "
71 + "partly cloudly, with highs in the 90's."
72 + )
73
74
75 tool_functions = {"get_current_weather": get_current_weather}
76
73 -tools = [{
74 - "type": "function",
75 - "function": {
76 - "name": "get_current_weather",
77 - "description": "Get the current weather in a given location",
78 - "parameters": {
79 - "type": "object",
80 - "properties": {
81 - "city": {
82 - "type":
83 - "string",
84 - "description":
85 - "The city to find the weather for, e.g. 'San Francisco'"
77 +tools = [
78 + {
79 + "type": "function",
80 + "function": {
81 + "name": "get_current_weather",
82 + "description": "Get the current weather in a given location",
83 + "parameters": {
84 + "type": "object",
85 + "properties": {
86 + "city": {
87 + "type": "string",
88 + "description": "The city to find the weather for, e.g. 'San Francisco'",
89 + },
90 + "state": {
91 + "type": "string",
92 + "description": "the two-letter abbreviation for the state that the city is"
93 + " in, e.g. 'CA' which would mean 'California'",
94 + },
95 + "unit": {
96 + "type": "string",
97 + "description": "The unit to fetch the temperature in",
98 + "enum": ["celsius", "fahrenheit"],
99 + },
100 },
87 - "state": {
88 - "type":
89 - "string",
90 - "description":
91 - "the two-letter abbreviation for the state that the city is"
92 - " in, e.g. 'CA' which would mean 'California'"
93 - },
94 - "unit": {
95 - "type": "string",
96 - "description": "The unit to fetch the temperature in",
97 - "enum": ["celsius", "fahrenheit"]
98 - }
101 + "required": ["city", "state", "unit"],
102 },
100 - "required": ["city", "state", "unit"]
101 - }
103 + },
104 }
103 -}]
105 +]
106
105 -messages = [{
106 - "role":
107 - "user",
108 - "content":
109 - "Can you tell me what the temperate will be in Dallas, in fahrenheit?"
110 -}]
107 +messages = [
108 + {
109 + "role": "user",
110 + "content": "Can you tell me what the temperate will be in Dallas, in fahrenheit?",
111 + }
112 +]
113
114 outputs = llm.chat(messages, sampling_params=sampling_params, tools=tools)
115 output = outputs[0].outputs[0].text.strip()
116
117 # append the assistant message
116 -messages.append({
117 - "role": "assistant",
118 - "content": output,
119 -})
118 +messages.append(
119 + {
120 + "role": "assistant",
121 + "content": output,
122 + }
123 +)
124
125 # let's now actually parse and execute the model's output simulating an API call by using the
126 # above defined function
127 tool_calls = json.loads(output)
128 tool_answers = [
125 - tool_functions[call['name']](**call['arguments']) for call in tool_calls
129 + tool_functions[call["name"]](**call["arguments"]) for call in tool_calls
130 ]
131
132 # append the answer as a tool message and let the LLM give you an answer
129 -messages.append({
130 - "role": "tool",
131 - "content": "\n\n".join(tool_answers),
132 - "tool_call_id": generate_random_id(),
133 -})
133 +messages.append(
134 + {
135 + "role": "tool",
136 + "content": "\n\n".join(tool_answers),
137 + "tool_call_id": generate_random_id(),
138 + }
139 +)
140
141 outputs = llm.chat(messages, sampling_params, tools=tools)
142
examples/offline_inference/data_parallel.py
+61 -50
@@ -27,6 +27,7 @@ Multi-node:
27 --master-addr=10.99.48.128 \
28 --master-port=13345
29 """
30 +
31 import os
32 from time import sleep
33
@@ -36,46 +37,46 @@ from vllm.utils import get_open_port
37
38 def parse_args():
39 import argparse
40 +
41 parser = argparse.ArgumentParser(description="Data Parallel Inference")
40 - parser.add_argument("--model",
41 - type=str,
42 - default="ibm-research/PowerMoE-3b",
43 - help="Model name or path")
44 - parser.add_argument("--dp-size",
45 - type=int,
46 - default=2,
47 - help="Data parallel size")
48 - parser.add_argument("--tp-size",
49 - type=int,
50 - default=2,
51 - help="Tensor parallel size")
52 - parser.add_argument("--node-size",
53 - type=int,
54 - default=1,
55 - help="Total number of nodes")
56 - parser.add_argument("--node-rank",
57 - type=int,
58 - default=0,
59 - help="Rank of the current node")
60 - parser.add_argument("--master-addr",
61 - type=str,
62 - default="",
63 - help="Master node IP address")
64 - parser.add_argument("--master-port",
65 - type=int,
66 - default=0,
67 - help="Master node port")
68 - parser.add_argument("--enforce-eager",
69 - action='store_true',
70 - help="Enforce eager mode execution.")
71 - parser.add_argument("--trust-remote-code",
72 - action='store_true',
73 - help="Trust remote code.")
42 + parser.add_argument(
43 + "--model",
44 + type=str,
45 + default="ibm-research/PowerMoE-3b",
46 + help="Model name or path",
47 + )
48 + parser.add_argument("--dp-size", type=int, default=2, help="Data parallel size")
49 + parser.add_argument("--tp-size", type=int, default=2, help="Tensor parallel size")
50 + parser.add_argument(
51 + "--node-size", type=int, default=1, help="Total number of nodes"
52 + )
53 + parser.add_argument(
54 + "--node-rank", type=int, default=0, help="Rank of the current node"
55 + )
56 + parser.add_argument(
57 + "--master-addr", type=str, default="", help="Master node IP address"
58 + )
59 + parser.add_argument("--master-port", type=int, default=0, help="Master node port")
60 + parser.add_argument(
61 + "--enforce-eager", action="store_true", help="Enforce eager mode execution."
62 + )
63 + parser.add_argument(
64 + "--trust-remote-code", action="store_true", help="Trust remote code."
65 + )
66 return parser.parse_args()
67
68
77 -def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,
78 - dp_master_port, GPUs_per_dp_rank, enforce_eager, trust_remote_code):
69 +def main(
70 + model,
71 + dp_size,
72 + local_dp_rank,
73 + global_dp_rank,
74 + dp_master_ip,
75 + dp_master_port,
76 + GPUs_per_dp_rank,
77 + enforce_eager,
78 + trust_remote_code,
79 +):
80 os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
81 os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
82 os.environ["VLLM_DP_SIZE"] = str(dp_size)
@@ -110,9 +111,9 @@ def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,
111 # since we are doing data parallel, every rank can have different
112 # sampling params. here we set different max_tokens for different
113 # ranks for demonstration.
113 - sampling_params = SamplingParams(temperature=0.8,
114 - top_p=0.95,
115 - max_tokens=[16, 20][global_dp_rank % 2])
114 + sampling_params = SamplingParams(
115 + temperature=0.8, top_p=0.95, max_tokens=[16, 20][global_dp_rank % 2]
116 + )
117
118 # Create an LLM.
119 llm = LLM(
@@ -130,15 +131,16 @@ def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,
131 break
132 prompt = output.prompt
133 generated_text = output.outputs[0].text
133 - print(f"DP rank {global_dp_rank}, Prompt: {prompt!r}, "
134 - f"Generated text: {generated_text!r}")
134 + print(
135 + f"DP rank {global_dp_rank}, Prompt: {prompt!r}, "
136 + f"Generated text: {generated_text!r}"
137 + )
138
139 # Give engines time to pause their processing loops before exiting.
140 sleep(1)
141
142
143 if __name__ == "__main__":
141 -
144 args = parse_args()
145
146 dp_size = args.dp_size
@@ -160,20 +162,29 @@ if __name__ == "__main__":
162
163 procs = []
164 for local_dp_rank, global_dp_rank in enumerate(
163 - range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)):
164 - proc = Process(target=main,
165 - args=(args.model, dp_size, local_dp_rank,
166 - global_dp_rank, dp_master_ip, dp_master_port,
167 - tp_size, args.enforce_eager,
168 - args.trust_remote_code))
165 + range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)
166 + ):
167 + proc = Process(
168 + target=main,
169 + args=(
170 + args.model,
171 + dp_size,
172 + local_dp_rank,
173 + global_dp_rank,
174 + dp_master_ip,
175 + dp_master_port,
176 + tp_size,
177 + args.enforce_eager,
178 + args.trust_remote_code,
179 + ),
180 + )
181 proc.start()
182 procs.append(proc)
183 exit_code = 0
184 for proc in procs:
185 proc.join(timeout=300)
186 if proc.exitcode is None:
175 - print(f"Killing process {proc.pid} that "
176 - f"didn't stop within 5 minutes.")
187 + print(f"Killing process {proc.pid} that didn't stop within 5 minutes.")
188 proc.kill()
189 exit_code = 1
190 elif proc.exitcode:
examples/offline_inference/disaggregated-prefill-v1/decode_example.py
+12 -11
@@ -22,17 +22,18 @@ def main():
22 prompts = read_prompts()
23 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
24
25 - llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct",
26 - enforce_eager=True,
27 - gpu_memory_utilization=0.8,
28 - max_num_batched_tokens=64,
29 - max_num_seqs=16,
30 - kv_transfer_config=KVTransferConfig(
31 - kv_connector="SharedStorageConnector",
32 - kv_role="kv_both",
33 - kv_connector_extra_config={
34 - "shared_storage_path": "local_storage"
35 - })) #, max_model_len=2048, max_num_batched_tokens=2048)
25 + llm = LLM(
26 + model="meta-llama/Llama-3.2-1B-Instruct",
27 + enforce_eager=True,
28 + gpu_memory_utilization=0.8,
29 + max_num_batched_tokens=64,
30 + max_num_seqs=16,
31 + kv_transfer_config=KVTransferConfig(
32 + kv_connector="SharedStorageConnector",
33 + kv_role="kv_both",
34 + kv_connector_extra_config={"shared_storage_path": "local_storage"},
35 + ),
36 + ) # , max_model_len=2048, max_num_batched_tokens=2048)
37
38 # 1ST generation (prefill instance)
39 outputs = llm.generate(prompts, sampling_params)
examples/offline_inference/disaggregated-prefill-v1/prefill_example.py
+10 -9
@@ -20,15 +20,16 @@ def main():
20
21 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
22
23 - llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct",
24 - enforce_eager=True,
25 - gpu_memory_utilization=0.8,
26 - kv_transfer_config=KVTransferConfig(
27 - kv_connector="SharedStorageConnector",
28 - kv_role="kv_both",
29 - kv_connector_extra_config={
30 - "shared_storage_path": "local_storage"
31 - })) #, max_model_len=2048, max_num_batched_tokens=2048)
23 + llm = LLM(
24 + model="meta-llama/Llama-3.2-1B-Instruct",
25 + enforce_eager=True,
26 + gpu_memory_utilization=0.8,
27 + kv_transfer_config=KVTransferConfig(
28 + kv_connector="SharedStorageConnector",
29 + kv_role="kv_both",
30 + kv_connector_extra_config={"shared_storage_path": "local_storage"},
31 + ),
32 + ) # , max_model_len=2048, max_num_batched_tokens=2048)
33
34 # 1ST generation (prefill instance)
35 outputs = llm.generate(
examples/offline_inference/disaggregated_prefill.py
+27 -18
@@ -4,6 +4,7 @@ This file demonstrates the example usage of disaggregated prefilling
4 We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
5 and then transfer the KV cache between them.
6 """
7 +
8 import os
9 import time
10 from multiprocessing import Event, Process
@@ -32,17 +33,21 @@ def run_prefill(prefill_done):
33 # This instance is the prefill node (kv_producer, rank 0).
34 # The number of parallel instances for KV cache transfer is set to 2,
35 # as required for PyNcclConnector.
35 - ktc = KVTransferConfig(kv_connector="PyNcclConnector",
36 - kv_role="kv_producer",
37 - kv_rank=0,
38 - kv_parallel_size=2)
36 + ktc = KVTransferConfig(
37 + kv_connector="PyNcclConnector",
38 + kv_role="kv_producer",
39 + kv_rank=0,
40 + kv_parallel_size=2,
41 + )
42
43 # Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
44 # memory. You may need to adjust the value to fit your GPU.
42 - llm = LLM(model="meta-llama/Meta-Llama-3.1-8B-Instruct",
43 - kv_transfer_config=ktc,
44 - max_model_len=2000,
45 - gpu_memory_utilization=0.8)
45 + llm = LLM(
46 + model="meta-llama/Meta-Llama-3.1-8B-Instruct",
47 + kv_transfer_config=ktc,
48 + max_model_len=2000,
49 + gpu_memory_utilization=0.8,
50 + )
51
52 llm.generate(prompts, sampling_params)
53 print("Prefill node is finished.")
@@ -72,17 +77,21 @@ def run_decode(prefill_done):
77 # This instance is the decode node (kv_consumer, rank 1).
78 # The number of parallel instances for KV cache transfer is set to 2,
79 # as required for PyNcclConnector.
75 - ktc = KVTransferConfig(kv_connector="PyNcclConnector",
76 - kv_role="kv_consumer",
77 - kv_rank=1,
78 - kv_parallel_size=2)
80 + ktc = KVTransferConfig(
81 + kv_connector="PyNcclConnector",
82 + kv_role="kv_consumer",
83 + kv_rank=1,
84 + kv_parallel_size=2,
85 + )
86
87 # Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
88 # memory. You may need to adjust the value to fit your GPU.
82 - llm = LLM(model="meta-llama/Meta-Llama-3.1-8B-Instruct",
83 - kv_transfer_config=ktc,
84 - max_model_len=2000,
85 - gpu_memory_utilization=0.8)
89 + llm = LLM(
90 + model="meta-llama/Meta-Llama-3.1-8B-Instruct",
91 + kv_transfer_config=ktc,
92 + max_model_len=2000,
93 + gpu_memory_utilization=0.8,
94 + )
95
96 # Wait for the producer to start the pipe
97 print("Waiting for prefill node to finish...")
@@ -99,8 +108,8 @@ def run_decode(prefill_done):
108
109 def main():
110 prefill_done = Event()
102 - prefill_process = Process(target=run_prefill, args=(prefill_done, ))
103 - decode_process = Process(target=run_decode, args=(prefill_done, ))
111 + prefill_process = Process(target=run_prefill, args=(prefill_done,))
112 + decode_process = Process(target=run_decode, args=(prefill_done,))
113
114 # Start prefill node
115 prefill_process.start()
examples/offline_inference/eagle.py
+23 -27
@@ -20,9 +20,7 @@ def load_prompts(dataset_path, num_prompts):
20 print(f"Error reading dataset: {e}")
21 return []
22 else:
23 - prompts = [
24 - "The future of AI is", "The president of the United States is"
25 - ]
23 + prompts = ["The future of AI is", "The president of the United States is"]
24
25 return prompts[:num_prompts]
26
@@ -33,34 +31,32 @@ def parse_args():
31 "--dataset",
32 type=str,
33 default="./examples/data/gsm8k.jsonl",
36 - help="downloaded from the eagle repo " \
37 - "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/"
34 + help="downloaded from the eagle repo "
35 + "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/",
36 + )
37 + parser.add_argument(
38 + "--method", type=str, default="eagle", choices=["eagle", "eagle3"]
39 )
39 - parser.add_argument("--method",
40 - type=str,
41 - default='eagle',
42 - choices=['eagle', 'eagle3'])
40 parser.add_argument("--max_num_seqs", type=int, default=8)
41 parser.add_argument("--num_prompts", type=int, default=80)
42 parser.add_argument("--num_spec_tokens", type=int, default=2)
43 parser.add_argument("--tp", type=int, default=1)
44 parser.add_argument("--draft_tp", type=int, default=1)
48 - parser.add_argument("--enforce_eager", action='store_true')
49 - parser.add_argument("--enable_chunked_prefill", action='store_true')
45 + parser.add_argument("--enforce_eager", action="store_true")
46 + parser.add_argument("--enable_chunked_prefill", action="store_true")
47 parser.add_argument("--max_num_batched_tokens", type=int, default=2048)
48 parser.add_argument("--temp", type=float, default=0)
49 return parser.parse_args()
50
51
52 def main():
56 -
53 args = parse_args()
54
55 model_dir = "meta-llama/Llama-3.1-8B-Instruct"
56
61 - if args.method == 'eagle':
57 + if args.method == "eagle":
58 eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
63 - elif args.method == 'eagle3':
59 + elif args.method == "eagle3":
60 eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
61 else:
62 raise ValueError(f"unknown method: {args.method}")
@@ -72,11 +68,9 @@ def main():
68 prompts = load_prompts(args.dataset, args.num_prompts)
69
70 prompt_ids = [
75 - tokenizer.apply_chat_template([{
76 - "role": "user",
77 - "content": prompt
78 - }],
79 - add_generation_prompt=True)
71 + tokenizer.apply_chat_template(
72 + [{"role": "user", "content": prompt}], add_generation_prompt=True
73 + )
74 for prompt in prompts
75 ]
76
@@ -102,8 +96,7 @@ def main():
96
97 sampling_params = SamplingParams(temperature=args.temp, max_tokens=256)
98
105 - outputs = llm.generate(prompt_token_ids=prompt_ids,
106 - sampling_params=sampling_params)
99 + outputs = llm.generate(prompt_token_ids=prompt_ids, sampling_params=sampling_params)
100
101 # print the generated text
102 for output in outputs:
@@ -120,19 +113,22 @@ def main():
113 # accepted
114 acceptance_counts = [0] * (args.num_spec_tokens + 1)
115 for output in outputs:
123 - for step, count in enumerate(
124 - output.metrics.spec_token_acceptance_counts):
116 + for step, count in enumerate(output.metrics.spec_token_acceptance_counts):
117 acceptance_counts[step] += count
118
119 print("-" * 50)
128 - print(f"mean acceptance length (including bonus tokens): \
129 - {1 + (sum(acceptance_counts) / acceptance_counts[0]):.2f}")
120 + print(
121 + f"mean acceptance length (including bonus tokens): \
122 + {1 + (sum(acceptance_counts) / acceptance_counts[0]):.2f}"
123 + )
124 print("-" * 50)
125
126 # print acceptance at each token position
127 for i in range(len(acceptance_counts)):
134 - print(f"acceptance at token {i}:"
135 - f"{acceptance_counts[i] / (acceptance_counts[0]):.2f}")
128 + print(
129 + f"acceptance at token {i}:"
130 + f"{acceptance_counts[i] / (acceptance_counts[0]):.2f}"
131 + )
132
133
134 if __name__ == "__main__":
examples/offline_inference/embed_jina_embeddings_v3.py
+11 -8
@@ -10,9 +10,9 @@ def parse_args():
10 parser = FlexibleArgumentParser()
11 parser = EngineArgs.add_cli_args(parser)
12 # Set example specific arguments
13 - parser.set_defaults(model="jinaai/jina-embeddings-v3",
14 - task="embed",
15 - trust_remote_code=True)
13 + parser.set_defaults(
14 + model="jinaai/jina-embeddings-v3", task="embed", trust_remote_code=True
15 + )
16 return parser.parse_args()
17
18
@@ -41,11 +41,14 @@ def main(args: Namespace):
41 print("-" * 60)
42 for prompt, output in zip(prompts, outputs):
43 embeds = output.outputs.embedding
44 - embeds_trimmed = ((str(embeds[:16])[:-1] +
45 - ", ...]") if len(embeds) > 16 else embeds)
46 - print(f"Prompt: {prompt!r} \n"
47 - f"Embeddings for text matching: {embeds_trimmed} "
48 - f"(size={len(embeds)})")
44 + embeds_trimmed = (
45 + (str(embeds[:16])[:-1] + ", ...]") if len(embeds) > 16 else embeds
46 + )
47 + print(
48 + f"Prompt: {prompt!r} \n"
49 + f"Embeddings for text matching: {embeds_trimmed} "
50 + f"(size={len(embeds)})"
51 + )
52 print("-" * 60)
53
54
examples/offline_inference/embed_matryoshka_fy.py
+7 -8
@@ -10,9 +10,9 @@ def parse_args():
10 parser = FlexibleArgumentParser()
11 parser = EngineArgs.add_cli_args(parser)
12 # Set example specific arguments
13 - parser.set_defaults(model="jinaai/jina-embeddings-v3",
14 - task="embed",
15 - trust_remote_code=True)
13 + parser.set_defaults(
14 + model="jinaai/jina-embeddings-v3", task="embed", trust_remote_code=True
15 + )
16 return parser.parse_args()
17
18
@@ -39,11 +39,10 @@ def main(args: Namespace):
39 print("-" * 60)
40 for prompt, output in zip(prompts, outputs):
41 embeds = output.outputs.embedding
42 - embeds_trimmed = ((str(embeds[:16])[:-1] +
43 - ", ...]") if len(embeds) > 16 else embeds)
44 - print(f"Prompt: {prompt!r} \n"
45 - f"Embeddings: {embeds_trimmed} "
46 - f"(size={len(embeds)})")
42 + embeds_trimmed = (
43 + (str(embeds[:16])[:-1] + ", ...]") if len(embeds) > 16 else embeds
44 + )
45 + print(f"Prompt: {prompt!r} \nEmbeddings: {embeds_trimmed} (size={len(embeds)})")
46 print("-" * 60)
47
48
examples/offline_inference/encoder_decoder.py
+26 -14
@@ -1,12 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -'''
2 +"""
3 Demonstrate prompting of text-to-text
4 encoder/decoder models, specifically BART
5 -'''
5 +"""
6
7 from vllm import LLM, SamplingParams
8 -from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
9 - TokensPrompt, zip_enc_dec_prompts)
8 +from vllm.inputs import (
9 + ExplicitEncoderDecoderPrompt,
10 + TextPrompt,
11 + TokensPrompt,
12 + zip_enc_dec_prompts,
13 +)
14
15
16 def create_prompts(tokenizer):
@@ -18,8 +22,9 @@ def create_prompts(tokenizer):
22 # - Helpers for building prompts
23 text_prompt_raw = "Hello, my name is"
24 text_prompt = TextPrompt(prompt="The president of the United States is")
21 - tokens_prompt = TokensPrompt(prompt_token_ids=tokenizer.encode(
22 - prompt="The capital of France is"))
25 + tokens_prompt = TokensPrompt(
26 + prompt_token_ids=tokenizer.encode(prompt="The capital of France is")
27 + )
28 # - Pass a single prompt to encoder/decoder model
29 # (implicitly encoder input prompt);
30 # decoder input prompt is assumed to be None
@@ -57,14 +62,19 @@ def create_prompts(tokenizer):
62 # decoder prompts together into a list of ExplicitEncoderDecoderPrompt
63 # instances
64 zipped_prompt_list = zip_enc_dec_prompts(
60 - ['An encoder prompt', 'Another encoder prompt'],
61 - ['A decoder prompt', 'Another decoder prompt'])
65 + ["An encoder prompt", "Another encoder prompt"],
66 + ["A decoder prompt", "Another decoder prompt"],
67 + )
68
69 # - Let's put all of the above example prompts together into one list
70 # which we will pass to the encoder/decoder LLM.
71 return [
66 - single_text_prompt_raw, single_text_prompt, single_tokens_prompt,
67 - enc_dec_prompt1, enc_dec_prompt2, enc_dec_prompt3
72 + single_text_prompt_raw,
73 + single_text_prompt,
74 + single_tokens_prompt,
75 + enc_dec_prompt1,
76 + enc_dec_prompt2,
77 + enc_dec_prompt3,
78 ] + zipped_prompt_list
79
80
@@ -85,10 +95,12 @@ def print_outputs(outputs):
95 prompt = output.prompt
96 encoder_prompt = output.encoder_prompt
97 generated_text = output.outputs[0].text
88 - print(f"Output {i+1}:")
89 - print(f"Encoder prompt: {encoder_prompt!r}\n"
90 - f"Decoder prompt: {prompt!r}\n"
91 - f"Generated text: {generated_text!r}")
98 + print(f"Output {i + 1}:")
99 + print(
100 + f"Encoder prompt: {encoder_prompt!r}\n"
101 + f"Decoder prompt: {prompt!r}\n"
102 + f"Generated text: {generated_text!r}"
103 + )
104 print("-" * 50)
105
106
examples/offline_inference/encoder_decoder_multimodal.py
+32 -30
@@ -3,6 +3,7 @@
3 This example shows how to use vLLM for running offline inference with
4 the explicit/implicit prompt format on enc-dec LMMs for text generation.
5 """
6 +
7 import time
8 from collections.abc import Sequence
9 from dataclasses import asdict
@@ -30,18 +31,14 @@ def run_florence2():
31 )
32
33 prompts = [
33 - { # implicit prompt with task token
34 + { # implicit prompt with task token
35 "prompt": "<DETAILED_CAPTION>",
35 - "multi_modal_data": {
36 - "image": ImageAsset("stop_sign").pil_image
37 - },
36 + "multi_modal_data": {"image": ImageAsset("stop_sign").pil_image},
37 },
39 - { # explicit encoder/decoder prompt
38 + { # explicit encoder/decoder prompt
39 "encoder_prompt": {
40 "prompt": "Describe in detail what is shown in the image.",
42 - "multi_modal_data": {
43 - "image": ImageAsset("cherry_blossom").pil_image
44 - },
41 + "multi_modal_data": {"image": ImageAsset("cherry_blossom").pil_image},
42 },
43 "decoder_prompt": "",
44 },
@@ -63,20 +60,20 @@ def run_mllama():
60 )
61
62 prompts = [
66 - { # Implicit prompt
67 - "prompt": "<|image|><|begin_of_text|>What is the content of this image?", # noqa: E501
63 + { # Implicit prompt
64 + "prompt": "<|image|><|begin_of_text|>What is the content of this image?", # noqa: E501
65 "multi_modal_data": {
66 "image": ImageAsset("stop_sign").pil_image,
67 },
68 },
72 - { # Explicit prompt
69 + { # Explicit prompt
70 "encoder_prompt": {
71 "prompt": "<|image|>",
72 "multi_modal_data": {
73 "image": ImageAsset("stop_sign").pil_image,
74 },
75 },
79 - "decoder_prompt": "<|image|><|begin_of_text|>Please describe the image.", # noqa: E501
76 + "decoder_prompt": "<|image|><|begin_of_text|>Please describe the image.", # noqa: E501
77 },
78 ]
79
@@ -96,13 +93,13 @@ def run_whisper():
93 )
94
95 prompts = [
99 - { # Test implicit prompt
96 + { # Test implicit prompt
97 "prompt": "<|startoftranscript|>",
98 "multi_modal_data": {
99 "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
100 },
101 },
105 - { # Test explicit encoder/decoder prompt
102 + { # Test explicit encoder/decoder prompt
103 "encoder_prompt": {
104 "prompt": "",
105 "multi_modal_data": {
@@ -110,7 +107,7 @@ def run_whisper():
107 },
108 },
109 "decoder_prompt": "<|startoftranscript|>",
113 - }
110 + },
111 ]
112
113 return ModelRequestData(
@@ -128,18 +125,23 @@ model_example_map = {
125
126 def parse_args():
127 parser = FlexibleArgumentParser(
131 - description='Demo on using vLLM for offline inference with '
132 - 'vision language models for text generation')
133 - parser.add_argument('--model-type',
134 - '-m',
135 - type=str,
136 - default="mllama",
137 - choices=model_example_map.keys(),
138 - help='Huggingface "model_type".')
139 - parser.add_argument("--seed",
140 - type=int,
141 - default=None,
142 - help="Set the seed when initializing `vllm.LLM`.")
128 + description="Demo on using vLLM for offline inference with "
129 + "vision language models for text generation"
130 + )
131 + parser.add_argument(
132 + "--model-type",
133 + "-m",
134 + type=str,
135 + default="mllama",
136 + choices=model_example_map.keys(),
137 + help='Huggingface "model_type".',
138 + )
139 + parser.add_argument(
140 + "--seed",
141 + type=int,
142 + default=None,
143 + help="Set the seed when initializing `vllm.LLM`.",
144 + )
145 return parser.parse_args()
146
147
@@ -153,7 +155,8 @@ def main(args):
155 # Disable other modalities to save memory
156 default_limits = {"image": 0, "video": 0, "audio": 0}
157 req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
156 - req_data.engine_args.limit_mm_per_prompt or {})
158 + req_data.engine_args.limit_mm_per_prompt or {}
159 + )
160
161 engine_args = asdict(req_data.engine_args) | {"seed": args.seed}
162 llm = LLM(**engine_args)
@@ -179,8 +182,7 @@ def main(args):
182 for output in outputs:
183 prompt = output.prompt
184 generated_text = output.outputs[0].text
182 - print(f"Decoder prompt: {prompt!r}, "
183 - f"Generated text: {generated_text!r}")
185 + print(f"Decoder prompt: {prompt!r}, Generated text: {generated_text!r}")
186
187 duration = time.time() - start
188
examples/offline_inference/llm_engine_example.py
+19 -15
@@ -3,6 +3,7 @@
3 This file demonstrates using the `LLMEngine`
4 for processing prompts with various sampling parameters.
5 """
6 +
7 import argparse
8
9 from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
@@ -12,24 +13,26 @@ from vllm.utils import FlexibleArgumentParser
13 def create_test_prompts() -> list[tuple[str, SamplingParams]]:
14 """Create a list of test prompts with their sampling parameters."""
15 return [
15 - ("A robot may not injure a human being",
16 - SamplingParams(temperature=0.0, logprobs=1, prompt_logprobs=1)),
17 - ("To be or not to be,",
18 - SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2)),
19 - ("What is the meaning of life?",
20 - SamplingParams(n=2,
21 - temperature=0.8,
22 - top_p=0.95,
23 - frequency_penalty=0.1)),
16 + (
17 + "A robot may not injure a human being",
18 + SamplingParams(temperature=0.0, logprobs=1, prompt_logprobs=1),
19 + ),
20 + (
21 + "To be or not to be,",
22 + SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2),
23 + ),
24 + (
25 + "What is the meaning of life?",
26 + SamplingParams(n=2, temperature=0.8, top_p=0.95, frequency_penalty=0.1),
27 + ),
28 ]
29
30
27 -def process_requests(engine: LLMEngine,
28 - test_prompts: list[tuple[str, SamplingParams]]):
31 +def process_requests(engine: LLMEngine, test_prompts: list[tuple[str, SamplingParams]]):
32 """Continuously process a list of prompts and handle the outputs."""
33 request_id = 0
34
32 - print('-' * 50)
35 + print("-" * 50)
36 while test_prompts or engine.has_unfinished_requests():
37 if test_prompts:
38 prompt, sampling_params = test_prompts.pop(0)
@@ -41,7 +44,7 @@ def process_requests(engine: LLMEngine,
44 for request_output in request_outputs:
45 if request_output.finished:
46 print(request_output)
44 - print('-' * 50)
47 + print("-" * 50)
48
49
50 def initialize_engine(args: argparse.Namespace) -> LLMEngine:
@@ -52,7 +55,8 @@ def initialize_engine(args: argparse.Namespace) -> LLMEngine:
55
56 def parse_args():
57 parser = FlexibleArgumentParser(
55 - description='Demo on using the LLMEngine class directly')
58 + description="Demo on using the LLMEngine class directly"
59 + )
60 parser = EngineArgs.add_cli_args(parser)
61 return parser.parse_args()
62
@@ -64,6 +68,6 @@ def main(args: argparse.Namespace):
68 process_requests(engine, test_prompts)
69
70
67 -if __name__ == '__main__':
71 +if __name__ == "__main__":
72 args = parse_args()
73 main(args)
examples/offline_inference/load_sharded_state.py
+19 -19
@@ -36,22 +36,21 @@ def parse_args():
36 parser.set_defaults(load_format="sharded_state")
37
38 # Add validation arguments
39 - parser.add_argument("--prompt",
40 - type=str,
41 - default="Hello, world!",
42 - help="Prompt for validation")
43 - parser.add_argument("--max-tokens",
44 - type=int,
45 - default=100,
46 - help="Maximum number of tokens to generate")
47 - parser.add_argument("--temperature",
48 - type=float,
49 - default=0.7,
50 - help="Sampling temperature")
51 - parser.add_argument("--top-p",
52 - type=float,
53 - default=1.0,
54 - help="Top-p sampling parameter")
39 + parser.add_argument(
40 + "--prompt", type=str, default="Hello, world!", help="Prompt for validation"
41 + )
42 + parser.add_argument(
43 + "--max-tokens",
44 + type=int,
45 + default=100,
46 + help="Maximum number of tokens to generate",
47 + )
48 + parser.add_argument(
49 + "--temperature", type=float, default=0.7, help="Sampling temperature"
50 + )
51 + parser.add_argument(
52 + "--top-p", type=float, default=1.0, help="Top-p sampling parameter"
53 + )
54
55 return parser.parse_args()
56
@@ -60,8 +59,9 @@ def main():
59 args = parse_args()
60 engine_args = EngineArgs.from_cli_args(args)
61
63 - print(f"Loading model from {engine_args.model} "
64 - f"using format {engine_args.load_format}")
62 + print(
63 + f"Loading model from {engine_args.model} using format {engine_args.load_format}"
64 + )
65 print(f"Tensor parallel size: {engine_args.tensor_parallel_size}")
66
67 # Load the model using engine args
@@ -90,4 +90,4 @@ def main():
90
91
92 if __name__ == "__main__":
93 - main()
\ No newline at end of file
93 + main()
examples/offline_inference/lora_with_quantization_inference.py
+61 -55
@@ -17,50 +17,55 @@ from vllm.lora.request import LoRARequest
17
18
19 def create_test_prompts(
20 - lora_path: str
20 + lora_path: str,
21 ) -> list[tuple[str, SamplingParams, Optional[LoRARequest]]]:
22 return [
23 # this is an example of using quantization without LoRA
24 - ("My name is",
25 - SamplingParams(temperature=0.0,
26 - logprobs=1,
27 - prompt_logprobs=1,
28 - max_tokens=128), None),
24 + (
25 + "My name is",
26 + SamplingParams(
27 + temperature=0.0, logprobs=1, prompt_logprobs=1, max_tokens=128
28 + ),
29 + None,
30 + ),
31 # the next three examples use quantization with LoRA
30 - ("my name is",
31 - SamplingParams(temperature=0.0,
32 - logprobs=1,
33 - prompt_logprobs=1,
34 - max_tokens=128),
35 - LoRARequest("lora-test-1", 1, lora_path)),
36 - ("The capital of USA is",
37 - SamplingParams(temperature=0.0,
38 - logprobs=1,
39 - prompt_logprobs=1,
40 - max_tokens=128),
41 - LoRARequest("lora-test-2", 1, lora_path)),
42 - ("The capital of France is",
43 - SamplingParams(temperature=0.0,
44 - logprobs=1,
45 - prompt_logprobs=1,
46 - max_tokens=128),
47 - LoRARequest("lora-test-3", 1, lora_path)),
32 + (
33 + "my name is",
34 + SamplingParams(
35 + temperature=0.0, logprobs=1, prompt_logprobs=1, max_tokens=128
36 + ),
37 + LoRARequest("lora-test-1", 1, lora_path),
38 + ),
39 + (
40 + "The capital of USA is",
41 + SamplingParams(
42 + temperature=0.0, logprobs=1, prompt_logprobs=1, max_tokens=128
43 + ),
44 + LoRARequest("lora-test-2", 1, lora_path),
45 + ),
46 + (
47 + "The capital of France is",
48 + SamplingParams(
49 + temperature=0.0, logprobs=1, prompt_logprobs=1, max_tokens=128
50 + ),
51 + LoRARequest("lora-test-3", 1, lora_path),
52 + ),
53 ]
54
55
51 -def process_requests(engine: LLMEngine,
52 - test_prompts: list[tuple[str, SamplingParams,
53 - Optional[LoRARequest]]]):
56 +def process_requests(
57 + engine: LLMEngine,
58 + test_prompts: list[tuple[str, SamplingParams, Optional[LoRARequest]]],
59 +):
60 """Continuously process a list of prompts and handle the outputs."""
61 request_id = 0
62
63 while test_prompts or engine.has_unfinished_requests():
64 if test_prompts:
65 prompt, sampling_params, lora_request = test_prompts.pop(0)
60 - engine.add_request(str(request_id),
61 - prompt,
62 - sampling_params,
63 - lora_request=lora_request)
66 + engine.add_request(
67 + str(request_id), prompt, sampling_params, lora_request=lora_request
68 + )
69 request_id += 1
70
71 request_outputs: list[RequestOutput] = engine.step()
@@ -71,15 +76,18 @@ def process_requests(engine: LLMEngine,
76 print(f"Output: {request_output.outputs[0].text}")
77
78
74 -def initialize_engine(model: str, quantization: str,
75 - lora_repo: Optional[str]) -> LLMEngine:
79 +def initialize_engine(
80 + model: str, quantization: str, lora_repo: Optional[str]
81 +) -> LLMEngine:
82 """Initialize the LLMEngine."""
83
78 - engine_args = EngineArgs(model=model,
79 - quantization=quantization,
80 - enable_lora=True,
81 - max_lora_rank=64,
82 - max_loras=4)
84 + engine_args = EngineArgs(
85 + model=model,
86 + quantization=quantization,
87 + enable_lora=True,
88 + max_lora_rank=64,
89 + max_loras=4,
90 + )
91 return LLMEngine.from_engine_args(engine_args)
92
93
@@ -90,32 +98,30 @@ def main():
98 # QLoRA (https://arxiv.org/abs/2305.14314)
99 {
100 "name": "qlora_inference_example",
93 - 'model': "huggyllama/llama-7b",
94 - 'quantization': "bitsandbytes",
95 - 'lora_repo': 'timdettmers/qlora-flan-7b'
101 + "model": "huggyllama/llama-7b",
102 + "quantization": "bitsandbytes",
103 + "lora_repo": "timdettmers/qlora-flan-7b",
104 },
105 {
106 "name": "AWQ_inference_with_lora_example",
99 - 'model': 'TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ',
100 - 'quantization': "awq",
101 - 'lora_repo': 'jashing/tinyllama-colorist-lora'
107 + "model": "TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ",
108 + "quantization": "awq",
109 + "lora_repo": "jashing/tinyllama-colorist-lora",
110 },
111 {
112 "name": "GPTQ_inference_with_lora_example",
105 - 'model': 'TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ',
106 - 'quantization': "gptq",
107 - 'lora_repo': 'jashing/tinyllama-colorist-lora'
108 - }
113 + "model": "TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ",
114 + "quantization": "gptq",
115 + "lora_repo": "jashing/tinyllama-colorist-lora",
116 + },
117 ]
118
119 for test_config in test_configs:
112 - print(
113 - f"~~~~~~~~~~~~~~~~ Running: {test_config['name']} ~~~~~~~~~~~~~~~~"
120 + print(f"~~~~~~~~~~~~~~~~ Running: {test_config['name']} ~~~~~~~~~~~~~~~~")
121 + engine = initialize_engine(
122 + test_config["model"], test_config["quantization"], test_config["lora_repo"]
123 )
115 - engine = initialize_engine(test_config['model'],
116 - test_config['quantization'],
117 - test_config['lora_repo'])
118 - lora_path = snapshot_download(repo_id=test_config['lora_repo'])
124 + lora_path = snapshot_download(repo_id=test_config["lora_repo"])
125 test_prompts = create_test_prompts(lora_path)
126 process_requests(engine, test_prompts)
127
@@ -125,5 +131,5 @@ def main():
131 torch.cuda.empty_cache()
132
133
128 -if __name__ == '__main__':
134 +if __name__ == "__main__":
135 main()
examples/offline_inference/mistral-small.py
+20 -44
@@ -74,19 +74,10 @@ def run_simple_demo(args: argparse.Namespace):
74
75 messages = [
76 {
77 - "role":
78 - "user",
77 + "role": "user",
78 "content": [
80 - {
81 - "type": "text",
82 - "text": prompt
83 - },
84 - {
85 - "type": "image_url",
86 - "image_url": {
87 - "url": image_url
88 - }
89 - },
79 + {"type": "text", "text": prompt},
80 + {"type": "image_url", "image_url": {"url": image_url}},
81 ],
82 },
83 ]
@@ -121,25 +112,11 @@ def run_advanced_demo(args: argparse.Namespace):
112
113 messages = [
114 {
124 - "role":
125 - "user",
115 + "role": "user",
116 "content": [
127 - {
128 - "type": "text",
129 - "text": prompt
130 - },
131 - {
132 - "type": "image_url",
133 - "image_url": {
134 - "url": url_1
135 - }
136 - },
137 - {
138 - "type": "image_url",
139 - "image_url": {
140 - "url": url_2
141 - }
142 - },
117 + {"type": "text", "text": prompt},
118 + {"type": "image_url", "image_url": {"url": url_1}},
119 + {"type": "image_url", "image_url": {"url": url_2}},
120 ],
121 },
122 {
@@ -153,12 +130,7 @@ def run_advanced_demo(args: argparse.Namespace):
130 {
131 "role": "user",
132 "content": [
156 - {
157 - "type": "image_url",
158 - "image_url": {
159 - "url": url_3
160 - }
161 - },
133 + {"type": "image_url", "image_url": {"url": url_3}},
134 ],
135 },
136 ]
@@ -171,7 +143,8 @@ def run_advanced_demo(args: argparse.Namespace):
143
144 def parse_args():
145 parser = argparse.ArgumentParser(
174 - description="Run a demo in simple or advanced mode.")
146 + description="Run a demo in simple or advanced mode."
147 + )
148
149 parser.add_argument(
150 "mode",
@@ -179,15 +152,18 @@ def parse_args():
152 help="Specify the demo mode: 'simple' or 'advanced'",
153 )
154
182 - parser.add_argument('--format',
183 - choices=["mistral", "hf"],
184 - default="mistral",
185 - help='Specify the format of the model to load.')
155 + parser.add_argument(
156 + "--format",
157 + choices=["mistral", "hf"],
158 + default="mistral",
159 + help="Specify the format of the model to load.",
160 + )
161
162 parser.add_argument(
188 - '--disable-mm-preprocessor-cache',
189 - action='store_true',
190 - help='If True, disables caching of multi-modal preprocessor/mapper.')
163 + "--disable-mm-preprocessor-cache",
164 + action="store_true",
165 + help="If True, disables caching of multi-modal preprocessor/mapper.",
166 + )
167 return parser.parse_args()
168
169
examples/offline_inference/mlpspeculator.py
+6 -5
@@ -13,8 +13,9 @@ import time
13 from vllm import LLM, SamplingParams
14
15
16 -def time_generation(llm: LLM, prompts: list[str],
17 - sampling_params: SamplingParams, title: str):
16 +def time_generation(
17 + llm: LLM, prompts: list[str], sampling_params: SamplingParams, title: str
18 +):
19 # Generate texts from the prompts. The output is a list of RequestOutput
20 # objects that contain the prompt, generated text, and other information.
21 # Warmup first
@@ -25,8 +26,7 @@ def time_generation(llm: LLM, prompts: list[str],
26 end = time.time()
27 print("-" * 50)
28 print(title)
28 - print("time: ",
29 - (end - start) / sum(len(o.outputs[0].token_ids) for o in outputs))
29 + print("time: ", (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
@@ -38,7 +38,8 @@ def main():
38 template = (
39 "Below is an instruction that describes a task. Write a response "
40 "that appropriately completes the request.\n\n### Instruction:\n{}"
41 - "\n\n### Response:\n")
41 + "\n\n### Response:\n"
42 + )
43
44 # Sample prompts.
45 prompts = [
examples/offline_inference/multilora_inference.py
+49 -37
@@ -15,7 +15,7 @@ from vllm.lora.request import LoRARequest
15
16
17 def create_test_prompts(
18 - lora_path: str
18 + lora_path: str,
19 ) -> list[tuple[str, SamplingParams, Optional[LoRARequest]]]:
20 """Create a list of test prompts with their sampling parameters.
21
@@ -26,38 +26,49 @@ def create_test_prompts(
26 first adapter have finished.
27 """
28 return [
29 - ("A robot may not injure a human being",
30 - SamplingParams(temperature=0.0,
31 - logprobs=1,
32 - prompt_logprobs=1,
33 - max_tokens=128), None),
34 - ("To be or not to be,",
35 - SamplingParams(temperature=0.8,
36 - top_k=5,
37 - presence_penalty=0.2,
38 - max_tokens=128), None),
29 + (
30 + "A robot may not injure a human being",
31 + SamplingParams(
32 + temperature=0.0, logprobs=1, prompt_logprobs=1, max_tokens=128
33 + ),
34 + None,
35 + ),
36 + (
37 + "To be or not to be,",
38 + SamplingParams(
39 + temperature=0.8, top_k=5, presence_penalty=0.2, max_tokens=128
40 + ),
41 + None,
42 + ),
43 (
44 "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501
41 - SamplingParams(temperature=0.0,
42 - logprobs=1,
43 - prompt_logprobs=1,
44 - max_tokens=128,
45 - stop_token_ids=[32003]),
46 - LoRARequest("sql-lora", 1, lora_path)),
45 + SamplingParams(
46 + temperature=0.0,
47 + logprobs=1,
48 + prompt_logprobs=1,
49 + max_tokens=128,
50 + stop_token_ids=[32003],
51 + ),
52 + LoRARequest("sql-lora", 1, lora_path),
53 + ),
54 (
55 "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501
49 - SamplingParams(temperature=0.0,
50 - logprobs=1,
51 - prompt_logprobs=1,
52 - max_tokens=128,
53 - stop_token_ids=[32003]),
54 - LoRARequest("sql-lora2", 2, lora_path)),
56 + SamplingParams(
57 + temperature=0.0,
58 + logprobs=1,
59 + prompt_logprobs=1,
60 + max_tokens=128,
61 + stop_token_ids=[32003],
62 + ),
63 + LoRARequest("sql-lora2", 2, lora_path),
64 + ),
65 ]
66
67
58 -def process_requests(engine: LLMEngine,
59 - test_prompts: list[tuple[str, SamplingParams,
60 - Optional[LoRARequest]]]):
68 +def process_requests(
69 + engine: LLMEngine,
70 + test_prompts: list[tuple[str, SamplingParams, Optional[LoRARequest]]],
71 +):
72 """Continuously process a list of prompts and handle the outputs."""
73 request_id = 0
74
@@ -65,10 +76,9 @@ def process_requests(engine: LLMEngine,
76 while test_prompts or engine.has_unfinished_requests():
77 if test_prompts:
78 prompt, sampling_params, lora_request = test_prompts.pop(0)
68 - engine.add_request(str(request_id),
69 - prompt,
70 - sampling_params,
71 - lora_request=lora_request)
79 + engine.add_request(
80 + str(request_id), prompt, sampling_params, lora_request=lora_request
81 + )
82 request_id += 1
83
84 request_outputs: list[RequestOutput] = engine.step()
@@ -88,12 +98,14 @@ def initialize_engine() -> LLMEngine:
98 # numbers will cause higher memory usage. If you know that all LoRAs will
99 # use the same rank, it is recommended to set this as low as possible.
100 # max_cpu_loras: controls the size of the CPU LoRA cache.
91 - engine_args = EngineArgs(model="meta-llama/Llama-2-7b-hf",
92 - enable_lora=True,
93 - max_loras=1,
94 - max_lora_rank=8,
95 - max_cpu_loras=2,
96 - max_num_seqs=256)
101 + engine_args = EngineArgs(
102 + model="meta-llama/Llama-2-7b-hf",
103 + enable_lora=True,
104 + max_loras=1,
105 + max_lora_rank=8,
106 + max_cpu_loras=2,
107 + max_num_seqs=256,
108 + )
109 return LLMEngine.from_engine_args(engine_args)
110
111
@@ -105,5 +117,5 @@ def main():
117 process_requests(engine, test_prompts)
118
119
108 -if __name__ == '__main__':
120 +if __name__ == "__main__":
121 main()
examples/offline_inference/neuron.py
+2 -1
@@ -30,7 +30,8 @@ def main():
30 # The device argument can be either unspecified for automated detection,
31 # or explicitly assigned.
32 device="neuron",
33 - tensor_parallel_size=2)
33 + tensor_parallel_size=2,
34 + )
35 # Generate texts from the prompts. The output is a list of RequestOutput objects
36 # that contain the prompt, generated text, and other information.
37 outputs = llm.generate(prompts, sampling_params)
examples/offline_inference/neuron_eagle.py
+3 -3
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 -This example shows how to run offline inference with an EAGLE speculative
3 +This example shows how to run offline inference with an EAGLE speculative
4 decoding model on neuron. To use EAGLE speculative decoding, you must use
5 a draft model that is specifically fine-tuned for EAGLE speculation.
6 Additionally, to use EAGLE with NxD Inference, the draft model must include
@@ -24,7 +24,7 @@ llm = LLM(
24 speculative_config={
25 "model": "/home/ubuntu/model_hf/Llama-3.1-70B-Instruct-EAGLE-Draft",
26 "num_speculative_tokens": 5,
27 - "max_model_len": 2048
27 + "max_model_len": 2048,
28 },
29 max_num_seqs=4,
30 # The max_model_len and block_size arguments are required to be same as
@@ -40,7 +40,7 @@ llm = LLM(
40 tensor_parallel_size=32,
41 override_neuron_config={
42 "enable_eagle_speculation": True,
43 - "enable_fused_speculation": True
43 + "enable_fused_speculation": True,
44 },
45 )
46
examples/offline_inference/neuron_int8_quantization.py
+5 -4
@@ -5,12 +5,12 @@ import os
5 from vllm import LLM, SamplingParams
6
7 # creates XLA hlo graphs for all the context length buckets.
8 -os.environ['NEURON_CONTEXT_LENGTH_BUCKETS'] = "128,512,1024,2048"
8 +os.environ["NEURON_CONTEXT_LENGTH_BUCKETS"] = "128,512,1024,2048"
9 # creates XLA hlo graphs for all the token gen buckets.
10 -os.environ['NEURON_TOKEN_GEN_BUCKETS'] = "128,512,1024,2048"
10 +os.environ["NEURON_TOKEN_GEN_BUCKETS"] = "128,512,1024,2048"
11 # Quantizes neuron model weight to int8 ,
12 # The default config for quantization is int8 dtype.
13 -os.environ['NEURON_QUANT_DTYPE'] = "s8"
13 +os.environ["NEURON_QUANT_DTYPE"] = "s8"
14
15 # Sample prompts.
16 prompts = [
@@ -44,7 +44,8 @@ def main():
44 override_neuron_config={
45 "cast_logits_dtype": "bfloat16",
46 },
47 - tensor_parallel_size=2)
47 + tensor_parallel_size=2,
48 + )
49 # Generate texts from the prompts. The output is a list of RequestOutput objects
50 # that contain the prompt, generated text, and other information.
51 outputs = llm.generate(prompts, sampling_params)
examples/offline_inference/neuron_speculation.py
+5 -5
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 -This example shows how to run offline inference with a speculative
3 +This example shows how to run offline inference with a speculative
4 decoding model on neuron.
5 """
6
@@ -19,9 +19,9 @@ prompts = [
19 def config_buckets():
20 """Configure context length and token gen buckets."""
21 # creates XLA hlo graphs for all the context length buckets.
22 - os.environ['NEURON_CONTEXT_LENGTH_BUCKETS'] = "128,512,1024,2048"
22 + os.environ["NEURON_CONTEXT_LENGTH_BUCKETS"] = "128,512,1024,2048"
23 # creates XLA hlo graphs for all the token gen buckets.
24 - os.environ['NEURON_TOKEN_GEN_BUCKETS'] = "128,512,1024,2048"
24 + os.environ["NEURON_TOKEN_GEN_BUCKETS"] = "128,512,1024,2048"
25
26
27 def initialize_model():
@@ -31,7 +31,7 @@ def initialize_model():
31 speculative_config={
32 "model": "openlm-research/open_llama_3b",
33 "num_speculative_tokens": 4,
34 - "max_model_len": 2048
34 + "max_model_len": 2048,
35 },
36 max_num_seqs=4,
37 max_model_len=2048,
@@ -60,5 +60,5 @@ def main():
60 process_requests(model, sampling_params)
61
62
63 -if __name__ == '__main__':
63 +if __name__ == "__main__":
64 main()
examples/offline_inference/prefix_caching.py
+13 -8
@@ -16,7 +16,8 @@ prefix = (
16 "teaching role. They have 5 years of previous teaching experience "
17 "as an assistant teacher at a co-ed, public school with experience "
18 "in middle school math teaching. Based on these information, fulfill "
19 - "the following paragraph: ")
19 + "the following paragraph: "
20 +)
21
22 # Sample prompts.
23 prompts = [
@@ -58,9 +59,11 @@ def main():
59 cleanup_dist_env_and_memory()
60
61 # Create an LLM with prefix caching enabled.
61 - prefix_cached_llm = LLM(model="facebook/opt-125m",
62 - enable_prefix_caching=True,
63 - gpu_memory_utilization=0.4)
62 + prefix_cached_llm = LLM(
63 + model="facebook/opt-125m",
64 + enable_prefix_caching=True,
65 + gpu_memory_utilization=0.4,
66 + )
67
68 # Warmup so that the shared prompt's KV cache is computed.
69 prefix_cached_llm.generate(generating_prompts[0], sampling_params)
@@ -81,10 +84,12 @@ def main():
84 print("-" * 50)
85
86 # Compare the results and display the speedup
84 - generated_same = all([
85 - regular_generated_texts[i] == cached_generated_texts[i]
86 - for i in range(len(prompts))
87 - ])
87 + generated_same = all(
88 + [
89 + regular_generated_texts[i] == cached_generated_texts[i]
90 + for i in range(len(prompts))
91 + ]
92 + )
93 print(f"Generated answers are the same: {generated_same}")
94
95
examples/offline_inference/prithvi_geospatial_mae.py
+82 -97
@@ -16,7 +16,8 @@ The requirements for running this script are:
16 Run the example:
17 python prithvi_geospatial_mae.py
18
19 -""" # noqa: E501
19 +""" # noqa: E501
20 +
21 import argparse
22 import datetime
23 import os
@@ -110,77 +111,67 @@ model_config = """{
111
112 # Temporarily creating the "config.json" for the model.
113 # This is going to disappear once the correct config.json is available on HF
113 -with open(os.path.join(os.path.dirname(__file__), "./model/config.json"),
114 - 'w') as config_file:
114 +with open(
115 + os.path.join(os.path.dirname(__file__), "./model/config.json"), "w"
116 +) as config_file:
117 config_file.write(model_config)
118
119 datamodule_config = {
118 - 'bands': ['BLUE', 'GREEN', 'RED', 'NIR_NARROW', 'SWIR_1', 'SWIR_2'],
119 - 'batch_size':
120 - 16,
121 - 'constant_scale':
122 - 0.0001,
123 - 'data_root':
124 - '/dccstor/geofm-finetuning/datasets/sen1floods11',
125 - 'drop_last':
126 - True,
127 - 'no_data_replace':
128 - 0.0,
129 - 'no_label_replace':
130 - -1,
131 - 'num_workers':
132 - 8,
133 - 'test_transform': [
134 - albumentations.Resize(always_apply=False,
135 - height=448,
136 - interpolation=1,
137 - p=1,
138 - width=448),
139 - albumentations.pytorch.ToTensorV2(transpose_mask=False,
140 - always_apply=True,
141 - p=1.0)
120 + "bands": ["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"],
121 + "batch_size": 16,
122 + "constant_scale": 0.0001,
123 + "data_root": "/dccstor/geofm-finetuning/datasets/sen1floods11",
124 + "drop_last": True,
125 + "no_data_replace": 0.0,
126 + "no_label_replace": -1,
127 + "num_workers": 8,
128 + "test_transform": [
129 + albumentations.Resize(
130 + always_apply=False, height=448, interpolation=1, p=1, width=448
131 + ),
132 + albumentations.pytorch.ToTensorV2(
133 + transpose_mask=False, always_apply=True, p=1.0
134 + ),
135 ],
136 }
137
138
139 class PrithviMAE:
147 -
140 def __init__(self):
141 print("Initializing PrithviMAE model")
150 - self.model = LLM(model=os.path.join(os.path.dirname(__file__),
151 - "./model"),
152 - skip_tokenizer_init=True,
153 - dtype="float32")
142 + self.model = LLM(
143 + model=os.path.join(os.path.dirname(__file__), "./model"),
144 + skip_tokenizer_init=True,
145 + dtype="float32",
146 + )
147
148 def run(self, input_data, location_coords):
149 print("################ Running inference on vLLM ##############")
150 # merge the inputs into one data structure
151 mm_data = {
159 - "pixel_values":
160 - torch.empty(0) if input_data is None else input_data,
161 - "location_coords":
162 - torch.empty(0) if location_coords is None else location_coords
152 + "pixel_values": torch.empty(0) if input_data is None else input_data,
153 + "location_coords": torch.empty(0)
154 + if location_coords is None
155 + else location_coords,
156 }
157
158 prompt = {"prompt_token_ids": [1], "multi_modal_data": mm_data}
159
160 outputs = self.model.encode(prompt, use_tqdm=False)
168 - print(
169 - "################ Inference done (it took seconds) ##############"
170 - )
161 + print("################ Inference done (it took seconds) ##############")
162
163 return outputs[0].outputs.data
164
165
166 def generate_datamodule():
167 datamodule = Sen1Floods11NonGeoDataModule(
177 - data_root=datamodule_config['data_root'],
168 + data_root=datamodule_config["data_root"],
169 batch_size=datamodule_config["batch_size"],
170 num_workers=datamodule_config["num_workers"],
171 bands=datamodule_config["bands"],
172 drop_last=datamodule_config["drop_last"],
182 - test_transform=datamodule_config["test_transform"
183 - ""])
173 + test_transform=datamodule_config["test_transform"],
174 + )
175
176 return datamodule
177
@@ -204,8 +195,7 @@ def process_channel_group(orig_img, channels):
195 max_value = max(3000, np.percentile(orig_img[valid_mask], PERCENTILE))
196 min_value = OFFSET
197
207 - orig_img = torch.clamp((orig_img - min_value) / (max_value - min_value), 0,
208 - 1)
198 + orig_img = torch.clamp((orig_img - min_value) / (max_value - min_value), 0, 1)
199
200 # No data as zeros
201 orig_img[~valid_mask] = 0
@@ -300,18 +290,21 @@ def load_example(
290 location_coords.append(coords)
291
292 try:
303 - match = re.search(r'(\d{7,8}T\d{6})', file)
293 + match = re.search(r"(\d{7,8}T\d{6})", file)
294 if match:
295 year = int(match.group(1)[:4])
306 - julian_day = match.group(1).split('T')[0][4:]
296 + julian_day = match.group(1).split("T")[0][4:]
297 if len(julian_day) == 3:
298 julian_day = int(julian_day)
299 else:
310 - julian_day = datetime.datetime.strptime(
311 - julian_day, '%m%d').timetuple().tm_yday
300 + julian_day = (
301 + datetime.datetime.strptime(julian_day, "%m%d")
302 + .timetuple()
303 + .tm_yday
304 + )
305 temporal_coords.append([year, julian_day])
306 except Exception as e:
314 - print(f'Could not extract timestamp for {file} ({e})')
307 + print(f"Could not extract timestamp for {file} ({e})")
308
309 imgs = np.stack(imgs, axis=0) # num_frames, H, W, C
310 imgs = np.moveaxis(imgs, -1, 0).astype("float32")
@@ -320,50 +313,44 @@ def load_example(
313 return imgs, temporal_coords, location_coords, metas
314
315
323 -def run_model(input_data,
324 - temporal_coords,
325 - location_coords,
326 - model,
327 - datamodule,
328 - img_size,
329 - lightning_model=None):
316 +def run_model(
317 + input_data,
318 + temporal_coords,
319 + location_coords,
320 + model,
321 + datamodule,
322 + img_size,
323 + lightning_model=None,
324 +):
325 # Reflect pad if not divisible by img_size
326 original_h, original_w = input_data.shape[-2:]
327 pad_h = (img_size - (original_h % img_size)) % img_size
328 pad_w = (img_size - (original_w % img_size)) % img_size
334 - input_data = np.pad(input_data,
335 - ((0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)),
336 - mode="reflect")
329 + input_data = np.pad(
330 + input_data, ((0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)), mode="reflect"
331 + )
332
333 # Build sliding window
334 batch_size = 1
335 batch = torch.tensor(input_data, device="cpu")
341 - windows = (batch.unfold(3, img_size,
342 - img_size).unfold(4, img_size, img_size))
336 + windows = batch.unfold(3, img_size, img_size).unfold(4, img_size, img_size)
337 h1, w1 = windows.shape[3:5]
344 - windows = rearrange(windows,
345 - "b c t h1 w1 h w -> (b h1 w1) c t h w",
346 - h=img_size,
347 - w=img_size)
338 + windows = rearrange(
339 + windows, "b c t h1 w1 h w -> (b h1 w1) c t h w", h=img_size, w=img_size
340 + )
341
342 # Split into batches if number of windows > batch_size
350 - num_batches = windows.shape[0] // batch_size if windows.shape[
351 - 0] > batch_size else 1
343 + num_batches = windows.shape[0] // batch_size if windows.shape[0] > batch_size else 1
344 windows = torch.tensor_split(windows, num_batches, dim=0)
345
354 - if torch.cuda.is_available():
355 - device = torch.device('cuda')
356 - else:
357 - device = torch.device('cpu')
346 + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
347
348 if temporal_coords:
360 - temporal_coords = torch.tensor(temporal_coords,
361 - device=device).unsqueeze(0)
349 + temporal_coords = torch.tensor(temporal_coords, device=device).unsqueeze(0)
350 else:
351 temporal_coords = None
352 if location_coords:
365 - location_coords = torch.tensor(location_coords[0],
366 - device=device).unsqueeze(0)
353 + location_coords = torch.tensor(location_coords[0], device=device).unsqueeze(0)
354 else:
355 location_coords = None
356
@@ -371,26 +358,24 @@ def run_model(input_data,
358 pred_imgs = []
359 for x in windows:
360 # Apply standardization
374 - x = datamodule.test_transform(
375 - image=x.squeeze().numpy().transpose(1, 2, 0))
376 - x = datamodule.aug(x)['image']
361 + x = datamodule.test_transform(image=x.squeeze().numpy().transpose(1, 2, 0))
362 + x = datamodule.aug(x)["image"]
363
364 with torch.no_grad():
365 x = x.to(device)
366 pred = model.run(x, location_coords=location_coords)
367 if lightning_model:
368 pred_lightning = lightning_model(
383 - x,
384 - temporal_coords=temporal_coords,
385 - location_coords=location_coords)
369 + x, temporal_coords=temporal_coords, location_coords=location_coords
370 + )
371 pred_lightning = pred_lightning.output.detach().cpu()
372 if not torch.equal(pred, pred_lightning):
373 print("Inference output is not equal")
374 y_hat = pred.argmax(dim=1)
375
391 - y_hat = torch.nn.functional.interpolate(y_hat.unsqueeze(1).float(),
392 - size=img_size,
393 - mode="nearest")
376 + y_hat = torch.nn.functional.interpolate(
377 + y_hat.unsqueeze(1).float(), size=img_size, mode="nearest"
378 + )
379
380 pred_imgs.append(y_hat)
381
@@ -437,8 +422,7 @@ def parse_args():
422 default=[1, 2, 3, 8, 11, 12],
423 type=int,
424 nargs="+",
440 - help=
441 - "0-based indices of the six Prithvi channels to be selected from the "
425 + help="0-based indices of the six Prithvi channels to be selected from the "
426 "input. By default selects [1,2,3,8,11,12] for S2L1C data.",
427 )
428 parser.add_argument(
@@ -478,17 +462,18 @@ def main(
462 # Running model ------------------------------------------------------------
463
464 channels = [
481 - datamodule_config['bands'].index(b) for b in ["RED", "GREEN", "BLUE"]
465 + datamodule_config["bands"].index(b) for b in ["RED", "GREEN", "BLUE"]
466 ] # BGR -> RGB
467
484 - pred = run_model(input_data, temporal_coords, location_coords, model_obj,
485 - datamodule, img_size)
468 + pred = run_model(
469 + input_data, temporal_coords, location_coords, model_obj, datamodule, img_size
470 + )
471
472 # Save pred
473 meta_data.update(count=1, dtype="uint8", compress="lzw", nodata=0)
474 pred_file = os.path.join(
490 - output_dir,
491 - f"pred_{os.path.splitext(os.path.basename(data_file))[0]}.tiff")
475 + output_dir, f"pred_{os.path.splitext(os.path.basename(data_file))[0]}.tiff"
476 + )
477 save_geotiff(_convert_np_uint8(pred), pred_file, meta_data)
478
479 # Save image + pred
@@ -502,13 +487,13 @@ def main(
487 channels=channels,
488 )
489
505 - pred[pred == 0.] = np.nan
490 + pred[pred == 0.0] = np.nan
491 img_pred = rgb_orig * 0.7 + pred * 0.3
492 img_pred[img_pred.isnan()] = rgb_orig[img_pred.isnan()]
493
494 img_pred_file = os.path.join(
510 - output_dir,
511 - f"rgb_pred_{os.path.splitext(os.path.basename(data_file))[0]}.tiff")
495 + output_dir, f"rgb_pred_{os.path.splitext(os.path.basename(data_file))[0]}.tiff"
496 + )
497 save_geotiff(
498 image=_convert_np_uint8(img_pred),
499 output_path=img_pred_file,
@@ -518,8 +503,9 @@ def main(
503 # Save image rgb
504 if rgb_outputs:
505 rgb_file = os.path.join(
521 - output_dir, "original_rgb_"
522 - f"{os.path.splitext(os.path.basename(data_file))[0]}.tiff")
506 + output_dir,
507 + f"original_rgb_{os.path.splitext(os.path.basename(data_file))[0]}.tiff",
508 + )
509 save_geotiff(
510 image=_convert_np_uint8(rgb_orig),
511 output_path=rgb_file,
@@ -528,7 +514,6 @@ def main(
514
515
516 if __name__ == "__main__":
531 -
517 args = parse_args()
518
519 main(**vars(args))
examples/offline_inference/profiling.py
+133 -92
@@ -44,14 +44,17 @@ def get_dtype(dtype: str):
44
45
46 OutputLen_NumReqs_Map: TypeAlias = dict[int, int]
47 -def compute_request_output_lengths(batch_size: int, step_requests: list[int]) \
48 - -> OutputLen_NumReqs_Map:
47 +
48 +
49 +def compute_request_output_lengths(
50 + batch_size: int, step_requests: list[int]
51 +) -> OutputLen_NumReqs_Map:
52 """
53 Given the number of requests, batch_size, and the number of requests
54 that each engine-step should process, step_requests, determine the
55 output lengths of the requests such that step_request is honoured.
56
54 - Example:
57 + Example:
58 if batch size = 128 and step_request = [128, 128, 96, 64, 32, 1]
59 then return,
60 {2 : 32, 3 : 32, 4 : 32, 5 : 31, 6 : 1}, meaning,
@@ -100,17 +103,19 @@ def compute_request_output_lengths(batch_size: int, step_requests: list[int]) \
103 output_length -= 1
104
105 # sanity checks.
103 - assert sum(ol_nr.values()) == batch_size, \
104 - ("Number of requests in output-length assignment does not match "
105 - f"batch-size.\n batch size {batch_size} - "
106 - f"step requests {step_requests} - assignments {ol_nr}")
106 + assert sum(ol_nr.values()) == batch_size, (
107 + "Number of requests in output-length assignment does not match "
108 + f"batch-size.\n batch size {batch_size} - "
109 + f"step requests {step_requests} - assignments {ol_nr}"
110 + )
111
112 # Check that the output-length is in [1, num-steps]. Output length must be
113 # at least 1 as all requests must participate in the prefill-step.
110 - assert all(ol >= 1 and ol <= num_steps for ol in ol_nr), \
111 - ("Output lengths of requests should be in range "
112 - f"[1, num-engine-steps].\n batch size {batch_size} - "
113 - f"step requests {step_requests} - assignments {ol_nr}")
114 + assert all(ol >= 1 and ol <= num_steps for ol in ol_nr), (
115 + "Output lengths of requests should be in range "
116 + f"[1, num-engine-steps].\n batch size {batch_size} - "
117 + f"step requests {step_requests} - assignments {ol_nr}"
118 + )
119
120 return ol_nr
121
@@ -131,7 +136,7 @@ def determine_requests_per_step(context: ProfileContext) -> list[int]:
136 context: ProfileContext object.
137
138 Returns:
134 - list[int]: Number of requests to process for all engine-steps.
139 + list[int]: Number of requests to process for all engine-steps.
140 output[i], contains the number of requests that the ith step
141 should process.
142 """
@@ -140,10 +145,13 @@ def determine_requests_per_step(context: ProfileContext) -> list[int]:
145 # that their output lengths must be equal to num_engine_steps.
146 return [context.batch_size] * context.num_steps
147
143 - assert context.complete_num_requests_per_step and \
144 - context.complete_num_requests_per_step > 0, \
145 - (f"Expected a positive complete_num_requests_per_step argument."
146 - f"Instead got {context.complete_num_requests_per_step}")
148 + assert (
149 + context.complete_num_requests_per_step
150 + and context.complete_num_requests_per_step > 0
151 + ), (
152 + f"Expected a positive complete_num_requests_per_step argument."
153 + f"Instead got {context.complete_num_requests_per_step}"
154 + )
155
156 # We start dropping after the first decode step.
157 step_requests = [
@@ -165,8 +173,9 @@ def determine_requests_per_step(context: ProfileContext) -> list[int]:
173 return step_requests
174
175
168 -def run_profile(context: ProfileContext, csv_output: Optional[str],
169 - json_output: Optional[str]):
176 +def run_profile(
177 + context: ProfileContext, csv_output: Optional[str], json_output: Optional[str]
178 +):
179 print("Run profile with:")
180 for key, value in asdict(context).items():
181 print(f" {key} = {value}")
@@ -174,7 +183,8 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
183 requests_per_step: list[int] = determine_requests_per_step(context)
184
185 ol_nr: OutputLen_NumReqs_Map = compute_request_output_lengths(
177 - context.batch_size, requests_per_step)
186 + context.batch_size, requests_per_step
187 + )
188
189 num_steps_to_profile: int = len(requests_per_step)
190 max_output_len: int = max(ol_nr.keys())
@@ -186,7 +196,8 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
196 top_p=0.95,
197 # max_tokens is set on a per-request basis.
198 max_tokens=None,
189 - ignore_eos=True)
199 + ignore_eos=True,
200 + )
201
202 # Create LLM
203 llm = LLM(**asdict(context.engine_args))
@@ -199,31 +210,37 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
210 max_num_seqs = scheduler_config.max_num_seqs
211
212 if batch_size * prompt_len > max_num_batched_tokens:
202 - print(f"ERROR: chosen batch_size * prompt_len "
203 - f"({batch_size} * {prompt_len} = {batch_size * prompt_len}) is "
204 - f"larger than max_num_batched_tokens ({max_num_batched_tokens}) "
205 - f"and therefore cannot be run in a single profile step, please "
206 - f"choose a smaller batch size or prompt length, or increase "
207 - f"--max-num-batched-tokens")
213 + print(
214 + f"ERROR: chosen batch_size * prompt_len "
215 + f"({batch_size} * {prompt_len} = {batch_size * prompt_len}) is "
216 + f"larger than max_num_batched_tokens ({max_num_batched_tokens}) "
217 + f"and therefore cannot be run in a single profile step, please "
218 + f"choose a smaller batch size or prompt length, or increase "
219 + f"--max-num-batched-tokens"
220 + )
221 sys.exit(-1)
222 if batch_size > max_num_seqs:
223 print(
224 f"ERROR: chosen batch_size ({batch_size}) is larger than "
225 f"max_num_seqs ({max_num_seqs}) and therefore cannot be run in a "
213 - f"single profile step, please choose a smaller batch size")
226 + f"single profile step, please choose a smaller batch size"
227 + )
228 sys.exit(-1)
215 - print("llm.llm_engine.model_config.max_model_len: ",
216 - llm.llm_engine.model_config.max_model_len)
229 + print(
230 + "llm.llm_engine.model_config.max_model_len: ",
231 + llm.llm_engine.model_config.max_model_len,
232 + )
233 if prompt_len + max_output_len > llm.llm_engine.model_config.max_model_len:
218 - print(f"ERROR: chosen prompt_len + max_output_len ({prompt_len} + "
219 - f"{max_output_len} = {prompt_len + max_output_len}) is larger "
220 - f"than the model's max_model_len ({max_model_len}), please "
221 - f"choose a smaller prompt_len or max_output_len, or increase "
222 - f"--max-model-len")
234 + print(
235 + f"ERROR: chosen prompt_len + max_output_len ({prompt_len} + "
236 + f"{max_output_len} = {prompt_len + max_output_len}) is larger "
237 + f"than the model's max_model_len ({max_model_len}), please "
238 + f"choose a smaller prompt_len or max_output_len, or increase "
239 + f"--max-model-len"
240 + )
241 sys.exit(-1)
242
243 def add_requests():
226 -
244 def get_output_len_generator() -> Generator[int, Any, Any]:
245 for output_len, num_reqs in ol_nr.items():
246 for _ in range(num_reqs):
@@ -234,13 +251,15 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
251 sampling_params.max_tokens = next(output_len_generator)
252 assert isinstance(sampling_params.max_tokens, int)
253
237 - prompt_token_ids = torch.randint(llm.get_tokenizer().vocab_size,
238 - size=(prompt_len, )).tolist()
254 + prompt_token_ids = torch.randint(
255 + llm.get_tokenizer().vocab_size, size=(prompt_len,)
256 + ).tolist()
257
258 llm.llm_engine.add_request(
259 request_id=f"seq{i}",
242 - prompt={'prompt_token_ids': prompt_token_ids},
243 - params=sampling_params)
260 + prompt={"prompt_token_ids": prompt_token_ids},
261 + params=sampling_params,
262 + )
263
264 def abort_requests():
265 for i in range(batch_size):
@@ -261,10 +280,8 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
280
281 decode_profs = []
282 for _ in tqdm.tqdm(range(num_steps_to_profile - 1)):
264 - num_running_seqs = llm.llm_engine.scheduler[
265 - 0].get_num_unfinished_seq_groups()
266 - with layerwise_profile(
267 - num_running_seqs=num_running_seqs) as decode_prof:
283 + num_running_seqs = llm.llm_engine.scheduler[0].get_num_unfinished_seq_groups()
284 + with layerwise_profile(num_running_seqs=num_running_seqs) as decode_prof:
285 llm.llm_engine.step()
286 decode_profs.append(decode_prof)
287
@@ -274,8 +291,7 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
291
292 LINE_WIDTH = 80
293 print("=" * LINE_WIDTH)
277 - print(f"= Prefill Model Table "
278 - f"(prompt_len={prompt_len}, batch_size={batch_size})")
294 + print(f"= Prefill Model Table (prompt_len={prompt_len}, batch_size={batch_size})")
295 print("=" * LINE_WIDTH)
296 print()
297 prefill_results.print_model_table()
@@ -283,16 +299,17 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
299 if has_decode:
300 print()
301 print("=" * LINE_WIDTH)
286 - print(f"= First Decode Step Model Table "
287 - f"(prompt_len={prompt_len}, batch_size={batch_size})")
302 + print(
303 + f"= First Decode Step Model Table "
304 + f"(prompt_len={prompt_len}, batch_size={batch_size})"
305 + )
306 print("=" * LINE_WIDTH)
307 print()
308 decode_results_list[0].print_model_table()
309
310 print()
311 print("=" * LINE_WIDTH)
294 - print(f"= Prefill Summary Table "
295 - f"(prompt_len={prompt_len}, batch_size={batch_size})")
312 + print(f"= Prefill Summary Table (prompt_len={prompt_len}, batch_size={batch_size})")
313 print("=" * LINE_WIDTH)
314 print()
315 prefill_results.print_summary_table()
@@ -300,25 +317,32 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
317 if has_decode:
318 print()
319 print("=" * LINE_WIDTH)
303 - print(f"= First Decode Step Summary Table "
304 - f"(prompt_len={prompt_len}, batch_size={batch_size})")
320 + print(
321 + f"= First Decode Step Summary Table "
322 + f"(prompt_len={prompt_len}, batch_size={batch_size})"
323 + )
324 print("=" * LINE_WIDTH)
325 print()
326 decode_results_list[0].print_summary_table()
327
328 if csv_output:
310 - csv_filename_base = csv_output[:-4] \
311 - if csv_output.endswith('.csv') else csv_output
329 + csv_filename_base = (
330 + csv_output[:-4] if csv_output.endswith(".csv") else csv_output
331 + )
332 prefill_results.export_model_stats_table_csv(
313 - csv_filename_base + "_prefill_model_table.csv")
333 + csv_filename_base + "_prefill_model_table.csv"
334 + )
335 prefill_results.export_summary_stats_table_csv(
315 - csv_filename_base + "_prefill_summary_table.csv")
336 + csv_filename_base + "_prefill_summary_table.csv"
337 + )
338
339 if has_decode:
318 - decode_results_list[0].export_model_stats_table_csv(\
319 - csv_filename_base + "_decode_model_table.csv")
340 + decode_results_list[0].export_model_stats_table_csv(
341 + csv_filename_base + "_decode_model_table.csv"
342 + )
343 decode_results_list[0].export_summary_stats_table_csv(
321 - csv_filename_base + "_decode_summary_table.csv")
344 + csv_filename_base + "_decode_summary_table.csv"
345 + )
346
347 if json_output:
348 cuda_devices = [
@@ -332,7 +356,7 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
356 "torch_version": f"{torch.__version__}",
357 "torch_cuda_version": f"{torch.version.cuda}",
358 "cuda_devices": f"{cuda_devices}",
335 - **asdict(context)
359 + **asdict(context),
360 },
361 "prefill": prefill_results.convert_stats_to_dict(),
362 }
@@ -342,8 +366,9 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
366 json_dict[f"decode_{idx + 1}"] = dr.convert_stats_to_dict()
367
368 # Add .json to json_output filename if it doesn't exist already.
345 - json_output_file = json_output if json_output.endswith(
346 - '.json') else json_output + '.json'
369 + json_output_file = (
370 + json_output if json_output.endswith(".json") else json_output + ".json"
371 + )
372 with open(json_output_file, "w+") as f:
373 json.dump(json_dict, f, indent=2)
374 pass
@@ -351,16 +376,21 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
376 if context.save_chrome_traces_folder is not None:
377 os.makedirs(context.save_chrome_traces_folder, exist_ok=True)
378 prefill_prof.profiler.export_chrome_trace(
354 - context.save_chrome_traces_folder + "/prefill.json")
379 + context.save_chrome_traces_folder + "/prefill.json"
380 + )
381 for idx, decode_prof in enumerate(decode_profs):
382 decode_prof.profiler.export_chrome_trace(
357 - context.save_chrome_traces_folder + f"/decode_{idx + 1}.json")
358 - print("Traces saved as prefill.json and decode_1.json, etc."
359 - f" in folder {context.save_chrome_traces_folder}")
383 + context.save_chrome_traces_folder + f"/decode_{idx + 1}.json"
384 + )
385 + print(
386 + "Traces saved as prefill.json and decode_1.json, etc."
387 + f" in folder {context.save_chrome_traces_folder}"
388 + )
389
390
391 def parse_args():
363 - parser = FlexibleArgumentParser(description="""
392 + parser = FlexibleArgumentParser(
393 + description="""
394 Profile a model
395
396 example:
@@ -384,7 +414,8 @@ Profile a model
414 --output-directory profile_breakdown --plot-metric pct_cuda_time
415 ```
416 """,
387 - formatter_class=RawTextHelpFormatter)
417 + formatter_class=RawTextHelpFormatter,
418 + )
419 parser.add_argument(
420 "--csv",
421 type=str,
@@ -393,59 +424,68 @@ Profile a model
424 "filename, will create <filename>_prefill_model_table.csv, "
425 "<filename>_prefill_summary_table.csv, "
426 "<filename>_decode_model_table.csv, and "
396 - "<filename>_decode_summary_table.csv")
427 + "<filename>_decode_summary_table.csv",
428 + )
429 parser.add_argument(
430 "--json",
431 type=str,
432 default=None,
401 - help="Export the results as a json file. This should be the filename")
402 - parser.add_argument("--save-chrome-traces-folder",
403 - type=str,
404 - help="Save chrome traces for the prefill and decode "
405 - "will save traces as prefill.json and decode_1.json, "
406 - "etc. inside this folder")
433 + help="Export the results as a json file. This should be the filename",
434 + )
435 + parser.add_argument(
436 + "--save-chrome-traces-folder",
437 + type=str,
438 + help="Save chrome traces for the prefill and decode "
439 + "will save traces as prefill.json and decode_1.json, "
440 + "etc. inside this folder",
441 + )
442 parser.add_argument(
443 "--prompt-len",
444 type=int,
445 default=PROMPT_LEN_DEFAULT,
446 help=f"Length of the random prompt to use when profiling, all batched "
412 - f"requests use the same prompt_len, default={PROMPT_LEN_DEFAULT}")
413 - parser.add_argument("--batch-size",
414 - type=int,
415 - default=BATCH_SIZE_DEFAULT,
416 - help=f"Number of requests to run as a single batch, "
417 - f"default={BATCH_SIZE_DEFAULT}")
447 + f"requests use the same prompt_len, default={PROMPT_LEN_DEFAULT}",
448 + )
449 + parser.add_argument(
450 + "--batch-size",
451 + type=int,
452 + default=BATCH_SIZE_DEFAULT,
453 + help=f"Number of requests to run as a single batch, "
454 + f"default={BATCH_SIZE_DEFAULT}",
455 + )
456
457 subparsers = parser.add_subparsers(dest="cmd")
458
459 run_num_steps_parser = subparsers.add_parser(
422 - "run_num_steps",
423 - help="This variation profiles n engine.step() invocations.")
460 + "run_num_steps", help="This variation profiles n engine.step() invocations."
461 + )
462 run_num_steps_parser.add_argument(
425 - '-n',
426 - '--num-steps',
463 + "-n",
464 + "--num-steps",
465 type=int,
466 help="Number of engine steps to profile.\n"
467 "Setting it to 1, profiles only the prefill step.\n"
468 "Setting it to 2, profiles the prefill and first decode step\n"
469 "Setting it to 3, profiles the prefill, 1st and 2nd decode steps\n"
432 - "and so on ...")
470 + "and so on ...",
471 + )
472
473 run_to_completion_parser = subparsers.add_parser(
474 "run_to_completion",
475 help="This variation profiles all the engine.step() invocations"
437 - "until the engine exhausts all submitted requests.")
476 + "until the engine exhausts all submitted requests.",
477 + )
478 run_to_completion_parser.add_argument(
439 - '-n',
440 - '--complete-num-requests-per-step',
479 + "-n",
480 + "--complete-num-requests-per-step",
481 type=int,
442 - help=
443 - "Complete complete_num_requests_per_step requests every decode step."
482 + help="Complete complete_num_requests_per_step requests every decode step."
483 "For e.g., with batch_size 128 and complete_num_requests_per_step 32,"
484 "the profiler is run for 6 engine steps, with the steps processing, "
485 "128, 128, 96, 64, 32, 1 requests respectively.\n"
486 "Note that we tack-on a one-request step at the end as it is often "
448 - "useful.")
487 + "useful.",
488 + )
489
490 EngineArgs.add_cli_args(parser)
491
@@ -459,7 +499,8 @@ def main(args):
499 k: v
500 for k, v in vars(args).items()
501 if k in inspect.signature(ProfileContext).parameters
462 - })
502 + },
503 + )
504 run_profile(context, csv_output=args.csv, json_output=args.json)
505
506
examples/offline_inference/profiling_tpu/profiling.py
+36 -33
@@ -31,18 +31,16 @@ def main(args: argparse.Namespace):
31 max_tokens=args.output_len,
32 )
33 print(sampling_params)
34 - dummy_prompt_token_ids = np.random.randint(10000,
35 - size=(args.batch_size,
36 - args.input_len))
37 - dummy_prompts: list[PromptType] = [{
38 - "prompt_token_ids": batch
39 - } for batch in dummy_prompt_token_ids.tolist()]
34 + dummy_prompt_token_ids = np.random.randint(
35 + 10000, size=(args.batch_size, args.input_len)
36 + )
37 + dummy_prompts: list[PromptType] = [
38 + {"prompt_token_ids": batch} for batch in dummy_prompt_token_ids.tolist()
39 + ]
40
41 def run_to_completion():
42 start_time = time.perf_counter()
43 - llm.generate(dummy_prompts,
44 - sampling_params=sampling_params,
45 - use_tqdm=False)
43 + llm.generate(dummy_prompts, sampling_params=sampling_params, use_tqdm=False)
44 end_time = time.perf_counter()
45 latency = end_time - start_time
46 return latency
@@ -58,10 +56,9 @@ def main(args: argparse.Namespace):
56 profile_dir = args.profile_result_dir
57 print(f"Profiling (results will be saved to '{profile_dir}')...")
58 # Enable tracing on server
61 - xp.trace_detached("localhost:9012",
62 - profile_dir,
63 - delay_ms=DELAY_MS,
64 - duration_ms=DURATION_MS)
59 + xp.trace_detached(
60 + "localhost:9012", profile_dir, delay_ms=DELAY_MS, duration_ms=DURATION_MS
61 + )
62 if DELAY_MS == 0:
63 time.sleep(1.0)
64 profile_latencies = []
@@ -72,30 +69,36 @@ def main(args: argparse.Namespace):
69 return
70
71
75 -if __name__ == '__main__':
72 +if __name__ == "__main__":
73 parser = FlexibleArgumentParser(
77 - description='Benchmark the latency of processing a single batch of '
78 - 'requests till completion.')
79 - parser.add_argument('--input-len', type=int, default=32)
80 - parser.add_argument('--output-len', type=int, default=128)
81 - parser.add_argument('--batch-size', type=int, default=8)
82 - parser.add_argument('--num-iters-warmup',
83 - type=int,
84 - default=5,
85 - help='Number of iterations to run for warmup.')
86 - parser.add_argument('--num-iters',
87 - type=int,
88 - default=1,
89 - help='Number of iterations to run for profiling.')
74 + description="Benchmark the latency of processing a single batch of "
75 + "requests till completion."
76 + )
77 + parser.add_argument("--input-len", type=int, default=32)
78 + parser.add_argument("--output-len", type=int, default=128)
79 + parser.add_argument("--batch-size", type=int, default=8)
80 + parser.add_argument(
81 + "--num-iters-warmup",
82 + type=int,
83 + default=5,
84 + help="Number of iterations to run for warmup.",
85 + )
86 + parser.add_argument(
87 + "--num-iters",
88 + type=int,
89 + default=1,
90 + help="Number of iterations to run for profiling.",
91 + )
92 parser.add_argument(
91 - '--profile-result-dir',
93 + "--profile-result-dir",
94 type=str,
95 default="profiles",
94 - help=
95 - ('path to save the pytorch profiler output. Can be visualized '
96 - 'with ui.perfetto.dev or Tensorboard '
97 - '(https://cloud.google.com/tpu/docs/pytorch-xla-performance-profiling-tpu-vm).'
98 - ))
96 + help=(
97 + "path to save the pytorch profiler output. Can be visualized "
98 + "with ui.perfetto.dev or Tensorboard "
99 + "(https://cloud.google.com/tpu/docs/pytorch-xla-performance-profiling-tpu-vm)."
100 + ),
101 + )
102
103 parser = EngineArgs.add_cli_args(parser)
104 args = parser.parse_args()
examples/offline_inference/prompt_embed_inference.py
+29 -36
@@ -18,8 +18,7 @@ Run:
18 """
19
20 import torch
21 -from transformers import (AutoModelForCausalLM, AutoTokenizer,
22 - PreTrainedTokenizer)
21 +from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer
22
23 from vllm import LLM
24
@@ -32,27 +31,29 @@ def init_tokenizer_and_llm(model_name: str):
31 return tokenizer, embedding_layer, llm
32
33
35 -def get_prompt_embeds(chat: list[dict[str,
36 - str]], tokenizer: PreTrainedTokenizer,
37 - embedding_layer: torch.nn.Module):
38 - token_ids = tokenizer.apply_chat_template(chat,
39 - add_generation_prompt=True,
40 - return_tensors='pt')
34 +def get_prompt_embeds(
35 + chat: list[dict[str, str]],
36 + tokenizer: PreTrainedTokenizer,
37 + embedding_layer: torch.nn.Module,
38 +):
39 + token_ids = tokenizer.apply_chat_template(
40 + chat, add_generation_prompt=True, return_tensors="pt"
41 + )
42 prompt_embeds = embedding_layer(token_ids).squeeze(0)
43 return prompt_embeds
44
45
45 -def single_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
46 - embedding_layer: torch.nn.Module):
47 - chat = [{
48 - "role": "user",
49 - "content": "Please tell me about the capital of France."
50 - }]
46 +def single_prompt_inference(
47 + llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module
48 +):
49 + chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
50 prompt_embeds = get_prompt_embeds(chat, tokenizer, embedding_layer)
51
53 - outputs = llm.generate({
54 - "prompt_embeds": prompt_embeds,
55 - })
52 + outputs = llm.generate(
53 + {
54 + "prompt_embeds": prompt_embeds,
55 + }
56 + )
57
58 print("\n[Single Inference Output]")
59 print("-" * 30)
@@ -61,34 +62,26 @@ def single_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
62 print("-" * 30)
63
64
64 -def batch_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
65 - embedding_layer: torch.nn.Module):
66 - chats = [[{
67 - "role": "user",
68 - "content": "Please tell me about the capital of France."
69 - }],
70 - [{
71 - "role": "user",
72 - "content": "When is the day longest during the year?"
73 - }],
74 - [{
75 - "role": "user",
76 - "content": "Where is bigger, the moon or the sun?"
77 - }]]
65 +def batch_prompt_inference(
66 + llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module
67 +):
68 + chats = [
69 + [{"role": "user", "content": "Please tell me about the capital of France."}],
70 + [{"role": "user", "content": "When is the day longest during the year?"}],
71 + [{"role": "user", "content": "Where is bigger, the moon or the sun?"}],
72 + ]
73
74 prompt_embeds_list = [
75 get_prompt_embeds(chat, tokenizer, embedding_layer) for chat in chats
76 ]
77
83 - outputs = llm.generate([{
84 - "prompt_embeds": embeds
85 - } for embeds in prompt_embeds_list])
78 + outputs = llm.generate([{"prompt_embeds": embeds} for embeds in prompt_embeds_list])
79
80 print("\n[Batch Inference Outputs]")
81 print("-" * 30)
82 for i, o in enumerate(outputs):
90 - print(f"Q{i+1}: {chats[i][0]['content']}")
91 - print(f"A{i+1}: {o.outputs[0].text}\n")
83 + print(f"Q{i + 1}: {chats[i][0]['content']}")
84 + print(f"A{i + 1}: {o.outputs[0].text}\n")
85 print("-" * 30)
86
87
examples/offline_inference/qwen2_5_omni/only_thinker.py
+68 -59
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 -This example shows how to use vLLM for running offline inference
3 +This example shows how to use vLLM for running offline inference
4 with the correct prompt format on Qwen2.5-Omni (thinker only).
5 """
6
@@ -27,51 +27,55 @@ class QueryResult(NamedTuple):
27 default_system = (
28 "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
29 "Group, capable of perceiving auditory and visual inputs, as well as "
30 - "generating text and speech.")
30 + "generating text and speech."
31 +)
32
33
34 def get_mixed_modalities_query() -> QueryResult:
34 - question = ("What is recited in the audio? "
35 - "What is the content of this image? Why is this video funny?")
36 - prompt = (f"<|im_start|>system\n{default_system}<|im_end|>\n"
37 - "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
38 - "<|vision_bos|><|IMAGE|><|vision_eos|>"
39 - "<|vision_bos|><|VIDEO|><|vision_eos|>"
40 - f"{question}<|im_end|>\n"
41 - f"<|im_start|>assistant\n")
35 + question = (
36 + "What is recited in the audio? "
37 + "What is the content of this image? Why is this video funny?"
38 + )
39 + prompt = (
40 + f"<|im_start|>system\n{default_system}<|im_end|>\n"
41 + "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
42 + "<|vision_bos|><|IMAGE|><|vision_eos|>"
43 + "<|vision_bos|><|VIDEO|><|vision_eos|>"
44 + f"{question}<|im_end|>\n"
45 + f"<|im_start|>assistant\n"
46 + )
47 return QueryResult(
48 inputs={
49 "prompt": prompt,
50 "multi_modal_data": {
46 - "audio":
47 - AudioAsset("mary_had_lamb").audio_and_sample_rate,
48 - "image":
49 - convert_image_mode(
50 - ImageAsset("cherry_blossom").pil_image, "RGB"),
51 - "video":
52 - VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
51 + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
52 + "image": convert_image_mode(
53 + ImageAsset("cherry_blossom").pil_image, "RGB"
54 + ),
55 + "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
56 },
57 },
55 - limit_mm_per_prompt={
56 - "audio": 1,
57 - "image": 1,
58 - "video": 1
59 - },
58 + limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1},
59 )
60
61
62 def get_use_audio_in_video_query() -> QueryResult:
64 - question = ("Describe the content of the video, "
65 - "then convert what the baby say into text.")
66 - prompt = (f"<|im_start|>system\n{default_system}<|im_end|>\n"
67 - "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
68 - f"{question}<|im_end|>\n"
69 - f"<|im_start|>assistant\n")
63 + question = (
64 + "Describe the content of the video, then convert what the baby say into text."
65 + )
66 + prompt = (
67 + f"<|im_start|>system\n{default_system}<|im_end|>\n"
68 + "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
69 + f"{question}<|im_end|>\n"
70 + f"<|im_start|>assistant\n"
71 + )
72 asset = VideoAsset(name="baby_reading", num_frames=16)
73 audio = asset.get_audio(sampling_rate=16000)
72 - assert not envs.VLLM_USE_V1, ("V1 does not support use_audio_in_video. "
73 - "Please launch this example with "
74 - "`VLLM_USE_V1=0`.")
74 + assert not envs.VLLM_USE_V1, (
75 + "V1 does not support use_audio_in_video. "
76 + "Please launch this example with "
77 + "`VLLM_USE_V1=0`."
78 + )
79 return QueryResult(
80 inputs={
81 "prompt": prompt,
@@ -83,20 +87,19 @@ def get_use_audio_in_video_query() -> QueryResult:
87 "use_audio_in_video": True,
88 },
89 },
86 - limit_mm_per_prompt={
87 - "audio": 1,
88 - "video": 1
89 - },
90 + limit_mm_per_prompt={"audio": 1, "video": 1},
91 )
92
93
94 def get_multi_audios_query() -> QueryResult:
95 question = "Are these two audio clips the same?"
95 - prompt = (f"<|im_start|>system\n{default_system}<|im_end|>\n"
96 - "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
97 - "<|audio_bos|><|AUDIO|><|audio_eos|>"
98 - f"{question}<|im_end|>\n"
99 - f"<|im_start|>assistant\n")
96 + prompt = (
97 + f"<|im_start|>system\n{default_system}<|im_end|>\n"
98 + "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
99 + "<|audio_bos|><|AUDIO|><|audio_eos|>"
100 + f"{question}<|im_end|>\n"
101 + f"<|im_start|>assistant\n"
102 + )
103 return QueryResult(
104 inputs={
105 "prompt": prompt,
@@ -124,18 +127,19 @@ def main(args):
127 model_name = "Qwen/Qwen2.5-Omni-7B"
128 query_result = query_map[args.query_type]()
129
127 - llm = LLM(model=model_name,
128 - max_model_len=5632,
129 - max_num_seqs=5,
130 - limit_mm_per_prompt=query_result.limit_mm_per_prompt,
131 - seed=args.seed)
130 + llm = LLM(
131 + model=model_name,
132 + max_model_len=5632,
133 + max_num_seqs=5,
134 + limit_mm_per_prompt=query_result.limit_mm_per_prompt,
135 + seed=args.seed,
136 + )
137
138 # We set temperature to 0.2 so that outputs can be different
139 # even when all prompts are identical when running batch inference.
140 sampling_params = SamplingParams(temperature=0.2, max_tokens=64)
141
137 - outputs = llm.generate(query_result.inputs,
138 - sampling_params=sampling_params)
142 + outputs = llm.generate(query_result.inputs, sampling_params=sampling_params)
143
144 for o in outputs:
145 generated_text = o.outputs[0].text
@@ -144,18 +148,23 @@ def main(args):
148
149 def parse_args():
150 parser = FlexibleArgumentParser(
147 - description='Demo on using vLLM for offline inference with '
148 - 'audio language models')
149 - parser.add_argument('--query-type',
150 - '-q',
151 - type=str,
152 - default="mixed_modalities",
153 - choices=query_map.keys(),
154 - help='Query type.')
155 - parser.add_argument("--seed",
156 - type=int,
157 - default=None,
158 - help="Set the seed when initializing `vllm.LLM`.")
151 + description="Demo on using vLLM for offline inference with "
152 + "audio language models"
153 + )
154 + parser.add_argument(
155 + "--query-type",
156 + "-q",
157 + type=str,
158 + default="mixed_modalities",
159 + choices=query_map.keys(),
160 + help="Query type.",
161 + )
162 + parser.add_argument(
163 + "--seed",
164 + type=int,
165 + default=None,
166 + help="Set the seed when initializing `vllm.LLM`.",
167 + )
168
169 return parser.parse_args()
170
examples/offline_inference/qwen_1m.py
+17 -13
@@ -17,10 +17,10 @@ def load_prompt() -> str:
17 # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt
18
19 with urlopen(
20 - "https://qianwen-res.oss-cn-beijing.aliyuncs.com"
21 - "/Qwen2.5-1M/test-data/600k.txt",
22 - timeout=5) as response:
23 - prompt = response.read().decode('utf-8')
20 + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/600k.txt",
21 + timeout=5,
22 + ) as response:
23 + prompt = response.read().decode("utf-8")
24 return prompt
25
26
@@ -41,18 +41,22 @@ def process_requests(llm: LLM, prompts: list[str]) -> None:
41 for output in outputs:
42 prompt_token_ids = output.prompt_token_ids
43 generated_text = output.outputs[0].text
44 - print(f"Prompt length: {len(prompt_token_ids)}, "
45 - f"Generated text: {generated_text!r}")
44 + print(
45 + f"Prompt length: {len(prompt_token_ids)}, "
46 + f"Generated text: {generated_text!r}"
47 + )
48
49
50 # Create an LLM.
51 def initialize_engine() -> LLM:
50 - llm = LLM(model="Qwen/Qwen2.5-7B-Instruct-1M",
51 - max_model_len=1048576,
52 - tensor_parallel_size=4,
53 - enforce_eager=True,
54 - enable_chunked_prefill=True,
55 - max_num_batched_tokens=131072)
52 + llm = LLM(
53 + model="Qwen/Qwen2.5-7B-Instruct-1M",
54 + max_model_len=1048576,
55 + tensor_parallel_size=4,
56 + enforce_eager=True,
57 + enable_chunked_prefill=True,
58 + max_num_batched_tokens=131072,
59 + )
60 return llm
61
62
@@ -62,5 +66,5 @@ def main():
66 process_requests(llm, [prompt])
67
68
65 -if __name__ == '__main__':
69 +if __name__ == "__main__":
70 main()
examples/offline_inference/rlhf.py
+10 -11
@@ -12,6 +12,7 @@ inference instance. In practice, there could be multiple training instances
12 and multiple inference instances. For the full implementation, please refer
13 to the OpenRLHF framework.
14 """
15 +
16 import os
17
18 import ray
@@ -26,7 +27,6 @@ from vllm.utils import get_ip, get_open_port
27
28
29 class MyLLM(LLM):
29 -
30 def __init__(self, *args, **kwargs):
31 # a hack to make the script work.
32 # stop ray from manipulating CUDA_VISIBLE_DEVICES
@@ -89,8 +89,7 @@ print("-" * 50)
89 for output in outputs:
90 prompt = output.prompt
91 generated_text = output.outputs[0].text
92 - print(f"Prompt: {prompt!r}\n"
93 - f"Generated text: {generated_text!r}")
92 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
93 print("-" * 50)
94
95 # set up the communication between the training process
@@ -98,11 +97,13 @@ for output in outputs:
97 master_address = get_ip()
98 master_port = get_open_port()
99
101 -handle = llm.collective_rpc.remote("init_weight_update_group",
102 - args=(master_address, master_port, 1, 3))
100 +handle = llm.collective_rpc.remote(
101 + "init_weight_update_group", args=(master_address, master_port, 1, 3)
102 +)
103
104 -model_update_group = stateless_init_process_group(master_address, master_port,
105 - 0, 3, torch.device("cuda:0"))
104 +model_update_group = stateless_init_process_group(
105 + master_address, master_port, 0, 3, torch.device("cuda:0")
106 +)
107 ray.get(handle)
108
109 # simulate training, modify the weights of the model.
@@ -111,8 +112,7 @@ for name, p in train_model.named_parameters():
112
113 # sync weight from the training process to the inference engine.
114 for name, p in train_model.named_parameters():
114 - handle = llm.collective_rpc.remote("update_weight",
115 - args=(name, p.dtype, p.shape))
115 + handle = llm.collective_rpc.remote("update_weight", args=(name, p.dtype, p.shape))
116 model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
117 ray.get(handle)
118
@@ -126,6 +126,5 @@ print("-" * 50)
126 for output in outputs_updated:
127 prompt = output.prompt
128 generated_text = output.outputs[0].text
129 - print(f"Prompt: {prompt!r}\n"
130 - f"Generated text: {generated_text!r}")
129 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
130 print("-" * 50)
examples/offline_inference/rlhf_colocate.py
+13 -10
@@ -9,6 +9,7 @@ The key points:
9 - Use cuda-ipc to pass tensors, since NCCL does not work when we have
10 multiple processes on the same GPU.
11 """
12 +
13 import os
14
15 import ray
@@ -20,7 +21,6 @@ from vllm import LLM
21
22
23 class MyLLM(LLM):
23 -
24 def __init__(self, *args, bundle_indices: list, **kwargs):
25 # a hack to make the script work.
26 # stop ray from manipulating CUDA_VISIBLE_DEVICES
@@ -29,17 +29,16 @@ class MyLLM(LLM):
29 # every worker will use 0.4 GPU, so that we can schedule
30 # 2 instances on the same GPUs.
31 os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
32 - os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(
33 - map(str, bundle_indices))
32 + os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(map(str, bundle_indices))
33 print(f"creating LLM with bundle_indices={bundle_indices}")
34 super().__init__(*args, **kwargs)
35
36
37 class RayTrainingActor:
39 -
38 def __init__(self):
39 # ray will set CUDA_VISIBLE_DEVICES to the assigned GPUs
40 from transformers import AutoModelForCausalLM
41 +
42 self.model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
43 self.model.to("cuda:0")
44 for name, p in self.model.named_parameters():
@@ -48,6 +47,7 @@ class RayTrainingActor:
47 # the argument for get_device_uuid is the index
48 # of the GPU in the visible devices.
49 from vllm.platforms import current_platform
50 +
51 self.device_uuid = current_platform.get_device_uuid(0)
52
53 def report_device_id(self) -> str:
@@ -55,6 +55,7 @@ class RayTrainingActor:
55
56 def get_weight_ipc_handles(self):
57 from torch.multiprocessing.reductions import reduce_tensor
58 +
59 data = {}
60 for name, p in self.model.named_parameters():
61 # the training actor might only have a subset of the weights
@@ -101,7 +102,7 @@ for bundle_index, training_actor in enumerate(training_actors):
102 print(f"training actor {bundle_index} is on {device_id}")
103 training_actor_device_ids.append(device_id)
104
104 -for (i, bundle_indices) in enumerate([[0, 1], [2, 3]]):
105 +for i, bundle_indices in enumerate([[0, 1], [2, 3]]):
106 # IMPORTANT: when creating vLLM instances, we need to
107 # make sure there are no GPU activities on the target GPUs,
108 # otherwise, they will interfere with the vLLM memory profiling,
@@ -128,7 +129,8 @@ for (i, bundle_indices) in enumerate([[0, 1], [2, 3]]):
129
130 for i, llm in enumerate(inference_engines):
131 inference_engine_device_ids.append(
131 - ray.get(llm.collective_rpc.remote("report_device_id", args=tuple())))
132 + ray.get(llm.collective_rpc.remote("report_device_id", args=tuple()))
133 + )
134 print(f"inference engine {i} is on {inference_engine_device_ids[-1]}")
135
136 # check the placement
@@ -147,9 +149,10 @@ for actor in training_actors:
149 print("update the weights of the inference engines")
150 for llm in inference_engines:
151 ray.get(
150 - llm.collective_rpc.remote("update_weights_from_ipc_handles",
151 - args=(ipc_handles, )))
152 + llm.collective_rpc.remote(
153 + "update_weights_from_ipc_handles", args=(ipc_handles,)
154 + )
155 + )
156 print("check if the weights are updated")
157 for llm in inference_engines:
154 - assert ray.get(
155 - llm.collective_rpc.remote("check_weights_changed", args=tuple()))
158 + assert ray.get(llm.collective_rpc.remote("check_weights_changed", args=tuple()))
examples/offline_inference/rlhf_utils.py
+16 -16
@@ -2,21 +2,20 @@
2 import torch
3
4
5 -def stateless_init_process_group(master_address, master_port, rank, world_size,
6 - device):
5 +def stateless_init_process_group(master_address, master_port, rank, world_size, device):
6 """
7 vLLM provides `StatelessProcessGroup` to create a process group
8 without considering the global process group in torch.distributed.
9 It is recommended to create `StatelessProcessGroup`, and then initialize
11 - the data-plane communication (NCCL) between external (train processes)
10 + the data-plane communication (NCCL) between external (train processes)
11 and vLLM workers.
12 """
13 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
14 from vllm.distributed.utils import StatelessProcessGroup
16 - pg = StatelessProcessGroup.create(host=master_address,
17 - port=master_port,
18 - rank=rank,
19 - world_size=world_size)
15 +
16 + pg = StatelessProcessGroup.create(
17 + host=master_address, port=master_port, rank=rank, world_size=world_size
18 + )
19 pynccl = PyNcclCommunicator(pg, device=device)
20 return pynccl
21
@@ -31,9 +30,11 @@ class WorkerExtension:
30 should pass the full qualified name as `worker_extension_cls` argument.
31 """
32
34 - def init_weight_update_group(self, master_address, master_port,
35 - rank_offset, world_size):
33 + def init_weight_update_group(
34 + self, master_address, master_port, rank_offset, world_size
35 + ):
36 from vllm.distributed.parallel_state import get_world_group
37 +
38 rank = get_world_group().rank + rank_offset
39 self.model_update_group = stateless_init_process_group(
40 master_address,
@@ -45,9 +46,9 @@ class WorkerExtension:
46
47 def update_weight(self, name, dtype, shape):
48 weight = torch.empty(shape, dtype=dtype, device="cuda")
48 - self.model_update_group.broadcast(weight,
49 - src=0,
50 - stream=torch.cuda.current_stream())
49 + self.model_update_group.broadcast(
50 + weight, src=0, stream=torch.cuda.current_stream()
51 + )
52
53 self.model_runner.model.load_weights(weights=[(name, weight)])
54
@@ -59,8 +60,7 @@ class WorkerExtension:
60 """
61 weights_updated = True
62 for name, p in self.model_runner.model.named_parameters():
62 - weights_updated = weights_updated and torch.allclose(
63 - p, torch.zeros_like(p))
63 + weights_updated = weights_updated and torch.allclose(p, torch.zeros_like(p))
64 return weights_updated
65
66
@@ -76,6 +76,7 @@ class ColocateWorkerExtension:
76
77 def report_device_id(self) -> str:
78 from vllm.platforms import current_platform
79 +
80 self.device_uuid = current_platform.get_device_uuid(self.device.index)
81 return self.device_uuid
82
@@ -100,6 +101,5 @@ class ColocateWorkerExtension:
101 """
102 weights_updated = True
103 for name, p in self.model_runner.model.named_parameters():
103 - weights_updated = weights_updated and torch.allclose(
104 - p, torch.zeros_like(p))
104 + weights_updated = weights_updated and torch.allclose(p, torch.zeros_like(p))
105 return weights_updated
examples/offline_inference/save_sharded_state.py
+21 -20
@@ -21,6 +21,7 @@ llm = LLM(
21 tensor_parallel_size=8,
22 )
23 """
24 +
25 import dataclasses
26 import os
27 import shutil
@@ -33,18 +34,18 @@ from vllm.utils import FlexibleArgumentParser
34 def parse_args():
35 parser = FlexibleArgumentParser()
36 EngineArgs.add_cli_args(parser)
36 - parser.add_argument("--output",
37 - "-o",
38 - required=True,
39 - type=str,
40 - help="path to output checkpoint")
41 - parser.add_argument("--file-pattern",
42 - type=str,
43 - help="string pattern of saved filenames")
44 - parser.add_argument("--max-file-size",
45 - type=str,
46 - default=5 * 1024**3,
47 - help="max size (in bytes) of each safetensors file")
37 + parser.add_argument(
38 + "--output", "-o", required=True, type=str, help="path to output checkpoint"
39 + )
40 + parser.add_argument(
41 + "--file-pattern", type=str, help="string pattern of saved filenames"
42 + )
43 + parser.add_argument(
44 + "--max-file-size",
45 + type=str,
46 + default=5 * 1024**3,
47 + help="max size (in bytes) of each safetensors file",
48 + )
49 return parser.parse_args()
50
51
@@ -68,23 +69,23 @@ def main(args):
69 # For V1 engine, we need to use engine_core.save_sharded_state
70 print("Using V1 engine save path")
71 llm.llm_engine.engine_core.save_sharded_state(
71 - path=args.output,
72 - pattern=args.file_pattern,
73 - max_size=args.max_file_size)
72 + path=args.output, pattern=args.file_pattern, max_size=args.max_file_size
73 + )
74 else:
75 # For V0 engine
76 print("Using V0 engine save path")
77 model_executor = llm.llm_engine.model_executor
78 - model_executor.save_sharded_state(path=args.output,
79 - pattern=args.file_pattern,
80 - max_size=args.max_file_size)
78 + model_executor.save_sharded_state(
79 + path=args.output, pattern=args.file_pattern, max_size=args.max_file_size
80 + )
81
82 # Copy metadata files to output directory
83 for file in os.listdir(model_path):
84 if os.path.splitext(file)[1] not in (".bin", ".pt", ".safetensors"):
85 if os.path.isdir(os.path.join(model_path, file)):
86 - shutil.copytree(os.path.join(model_path, file),
87 - os.path.join(args.output, file))
86 + shutil.copytree(
87 + os.path.join(model_path, file), os.path.join(args.output, file)
88 + )
89 else:
90 shutil.copy(os.path.join(model_path, file), args.output)
91
examples/offline_inference/structured_outputs.py
+21 -22
@@ -1,9 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 -This file demonstrates the example usage of guided decoding
4 -to generate structured outputs using vLLM. It shows how to apply
5 -different guided decoding techniques such as Choice, Regex, JSON schema,
6 -and Grammar to produce structured and formatted results
3 +This file demonstrates the example usage of guided decoding
4 +to generate structured outputs using vLLM. It shows how to apply
5 +different guided decoding techniques such as Choice, Regex, JSON schema,
6 +and Grammar to produce structured and formatted results
7 based on specific prompts.
8 """
9
@@ -15,20 +15,20 @@ from vllm import LLM, SamplingParams
15 from vllm.sampling_params import GuidedDecodingParams
16
17 # Guided decoding by Choice (list of possible options)
18 -guided_decoding_params_choice = GuidedDecodingParams(
19 - choice=["Positive", "Negative"])
20 -sampling_params_choice = SamplingParams(
21 - guided_decoding=guided_decoding_params_choice)
18 +guided_decoding_params_choice = GuidedDecodingParams(choice=["Positive", "Negative"])
19 +sampling_params_choice = SamplingParams(guided_decoding=guided_decoding_params_choice)
20 prompt_choice = "Classify this sentiment: vLLM is wonderful!"
21
22 # Guided decoding by Regex
23 guided_decoding_params_regex = GuidedDecodingParams(regex=r"\w+@\w+\.com\n")
24 sampling_params_regex = SamplingParams(
27 - guided_decoding=guided_decoding_params_regex, stop=["\n"])
25 + guided_decoding=guided_decoding_params_regex, stop=["\n"]
26 +)
27 prompt_regex = (
28 "Generate an email address for Alan Turing, who works in Enigma."
29 "End in .com and new line. Example result:"
31 - "alan.turing@enigma.com\n")
30 + "alan.turing@enigma.com\n"
31 +)
32
33
34 # Guided decoding by JSON using Pydantic schema
@@ -47,10 +47,11 @@ class CarDescription(BaseModel):
47
48 json_schema = CarDescription.model_json_schema()
49 guided_decoding_params_json = GuidedDecodingParams(json=json_schema)
50 -sampling_params_json = SamplingParams(
51 - guided_decoding=guided_decoding_params_json)
52 -prompt_json = ("Generate a JSON with the brand, model and car_type of"
53 - "the most iconic car from the 90's")
50 +sampling_params_json = SamplingParams(guided_decoding=guided_decoding_params_json)
51 +prompt_json = (
52 + "Generate a JSON with the brand, model and car_type of"
53 + "the most iconic car from the 90's"
54 +)
55
56 # Guided decoding by Grammar
57 simplified_sql_grammar = """
@@ -61,12 +62,11 @@ table ::= "table_1 " | "table_2 "
62 condition ::= column "= " number
63 number ::= "1 " | "2 "
64 """
64 -guided_decoding_params_grammar = GuidedDecodingParams(
65 - grammar=simplified_sql_grammar)
66 -sampling_params_grammar = SamplingParams(
67 - guided_decoding=guided_decoding_params_grammar)
68 -prompt_grammar = ("Generate an SQL query to show the 'username' and 'email'"
69 - "from the 'users' table.")
65 +guided_decoding_params_grammar = GuidedDecodingParams(grammar=simplified_sql_grammar)
66 +sampling_params_grammar = SamplingParams(guided_decoding=guided_decoding_params_grammar)
67 +prompt_grammar = (
68 + "Generate an SQL query to show the 'username' and 'email'from the 'users' table."
69 +)
70
71
72 def format_output(title: str, output: str):
@@ -90,8 +90,7 @@ def main():
90 json_output = generate_output(prompt_json, sampling_params_json, llm)
91 format_output("Guided decoding by JSON", json_output)
92
93 - grammar_output = generate_output(prompt_grammar, sampling_params_grammar,
94 - llm)
93 + grammar_output = generate_output(prompt_grammar, sampling_params_grammar, llm)
94 format_output("Guided decoding by Grammar", grammar_output)
95
96
examples/offline_inference/torchrun_example.py
+1 -2
@@ -45,8 +45,7 @@ if dist.get_rank() == 0:
45 for output in outputs:
46 prompt = output.prompt
47 generated_text = output.outputs[0].text
48 - print(f"Prompt: {prompt!r}\n"
49 - f"Generated text: {generated_text!r}\n")
48 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}\n")
49 print("-" * 50)
50 """
51 Further tips:
examples/offline_inference/tpu.py
+6 -4
@@ -20,10 +20,12 @@ sampling_params = SamplingParams(temperature=0, top_p=1.0, n=N, max_tokens=16)
20 def main():
21 # Set `enforce_eager=True` to avoid ahead-of-time compilation.
22 # In real workloads, `enforace_eager` should be `False`.
23 - llm = LLM(model="Qwen/Qwen2-1.5B-Instruct",
24 - max_num_batched_tokens=64,
25 - max_num_seqs=4,
26 - max_model_len=128)
23 + llm = LLM(
24 + model="Qwen/Qwen2-1.5B-Instruct",
25 + max_num_batched_tokens=64,
26 + max_num_seqs=4,
27 + max_model_len=128,
28 + )
29 outputs = llm.generate(prompts, sampling_params)
30 print("-" * 50)
31 for output, answer in zip(outputs, answers):
examples/offline_inference/vision_language.py
+223 -210
@@ -6,6 +6,7 @@ the correct prompt format on vision language models for text generation.
6 For most models, the prompt format should follow corresponding examples
7 on HuggingFace model repository.
8 """
9 +
10 import os
11 import random
12 from contextlib import contextmanager
@@ -49,9 +50,13 @@ def run_aria(questions: list[str], modality: str) -> ModelRequestData:
50 limit_mm_per_prompt={modality: 1},
51 )
52
52 - prompts = [(f"<|im_start|>user\n<fim_prefix><|img|><fim_suffix>{question}"
53 - "<|im_end|>\n<|im_start|>assistant\n")
54 - for question in questions]
53 + prompts = [
54 + (
55 + f"<|im_start|>user\n<fim_prefix><|img|><fim_suffix>{question}"
56 + "<|im_end|>\n<|im_start|>assistant\n"
57 + )
58 + for question in questions
59 + ]
60
61 stop_token_ids = [93532, 93653, 944, 93421, 1019, 93653, 93519]
62
@@ -135,8 +140,7 @@ def run_deepseek_vl2(questions: list[str], modality: str) -> ModelRequestData:
140 )
141
142 prompts = [
138 - f"<|User|>: <image>\n{question}\n\n<|Assistant|>:"
139 - for question in questions
143 + f"<|User|>: <image>\n{question}\n\n<|Assistant|>:" for question in questions
144 ]
145
146 return ModelRequestData(
@@ -198,9 +202,14 @@ def run_gemma3(questions: list[str], modality: str) -> ModelRequestData:
202 limit_mm_per_prompt={modality: 1},
203 )
204
201 - prompts = [("<bos><start_of_turn>user\n"
202 - f"<start_of_image>{question}<end_of_turn>\n"
203 - "<start_of_turn>model\n") for question in questions]
205 + prompts = [
206 + (
207 + "<bos><start_of_turn>user\n"
208 + f"<start_of_image>{question}<end_of_turn>\n"
209 + "<start_of_turn>model\n"
210 + )
211 + for question in questions
212 + ]
213
214 return ModelRequestData(
215 engine_args=engine_args,
@@ -225,7 +234,8 @@ def run_glm4v(questions: list[str], modality: str) -> ModelRequestData:
234
235 prompts = [
236 f"<|user|>\n<|begin_of_image|><|endoftext|><|end_of_image|>\
228 - {question}<|assistant|>" for question in questions
237 + {question}<|assistant|>"
238 + for question in questions
239 ]
240
241 stop_token_ids = [151329, 151336, 151338]
@@ -250,15 +260,13 @@ def run_h2ovl(questions: list[str], modality: str) -> ModelRequestData:
260 limit_mm_per_prompt={modality: 1},
261 )
262
253 - tokenizer = AutoTokenizer.from_pretrained(model_name,
254 - trust_remote_code=True)
255 - messages = [[{
256 - 'role': 'user',
257 - 'content': f"<image>\n{question}"
258 - }] for question in questions]
259 - prompts = tokenizer.apply_chat_template(messages,
260 - tokenize=False,
261 - add_generation_prompt=True)
263 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
264 + messages = [
265 + [{"role": "user", "content": f"<image>\n{question}"}] for question in questions
266 + ]
267 + prompts = tokenizer.apply_chat_template(
268 + messages, tokenize=False, add_generation_prompt=True
269 + )
270
271 # Stop tokens for H2OVL-Mississippi
272 # https://huggingface.co/h2oai/h2ovl-mississippi-800m
@@ -284,15 +292,14 @@ def run_idefics3(questions: list[str], modality: str) -> ModelRequestData:
292 # if you are running out of memory, you can reduce the "longest_edge".
293 # see: https://huggingface.co/HuggingFaceM4/Idefics3-8B-Llama3#model-optimizations
294 mm_processor_kwargs={
287 - "size": {
288 - "longest_edge": 3 * 364
289 - },
295 + "size": {"longest_edge": 3 * 364},
296 },
297 limit_mm_per_prompt={modality: 1},
298 )
293 - prompts = [(
294 - f"<|begin_of_text|>User:<image>{question}<end_of_utterance>\nAssistant:"
295 - ) for question in questions]
299 + prompts = [
300 + (f"<|begin_of_text|>User:<image>{question}<end_of_utterance>\nAssistant:")
301 + for question in questions
302 + ]
303
304 return ModelRequestData(
305 engine_args=engine_args,
@@ -311,9 +318,7 @@ def run_smolvlm(questions: list[str], modality: str) -> ModelRequestData:
318 max_num_seqs=2,
319 enforce_eager=True,
320 mm_processor_kwargs={
314 - "max_image_size": {
315 - "longest_edge": 384
316 - },
321 + "max_image_size": {"longest_edge": 384},
322 },
323 limit_mm_per_prompt={modality: 1},
324 )
@@ -330,7 +335,6 @@ def run_smolvlm(questions: list[str], modality: str) -> ModelRequestData:
335
336 # InternVL
337 def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
333 -
338 model_name = "OpenGVLab/InternVL3-2B"
339
340 engine_args = EngineArgs(
@@ -345,15 +349,14 @@ def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
349 elif modality == "video":
350 placeholder = "<video>"
351
348 - tokenizer = AutoTokenizer.from_pretrained(model_name,
349 - trust_remote_code=True)
350 - messages = [[{
351 - 'role': 'user',
352 - 'content': f"{placeholder}\n{question}"
353 - }] for question in questions]
354 - prompts = tokenizer.apply_chat_template(messages,
355 - tokenize=False,
356 - add_generation_prompt=True)
352 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
353 + messages = [
354 + [{"role": "user", "content": f"{placeholder}\n{question}"}]
355 + for question in questions
356 + ]
357 + prompts = tokenizer.apply_chat_template(
358 + messages, tokenize=False, add_generation_prompt=True
359 + )
360
361 # Stop tokens for InternVL
362 # models variants may have different stop tokens
@@ -361,9 +364,7 @@ def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
364 # https://huggingface.co/OpenGVLab/InternVL2-2B/blob/main/conversation.py
365 stop_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|end|>"]
366 stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
364 - stop_token_ids = [
365 - token_id for token_id in stop_token_ids if token_id is not None
366 - ]
367 + stop_token_ids = [token_id for token_id in stop_token_ids if token_id is not None]
368
369 return ModelRequestData(
370 engine_args=engine_args,
@@ -379,7 +380,8 @@ def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData:
380 prompts = [
381 "<|im_user|>user<|im_middle|><|media_start|>image<|media_content|>"
382 f"<|media_pad|><|media_end|>{question}<|im_end|>"
382 - "<|im_assistant|>assistant<|im_middle|>" for question in questions
383 + "<|im_assistant|>assistant<|im_middle|>"
384 + for question in questions
385 ]
386
387 engine_args = EngineArgs(
@@ -399,9 +401,7 @@ def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData:
401 def run_llava(questions: list[str], modality: str) -> ModelRequestData:
402 assert modality == "image"
403
402 - prompts = [
403 - f"USER: <image>\n{question}\nASSISTANT:" for question in questions
404 - ]
404 + prompts = [f"USER: <image>\n{question}\nASSISTANT:" for question in questions]
405
406 engine_args = EngineArgs(
407 model="llava-hf/llava-1.5-7b-hf",
@@ -434,13 +434,10 @@ def run_llava_next(questions: list[str], modality: str) -> ModelRequestData:
434
435 # LlaVA-NeXT-Video
436 # Currently only support for video input
437 -def run_llava_next_video(questions: list[str],
438 - modality: str) -> ModelRequestData:
437 +def run_llava_next_video(questions: list[str], modality: str) -> ModelRequestData:
438 assert modality == "video"
439
441 - prompts = [
442 - f"USER: <video>\n{question} ASSISTANT:" for question in questions
443 - ]
440 + prompts = [f"USER: <video>\n{question} ASSISTANT:" for question in questions]
441 engine_args = EngineArgs(
442 model="llava-hf/LLaVA-NeXT-Video-7B-hf",
443 max_model_len=8192,
@@ -455,19 +452,19 @@ def run_llava_next_video(questions: list[str],
452
453
454 # LLaVA-OneVision
458 -def run_llava_onevision(questions: list[str],
459 - modality: str) -> ModelRequestData:
460 -
455 +def run_llava_onevision(questions: list[str], modality: str) -> ModelRequestData:
456 if modality == "video":
457 prompts = [
458 f"<|im_start|>user <video>\n{question}<|im_end|> \
464 - <|im_start|>assistant\n" for question in questions
459 + <|im_start|>assistant\n"
460 + for question in questions
461 ]
462
463 elif modality == "image":
464 prompts = [
465 f"<|im_start|>user <image>\n{question}<|im_end|> \
470 - <|im_start|>assistant\n" for question in questions
466 + <|im_start|>assistant\n"
467 + for question in questions
468 ]
469
470 engine_args = EngineArgs(
@@ -486,11 +483,8 @@ def run_llava_onevision(questions: list[str],
483 def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
484 assert modality == "image"
485
489 - llama3_template = '<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n' # noqa: E501
490 - prompts = [
491 - llama3_template.format(f"{question}\n<image>")
492 - for question in questions
493 - ]
486 + llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" # noqa: E501
487 + prompts = [llama3_template.format(f"{question}\n<image>") for question in questions]
488
489 engine_args = EngineArgs(
490 model="TIGER-Lab/Mantis-8B-siglip-llama3",
@@ -530,8 +524,7 @@ def run_minicpmv_base(questions: list[str], modality: str, model_name):
524 # 2.6: image, video
525 # o2.6: image, video, audio
526 # model_name = "openbmb/MiniCPM-o-2_6"
533 - tokenizer = AutoTokenizer.from_pretrained(model_name,
534 - trust_remote_code=True)
527 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
528 engine_args = EngineArgs(
529 model=model_name,
530 max_model_len=4096,
@@ -547,7 +540,7 @@ def run_minicpmv_base(questions: list[str], modality: str, model_name):
540 # stop_token_ids = [tokenizer.eos_id, tokenizer.eot_id]
541
542 # 2.6 / o2.6
550 - stop_tokens = ['<|im_end|>', '<|endoftext|>']
543 + stop_tokens = ["<|im_end|>", "<|endoftext|>"]
544 stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
545
546 modality_placeholder = {
@@ -557,12 +550,16 @@ def run_minicpmv_base(questions: list[str], modality: str, model_name):
550
551 prompts = [
552 tokenizer.apply_chat_template(
560 - [{
561 - 'role': 'user',
562 - 'content': f"{modality_placeholder[modality]}\n{question}"
563 - }],
553 + [
554 + {
555 + "role": "user",
556 + "content": f"{modality_placeholder[modality]}\n{question}",
557 + }
558 + ],
559 tokenize=False,
565 - add_generation_prompt=True) for question in questions
560 + add_generation_prompt=True,
561 + )
562 + for question in questions
563 ]
564
565 return ModelRequestData(
@@ -622,19 +619,18 @@ def run_mllama(questions: list[str], modality: str) -> ModelRequestData:
619 )
620
621 tokenizer = AutoTokenizer.from_pretrained(model_name)
625 - messages = [[{
626 - "role":
627 - "user",
628 - "content": [{
629 - "type": "image"
630 - }, {
631 - "type": "text",
632 - "text": question
633 - }]
634 - }] for question in questions]
635 - prompts = tokenizer.apply_chat_template(messages,
636 - add_generation_prompt=True,
637 - tokenize=False)
622 + messages = [
623 + [
624 + {
625 + "role": "user",
626 + "content": [{"type": "image"}, {"type": "text", "text": question}],
627 + }
628 + ]
629 + for question in questions
630 + ]
631 + prompts = tokenizer.apply_chat_template(
632 + messages, add_generation_prompt=True, tokenize=False
633 + )
634
635 return ModelRequestData(
636 engine_args=engine_args,
@@ -657,19 +653,18 @@ def run_llama4(questions: list[str], modality: str) -> ModelRequestData:
653 )
654
655 tokenizer = AutoTokenizer.from_pretrained(model_name)
660 - messages = [[{
661 - "role":
662 - "user",
663 - "content": [{
664 - "type": "image"
665 - }, {
666 - "type": "text",
667 - "text": f"{question}"
668 - }]
669 - }] for question in questions]
670 - prompts = tokenizer.apply_chat_template(messages,
671 - add_generation_prompt=True,
672 - tokenize=False)
656 + messages = [
657 + [
658 + {
659 + "role": "user",
660 + "content": [{"type": "image"}, {"type": "text", "text": f"{question}"}],
661 + }
662 + ]
663 + for question in questions
664 + ]
665 + prompts = tokenizer.apply_chat_template(
666 + messages, add_generation_prompt=True, tokenize=False
667 + )
668 stop_token_ids = None
669 return ModelRequestData(
670 engine_args=engine_args,
@@ -693,7 +688,8 @@ def run_molmo(questions: list[str], modality: str) -> ModelRequestData:
688
689 prompts = [
690 f"<|im_start|>user <image>\n{question}<|im_end|> \
696 - <|im_start|>assistant\n" for question in questions
691 + <|im_start|>assistant\n"
692 + for question in questions
693 ]
694
695 return ModelRequestData(
@@ -717,15 +713,13 @@ def run_nvlm_d(questions: list[str], modality: str) -> ModelRequestData:
713 limit_mm_per_prompt={modality: 1},
714 )
715
720 - tokenizer = AutoTokenizer.from_pretrained(model_name,
721 - trust_remote_code=True)
722 - messages = [[{
723 - 'role': 'user',
724 - 'content': f"<image>\n{question}"
725 - }] for question in questions]
726 - prompts = tokenizer.apply_chat_template(messages,
727 - tokenize=False,
728 - add_generation_prompt=True)
716 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
717 + messages = [
718 + [{"role": "user", "content": f"<image>\n{question}"}] for question in questions
719 + ]
720 + prompts = tokenizer.apply_chat_template(
721 + messages, tokenize=False, add_generation_prompt=True
722 + )
723
724 return ModelRequestData(
725 engine_args=engine_args,
@@ -748,15 +742,13 @@ def run_ovis(questions: list[str], modality: str) -> ModelRequestData:
742 limit_mm_per_prompt={modality: 1},
743 )
744
751 - tokenizer = AutoTokenizer.from_pretrained(model_name,
752 - trust_remote_code=True)
753 - messages = [[{
754 - 'role': 'user',
755 - 'content': f"<image>\n{question}"
756 - }] for question in questions]
757 - prompts = tokenizer.apply_chat_template(messages,
758 - tokenize=False,
759 - add_generation_prompt=True)
745 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
746 + messages = [
747 + [{"role": "user", "content": f"<image>\n{question}"}] for question in questions
748 + ]
749 + prompts = tokenizer.apply_chat_template(
750 + messages, tokenize=False, add_generation_prompt=True
751 + )
752
753 return ModelRequestData(
754 engine_args=engine_args,
@@ -847,8 +839,7 @@ def run_phi4mm(questions: list[str], modality: str) -> ModelRequestData:
839 # we have to manually specify the path of the lora weights.
840 vision_lora_path = os.path.join(model_path, "vision-lora")
841 prompts = [
850 - f"<|user|><|image_1|>{question}<|end|><|assistant|>"
851 - for question in questions
842 + f"<|user|><|image_1|>{question}<|end|><|assistant|>" for question in questions
843 ]
844 engine_args = EngineArgs(
845 model=model_path,
@@ -915,7 +906,6 @@ def run_qwen_vl(questions: list[str], modality: str) -> ModelRequestData:
906
907 # Qwen2-VL
908 def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
918 -
909 model_name = "Qwen/Qwen2-VL-7B-Instruct"
910
911 engine_args = EngineArgs(
@@ -936,10 +926,13 @@ def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
926 placeholder = "<|video_pad|>"
927
928 prompts = [
939 - ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
940 - f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
941 - f"{question}<|im_end|>\n"
942 - "<|im_start|>assistant\n") for question in questions
929 + (
930 + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
931 + f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
932 + f"{question}<|im_end|>\n"
933 + "<|im_start|>assistant\n"
934 + )
935 + for question in questions
936 ]
937
938 return ModelRequestData(
@@ -950,7 +943,6 @@ def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
943
944 # Qwen2.5-VL
945 def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
953 -
946 model_name = "Qwen/Qwen2.5-VL-3B-Instruct"
947
948 engine_args = EngineArgs(
@@ -971,10 +963,13 @@ def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
963 placeholder = "<|video_pad|>"
964
965 prompts = [
974 - ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
975 - f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
976 - f"{question}<|im_end|>\n"
977 - "<|im_start|>assistant\n") for question in questions
966 + (
967 + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
968 + f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
969 + f"{question}<|im_end|>\n"
970 + "<|im_start|>assistant\n"
971 + )
972 + for question in questions
973 ]
974
975 return ModelRequestData(
@@ -1007,12 +1002,18 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
1002 default_system = (
1003 "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
1004 "Group, capable of perceiving auditory and visual inputs, as well as "
1010 - "generating text and speech.")
1005 + "generating text and speech."
1006 + )
1007
1012 - prompts = [(f"<|im_start|>system\n{default_system}<|im_end|>\n"
1013 - f"<|im_start|>user\n<|vision_bos|>{placeholder}<|vision_eos|>"
1014 - f"{question}<|im_end|>\n"
1015 - "<|im_start|>assistant\n") for question in questions]
1008 + prompts = [
1009 + (
1010 + f"<|im_start|>system\n{default_system}<|im_end|>\n"
1011 + f"<|im_start|>user\n<|vision_bos|>{placeholder}<|vision_eos|>"
1012 + f"{question}<|im_end|>\n"
1013 + "<|im_start|>assistant\n"
1014 + )
1015 + for question in questions
1016 + ]
1017 return ModelRequestData(
1018 engine_args=engine_args,
1019 prompts=prompts,
@@ -1032,15 +1033,13 @@ def run_skyworkr1v(questions: list[str], modality: str) -> ModelRequestData:
1033 limit_mm_per_prompt={modality: 1},
1034 )
1035
1035 - tokenizer = AutoTokenizer.from_pretrained(model_name,
1036 - trust_remote_code=True)
1037 - messages = [[{
1038 - 'role': 'user',
1039 - 'content': f"<image>\n{question}"
1040 - }] for question in questions]
1041 - prompts = tokenizer.apply_chat_template(messages,
1042 - tokenize=False,
1043 - add_generation_prompt=True)
1036 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
1037 + messages = [
1038 + [{"role": "user", "content": f"<image>\n{question}"}] for question in questions
1039 + ]
1040 + prompts = tokenizer.apply_chat_template(
1041 + messages, tokenize=False, add_generation_prompt=True
1042 + )
1043
1044 # Stop tokens for SkyworkR1V
1045 # https://huggingface.co/Skywork/Skywork-R1V-38B/blob/main/conversation.py
@@ -1104,8 +1103,7 @@ def get_multi_modal_input(args):
1103 """
1104 if args.modality == "image":
1105 # Input image and question
1107 - image = convert_image_mode(
1108 - ImageAsset("cherry_blossom").pil_image, "RGB")
1106 + image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
1107 img_questions = [
1108 "What is the content of this image?",
1109 "Describe the content of this image in detail.",
@@ -1120,8 +1118,7 @@ def get_multi_modal_input(args):
1118
1119 if args.modality == "video":
1120 # Input video and question
1123 - video = VideoAsset(name="baby_reading",
1124 - num_frames=args.num_frames).np_ndarrays
1121 + video = VideoAsset(name="baby_reading", num_frames=args.num_frames).np_ndarrays
1122 vid_questions = ["Why is this video funny?"]
1123
1124 return {
@@ -1133,12 +1130,13 @@ def get_multi_modal_input(args):
1130 raise ValueError(msg)
1131
1132
1136 -def apply_image_repeat(image_repeat_prob, num_prompts, data,
1137 - prompts: list[str], modality):
1138 - """Repeats images with provided probability of "image_repeat_prob".
1133 +def apply_image_repeat(
1134 + image_repeat_prob, num_prompts, data, prompts: list[str], modality
1135 +):
1136 + """Repeats images with provided probability of "image_repeat_prob".
1137 Used to simulate hit/miss for the MM preprocessor cache.
1138 """
1141 - assert (image_repeat_prob <= 1.0 and image_repeat_prob >= 0)
1139 + assert image_repeat_prob <= 1.0 and image_repeat_prob >= 0
1140 no_yes = [0, 1]
1141 probs = [1.0 - image_repeat_prob, image_repeat_prob]
1142
@@ -1153,12 +1151,12 @@ def apply_image_repeat(image_repeat_prob, num_prompts, data,
1151 new_val = (i // 256 // 256, i // 256, i % 256)
1152 cur_image.putpixel((0, 0), new_val)
1153
1156 - inputs.append({
1157 - "prompt": prompts[i % len(prompts)],
1158 - "multi_modal_data": {
1159 - modality: cur_image
1154 + inputs.append(
1155 + {
1156 + "prompt": prompts[i % len(prompts)],
1157 + "multi_modal_data": {modality: cur_image},
1158 }
1161 - })
1159 + )
1160
1161 return inputs
1162
@@ -1167,6 +1165,7 @@ def apply_image_repeat(image_repeat_prob, num_prompts, data,
1165 def time_counter(enable: bool):
1166 if enable:
1167 import time
1168 +
1169 start_time = time.time()
1170 yield
1171 elapsed_time = time.time() - start_time
@@ -1179,54 +1178,65 @@ def time_counter(enable: bool):
1178
1179 def parse_args():
1180 parser = FlexibleArgumentParser(
1182 - description='Demo on using vLLM for offline inference with '
1183 - 'vision language models for text generation')
1184 - parser.add_argument('--model-type',
1185 - '-m',
1186 - type=str,
1187 - default="llava",
1188 - choices=model_example_map.keys(),
1189 - help='Huggingface "model_type".')
1190 - parser.add_argument('--num-prompts',
1191 - type=int,
1192 - default=4,
1193 - help='Number of prompts to run.')
1194 - parser.add_argument('--modality',
1195 - type=str,
1196 - default="image",
1197 - choices=['image', 'video'],
1198 - help='Modality of the input.')
1199 - parser.add_argument('--num-frames',
1200 - type=int,
1201 - default=16,
1202 - help='Number of frames to extract from the video.')
1203 - parser.add_argument("--seed",
1204 - type=int,
1205 - default=None,
1206 - help="Set the seed when initializing `vllm.LLM`.")
1181 + description="Demo on using vLLM for offline inference with "
1182 + "vision language models for text generation"
1183 + )
1184 + parser.add_argument(
1185 + "--model-type",
1186 + "-m",
1187 + type=str,
1188 + default="llava",
1189 + choices=model_example_map.keys(),
1190 + help='Huggingface "model_type".',
1191 + )
1192 + parser.add_argument(
1193 + "--num-prompts", type=int, default=4, help="Number of prompts to run."
1194 + )
1195 + parser.add_argument(
1196 + "--modality",
1197 + type=str,
1198 + default="image",
1199 + choices=["image", "video"],
1200 + help="Modality of the input.",
1201 + )
1202 + parser.add_argument(
1203 + "--num-frames",
1204 + type=int,
1205 + default=16,
1206 + help="Number of frames to extract from the video.",
1207 + )
1208 + parser.add_argument(
1209 + "--seed",
1210 + type=int,
1211 + default=None,
1212 + help="Set the seed when initializing `vllm.LLM`.",
1213 + )
1214
1215 parser.add_argument(
1209 - '--image-repeat-prob',
1216 + "--image-repeat-prob",
1217 type=float,
1218 default=None,
1212 - help='Simulates the hit-ratio for multi-modal preprocessor cache'
1213 - ' (if enabled)')
1219 + help="Simulates the hit-ratio for multi-modal preprocessor cache (if enabled)",
1220 + )
1221
1222 parser.add_argument(
1216 - '--disable-mm-preprocessor-cache',
1217 - action='store_true',
1218 - help='If True, disables caching of multi-modal preprocessor/mapper.')
1223 + "--disable-mm-preprocessor-cache",
1224 + action="store_true",
1225 + help="If True, disables caching of multi-modal preprocessor/mapper.",
1226 + )
1227
1228 parser.add_argument(
1221 - '--time-generate',
1222 - action='store_true',
1223 - help='If True, then print the total generate() call time')
1229 + "--time-generate",
1230 + action="store_true",
1231 + help="If True, then print the total generate() call time",
1232 + )
1233
1234 parser.add_argument(
1226 - '--use-different-prompt-per-request',
1227 - action='store_true',
1228 - help='If True, then use different prompt (with the same multi-modal '
1229 - 'data) for each request.')
1235 + "--use-different-prompt-per-request",
1236 + action="store_true",
1237 + help="If True, then use different prompt (with the same multi-modal "
1238 + "data) for each request.",
1239 + )
1240 return parser.parse_args()
1241
1242
@@ -1245,7 +1255,8 @@ def main(args):
1255 # Disable other modalities to save memory
1256 default_limits = {"image": 0, "video": 0, "audio": 0}
1257 req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
1248 - req_data.engine_args.limit_mm_per_prompt or {})
1258 + req_data.engine_args.limit_mm_per_prompt or {}
1259 + )
1260
1261 engine_args = asdict(req_data.engine_args) | {
1262 "seed": args.seed,
@@ -1254,44 +1265,46 @@ def main(args):
1265 llm = LLM(**engine_args)
1266
1267 # Don't want to check the flag multiple times, so just hijack `prompts`.
1257 - prompts = req_data.prompts if args.use_different_prompt_per_request else [
1258 - req_data.prompts[0]
1259 - ]
1268 + prompts = (
1269 + req_data.prompts
1270 + if args.use_different_prompt_per_request
1271 + else [req_data.prompts[0]]
1272 + )
1273
1274 # We set temperature to 0.2 so that outputs can be different
1275 # even when all prompts are identical when running batch inference.
1263 - sampling_params = SamplingParams(temperature=0.2,
1264 - max_tokens=64,
1265 - stop_token_ids=req_data.stop_token_ids)
1276 + sampling_params = SamplingParams(
1277 + temperature=0.2, max_tokens=64, stop_token_ids=req_data.stop_token_ids
1278 + )
1279
1280 assert args.num_prompts > 0
1281 if args.num_prompts == 1:
1282 # Single inference
1283 inputs = {
1284 "prompt": prompts[0],
1272 - "multi_modal_data": {
1273 - modality: data
1274 - },
1285 + "multi_modal_data": {modality: data},
1286 }
1287 else:
1288 # Batch inference
1289 if args.image_repeat_prob is not None:
1290 # Repeat images with specified probability of "image_repeat_prob"
1280 - inputs = apply_image_repeat(args.image_repeat_prob,
1281 - args.num_prompts, data, prompts,
1282 - modality)
1291 + inputs = apply_image_repeat(
1292 + args.image_repeat_prob, args.num_prompts, data, prompts, modality
1293 + )
1294 else:
1295 # Use the same image for all prompts
1285 - inputs = [{
1286 - "prompt": prompts[i % len(prompts)],
1287 - "multi_modal_data": {
1288 - modality: data
1289 - },
1290 - } for i in range(args.num_prompts)]
1296 + inputs = [
1297 + {
1298 + "prompt": prompts[i % len(prompts)],
1299 + "multi_modal_data": {modality: data},
1300 + }
1301 + for i in range(args.num_prompts)
1302 + ]
1303
1304 # Add LoRA request if applicable
1293 - lora_request = (req_data.lora_requests *
1294 - args.num_prompts if req_data.lora_requests else None)
1305 + lora_request = (
1306 + req_data.lora_requests * args.num_prompts if req_data.lora_requests else None
1307 + )
1308
1309 with time_counter(args.time_generate):
1310 outputs = llm.generate(
examples/offline_inference/vision_language_embedding.py
+41 -30
@@ -6,6 +6,7 @@ the correct prompt format on vision language models for multimodal embedding.
6 For most models, the prompt format should follow corresponding examples
7 on HuggingFace model repository.
8 """
9 +
10 from argparse import Namespace
11 from dataclasses import asdict
12 from typing import Literal, NamedTuple, Optional, TypedDict, Union, get_args
@@ -44,19 +45,17 @@ class ModelRequestData(NamedTuple):
45
46
47 def run_e5_v(query: Query) -> ModelRequestData:
47 - llama3_template = '<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n' # noqa: E501
48 + llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n" # noqa: E501
49
50 if query["modality"] == "text":
51 text = query["text"]
51 - prompt = llama3_template.format(
52 - f"{text}\nSummary above sentence in one word: ")
52 + prompt = llama3_template.format(f"{text}\nSummary above sentence in one word: ")
53 image = None
54 elif query["modality"] == "image":
55 - prompt = llama3_template.format(
56 - "<image>\nSummary above image in one word: ")
55 + prompt = llama3_template.format("<image>\nSummary above image in one word: ")
56 image = query["image"]
57 else:
59 - modality = query['modality']
58 + modality = query["modality"]
59 raise ValueError(f"Unsupported query modality: '{modality}'")
60
61 engine_args = EngineArgs(
@@ -83,10 +82,12 @@ def run_vlm2vec(query: Query) -> ModelRequestData:
82 image = query["image"]
83 elif query["modality"] == "text+image":
84 text = query["text"]
86 - prompt = f"<|image_1|> Represent the given image with the following question: {text}" # noqa: E501
85 + prompt = (
86 + f"<|image_1|> Represent the given image with the following question: {text}" # noqa: E501
87 + )
88 image = query["image"]
89 else:
89 - modality = query['modality']
90 + modality = query["modality"]
91 raise ValueError(f"Unsupported query modality: '{modality}'")
92
93 engine_args = EngineArgs(
@@ -136,7 +137,8 @@ def run_encode(model: str, modality: QueryModality, seed: Optional[int]):
137 # Disable other modalities to save memory
138 default_limits = {"image": 0, "video": 0, "audio": 0}
139 req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
139 - req_data.engine_args.limit_mm_per_prompt or {})
140 + req_data.engine_args.limit_mm_per_prompt or {}
141 + )
142
143 engine_args = asdict(req_data.engine_args) | {"seed": seed}
144 llm = LLM(**engine_args)
@@ -145,10 +147,12 @@ def run_encode(model: str, modality: QueryModality, seed: Optional[int]):
147 if req_data.image is not None:
148 mm_data["image"] = req_data.image
149
148 - outputs = llm.embed({
149 - "prompt": req_data.prompt,
150 - "multi_modal_data": mm_data,
151 - })
150 + outputs = llm.embed(
151 + {
152 + "prompt": req_data.prompt,
153 + "multi_modal_data": mm_data,
154 + }
155 + )
156
157 print("-" * 50)
158 for output in outputs:
@@ -164,23 +168,30 @@ model_example_map = {
168
169 def parse_args():
170 parser = FlexibleArgumentParser(
167 - description='Demo on using vLLM for offline inference with '
168 - 'vision language models for multimodal embedding')
169 - parser.add_argument('--model-name',
170 - '-m',
171 - type=str,
172 - default="vlm2vec",
173 - choices=model_example_map.keys(),
174 - help='The name of the embedding model.')
175 - parser.add_argument('--modality',
176 - type=str,
177 - default="image",
178 - choices=get_args(QueryModality),
179 - help='Modality of the input.')
180 - parser.add_argument("--seed",
181 - type=int,
182 - default=None,
183 - help="Set the seed when initializing `vllm.LLM`.")
171 + description="Demo on using vLLM for offline inference with "
172 + "vision language models for multimodal embedding"
173 + )
174 + parser.add_argument(
175 + "--model-name",
176 + "-m",
177 + type=str,
178 + default="vlm2vec",
179 + choices=model_example_map.keys(),
180 + help="The name of the embedding model.",
181 + )
182 + parser.add_argument(
183 + "--modality",
184 + type=str,
185 + default="image",
186 + choices=get_args(QueryModality),
187 + help="Modality of the input.",
188 + )
189 + parser.add_argument(
190 + "--seed",
191 + type=int,
192 + default=None,
193 + help="Set the seed when initializing `vllm.LLM`.",
194 + )
195 return parser.parse_args()
196
197
examples/offline_inference/vision_language_multi_image.py
+210 -215
@@ -4,6 +4,7 @@ This example shows how to use vLLM for running offline inference with
4 multi-image input on vision language models for text generation,
5 using the chat template defined by the model.
6 """
7 +
8 import os
9 from argparse import Namespace
10 from dataclasses import asdict
@@ -59,8 +60,9 @@ def load_aria(question: str, image_urls: list[str]) -> ModelRequestData:
60 limit_mm_per_prompt={"image": len(image_urls)},
61 )
62 placeholders = "<fim_prefix><|img|><fim_suffix>\n" * len(image_urls)
62 - prompt = (f"<|im_start|>user\n{placeholders}{question}<|im_end|>\n"
63 - "<|im_start|>assistant\n")
63 + prompt = (
64 + f"<|im_start|>user\n{placeholders}{question}<|im_end|>\n<|im_start|>assistant\n"
65 + )
66 stop_token_ids = [93532, 93653, 944, 93421, 1019, 93653, 93519]
67
68 return ModelRequestData(
@@ -81,23 +83,21 @@ def load_aya_vision(question: str, image_urls: list[str]) -> ModelRequestData:
83 )
84
85 placeholders = [{"type": "image", "image": url} for url in image_urls]
84 - messages = [{
85 - "role":
86 - "user",
87 - "content": [
88 - *placeholders,
89 - {
90 - "type": "text",
91 - "text": question
92 - },
93 - ],
94 - }]
86 + messages = [
87 + {
88 + "role": "user",
89 + "content": [
90 + *placeholders,
91 + {"type": "text", "text": question},
92 + ],
93 + }
94 + ]
95
96 processor = AutoProcessor.from_pretrained(model_name)
97
98 - prompt = processor.apply_chat_template(messages,
99 - tokenize=False,
100 - add_generation_prompt=True)
98 + prompt = processor.apply_chat_template(
99 + messages, tokenize=False, add_generation_prompt=True
100 + )
101
102 return ModelRequestData(
103 engine_args=engine_args,
@@ -106,8 +106,7 @@ def load_aya_vision(question: str, image_urls: list[str]) -> ModelRequestData:
106 )
107
108
109 -def load_deepseek_vl2(question: str,
110 - image_urls: list[str]) -> ModelRequestData:
109 +def load_deepseek_vl2(question: str, image_urls: list[str]) -> ModelRequestData:
110 model_name = "deepseek-ai/deepseek-vl2-tiny"
111
112 engine_args = EngineArgs(
@@ -118,8 +117,9 @@ def load_deepseek_vl2(question: str,
117 limit_mm_per_prompt={"image": len(image_urls)},
118 )
119
121 - placeholder = "".join(f"image_{i}:<image>\n"
122 - for i, _ in enumerate(image_urls, start=1))
120 + placeholder = "".join(
121 + f"image_{i}:<image>\n" for i, _ in enumerate(image_urls, start=1)
122 + )
123 prompt = f"<|User|>: {placeholder}{question}\n\n<|Assistant|>:"
124
125 return ModelRequestData(
@@ -140,23 +140,21 @@ def load_gemma3(question: str, image_urls: list[str]) -> ModelRequestData:
140 )
141
142 placeholders = [{"type": "image", "image": url} for url in image_urls]
143 - messages = [{
144 - "role":
145 - "user",
146 - "content": [
147 - *placeholders,
148 - {
149 - "type": "text",
150 - "text": question
151 - },
152 - ],
153 - }]
143 + messages = [
144 + {
145 + "role": "user",
146 + "content": [
147 + *placeholders,
148 + {"type": "text", "text": question},
149 + ],
150 + }
151 + ]
152
153 processor = AutoProcessor.from_pretrained(model_name)
154
157 - prompt = processor.apply_chat_template(messages,
158 - tokenize=False,
159 - add_generation_prompt=True)
155 + prompt = processor.apply_chat_template(
156 + messages, tokenize=False, add_generation_prompt=True
157 + )
158
159 return ModelRequestData(
160 engine_args=engine_args,
@@ -176,15 +174,15 @@ def load_h2ovl(question: str, image_urls: list[str]) -> ModelRequestData:
174 mm_processor_kwargs={"max_dynamic_patch": 4},
175 )
176
179 - placeholders = "\n".join(f"Image-{i}: <image>\n"
180 - for i, _ in enumerate(image_urls, start=1))
181 - messages = [{'role': 'user', 'content': f"{placeholders}\n{question}"}]
177 + placeholders = "\n".join(
178 + f"Image-{i}: <image>\n" for i, _ in enumerate(image_urls, start=1)
179 + )
180 + messages = [{"role": "user", "content": f"{placeholders}\n{question}"}]
181
183 - tokenizer = AutoTokenizer.from_pretrained(model_name,
184 - trust_remote_code=True)
185 - prompt = tokenizer.apply_chat_template(messages,
186 - tokenize=False,
187 - add_generation_prompt=True)
182 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
183 + prompt = tokenizer.apply_chat_template(
184 + messages, tokenize=False, add_generation_prompt=True
185 + )
186
187 # Stop tokens for H2OVL-Mississippi
188 # https://huggingface.co/h2oai/h2ovl-mississippi-800m
@@ -211,14 +209,13 @@ def load_idefics3(question: str, image_urls: list[str]) -> ModelRequestData:
209 # if you are running out of memory, you can reduce the "longest_edge".
210 # see: https://huggingface.co/HuggingFaceM4/Idefics3-8B-Llama3#model-optimizations
211 mm_processor_kwargs={
214 - "size": {
215 - "longest_edge": 2 * 364
216 - },
212 + "size": {"longest_edge": 2 * 364},
213 },
214 )
215
220 - placeholders = "\n".join(f"Image-{i}: <image>\n"
221 - for i, _ in enumerate(image_urls, start=1))
216 + placeholders = "\n".join(
217 + f"Image-{i}: <image>\n" for i, _ in enumerate(image_urls, start=1)
218 + )
219 prompt = f"<|begin_of_text|>User:{placeholders}\n{question}<end_of_utterance>\nAssistant:" # noqa: E501
220 return ModelRequestData(
221 engine_args=engine_args,
@@ -238,15 +235,16 @@ def load_smolvlm(question: str, image_urls: list[str]) -> ModelRequestData:
235 enforce_eager=True,
236 limit_mm_per_prompt={"image": len(image_urls)},
237 mm_processor_kwargs={
241 - "max_image_size": {
242 - "longest_edge": 384
243 - },
238 + "max_image_size": {"longest_edge": 384},
239 },
240 )
241
247 - placeholders = "\n".join(f"Image-{i}: <image>\n"
248 - for i, _ in enumerate(image_urls, start=1))
249 - prompt = f"<|im_start|>User:{placeholders}\n{question}<end_of_utterance>\nAssistant:" # noqa: E501
242 + placeholders = "\n".join(
243 + f"Image-{i}: <image>\n" for i, _ in enumerate(image_urls, start=1)
244 + )
245 + prompt = (
246 + f"<|im_start|>User:{placeholders}\n{question}<end_of_utterance>\nAssistant:" # noqa: E501
247 + )
248 return ModelRequestData(
249 engine_args=engine_args,
250 prompt=prompt,
@@ -265,15 +263,15 @@ def load_internvl(question: str, image_urls: list[str]) -> ModelRequestData:
263 mm_processor_kwargs={"max_dynamic_patch": 4},
264 )
265
268 - placeholders = "\n".join(f"Image-{i}: <image>\n"
269 - for i, _ in enumerate(image_urls, start=1))
270 - messages = [{'role': 'user', 'content': f"{placeholders}\n{question}"}]
266 + placeholders = "\n".join(
267 + f"Image-{i}: <image>\n" for i, _ in enumerate(image_urls, start=1)
268 + )
269 + messages = [{"role": "user", "content": f"{placeholders}\n{question}"}]
270
272 - tokenizer = AutoTokenizer.from_pretrained(model_name,
273 - trust_remote_code=True)
274 - prompt = tokenizer.apply_chat_template(messages,
275 - tokenize=False,
276 - add_generation_prompt=True)
271 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
272 + prompt = tokenizer.apply_chat_template(
273 + messages, tokenize=False, add_generation_prompt=True
274 + )
275
276 # Stop tokens for InternVL
277 # models variants may have different stop tokens
@@ -301,23 +299,21 @@ def load_llama4(question: str, image_urls: list[str]) -> ModelRequestData:
299 )
300
301 placeholders = [{"type": "image", "image": url} for url in image_urls]
304 - messages = [{
305 - "role":
306 - "user",
307 - "content": [
308 - *placeholders,
309 - {
310 - "type": "text",
311 - "text": question
312 - },
313 - ],
314 - }]
302 + messages = [
303 + {
304 + "role": "user",
305 + "content": [
306 + *placeholders,
307 + {"type": "text", "text": question},
308 + ],
309 + }
310 + ]
311
312 processor = AutoProcessor.from_pretrained(model_name)
313
318 - prompt = processor.apply_chat_template(messages,
319 - tokenize=False,
320 - add_generation_prompt=True)
314 + prompt = processor.apply_chat_template(
315 + messages, tokenize=False, add_generation_prompt=True
316 + )
317
318 return ModelRequestData(
319 engine_args=engine_args,
@@ -338,24 +334,21 @@ def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:
334 )
335
336 placeholders = [{"type": "image", "image": url} for url in image_urls]
341 - messages = [{
342 - "role":
343 - "user",
344 - "content": [
345 - *placeholders,
346 - {
347 - "type": "text",
348 - "text": question
349 - },
350 - ],
351 - }]
337 + messages = [
338 + {
339 + "role": "user",
340 + "content": [
341 + *placeholders,
342 + {"type": "text", "text": question},
343 + ],
344 + }
345 + ]
346
353 - processor = AutoProcessor.from_pretrained(model_name,
354 - trust_remote_code=True)
347 + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
348
356 - prompt = processor.apply_chat_template(messages,
357 - tokenize=False,
358 - add_generation_prompt=True)
349 + prompt = processor.apply_chat_template(
350 + messages, tokenize=False, add_generation_prompt=True
351 + )
352
353 return ModelRequestData(
354 engine_args=engine_args,
@@ -419,15 +412,15 @@ def load_nvlm_d(question: str, image_urls: list[str]) -> ModelRequestData:
412 mm_processor_kwargs={"max_dynamic_patch": 4},
413 )
414
422 - placeholders = "\n".join(f"Image-{i}: <image>\n"
423 - for i, _ in enumerate(image_urls, start=1))
424 - messages = [{'role': 'user', 'content': f"{placeholders}\n{question}"}]
415 + placeholders = "\n".join(
416 + f"Image-{i}: <image>\n" for i, _ in enumerate(image_urls, start=1)
417 + )
418 + messages = [{"role": "user", "content": f"{placeholders}\n{question}"}]
419
426 - tokenizer = AutoTokenizer.from_pretrained(model_name,
427 - trust_remote_code=True)
428 - prompt = tokenizer.apply_chat_template(messages,
429 - tokenize=False,
430 - add_generation_prompt=True)
420 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
421 + prompt = tokenizer.apply_chat_template(
422 + messages, tokenize=False, add_generation_prompt=True
423 + )
424
425 return ModelRequestData(
426 engine_args=engine_args,
@@ -449,15 +442,15 @@ def load_ovis(question: str, image_urls: list[str]) -> ModelRequestData:
442 limit_mm_per_prompt={"image": len(image_urls)},
443 )
444
452 - placeholders = "\n".join(f"Image-{i}: <image>\n"
453 - for i, _ in enumerate(image_urls, start=1))
454 - messages = [{'role': 'user', 'content': f"{placeholders}\n{question}"}]
445 + placeholders = "\n".join(
446 + f"Image-{i}: <image>\n" for i, _ in enumerate(image_urls, start=1)
447 + )
448 + messages = [{"role": "user", "content": f"{placeholders}\n{question}"}]
449
456 - tokenizer = AutoTokenizer.from_pretrained(model_name,
457 - trust_remote_code=True)
458 - prompt = tokenizer.apply_chat_template(messages,
459 - tokenize=False,
460 - add_generation_prompt=True)
450 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
451 + prompt = tokenizer.apply_chat_template(
452 + messages, tokenize=False, add_generation_prompt=True
453 + )
454
455 return ModelRequestData(
456 engine_args=engine_args,
@@ -509,8 +502,9 @@ def load_phi3v(question: str, image_urls: list[str]) -> ModelRequestData:
502 limit_mm_per_prompt={"image": len(image_urls)},
503 mm_processor_kwargs={"num_crops": 4},
504 )
512 - placeholders = "\n".join(f"<|image_{i}|>"
513 - for i, _ in enumerate(image_urls, start=1))
505 + placeholders = "\n".join(
506 + f"<|image_{i}|>" for i, _ in enumerate(image_urls, start=1)
507 + )
508 prompt = f"<|user|>\n{placeholders}\n{question}<|end|>\n<|assistant|>\n"
509
510 return ModelRequestData(
@@ -542,8 +536,7 @@ def load_phi4mm(question: str, image_urls: list[str]) -> ModelRequestData:
536 mm_processor_kwargs={"dynamic_hd": 4},
537 )
538
545 - placeholders = "".join(f"<|image_{i}|>"
546 - for i, _ in enumerate(image_urls, start=1))
539 + placeholders = "".join(f"<|image_{i}|>" for i, _ in enumerate(image_urls, start=1))
540 prompt = f"<|user|>{placeholders}{question}<|end|><|assistant|>"
541
542 return ModelRequestData(
@@ -554,8 +547,7 @@ def load_phi4mm(question: str, image_urls: list[str]) -> ModelRequestData:
547 )
548
549
557 -def load_qwen_vl_chat(question: str,
558 - image_urls: list[str]) -> ModelRequestData:
550 +def load_qwen_vl_chat(question: str, image_urls: list[str]) -> ModelRequestData:
551 model_name = "Qwen/Qwen-VL-Chat"
552 engine_args = EngineArgs(
553 model=model_name,
@@ -565,24 +557,26 @@ def load_qwen_vl_chat(question: str,
557 hf_overrides={"architectures": ["QwenVLForConditionalGeneration"]},
558 limit_mm_per_prompt={"image": len(image_urls)},
559 )
568 - placeholders = "".join(f"Picture {i}: <img></img>\n"
569 - for i, _ in enumerate(image_urls, start=1))
560 + placeholders = "".join(
561 + f"Picture {i}: <img></img>\n" for i, _ in enumerate(image_urls, start=1)
562 + )
563
564 # This model does not have a chat_template attribute on its tokenizer,
565 # so we need to explicitly pass it. We use ChatML since it's used in the
566 # generation utils of the model:
567 # https://huggingface.co/Qwen/Qwen-VL-Chat/blob/main/qwen_generation_utils.py#L265
575 - tokenizer = AutoTokenizer.from_pretrained(model_name,
576 - trust_remote_code=True)
568 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
569
570 # Copied from: https://huggingface.co/docs/transformers/main/en/chat_templating
571 chat_template = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" # noqa: E501
572
581 - messages = [{'role': 'user', 'content': f"{placeholders}\n{question}"}]
582 - prompt = tokenizer.apply_chat_template(messages,
583 - tokenize=False,
584 - add_generation_prompt=True,
585 - chat_template=chat_template)
573 + messages = [{"role": "user", "content": f"{placeholders}\n{question}"}]
574 + prompt = tokenizer.apply_chat_template(
575 + messages,
576 + tokenize=False,
577 + add_generation_prompt=True,
578 + chat_template=chat_template,
579 + )
580
581 stop_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]
582 stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
@@ -600,9 +594,11 @@ def load_qwen2_vl(question: str, image_urls: list[str]) -> ModelRequestData:
594 try:
595 from qwen_vl_utils import process_vision_info
596 except ModuleNotFoundError:
603 - print('WARNING: `qwen-vl-utils` not installed, input images will not '
604 - 'be automatically resized. You can enable this functionality by '
605 - '`pip install qwen-vl-utils`.')
597 + print(
598 + "WARNING: `qwen-vl-utils` not installed, input images will not "
599 + "be automatically resized. You can enable this functionality by "
600 + "`pip install qwen-vl-utils`."
601 + )
602 process_vision_info = None
603
604 model_name = "Qwen/Qwen2-VL-7B-Instruct"
@@ -616,26 +612,22 @@ def load_qwen2_vl(question: str, image_urls: list[str]) -> ModelRequestData:
612 )
613
614 placeholders = [{"type": "image", "image": url} for url in image_urls]
619 - messages = [{
620 - "role": "system",
621 - "content": "You are a helpful assistant."
622 - }, {
623 - "role":
624 - "user",
625 - "content": [
626 - *placeholders,
627 - {
628 - "type": "text",
629 - "text": question
630 - },
631 - ],
632 - }]
615 + messages = [
616 + {"role": "system", "content": "You are a helpful assistant."},
617 + {
618 + "role": "user",
619 + "content": [
620 + *placeholders,
621 + {"type": "text", "text": question},
622 + ],
623 + },
624 + ]
625
626 processor = AutoProcessor.from_pretrained(model_name)
627
636 - prompt = processor.apply_chat_template(messages,
637 - tokenize=False,
638 - add_generation_prompt=True)
628 + prompt = processor.apply_chat_template(
629 + messages, tokenize=False, add_generation_prompt=True
630 + )
631
632 if process_vision_info is None:
633 image_data = [fetch_image(url) for url in image_urls]
@@ -653,9 +645,11 @@ def load_qwen2_5_vl(question: str, image_urls: list[str]) -> ModelRequestData:
645 try:
646 from qwen_vl_utils import process_vision_info
647 except ModuleNotFoundError:
656 - print('WARNING: `qwen-vl-utils` not installed, input images will not '
657 - 'be automatically resized. You can enable this functionality by '
658 - '`pip install qwen-vl-utils`.')
648 + print(
649 + "WARNING: `qwen-vl-utils` not installed, input images will not "
650 + "be automatically resized. You can enable this functionality by "
651 + "`pip install qwen-vl-utils`."
652 + )
653 process_vision_info = None
654
655 model_name = "Qwen/Qwen2.5-VL-3B-Instruct"
@@ -668,32 +662,27 @@ def load_qwen2_5_vl(question: str, image_urls: list[str]) -> ModelRequestData:
662 )
663
664 placeholders = [{"type": "image", "image": url} for url in image_urls]
671 - messages = [{
672 - "role": "system",
673 - "content": "You are a helpful assistant."
674 - }, {
675 - "role":
676 - "user",
677 - "content": [
678 - *placeholders,
679 - {
680 - "type": "text",
681 - "text": question
682 - },
683 - ],
684 - }]
665 + messages = [
666 + {"role": "system", "content": "You are a helpful assistant."},
667 + {
668 + "role": "user",
669 + "content": [
670 + *placeholders,
671 + {"type": "text", "text": question},
672 + ],
673 + },
674 + ]
675
676 processor = AutoProcessor.from_pretrained(model_name)
677
688 - prompt = processor.apply_chat_template(messages,
689 - tokenize=False,
690 - add_generation_prompt=True)
678 + prompt = processor.apply_chat_template(
679 + messages, tokenize=False, add_generation_prompt=True
680 + )
681
682 if process_vision_info is None:
683 image_data = [fetch_image(url) for url in image_urls]
684 else:
695 - image_data, _ = process_vision_info(messages,
696 - return_video_kwargs=False)
685 + image_data, _ = process_vision_info(messages, return_video_kwargs=False)
686
687 return ModelRequestData(
688 engine_args=engine_args,
@@ -726,23 +715,20 @@ model_example_map = {
715 }
716
717
729 -def run_generate(model, question: str, image_urls: list[str],
730 - seed: Optional[int]):
718 +def run_generate(model, question: str, image_urls: list[str], seed: Optional[int]):
719 req_data = model_example_map[model](question, image_urls)
720
721 engine_args = asdict(req_data.engine_args) | {"seed": args.seed}
722 llm = LLM(**engine_args)
723
736 - sampling_params = SamplingParams(temperature=0.0,
737 - max_tokens=256,
738 - stop_token_ids=req_data.stop_token_ids)
724 + sampling_params = SamplingParams(
725 + temperature=0.0, max_tokens=256, stop_token_ids=req_data.stop_token_ids
726 + )
727
728 outputs = llm.generate(
729 {
730 "prompt": req_data.prompt,
743 - "multi_modal_data": {
744 - "image": req_data.image_data
745 - },
731 + "multi_modal_data": {"image": req_data.image_data},
732 },
733 sampling_params=sampling_params,
734 lora_request=req_data.lora_requests,
@@ -755,38 +741,40 @@ def run_generate(model, question: str, image_urls: list[str],
741 print("-" * 50)
742
743
758 -def run_chat(model: str, question: str, image_urls: list[str],
759 - seed: Optional[int]):
744 +def run_chat(model: str, question: str, image_urls: list[str], seed: Optional[int]):
745 req_data = model_example_map[model](question, image_urls)
746
747 # Disable other modalities to save memory
748 default_limits = {"image": 0, "video": 0, "audio": 0}
749 req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
765 - req_data.engine_args.limit_mm_per_prompt or {})
750 + req_data.engine_args.limit_mm_per_prompt or {}
751 + )
752
753 engine_args = asdict(req_data.engine_args) | {"seed": seed}
754 llm = LLM(**engine_args)
755
770 - sampling_params = SamplingParams(temperature=0.0,
771 - max_tokens=256,
772 - stop_token_ids=req_data.stop_token_ids)
756 + sampling_params = SamplingParams(
757 + temperature=0.0, max_tokens=256, stop_token_ids=req_data.stop_token_ids
758 + )
759 outputs = llm.chat(
774 - [{
775 - "role":
776 - "user",
777 - "content": [
778 - {
779 - "type": "text",
780 - "text": question,
781 - },
782 - *({
783 - "type": "image_url",
784 - "image_url": {
785 - "url": image_url
760 + [
761 + {
762 + "role": "user",
763 + "content": [
764 + {
765 + "type": "text",
766 + "text": question,
767 },
787 - } for image_url in image_urls),
788 - ],
789 - }],
768 + *(
769 + {
770 + "type": "image_url",
771 + "image_url": {"url": image_url},
772 + }
773 + for image_url in image_urls
774 + ),
775 + ],
776 + }
777 + ],
778 sampling_params=sampling_params,
779 chat_template=req_data.chat_template,
780 lora_request=req_data.lora_requests,
@@ -801,32 +789,39 @@ def run_chat(model: str, question: str, image_urls: list[str],
789
790 def parse_args():
791 parser = FlexibleArgumentParser(
804 - description='Demo on using vLLM for offline inference with '
805 - 'vision language models that support multi-image input for text '
806 - 'generation')
807 - parser.add_argument('--model-type',
808 - '-m',
809 - type=str,
810 - default="phi3_v",
811 - choices=model_example_map.keys(),
812 - help='Huggingface "model_type".')
813 - parser.add_argument("--method",
814 - type=str,
815 - default="generate",
816 - choices=["generate", "chat"],
817 - help="The method to run in `vllm.LLM`.")
818 - parser.add_argument("--seed",
819 - type=int,
820 - default=None,
821 - help="Set the seed when initializing `vllm.LLM`.")
792 + description="Demo on using vLLM for offline inference with "
793 + "vision language models that support multi-image input for text "
794 + "generation"
795 + )
796 + parser.add_argument(
797 + "--model-type",
798 + "-m",
799 + type=str,
800 + default="phi3_v",
801 + choices=model_example_map.keys(),
802 + help='Huggingface "model_type".',
803 + )
804 + parser.add_argument(
805 + "--method",
806 + type=str,
807 + default="generate",
808 + choices=["generate", "chat"],
809 + help="The method to run in `vllm.LLM`.",
810 + )
811 + parser.add_argument(
812 + "--seed",
813 + type=int,
814 + default=None,
815 + help="Set the seed when initializing `vllm.LLM`.",
816 + )
817 parser.add_argument(
818 "--num-images",
819 "-n",
820 type=int,
826 - choices=list(range(1,
827 - len(IMAGE_URLS) + 1)), # the max number of images
821 + choices=list(range(1, len(IMAGE_URLS) + 1)), # the max number of images
822 default=2,
829 - help="Number of images to use for the demo.")
823 + help="Number of images to use for the demo.",
824 + )
825 return parser.parse_args()
826
827
@@ -835,7 +830,7 @@ def main(args: Namespace):
830 method = args.method
831 seed = args.seed
832
838 - image_urls = IMAGE_URLS[:args.num_images]
833 + image_urls = IMAGE_URLS[: args.num_images]
834
835 if method == "generate":
836 run_generate(model, QUESTION, image_urls, seed)
examples/online_serving/api_client.py
+9 -13
@@ -17,16 +17,15 @@ import requests
17
18
19 def clear_line(n: int = 1) -> None:
20 - LINE_UP = '\033[1A'
21 - LINE_CLEAR = '\x1b[2K'
20 + LINE_UP = "\033[1A"
21 + LINE_CLEAR = "\x1b[2K"
22 for _ in range(n):
23 print(LINE_UP, end=LINE_CLEAR, flush=True)
24
25
26 -def post_http_request(prompt: str,
27 - api_url: str,
28 - n: int = 1,
29 - stream: bool = False) -> requests.Response:
26 +def post_http_request(
27 + prompt: str, api_url: str, n: int = 1, stream: bool = False
28 +) -> requests.Response:
29 headers = {"User-Agent": "Test Client"}
30 pload = {
31 "prompt": prompt,
@@ -35,17 +34,14 @@ def post_http_request(prompt: str,
34 "max_tokens": 16,
35 "stream": stream,
36 }
38 - response = requests.post(api_url,
39 - headers=headers,
40 - json=pload,
41 - stream=stream)
37 + response = requests.post(api_url, headers=headers, json=pload, stream=stream)
38 return response
39
40
41 def get_streaming_response(response: requests.Response) -> Iterable[list[str]]:
46 - for chunk in response.iter_lines(chunk_size=8192,
47 - decode_unicode=False,
48 - delimiter=b"\n"):
42 + for chunk in response.iter_lines(
43 + chunk_size=8192, decode_unicode=False, delimiter=b"\n"
44 + ):
45 if chunk:
46 data = json.loads(chunk.decode("utf-8"))
47 output = data["text"]
examples/online_serving/cohere_rerank_client.py
+9 -8
@@ -6,6 +6,7 @@ Note that `pip install cohere` is needed to run this example.
6
7 run: vllm serve BAAI/bge-reranker-base
8 """
9 +
10 from typing import Union
11
12 import cohere
@@ -16,28 +17,28 @@ model = "BAAI/bge-reranker-base"
17 query = "What is the capital of France?"
18
19 documents = [
19 - "The capital of France is Paris", "Reranking is fun!",
20 - "vLLM is an open-source framework for fast AI serving"
20 + "The capital of France is Paris",
21 + "Reranking is fun!",
22 + "vLLM is an open-source framework for fast AI serving",
23 ]
24
25
24 -def cohere_rerank(client: Union[Client, ClientV2], model: str, query: str,
25 - documents: list[str]) -> dict:
26 +def cohere_rerank(
27 + client: Union[Client, ClientV2], model: str, query: str, documents: list[str]
28 +) -> dict:
29 return client.rerank(model=model, query=query, documents=documents)
30
31
32 def main():
33 # cohere v1 client
31 - cohere_v1 = cohere.Client(base_url="http://localhost:8000",
32 - api_key="sk-fake-key")
34 + cohere_v1 = cohere.Client(base_url="http://localhost:8000", api_key="sk-fake-key")
35 rerank_v1_result = cohere_rerank(cohere_v1, model, query, documents)
36 print("-" * 50)
37 print("rerank_v1_result:\n", rerank_v1_result)
38 print("-" * 50)
39
40 # or the v2
39 - cohere_v2 = cohere.ClientV2("sk-fake-key",
40 - base_url="http://localhost:8000")
41 + cohere_v2 = cohere.ClientV2("sk-fake-key", base_url="http://localhost:8000")
42 rerank_v2_result = cohere_rerank(cohere_v2, model, query, documents)
43 print("rerank_v2_result:\n", rerank_v2_result)
44 print("-" * 50)
examples/online_serving/disaggregated_serving/disagg_proxy_demo.py
+93 -95
@@ -13,6 +13,7 @@ launch this proxy demo through:
13 Note: This demo will be removed once the PDController implemented in PR 15343
14 (https://github.com/vllm-project/vllm/pull/15343) supports XpYd.
15 """
16 +
17 import argparse
18 import ipaddress
19 import itertools
@@ -26,8 +27,7 @@ from typing import Callable, Optional
27 import aiohttp
28 import requests
29 import uvicorn
29 -from fastapi import (APIRouter, Depends, FastAPI, Header, HTTPException,
30 - Request, status)
30 +from fastapi import APIRouter, Depends, FastAPI, Header, HTTPException, Request, status
31 from fastapi.responses import JSONResponse, StreamingResponse
32
33 AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
@@ -36,24 +36,24 @@ logging.basicConfig(level=logging.INFO)
36
37
38 class SchedulingPolicy(ABC):
39 -
39 @abstractmethod
40 def schedule(self, cycler: itertools.cycle):
41 raise NotImplementedError("Scheduling Proxy is not set.")
42
43
44 class Proxy:
46 -
45 def __init__(
46 self,
47 prefill_instances: list[str],
48 decode_instances: list[str],
49 model: str,
50 scheduling_policy: SchedulingPolicy,
53 - custom_create_completion: Optional[Callable[[Request],
54 - StreamingResponse]] = None,
55 - custom_create_chat_completion: Optional[Callable[
56 - [Request], StreamingResponse]] = None,
51 + custom_create_completion: Optional[
52 + Callable[[Request], StreamingResponse]
53 + ] = None,
54 + custom_create_chat_completion: Optional[
55 + Callable[[Request], StreamingResponse]
56 + ] = None,
57 ):
58 self.prefill_instances = prefill_instances
59 self.decode_instances = decode_instances
@@ -68,30 +68,30 @@ class Proxy:
68
69 def setup_routes(self):
70 self.router.post(
71 - "/v1/completions",
72 - dependencies=[
73 - Depends(self.validate_json_request)
74 - ])(self.custom_create_completion if self.
75 - custom_create_completion else self.create_completion)
71 + "/v1/completions", dependencies=[Depends(self.validate_json_request)]
72 + )(
73 + self.custom_create_completion
74 + if self.custom_create_completion
75 + else self.create_completion
76 + )
77 self.router.post(
77 - "/v1/chat/completions",
78 - dependencies=[
79 - Depends(self.validate_json_request)
80 - ])(self.custom_create_chat_completion if self.
81 - custom_create_chat_completion else self.create_chat_completion)
82 - self.router.get("/status",
83 - response_class=JSONResponse)(self.get_status)
84 - self.router.post("/instances/add",
85 - dependencies=[Depends(self.api_key_authenticate)
86 - ])(self.add_instance_endpoint)
78 + "/v1/chat/completions", dependencies=[Depends(self.validate_json_request)]
79 + )(
80 + self.custom_create_chat_completion
81 + if self.custom_create_chat_completion
82 + else self.create_chat_completion
83 + )
84 + self.router.get("/status", response_class=JSONResponse)(self.get_status)
85 + self.router.post(
86 + "/instances/add", dependencies=[Depends(self.api_key_authenticate)]
87 + )(self.add_instance_endpoint)
88
89 async def validate_json_request(self, raw_request: Request):
90 content_type = raw_request.headers.get("content-type", "").lower()
91 if content_type != "application/json":
92 raise HTTPException(
93 status_code=415,
93 - detail=
94 - "Unsupported Media Type: Only 'application/json' is allowed",
94 + detail="Unsupported Media Type: Only 'application/json' is allowed",
95 )
96
97 def api_key_authenticate(self, x_api_key: str = Header(...)):
@@ -103,8 +103,7 @@ class Proxy:
103 detail="Server configuration error.",
104 )
105 if x_api_key != expected_api_key:
106 - logger.warning("Unauthorized access attempt with API Key: %s",
107 - x_api_key)
106 + logger.warning("Unauthorized access attempt with API Key: %s", x_api_key)
107 raise HTTPException(
108 status_code=status.HTTP_403_FORBIDDEN,
109 detail="Forbidden: Invalid API Key.",
@@ -113,8 +112,7 @@ class Proxy:
112 async def validate_instance(self, instance: str) -> bool:
113 url = f"http://{instance}/v1/models"
114 try:
116 - async with aiohttp.ClientSession(
117 - timeout=AIOHTTP_TIMEOUT) as client:
115 + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as client:
116 logger.info("Verifying %s ...", instance)
117 async with client.get(url) as response:
118 if response.status == 200:
@@ -122,12 +120,15 @@ class Proxy:
120 if "data" in data and len(data["data"]) > 0:
121 model_cur = data["data"][0].get("id", "")
122 if model_cur == self.model:
125 - logger.info("Instance: %s could be added.",
126 - instance)
123 + logger.info("Instance: %s could be added.", instance)
124 return True
125 else:
129 - logger.warning("Mismatch model %s : %s != %s",
130 - instance, model_cur, self.model)
126 + logger.warning(
127 + "Mismatch model %s : %s != %s",
128 + instance,
129 + model_cur,
130 + self.model,
131 + )
132 return False
133 else:
134 return False
@@ -147,48 +148,47 @@ class Proxy:
148 instance_type = data.get("type")
149 instance = data.get("instance")
150 if instance_type not in ["prefill", "decode"]:
150 - raise HTTPException(status_code=400,
151 - detail="Invalid instance type.")
151 + raise HTTPException(status_code=400, detail="Invalid instance type.")
152 if not instance or ":" not in instance:
153 - raise HTTPException(status_code=400,
154 - detail="Invalid instance format.")
153 + raise HTTPException(status_code=400, detail="Invalid instance format.")
154 host, port_str = instance.split(":")
155 try:
156 if host != "localhost":
157 ipaddress.ip_address(host)
158 port = int(port_str)
159 if not (0 < port < 65536):
161 - raise HTTPException(status_code=400,
162 - detail="Invalid port number.")
160 + raise HTTPException(status_code=400, detail="Invalid port number.")
161 except Exception as e:
164 - raise HTTPException(status_code=400,
165 - detail="Invalid instance address.") from e
162 + raise HTTPException(
163 + status_code=400, detail="Invalid instance address."
164 + ) from e
165
166 is_valid = await self.validate_instance(instance)
167 if not is_valid:
169 - raise HTTPException(status_code=400,
170 - detail="Instance validation failed.")
168 + raise HTTPException(
169 + status_code=400, detail="Instance validation failed."
170 + )
171
172 if instance_type == "prefill":
173 if instance not in self.prefill_instances:
174 self.prefill_instances.append(instance)
175 - self.prefill_cycler = itertools.cycle(
176 - self.prefill_instances)
175 + self.prefill_cycler = itertools.cycle(self.prefill_instances)
176 else:
178 - raise HTTPException(status_code=400,
179 - detail="Instance already exists.")
177 + raise HTTPException(
178 + status_code=400, detail="Instance already exists."
179 + )
180 else:
181 if instance not in self.decode_instances:
182 self.decode_instances.append(instance)
183 self.decode_cycler = itertools.cycle(self.decode_instances)
184 else:
185 - raise HTTPException(status_code=400,
186 - detail="Instance already exists.")
185 + raise HTTPException(
186 + status_code=400, detail="Instance already exists."
187 + )
188
188 - return JSONResponse(content={
189 - "message":
190 - f"Added {instance} to {instance_type}_instances."
191 - })
189 + return JSONResponse(
190 + content={"message": f"Added {instance} to {instance_type}_instances."}
191 + )
192 except HTTPException as http_exc:
193 raise http_exc
194 except Exception as e:
@@ -197,16 +197,16 @@ class Proxy:
197
198 async def forward_request(self, url, data, use_chunked=True):
199 async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
200 - headers = {
201 - "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
202 - }
200 + headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
201 try:
204 - async with session.post(url=url, json=data,
205 - headers=headers) as response:
202 + async with session.post(
203 + url=url, json=data, headers=headers
204 + ) as response:
205 if 200 <= response.status < 300 or 400 <= response.status < 500: # noqa: E501
206 if use_chunked:
207 async for chunk_bytes in response.content.iter_chunked( # noqa: E501
209 - 1024):
208 + 1024
209 + ):
210 yield chunk_bytes
211 else:
212 content = await response.read()
@@ -217,20 +217,21 @@ class Proxy:
217 error_content = json.loads(error_content)
218 except json.JSONDecodeError:
219 error_content = error_content
220 - logger.error("Request failed with status %s: %s",
221 - response.status, error_content)
220 + logger.error(
221 + "Request failed with status %s: %s",
222 + response.status,
223 + error_content,
224 + )
225 raise HTTPException(
226 status_code=response.status,
224 - detail=
225 - f"Request failed with status {response.status}: "
227 + detail=f"Request failed with status {response.status}: "
228 f"{error_content}",
229 )
230 except aiohttp.ClientError as e:
231 logger.error("ClientError occurred: %s", str(e))
232 raise HTTPException(
233 status_code=502,
232 - detail=
233 - "Bad Gateway: Error communicating with upstream server.",
234 + detail="Bad Gateway: Error communicating with upstream server.",
235 ) from e
236 except Exception as e:
237 logger.error("Unexpected error: %s", str(e))
@@ -258,8 +259,8 @@ class Proxy:
259 prefill_instance = self.schedule(self.prefill_cycler)
260 try:
261 async for _ in self.forward_request(
261 - f"http://{prefill_instance}/v1/completions",
262 - kv_prepare_request):
262 + f"http://{prefill_instance}/v1/completions", kv_prepare_request
263 + ):
264 continue
265 except HTTPException as http_exc:
266 self.remove_instance_endpoint("prefill", prefill_instance)
@@ -270,7 +271,8 @@ class Proxy:
271
272 try:
273 generator = self.forward_request(
273 - f"http://{decode_instance}/v1/completions", request)
274 + f"http://{decode_instance}/v1/completions", request
275 + )
276 except HTTPException as http_exc:
277 self.remove_instance_endpoint("decode", decode_instance)
278 raise http_exc
@@ -295,8 +297,8 @@ class Proxy:
297 prefill_instance = self.schedule(self.prefill_cycler)
298 try:
299 async for _ in self.forward_request(
298 - f"http://{prefill_instance}/v1/chat/completions",
299 - kv_prepare_request):
300 + f"http://{prefill_instance}/v1/chat/completions", kv_prepare_request
301 + ):
302 continue
303 except HTTPException as http_exc:
304 self.remove_instance_endpoint("prefill", prefill_instance)
@@ -306,8 +308,8 @@ class Proxy:
308
309 try:
310 generator = self.forward_request(
309 - "http://" + decode_instance + "/v1/chat/completions",
310 - request)
311 + "http://" + decode_instance + "/v1/chat/completions", request
312 + )
313 except HTTPException as http_exc:
314 self.remove_instance_endpoint("decode", decode_instance)
315 raise http_exc
@@ -318,20 +320,20 @@ class Proxy:
320 error_messages = [str(e) for e in exc_info if e]
321 print("Error occurred in disagg proxy server")
322 print(error_messages)
321 - return StreamingResponse(content=iter(error_messages),
322 - media_type="text/event-stream")
323 + return StreamingResponse(
324 + content=iter(error_messages), media_type="text/event-stream"
325 + )
326
327 def remove_instance_endpoint(self, instance_type, instance):
325 - if (instance_type == "decode" and instance in self.decode_instances):
328 + if instance_type == "decode" and instance in self.decode_instances:
329 self.decode_instances.remove(instance)
330 self.decode_cycler = itertools.cycle(self.decode_instances)
328 - if (instance_type == "prefill" and instance in self.decode_instances):
331 + if instance_type == "prefill" and instance in self.decode_instances:
332 self.prefill_instances.remove(instance)
333 self.prefill_cycler = itertools.cycle(self.decode_instances)
334
335
336 class RoundRobinSchedulingPolicy(SchedulingPolicy):
334 -
337 def __init__(self):
338 super().__init__()
339
@@ -340,15 +342,12 @@ class RoundRobinSchedulingPolicy(SchedulingPolicy):
342
343
344 class ProxyServer:
343 -
345 def __init__(
346 self,
347 args: argparse.Namespace,
348 scheduling_policy: Optional[SchedulingPolicy] = None,
348 - create_completion: Optional[Callable[[Request],
349 - StreamingResponse]] = None,
350 - create_chat_completion: Optional[Callable[[Request],
351 - StreamingResponse]] = None,
349 + create_completion: Optional[Callable[[Request], StreamingResponse]] = None,
350 + create_chat_completion: Optional[Callable[[Request], StreamingResponse]] = None,
351 ):
352 self.validate_parsed_serve_args(args)
353 self.port = args.port
@@ -356,8 +355,11 @@ class ProxyServer:
355 prefill_instances=[] if args.prefill is None else args.prefill,
356 decode_instances=[] if args.decode is None else args.decode,
357 model=args.model,
359 - scheduling_policy=(scheduling_policy if scheduling_policy
360 - is not None else RoundRobinSchedulingPolicy()),
358 + scheduling_policy=(
359 + scheduling_policy
360 + if scheduling_policy is not None
361 + else RoundRobinSchedulingPolicy()
362 + ),
363 custom_create_completion=create_completion,
364 custom_create_chat_completion=create_chat_completion,
365 )
@@ -382,11 +384,9 @@ class ProxyServer:
384 ipaddress.ip_address(host)
385 port = int(port)
386 if not (0 < port < 65536):
385 - raise ValueError(
386 - f"Invalid port number in instance: {instance}")
387 + raise ValueError(f"Invalid port number in instance: {instance}")
388 except Exception as e:
388 - raise ValueError(
389 - f"Invalid instance {instance}: {str(e)}") from e
389 + raise ValueError(f"Invalid instance {instance}: {str(e)}") from e
390
391 def verify_model_config(self, instances: list, model: str) -> None:
392 model_suffix = model.split("/")[-1]
@@ -399,12 +399,14 @@ class ProxyServer:
399 if model_cur_suffix != model_suffix:
400 raise ValueError(
401 f"{instance} serves a different model: "
402 - f"{model_cur} != {model}")
402 + f"{model_cur} != {model}"
403 + )
404 else:
405 raise ValueError(f"Cannot get model id from {instance}!")
406 except requests.RequestException as e:
407 raise ValueError(
407 - f"Error communicating with {instance}: {str(e)}") from e
408 + f"Error communicating with {instance}: {str(e)}"
409 + ) from e
410
411 def run_server(self):
412 app = FastAPI()
@@ -417,11 +419,7 @@ class ProxyServer:
419 def parse_args():
420 # Todo: allow more config
421 parser = argparse.ArgumentParser("vLLM disaggregated proxy server.")
420 - parser.add_argument("--model",
421 - "-m",
422 - type=str,
423 - required=True,
424 - help="Model name")
422 + parser.add_argument("--model", "-m", type=str, required=True, help="Model name")
423
424 parser.add_argument(
425 "--prefill",
examples/online_serving/gradio_openai_chatbot_webserver.py
+38 -44
@@ -17,6 +17,7 @@ you can install it manually by following these steps:
17 2. Rename the downloaded file to: frpc_linux_amd64_v0.3
18 3. Move the file to this location: /home/user/.cache/huggingface/gradio/frpc
19 """
20 +
21 import argparse
22
23 import gradio as gr
@@ -24,16 +25,12 @@ from openai import OpenAI
25
26
27 def format_history_to_openai(history):
27 - history_openai_format = [{
28 - "role": "system",
29 - "content": "You are a great AI assistant."
30 - }]
28 + history_openai_format = [
29 + {"role": "system", "content": "You are a great AI assistant."}
30 + ]
31 for human, assistant in history:
32 history_openai_format.append({"role": "user", "content": human})
33 - history_openai_format.append({
34 - "role": "assistant",
35 - "content": assistant
36 - })
33 + history_openai_format.append({"role": "assistant", "content": assistant})
34 return history_openai_format
35
36
@@ -49,17 +46,17 @@ def predict(message, history, client, model_name, temp, stop_token_ids):
46 temperature=temp,
47 stream=True,
48 extra_body={
52 - 'repetition_penalty':
53 - 1,
54 - 'stop_token_ids':
55 - [int(id.strip())
56 - for id in stop_token_ids.split(',')] if stop_token_ids else []
57 - })
49 + "repetition_penalty": 1,
50 + "stop_token_ids": [int(id.strip()) for id in stop_token_ids.split(",")]
51 + if stop_token_ids
52 + else [],
53 + },
54 + )
55
56 # Collect all chunks and concatenate them into a full message
57 full_message = ""
58 for chunk in stream:
62 - full_message += (chunk.choices[0].delta.content or "")
59 + full_message += chunk.choices[0].delta.content or ""
60
61 # Return the full message as a single response
62 return full_message
@@ -67,38 +64,34 @@ def predict(message, history, client, model_name, temp, stop_token_ids):
64
65 def parse_args():
66 parser = argparse.ArgumentParser(
70 - description='Chatbot Interface with Customizable Parameters')
71 - parser.add_argument('--model-url',
72 - type=str,
73 - default='http://localhost:8000/v1',
74 - help='Model URL')
75 - parser.add_argument('-m',
76 - '--model',
77 - type=str,
78 - required=True,
79 - help='Model name for the chatbot')
80 - parser.add_argument('--temp',
81 - type=float,
82 - default=0.8,
83 - help='Temperature for text generation')
84 - parser.add_argument('--stop-token-ids',
85 - type=str,
86 - default='',
87 - help='Comma-separated stop token IDs')
67 + description="Chatbot Interface with Customizable Parameters"
68 + )
69 + parser.add_argument(
70 + "--model-url", type=str, default="http://localhost:8000/v1", help="Model URL"
71 + )
72 + parser.add_argument(
73 + "-m", "--model", type=str, required=True, help="Model name for the chatbot"
74 + )
75 + parser.add_argument(
76 + "--temp", type=float, default=0.8, help="Temperature for text generation"
77 + )
78 + parser.add_argument(
79 + "--stop-token-ids", type=str, default="", help="Comma-separated stop token IDs"
80 + )
81 parser.add_argument("--host", type=str, default=None)
82 parser.add_argument("--port", type=int, default=8001)
83 return parser.parse_args()
84
85
86 def build_gradio_interface(client, model_name, temp, stop_token_ids):
94 -
87 def chat_predict(message, history):
96 - return predict(message, history, client, model_name, temp,
97 - stop_token_ids)
88 + return predict(message, history, client, model_name, temp, stop_token_ids)
89
99 - return gr.ChatInterface(fn=chat_predict,
100 - title="Chatbot Interface",
101 - description="A simple chatbot powered by vLLM")
90 + return gr.ChatInterface(
91 + fn=chat_predict,
92 + title="Chatbot Interface",
93 + description="A simple chatbot powered by vLLM",
94 + )
95
96
97 def main():
@@ -113,12 +106,13 @@ def main():
106 client = OpenAI(api_key=openai_api_key, base_url=openai_api_base)
107
108 # Define the Gradio chatbot interface using the predict function
116 - gradio_interface = build_gradio_interface(client, args.model, args.temp,
117 - args.stop_token_ids)
109 + gradio_interface = build_gradio_interface(
110 + client, args.model, args.temp, args.stop_token_ids
111 + )
112
119 - gradio_interface.queue().launch(server_name=args.host,
120 - server_port=args.port,
121 - share=True)
113 + gradio_interface.queue().launch(
114 + server_name=args.host, server_port=args.port, share=True
115 + )
116
117
118 if __name__ == "__main__":
examples/online_serving/gradio_webserver.py
+14 -18
@@ -17,6 +17,7 @@ you can install it manually by following these steps:
17 2. Rename the downloaded file to: frpc_linux_amd64_v0.3
18 3. Move the file to this location: /home/user/.cache/huggingface/gradio/frpc
19 """
20 +
21 import argparse
22 import json
23
@@ -31,14 +32,11 @@ def http_bot(prompt):
32 "stream": True,
33 "max_tokens": 128,
34 }
34 - response = requests.post(args.model_url,
35 - headers=headers,
36 - json=pload,
37 - stream=True)
38 -
39 - for chunk in response.iter_lines(chunk_size=8192,
40 - decode_unicode=False,
41 - delimiter=b"\n"):
35 + response = requests.post(args.model_url, headers=headers, json=pload, stream=True)
36 +
37 + for chunk in response.iter_lines(
38 + chunk_size=8192, decode_unicode=False, delimiter=b"\n"
39 + ):
40 if chunk:
41 data = json.loads(chunk.decode("utf-8"))
42 output = data["text"][0]
@@ -48,10 +46,10 @@ def http_bot(prompt):
46 def build_demo():
47 with gr.Blocks() as demo:
48 gr.Markdown("# vLLM text completion demo\n")
51 - inputbox = gr.Textbox(label="Input",
52 - placeholder="Enter text and press ENTER")
53 - outputbox = gr.Textbox(label="Output",
54 - placeholder="Generated result from the model")
49 + inputbox = gr.Textbox(label="Input", placeholder="Enter text and press ENTER")
50 + outputbox = gr.Textbox(
51 + label="Output", placeholder="Generated result from the model"
52 + )
53 inputbox.submit(http_bot, [inputbox], [outputbox])
54 return demo
55
@@ -60,17 +58,15 @@ def parse_args():
58 parser = argparse.ArgumentParser()
59 parser.add_argument("--host", type=str, default=None)
60 parser.add_argument("--port", type=int, default=8001)
63 - parser.add_argument("--model-url",
64 - type=str,
65 - default="http://localhost:8000/generate")
61 + parser.add_argument(
62 + "--model-url", type=str, default="http://localhost:8000/generate"
63 + )
64 return parser.parse_args()
65
66
67 def main(args):
68 demo = build_demo()
71 - demo.queue().launch(server_name=args.host,
72 - server_port=args.port,
73 - share=True)
69 + demo.queue().launch(server_name=args.host, server_port=args.port, share=True)
70
71
72 if __name__ == "__main__":
examples/online_serving/jinaai_rerank_client.py
+6 -6
@@ -5,6 +5,7 @@ Jina and Cohere https://jina.ai/reranker
5
6 run: vllm serve BAAI/bge-reranker-base
7 """
8 +
9 import json
10
11 import requests
@@ -14,14 +15,13 @@ url = "http://127.0.0.1:8000/rerank"
15 headers = {"accept": "application/json", "Content-Type": "application/json"}
16
17 data = {
17 - "model":
18 - "BAAI/bge-reranker-base",
19 - "query":
20 - "What is the capital of France?",
18 + "model": "BAAI/bge-reranker-base",
19 + "query": "What is the capital of France?",
20 "documents": [
21 "The capital of Brazil is Brasilia.",
23 - "The capital of France is Paris.", "Horses and cows are both animals"
24 - ]
22 + "The capital of France is Paris.",
23 + "Horses and cows are both animals",
24 + ],
25 }
26
27
examples/online_serving/kv_events_subscriber.py
+7 -9
@@ -9,17 +9,14 @@ from msgspec.msgpack import Decoder
9 #
10 # Types copied from vllm.distributed.kv_events
11 #
12 -class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True,
13 - gc=False):
12 +class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
13 ts: float
14 events: list[Any]
15
16
18 -class KVCacheEvent(msgspec.Struct,
19 - array_like=True,
20 - omit_defaults=True,
21 - gc=False,
22 - tag=True):
17 +class KVCacheEvent(
18 + msgspec.Struct, array_like=True, omit_defaults=True, gc=False, tag=True
19 +):
20 """Base class for all KV cache-related events"""
21
22
@@ -77,8 +74,9 @@ def main():
74
75 if last_seq >= 0 and seq > last_seq + 1:
76 missed = seq - last_seq - 1
80 - print(f"Missed {missed} messages"
81 - f" (last: {last_seq}, current: {seq})")
77 + print(
78 + f"Missed {missed} messages (last: {last_seq}, current: {seq})"
79 + )
80
81 replay.send((last_seq + 1).to_bytes(8, "big"))
82
examples/online_serving/openai_chat_completion_client.py
+12 -16
@@ -12,26 +12,22 @@ from openai import OpenAI
12 openai_api_key = "EMPTY"
13 openai_api_base = "http://localhost:8000/v1"
14
15 -messages = [{
16 - "role": "system",
17 - "content": "You are a helpful assistant."
18 -}, {
19 - "role": "user",
20 - "content": "Who won the world series in 2020?"
21 -}, {
22 - "role": "assistant",
23 - "content": "The Los Angeles Dodgers won the World Series in 2020."
24 -}, {
25 - "role": "user",
26 - "content": "Where was it played?"
27 -}]
15 +messages = [
16 + {"role": "system", "content": "You are a helpful assistant."},
17 + {"role": "user", "content": "Who won the world series in 2020?"},
18 + {
19 + "role": "assistant",
20 + "content": "The Los Angeles Dodgers won the World Series in 2020.",
21 + },
22 + {"role": "user", "content": "Where was it played?"},
23 +]
24
25
26 def parse_args():
27 parser = argparse.ArgumentParser(description="Client for vLLM API server")
32 - parser.add_argument("--stream",
33 - action="store_true",
34 - help="Enable streaming response")
28 + parser.add_argument(
29 + "--stream", action="store_true", help="Enable streaming response"
30 + )
31 return parser.parse_args()
32
33
examples/online_serving/openai_chat_completion_client_for_multimodal.py
+114 -143
@@ -43,7 +43,7 @@ def encode_base64_content_from_url(content_url: str) -> str:
43
44 with requests.get(content_url) as response:
45 response.raise_for_status()
46 - result = base64.b64encode(response.content).decode('utf-8')
46 + result = base64.b64encode(response.content).decode("utf-8")
47
48 return result
49
@@ -51,10 +51,7 @@ def encode_base64_content_from_url(content_url: str) -> str:
51 # Text-only inference
52 def run_text_only(model: str) -> None:
53 chat_completion = client.chat.completions.create(
54 - messages=[{
55 - "role": "user",
56 - "content": "What's the capital of France?"
57 - }],
54 + messages=[{"role": "user", "content": "What's the capital of France?"}],
55 model=model,
56 max_completion_tokens=64,
57 )
@@ -65,26 +62,21 @@ def run_text_only(model: str) -> None:
62
63 # Single-image input inference
64 def run_single_image(model: str) -> None:
68 -
65 ## Use image url in the payload
66 image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
67 chat_completion_from_url = client.chat.completions.create(
72 - messages=[{
73 - "role":
74 - "user",
75 - "content": [
76 - {
77 - "type": "text",
78 - "text": "What's in this image?"
79 - },
80 - {
81 - "type": "image_url",
82 - "image_url": {
83 - "url": image_url
68 + messages=[
69 + {
70 + "role": "user",
71 + "content": [
72 + {"type": "text", "text": "What's in this image?"},
73 + {
74 + "type": "image_url",
75 + "image_url": {"url": image_url},
76 },
85 - },
86 - ],
87 - }],
77 + ],
78 + }
79 + ],
80 model=model,
81 max_completion_tokens=64,
82 )
@@ -95,22 +87,18 @@ def run_single_image(model: str) -> None:
87 ## Use base64 encoded image in the payload
88 image_base64 = encode_base64_content_from_url(image_url)
89 chat_completion_from_base64 = client.chat.completions.create(
98 - messages=[{
99 - "role":
100 - "user",
101 - "content": [
102 - {
103 - "type": "text",
104 - "text": "What's in this image?"
105 - },
106 - {
107 - "type": "image_url",
108 - "image_url": {
109 - "url": f"data:image/jpeg;base64,{image_base64}"
90 + messages=[
91 + {
92 + "role": "user",
93 + "content": [
94 + {"type": "text", "text": "What's in this image?"},
95 + {
96 + "type": "image_url",
97 + "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"},
98 },
111 - },
112 - ],
113 - }],
99 + ],
100 + }
101 + ],
102 model=model,
103 max_completion_tokens=64,
104 )
@@ -124,28 +112,22 @@ def run_multi_image(model: str) -> None:
112 image_url_duck = "https://upload.wikimedia.org/wikipedia/commons/d/da/2015_Kaczka_krzy%C5%BCowka_w_wodzie_%28samiec%29.jpg"
113 image_url_lion = "https://upload.wikimedia.org/wikipedia/commons/7/77/002_The_lion_king_Snyggve_in_the_Serengeti_National_Park_Photo_by_Giles_Laurent.jpg"
114 chat_completion_from_url = client.chat.completions.create(
127 - messages=[{
128 - "role":
129 - "user",
130 - "content": [
131 - {
132 - "type": "text",
133 - "text": "What are the animals in these images?"
134 - },
135 - {
136 - "type": "image_url",
137 - "image_url": {
138 - "url": image_url_duck
115 + messages=[
116 + {
117 + "role": "user",
118 + "content": [
119 + {"type": "text", "text": "What are the animals in these images?"},
120 + {
121 + "type": "image_url",
122 + "image_url": {"url": image_url_duck},
123 },
140 - },
141 - {
142 - "type": "image_url",
143 - "image_url": {
144 - "url": image_url_lion
124 + {
125 + "type": "image_url",
126 + "image_url": {"url": image_url_lion},
127 },
146 - },
147 - ],
148 - }],
128 + ],
129 + }
130 + ],
131 model=model,
132 max_completion_tokens=64,
133 )
@@ -161,22 +143,18 @@ def run_video(model: str) -> None:
143
144 ## Use video url in the payload
145 chat_completion_from_url = client.chat.completions.create(
164 - messages=[{
165 - "role":
166 - "user",
167 - "content": [
168 - {
169 - "type": "text",
170 - "text": "What's in this video?"
171 - },
172 - {
173 - "type": "video_url",
174 - "video_url": {
175 - "url": video_url
146 + messages=[
147 + {
148 + "role": "user",
149 + "content": [
150 + {"type": "text", "text": "What's in this video?"},
151 + {
152 + "type": "video_url",
153 + "video_url": {"url": video_url},
154 },
177 - },
178 - ],
179 - }],
155 + ],
156 + }
157 + ],
158 model=model,
159 max_completion_tokens=64,
160 )
@@ -186,22 +164,18 @@ def run_video(model: str) -> None:
164
165 ## Use base64 encoded video in the payload
166 chat_completion_from_base64 = client.chat.completions.create(
189 - messages=[{
190 - "role":
191 - "user",
192 - "content": [
193 - {
194 - "type": "text",
195 - "text": "What's in this video?"
196 - },
197 - {
198 - "type": "video_url",
199 - "video_url": {
200 - "url": f"data:video/mp4;base64,{video_base64}"
167 + messages=[
168 + {
169 + "role": "user",
170 + "content": [
171 + {"type": "text", "text": "What's in this video?"},
172 + {
173 + "type": "video_url",
174 + "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
175 },
202 - },
203 - ],
204 - }],
176 + ],
177 + }
178 + ],
179 model=model,
180 max_completion_tokens=64,
181 )
@@ -219,24 +193,22 @@ def run_audio(model: str) -> None:
193
194 # OpenAI-compatible schema (`input_audio`)
195 chat_completion_from_base64 = client.chat.completions.create(
222 - messages=[{
223 - "role":
224 - "user",
225 - "content": [
226 - {
227 - "type": "text",
228 - "text": "What's in this audio?"
229 - },
230 - {
231 - "type": "input_audio",
232 - "input_audio": {
233 - # Any format supported by librosa is supported
234 - "data": audio_base64,
235 - "format": "wav"
196 + messages=[
197 + {
198 + "role": "user",
199 + "content": [
200 + {"type": "text", "text": "What's in this audio?"},
201 + {
202 + "type": "input_audio",
203 + "input_audio": {
204 + # Any format supported by librosa is supported
205 + "data": audio_base64,
206 + "format": "wav",
207 + },
208 },
237 - },
238 - ],
239 - }],
209 + ],
210 + }
211 + ],
212 model=model,
213 max_completion_tokens=64,
214 )
@@ -246,23 +218,21 @@ def run_audio(model: str) -> None:
218
219 # HTTP URL
220 chat_completion_from_url = client.chat.completions.create(
249 - messages=[{
250 - "role":
251 - "user",
252 - "content": [
253 - {
254 - "type": "text",
255 - "text": "What's in this audio?"
256 - },
257 - {
258 - "type": "audio_url",
259 - "audio_url": {
260 - # Any format supported by librosa is supported
261 - "url": audio_url
221 + messages=[
222 + {
223 + "role": "user",
224 + "content": [
225 + {"type": "text", "text": "What's in this audio?"},
226 + {
227 + "type": "audio_url",
228 + "audio_url": {
229 + # Any format supported by librosa is supported
230 + "url": audio_url
231 + },
232 },
263 - },
264 - ],
265 - }],
233 + ],
234 + }
235 + ],
236 model=model,
237 max_completion_tokens=64,
238 )
@@ -272,23 +242,21 @@ def run_audio(model: str) -> None:
242
243 # base64 URL
244 chat_completion_from_base64 = client.chat.completions.create(
275 - messages=[{
276 - "role":
277 - "user",
278 - "content": [
279 - {
280 - "type": "text",
281 - "text": "What's in this audio?"
282 - },
283 - {
284 - "type": "audio_url",
285 - "audio_url": {
286 - # Any format supported by librosa is supported
287 - "url": f"data:audio/ogg;base64,{audio_base64}"
245 + messages=[
246 + {
247 + "role": "user",
248 + "content": [
249 + {"type": "text", "text": "What's in this audio?"},
250 + {
251 + "type": "audio_url",
252 + "audio_url": {
253 + # Any format supported by librosa is supported
254 + "url": f"data:audio/ogg;base64,{audio_base64}"
255 + },
256 },
289 - },
290 - ],
291 - }],
257 + ],
258 + }
259 + ],
260 model=model,
261 max_completion_tokens=64,
262 )
@@ -308,14 +276,17 @@ example_function_map = {
276
277 def parse_args():
278 parser = FlexibleArgumentParser(
311 - description='Demo on using OpenAI client for online serving with '
312 - 'multimodal language models served with vLLM.')
313 - parser.add_argument('--chat-type',
314 - '-c',
315 - type=str,
316 - default="single-image",
317 - choices=list(example_function_map.keys()),
318 - help='Conversation type with multimodal data.')
279 + description="Demo on using OpenAI client for online serving with "
280 + "multimodal language models served with vLLM."
281 + )
282 + parser.add_argument(
283 + "--chat-type",
284 + "-c",
285 + type=str,
286 + default="single-image",
287 + choices=list(example_function_map.keys()),
288 + help="Conversation type with multimodal data.",
289 + )
290 return parser.parse_args()
291
292
examples/online_serving/openai_chat_completion_client_with_tools.py
+74 -73
@@ -16,6 +16,7 @@ vllm serve NousResearch/Hermes-2-Pro-Llama-3-8B \
16 --chat-template examples/tool_chat_template_hermes.jinja \
17 --enable-auto-tool-choice --tool-call-parser hermes
18 """
19 +
20 import json
21 from typing import Any
22
@@ -25,55 +26,55 @@ from openai import OpenAI
26 openai_api_key = "EMPTY"
27 openai_api_base = "http://localhost:8000/v1"
28
28 -tools = [{
29 - "type": "function",
30 - "function": {
31 - "name": "get_current_weather",
32 - "description": "Get the current weather in a given location",
33 - "parameters": {
34 - "type": "object",
35 - "properties": {
36 - "city": {
37 - "type":
38 - "string",
39 - "description":
40 - "The city to find the weather for, e.g. 'San Francisco'"
41 - },
42 - "state": {
43 - "type":
44 - "string",
45 - "description":
46 - "the two-letter abbreviation for the state that the city is"
47 - " in, e.g. 'CA' which would mean 'California'"
48 - },
49 - "unit": {
50 - "type": "string",
51 - "description": "The unit to fetch the temperature in",
52 - "enum": ["celsius", "fahrenheit"]
53 - }
29 +properties = {
30 + "city": {
31 + "type": "string",
32 + "description": "The city to find the weather for, e.g. 'San Francisco'",
33 + },
34 + "state": {
35 + "type": "string",
36 + "description": "the two-letter abbreviation for the state that the city is"
37 + " in, e.g. 'CA' which would mean 'California'",
38 + },
39 + "unit": {
40 + "type": "string",
41 + "description": "The unit to fetch the temperature in",
42 + "enum": ["celsius", "fahrenheit"],
43 + },
44 +}
45 +
46 +tools = [
47 + {
48 + "type": "function",
49 + "function": {
50 + "name": "get_current_weather",
51 + "description": "Get the current weather in a given location",
52 + "parameters": {
53 + "type": "object",
54 + "properties": properties,
55 + "required": ["city", "state", "unit"],
56 },
55 - "required": ["city", "state", "unit"]
56 - }
57 + },
58 }
58 -}]
59 -
60 -messages = [{
61 - "role": "user",
62 - "content": "Hi! How are you doing today?"
63 -}, {
64 - "role": "assistant",
65 - "content": "I'm doing well! How can I help you?"
66 -}, {
67 - "role":
68 - "user",
69 - "content":
70 - "Can you tell me what the temperate will be in Dallas, in fahrenheit?"
71 -}]
72 -
73 -
74 -def get_current_weather(city: str, state: str, unit: 'str'):
75 - return ("The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
76 - "partly cloudly, with highs in the 90's.")
59 +]
60 +
61 +messages = [
62 + {"role": "user", "content": "Hi! How are you doing today?"},
63 + {"role": "assistant", "content": "I'm doing well! How can I help you?"},
64 + {
65 + "role": "user",
66 + "content": (
67 + "Can you tell me what the temperate will be in Dallas, in fahrenheit?"
68 + ),
69 + },
70 +]
71 +
72 +
73 +def get_current_weather(city: str, state: str, unit: "str"):
74 + return (
75 + "The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
76 + "partly cloudly, with highs in the 90's."
77 + )
78
79
80 def handle_tool_calls_stream(
@@ -82,10 +83,9 @@ def handle_tool_calls_stream(
83 model: str,
84 tools: list[dict[str, Any]],
85 ) -> list[Any]:
85 - tool_calls_stream = client.chat.completions.create(messages=messages,
86 - model=model,
87 - tools=tools,
88 - stream=True)
86 + tool_calls_stream = client.chat.completions.create(
87 + messages=messages, model=model, tools=tools, stream=True
88 + )
89 chunks = []
90 print("chunks: ")
91 for chunk in tool_calls_stream:
@@ -106,8 +106,7 @@ def handle_tool_calls_arguments(chunks: list[Any]) -> list[str]:
106 tool_call = chunk.choices[0].delta.tool_calls[0]
107 if tool_call.index != tool_call_idx:
108 if tool_call_idx >= 0:
109 - print(f"streamed tool call arguments: "
110 - f"{arguments[tool_call_idx]}")
109 + print(f"streamed tool call arguments: {arguments[tool_call_idx]}")
110 tool_call_idx = chunk.choices[0].delta.tool_calls[0].index
111 arguments.append("")
112 if tool_call.id:
@@ -115,8 +114,7 @@ def handle_tool_calls_arguments(chunks: list[Any]) -> list[str]:
114
115 if tool_call.function:
116 if tool_call.function.name:
118 - print(
119 - f"streamed tool call name: {tool_call.function.name}")
117 + print(f"streamed tool call name: {tool_call.function.name}")
118
119 if tool_call.function.arguments:
120 arguments[tool_call_idx] += tool_call.function.arguments
@@ -136,9 +134,9 @@ def main():
134 models = client.models.list()
135 model = models.data[0].id
136
139 - chat_completion = client.chat.completions.create(messages=messages,
140 - model=model,
141 - tools=tools)
137 + chat_completion = client.chat.completions.create(
138 + messages=messages, model=model, tools=tools
139 + )
140
141 print("-" * 70)
142 print("Chat completion results:")
@@ -158,10 +156,12 @@ def main():
156 print("-" * 70)
157
158 # Add tool call results to the conversation
161 - messages.append({
162 - "role": "assistant",
163 - "tool_calls": chat_completion.choices[0].message.tool_calls
164 - })
159 + messages.append(
160 + {
161 + "role": "assistant",
162 + "tool_calls": chat_completion.choices[0].message.tool_calls,
163 + }
164 + )
165
166 # Now, simulate a tool call
167 available_tools = {"get_current_weather": get_current_weather}
@@ -172,17 +172,18 @@ def main():
172 args = json.loads(call.function.arguments)
173 result = tool_to_call(**args)
174 print("tool_to_call result: ", result)
175 - messages.append({
176 - "role": "tool",
177 - "content": result,
178 - "tool_call_id": call.id,
179 - "name": call.function.name
180 - })
181 -
182 - chat_completion_2 = client.chat.completions.create(messages=messages,
183 - model=model,
184 - tools=tools,
185 - stream=False)
175 + messages.append(
176 + {
177 + "role": "tool",
178 + "content": result,
179 + "tool_call_id": call.id,
180 + "name": call.function.name,
181 + }
182 + )
183 +
184 + chat_completion_2 = client.chat.completions.create(
185 + messages=messages, model=model, tools=tools, stream=False
186 + )
187 print("Chat completion2 results:")
188 print(chat_completion_2)
189 print("-" * 70)
examples/online_serving/openai_chat_completion_client_with_tools_required.py
+24 -37
@@ -28,18 +28,16 @@ tools = [
28 "type": "object",
29 "properties": {
30 "city": {
31 - "type":
32 - "string",
33 - "description":
34 - "The city to find the weather for"
31 + "type": "string",
32 + "description": "The city to find the weather for"
33 ", e.g. 'San Francisco'",
34 },
35 "state": {
38 - "type":
39 - "string",
40 - "description":
41 - "the two-letter abbreviation for the state that the "
42 - "city is in, e.g. 'CA' which would mean 'California'",
36 + "type": "string",
37 + "description": (
38 + "the two-letter abbreviation for the state that the "
39 + "city is in, e.g. 'CA' which would mean 'California'"
40 + ),
41 },
42 "unit": {
43 "type": "string",
@@ -60,22 +58,20 @@ tools = [
58 "type": "object",
59 "properties": {
60 "city": {
63 - "type":
64 - "string",
65 - "description":
66 - "The city to get the forecast for, e.g. 'New York'",
61 + "type": "string",
62 + "description": (
63 + "The city to get the forecast for, e.g. 'New York'"
64 + ),
65 },
66 "state": {
69 - "type":
70 - "string",
71 - "description":
72 - "The two-letter abbreviation for the state, e.g. 'NY'",
67 + "type": "string",
68 + "description": (
69 + "The two-letter abbreviation for the state, e.g. 'NY'"
70 + ),
71 },
72 "days": {
75 - "type":
76 - "integer",
77 - "description":
78 - "Number of days to get the forecast for (1-7)",
73 + "type": "integer",
74 + "description": "Number of days to get the forecast for (1-7)",
75 },
76 "unit": {
77 "type": "string",
@@ -90,19 +86,11 @@ tools = [
86 ]
87
88 messages = [
89 + {"role": "user", "content": "Hi! How are you doing today?"},
90 + {"role": "assistant", "content": "I'm doing well! How can I help you?"},
91 {
92 "role": "user",
95 - "content": "Hi! How are you doing today?"
96 - },
97 - {
98 - "role": "assistant",
99 - "content": "I'm doing well! How can I help you?"
100 - },
101 - {
102 - "role":
103 - "user",
104 - "content":
105 - "Can you tell me what the current weather is in Dallas \
93 + "content": "Can you tell me what the current weather is in Dallas \
94 and the forecast for the next 5 days, in fahrenheit?",
95 },
96 ]
@@ -123,17 +111,16 @@ def main():
111 model=model,
112 tools=tools,
113 tool_choice="required",
126 - stream=True # Enable streaming response
114 + stream=True, # Enable streaming response
115 )
116
117 for chunk in chat_completion:
118 if chunk.choices and chunk.choices[0].delta.tool_calls:
119 print(chunk.choices[0].delta.tool_calls)
120
133 - chat_completion = client.chat.completions.create(messages=messages,
134 - model=model,
135 - tools=tools,
136 - tool_choice="required")
121 + chat_completion = client.chat.completions.create(
122 + messages=messages, model=model, tools=tools, tool_choice="required"
123 + )
124
125 print(chat_completion.choices[0].message.tool_calls)
126
examples/online_serving/openai_chat_completion_structured_outputs.py
+46 -34
@@ -20,10 +20,9 @@ openai_api_base = "http://localhost:8000/v1"
20 def guided_choice_completion(client: OpenAI, model: str):
21 completion = client.chat.completions.create(
22 model=model,
23 - messages=[{
24 - "role": "user",
25 - "content": "Classify this sentiment: vLLM is wonderful!"
26 - }],
23 + messages=[
24 + {"role": "user", "content": "Classify this sentiment: vLLM is wonderful!"}
25 + ],
26 extra_body={"guided_choice": ["positive", "negative"]},
27 )
28 return completion.choices[0].message.content
@@ -31,20 +30,21 @@ def guided_choice_completion(client: OpenAI, model: str):
30
31 # Guided decoding by Regex
32 def guided_regex_completion(client: OpenAI, model: str):
34 - prompt = ("Generate an email address for Alan Turing, who works in Enigma."
35 - "End in .com and new line. Example result:"
36 - "alan.turing@enigma.com\n")
33 + prompt = (
34 + "Generate an email address for Alan Turing, who works in Enigma."
35 + "End in .com and new line. Example result:"
36 + "alan.turing@enigma.com\n"
37 + )
38
39 completion = client.chat.completions.create(
40 model=model,
40 - messages=[{
41 - "role": "user",
42 - "content": prompt,
43 - }],
44 - extra_body={
45 - "guided_regex": r"\w+@\w+\.com\n",
46 - "stop": ["\n"]
47 - },
41 + messages=[
42 + {
43 + "role": "user",
44 + "content": prompt,
45 + }
46 + ],
47 + extra_body={"guided_regex": r"\w+@\w+\.com\n", "stop": ["\n"]},
48 )
49 return completion.choices[0].message.content
50
@@ -66,14 +66,18 @@ class CarDescription(BaseModel):
66 def guided_json_completion(client: OpenAI, model: str):
67 json_schema = CarDescription.model_json_schema()
68
69 - prompt = ("Generate a JSON with the brand, model and car_type of"
70 - "the most iconic car from the 90's")
69 + prompt = (
70 + "Generate a JSON with the brand, model and car_type of"
71 + "the most iconic car from the 90's"
72 + )
73 completion = client.chat.completions.create(
74 model=model,
73 - messages=[{
74 - "role": "user",
75 - "content": prompt,
76 - }],
75 + messages=[
76 + {
77 + "role": "user",
78 + "content": prompt,
79 + }
80 + ],
81 extra_body={"guided_json": json_schema},
82 )
83 return completion.choices[0].message.content
@@ -95,14 +99,18 @@ def guided_grammar_completion(client: OpenAI, model: str):
99 number ::= "1 " | "2 "
100 """
101
98 - prompt = ("Generate an SQL query to show the 'username' and 'email'"
99 - "from the 'users' table.")
102 + prompt = (
103 + "Generate an SQL query to show the 'username' and 'email'"
104 + "from the 'users' table."
105 + )
106 completion = client.chat.completions.create(
107 model=model,
102 - messages=[{
103 - "role": "user",
104 - "content": prompt,
105 - }],
108 + messages=[
109 + {
110 + "role": "user",
111 + "content": prompt,
112 + }
113 + ],
114 extra_body={"guided_grammar": simplified_sql_grammar},
115 )
116 return completion.choices[0].message.content
@@ -110,19 +118,23 @@ def guided_grammar_completion(client: OpenAI, model: str):
118
119 # Extra backend options
120 def extra_backend_options_completion(client: OpenAI, model: str):
113 - prompt = ("Generate an email address for Alan Turing, who works in Enigma."
114 - "End in .com and new line. Example result:"
115 - "alan.turing@enigma.com\n")
121 + prompt = (
122 + "Generate an email address for Alan Turing, who works in Enigma."
123 + "End in .com and new line. Example result:"
124 + "alan.turing@enigma.com\n"
125 + )
126
127 try:
128 # The guided_decoding_disable_fallback option forces vLLM to use
129 # xgrammar, so when it fails you get a 400 with the reason why
130 completion = client.chat.completions.create(
131 model=model,
122 - messages=[{
123 - "role": "user",
124 - "content": prompt,
125 - }],
132 + messages=[
133 + {
134 + "role": "user",
135 + "content": prompt,
136 + }
137 + ],
138 extra_body={
139 "guided_regex": r"\w+@\w+\.com\n",
140 "stop": ["\n"],
examples/online_serving/openai_chat_completion_structured_outputs_structural_tag.py
+21 -23
@@ -17,11 +17,10 @@ def main():
17 api_key=openai_api_key,
18 )
19
20 - messages = [{
21 - "role":
22 - "user",
23 - "content":
24 - """
20 + messages = [
21 + {
22 + "role": "user",
23 + "content": """
24 You have access to the following function to retrieve the weather in a city:
25
26 {
@@ -58,29 +57,28 @@ You are a helpful assistant.
57
58 Given the previous instructions, what is the weather in New York City, Boston,
59 and San Francisco?
61 -"""
62 - }]
60 +""",
61 + }
62 + ]
63
64 response = client.chat.completions.create(
65 model=client.models.list().data[0].id,
66 messages=messages,
67 response_format={
68 - "type":
69 - "structural_tag",
70 - "structures": [{
71 - "begin": "<function=get_weather>",
72 - "schema": {
73 - "type": "object",
74 - "properties": {
75 - "city": {
76 - "type": "string"
77 - }
78 - }
79 - },
80 - "end": "</function>"
81 - }],
82 - "triggers": ["<function="]
83 - })
68 + "type": "structural_tag",
69 + "structures": [
70 + {
71 + "begin": "<function=get_weather>",
72 + "schema": {
73 + "type": "object",
74 + "properties": {"city": {"type": "string"}},
75 + },
76 + "end": "</function>",
77 + }
78 + ],
79 + "triggers": ["<function="],
80 + },
81 + )
82 print(response)
83
84
examples/online_serving/openai_chat_completion_structured_outputs_with_reasoning.py
+35 -24
@@ -27,21 +27,22 @@ openai_api_base = "http://localhost:8000/v1"
27
28
29 def print_completion_details(completion):
30 - print("reasoning_content: ",
31 - completion.choices[0].message.reasoning_content)
30 + print("reasoning_content: ", completion.choices[0].message.reasoning_content)
31 print("content: ", completion.choices[0].message.content)
32
33
34 # Guided decoding by Regex
35 def guided_regex_completion(client: OpenAI, model: str):
37 - prompt = ("What is the capital of France?")
36 + prompt = "What is the capital of France?"
37
38 completion = client.chat.completions.create(
39 model=model,
41 - messages=[{
42 - "role": "user",
43 - "content": prompt,
44 - }],
40 + messages=[
41 + {
42 + "role": "user",
43 + "content": prompt,
44 + }
45 + ],
46 extra_body={
47 "guided_regex": "(Paris|London)",
48 },
@@ -57,13 +58,15 @@ class People(BaseModel):
58 def guided_json_completion(client: OpenAI, model: str):
59 json_schema = People.model_json_schema()
60
60 - prompt = ("Generate a JSON with the name and age of one random person.")
61 + prompt = "Generate a JSON with the name and age of one random person."
62 completion = client.chat.completions.create(
63 model=model,
63 - messages=[{
64 - "role": "user",
65 - "content": prompt,
66 - }],
64 + messages=[
65 + {
66 + "role": "user",
67 + "content": prompt,
68 + }
69 + ],
70 extra_body={"guided_json": json_schema},
71 )
72 print_completion_details(completion)
@@ -86,14 +89,18 @@ class CarDescription(BaseModel):
89 def guided_car_json_completion(client: OpenAI, model: str):
90 json_schema = CarDescription.model_json_schema()
91
89 - prompt = ("Generate a JSON with the brand, model and car_type of"
90 - "the most iconic car from the 90's")
92 + prompt = (
93 + "Generate a JSON with the brand, model and car_type of"
94 + "the most iconic car from the 90's"
95 + )
96 completion = client.chat.completions.create(
97 model=model,
93 - messages=[{
94 - "role": "user",
95 - "content": prompt,
96 - }],
98 + messages=[
99 + {
100 + "role": "user",
101 + "content": prompt,
102 + }
103 + ],
104 extra_body={"guided_json": json_schema},
105 )
106 print_completion_details(completion)
@@ -116,14 +123,18 @@ def guided_grammar_completion(client: OpenAI, model: str):
123 """
124
125 # This may be very slow https://github.com/vllm-project/vllm/issues/12122
119 - prompt = ("Generate an SQL query to show the 'username' and 'email'"
120 - "from the 'users' table.")
126 + prompt = (
127 + "Generate an SQL query to show the 'username' and 'email'"
128 + "from the 'users' table."
129 + )
130 completion = client.chat.completions.create(
131 model=model,
123 - messages=[{
124 - "role": "user",
125 - "content": prompt,
126 - }],
132 + messages=[
133 + {
134 + "role": "user",
135 + "content": prompt,
136 + }
137 + ],
138 extra_body={"guided_grammar": simplified_sql_grammar},
139 )
140 print_completion_details(completion)
examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
+70 -88
@@ -20,9 +20,11 @@ from openai import OpenAI
20
21
22 # Now, simulate a tool call
23 -def get_current_weather(city: str, state: str, unit: 'str'):
24 - return ("The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
25 - "partly cloudly, with highs in the 90's.")
23 +def get_current_weather(city: str, state: str, unit: "str"):
24 + return (
25 + "The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
26 + "partly cloudly, with highs in the 90's."
27 + )
28
29
30 available_tools = {"get_current_weather": get_current_weather}
@@ -31,49 +33,47 @@ available_tools = {"get_current_weather": get_current_weather}
33 openai_api_key = "EMPTY"
34 openai_api_base = "http://localhost:8000/v1"
35
34 -tools = [{
35 - "type": "function",
36 - "function": {
37 - "name": "get_current_weather",
38 - "description": "Get the current weather in a given location",
39 - "parameters": {
40 - "type": "object",
41 - "properties": {
42 - "city": {
43 - "type":
44 - "string",
45 - "description":
46 - "The city to find the weather for, e.g. 'San Francisco'"
47 - },
48 - "state": {
49 - "type":
50 - "string",
51 - "description":
52 - "the two-letter abbreviation for the state that the city is"
53 - " in, e.g. 'CA' which would mean 'California'"
54 - },
55 - "unit": {
56 - "type": "string",
57 - "description": "The unit to fetch the temperature in",
58 - "enum": ["celsius", "fahrenheit"]
59 - }
36 +properties = {
37 + "city": {
38 + "type": "string",
39 + "description": "The city to find the weather for, e.g. 'San Francisco'",
40 + },
41 + "state": {
42 + "type": "string",
43 + "description": "the two-letter abbreviation for the state that the city is"
44 + " in, e.g. 'CA' which would mean 'California'",
45 + },
46 + "unit": {
47 + "type": "string",
48 + "description": "The unit to fetch the temperature in",
49 + "enum": ["celsius", "fahrenheit"],
50 + },
51 +}
52 +
53 +tools = [
54 + {
55 + "type": "function",
56 + "function": {
57 + "name": "get_current_weather",
58 + "description": "Get the current weather in a given location",
59 + "parameters": {
60 + "type": "object",
61 + "properties": properties,
62 + "required": ["city", "state", "unit"],
63 },
61 - "required": ["city", "state", "unit"]
62 - }
64 + },
65 }
64 -}]
65 -messages = [{
66 - "role": "user",
67 - "content": "Hi! How are you doing today?"
68 -}, {
69 - "role": "assistant",
70 - "content": "I'm doing well! How can I help you?"
71 -}, {
72 - "role":
73 - "user",
74 - "content":
75 - "Can you tell me what the temperate will be in Dallas, in fahrenheit?"
76 -}]
66 +]
67 +messages = [
68 + {"role": "user", "content": "Hi! How are you doing today?"},
69 + {"role": "assistant", "content": "I'm doing well! How can I help you?"},
70 + {
71 + "role": "user",
72 + "content": (
73 + "Can you tell me what the temperate will be in Dallas, in fahrenheit?"
74 + ),
75 + },
76 +]
77
78
79 def extract_reasoning_and_calls(chunks: list):
@@ -110,73 +110,55 @@ def main():
110 models = client.models.list()
111 model = models.data[0].id
112
113 + print("---------Full Generate With Automatic Function Calling-------------")
114 + tool_calls = client.chat.completions.create(
115 + messages=messages, model=model, tools=tools
116 + )
117 + print(f"reasoning_content: {tool_calls.choices[0].message.reasoning_content}")
118 + print(f"function name: {tool_calls.choices[0].message.tool_calls[0].function.name}")
119 print(
114 - "---------Full Generate With Automatic Function Calling-------------")
115 - tool_calls = client.chat.completions.create(messages=messages,
116 - model=model,
117 - tools=tools)
118 - print(
119 - f"reasoning_content: {tool_calls.choices[0].message.reasoning_content}"
120 + f"function arguments: "
121 + f"{tool_calls.choices[0].message.tool_calls[0].function.arguments}"
122 )
121 - print(f"function name: "
122 - f"{tool_calls.choices[0].message.tool_calls[0].function.name}")
123 - print(f"function arguments: "
124 - f"{tool_calls.choices[0].message.tool_calls[0].function.arguments}")
123
126 - print(
127 - "----------Stream Generate With Automatic Function Calling-----------")
128 - tool_calls_stream = client.chat.completions.create(messages=messages,
129 - model=model,
130 - tools=tools,
131 - stream=True)
124 + print("----------Stream Generate With Automatic Function Calling-----------")
125 + tool_calls_stream = client.chat.completions.create(
126 + messages=messages, model=model, tools=tools, stream=True
127 + )
128
129 chunks = list(tool_calls_stream)
130
135 - reasoning_content, arguments, function_names = extract_reasoning_and_calls(
136 - chunks)
131 + reasoning_content, arguments, function_names = extract_reasoning_and_calls(chunks)
132
133 print(f"reasoning_content: {reasoning_content}")
134 print(f"function name: {function_names[0]}")
135 print(f"function arguments: {arguments[0]}")
136
142 - print(
143 - "----------Full Generate With Named Function Calling-----------------")
144 - tool_calls = client.chat.completions.create(messages=messages,
145 - model=model,
146 - tools=tools,
147 - tool_choice={
148 - "type": "function",
149 - "function": {
150 - "name":
151 - "get_current_weather"
152 - }
153 - })
137 + print("----------Full Generate With Named Function Calling-----------------")
138 + tool_calls = client.chat.completions.create(
139 + messages=messages,
140 + model=model,
141 + tools=tools,
142 + tool_choice={"type": "function", "function": {"name": "get_current_weather"}},
143 + )
144
145 tool_call = tool_calls.choices[0].message.tool_calls[0].function
156 - print(
157 - f"reasoning_content: {tool_calls.choices[0].message.reasoning_content}"
158 - )
146 + print(f"reasoning_content: {tool_calls.choices[0].message.reasoning_content}")
147 print(f"function name: {tool_call.name}")
148 print(f"function arguments: {tool_call.arguments}")
161 - print(
162 - "----------Stream Generate With Named Function Calling--------------")
149 + print("----------Stream Generate With Named Function Calling--------------")
150
151 tool_calls_stream = client.chat.completions.create(
152 messages=messages,
153 model=model,
154 tools=tools,
168 - tool_choice={
169 - "type": "function",
170 - "function": {
171 - "name": "get_current_weather"
172 - }
173 - },
174 - stream=True)
155 + tool_choice={"type": "function", "function": {"name": "get_current_weather"}},
156 + stream=True,
157 + )
158
159 chunks = list(tool_calls_stream)
160
178 - reasoning_content, arguments, function_names = extract_reasoning_and_calls(
179 - chunks)
161 + reasoning_content, arguments, function_names = extract_reasoning_and_calls(chunks)
162 print(f"reasoning_content: {reasoning_content}")
163 print(f"function name: {function_names[0]}")
164 print(f"function arguments: {arguments[0]}")
examples/online_serving/openai_chat_completion_with_reasoning.py
+6 -6
@@ -45,12 +45,12 @@ def main():
45
46 # Round 2
47 messages.append({"role": "assistant", "content": content})
48 - messages.append({
49 - "role":
50 - "user",
51 - "content":
52 - "How many Rs are there in the word 'strawberry'?",
53 - })
48 + messages.append(
49 + {
50 + "role": "user",
51 + "content": "How many Rs are there in the word 'strawberry'?",
52 + }
53 + )
54 response = client.chat.completions.create(model=model, messages=messages)
55
56 reasoning_content = response.choices[0].message.reasoning_content
examples/online_serving/openai_chat_completion_with_reasoning_streaming.py
+1 -3
@@ -43,9 +43,7 @@ def main():
43
44 # ruff: noqa: E501
45 # For granite: add: `extra_body={"chat_template_kwargs": {"thinking": True}}`
46 - stream = client.chat.completions.create(model=model,
47 - messages=messages,
48 - stream=True)
46 + stream = client.chat.completions.create(model=model, messages=messages, stream=True)
47
48 print("client: Start streaming chat completions...")
49 printed_reasoning_content = False
examples/online_serving/openai_chat_embedding_client_for_multimodal.py
+62 -65
@@ -14,26 +14,17 @@ def vlm2vec():
14 response = requests.post(
15 "http://localhost:8000/v1/embeddings",
16 json={
17 - "model":
18 - "TIGER-Lab/VLM2Vec-Full",
19 - "messages": [{
20 - "role":
21 - "user",
22 - "content": [
23 - {
24 - "type": "image_url",
25 - "image_url": {
26 - "url": image_url
27 - }
28 - },
29 - {
30 - "type": "text",
31 - "text": "Represent the given image."
32 - },
33 - ],
34 - }],
35 - "encoding_format":
36 - "float",
17 + "model": "TIGER-Lab/VLM2Vec-Full",
18 + "messages": [
19 + {
20 + "role": "user",
21 + "content": [
22 + {"type": "image_url", "image_url": {"url": image_url}},
23 + {"type": "text", "text": "Represent the given image."},
24 + ],
25 + }
26 + ],
27 + "encoding_format": "float",
28 },
29 )
30 response.raise_for_status()
@@ -45,19 +36,20 @@ def vlm2vec():
36 def dse_qwen2_vl(inp: dict):
37 # Embedding an Image
38 if inp["type"] == "image":
48 - messages = [{
49 - "role":
50 - "user",
51 - "content": [{
52 - "type": "image_url",
53 - "image_url": {
54 - "url": inp["image_url"],
55 - }
56 - }, {
57 - "type": "text",
58 - "text": "What is shown in this image?"
59 - }]
60 - }]
39 + messages = [
40 + {
41 + "role": "user",
42 + "content": [
43 + {
44 + "type": "image_url",
45 + "image_url": {
46 + "url": inp["image_url"],
47 + },
48 + },
49 + {"type": "text", "text": "What is shown in this image?"},
50 + ],
51 + }
52 + ]
53 # Embedding a Text Query
54 else:
55 # MrLight/dse-qwen2-2b-mrl-v1 requires a placeholder image
@@ -66,23 +58,21 @@ def dse_qwen2_vl(inp: dict):
58 image_placeholder = Image.new("RGB", (56, 56))
59 image_placeholder.save(buffer, "png")
60 buffer.seek(0)
69 - image_placeholder = base64.b64encode(buffer.read()).decode('utf-8')
70 - messages = [{
71 - "role":
72 - "user",
73 - "content": [
74 - {
75 - "type": "image_url",
76 - "image_url": {
77 - "url": f"data:image/jpeg;base64,{image_placeholder}",
78 - }
79 - },
80 - {
81 - "type": "text",
82 - "text": f"Query: {inp['content']}"
83 - },
84 - ]
85 - }]
61 + image_placeholder = base64.b64encode(buffer.read()).decode("utf-8")
62 + messages = [
63 + {
64 + "role": "user",
65 + "content": [
66 + {
67 + "type": "image_url",
68 + "image_url": {
69 + "url": f"data:image/jpeg;base64,{image_placeholder}",
70 + },
71 + },
72 + {"type": "text", "text": f"Query: {inp['content']}"},
73 + ],
74 + }
75 + ]
76
77 response = requests.post(
78 "http://localhost:8000/v1/embeddings",
@@ -101,12 +91,15 @@ def dse_qwen2_vl(inp: dict):
91 def parse_args():
92 parser = argparse.ArgumentParser(
93 "Script to call a specified VLM through the API. Make sure to serve "
104 - "the model with --task embed before running this.")
105 - parser.add_argument("--model",
106 - type=str,
107 - choices=["vlm2vec", "dse_qwen2_vl"],
108 - required=True,
109 - help="Which model to call.")
94 + "the model with --task embed before running this."
95 + )
96 + parser.add_argument(
97 + "--model",
98 + type=str,
99 + choices=["vlm2vec", "dse_qwen2_vl"],
100 + required=True,
101 + help="Which model to call.",
102 + )
103 return parser.parse_args()
104
105
@@ -114,16 +107,20 @@ def main(args):
107 if args.model == "vlm2vec":
108 vlm2vec()
109 elif args.model == "dse_qwen2_vl":
117 - dse_qwen2_vl({
118 - "type": "image",
119 - "image_url": image_url,
120 - })
121 - dse_qwen2_vl({
122 - "type": "text",
123 - "content": "What is the weather like today?",
124 - })
110 + dse_qwen2_vl(
111 + {
112 + "type": "image",
113 + "image_url": image_url,
114 + }
115 + )
116 + dse_qwen2_vl(
117 + {
118 + "type": "text",
119 + "content": "What is the weather like today?",
120 + }
121 + )
122
123
127 -if __name__ == '__main__':
124 +if __name__ == "__main__":
125 args = parse_args()
126 main(args)
examples/online_serving/openai_classification_client.py
+1 -3
@@ -16,9 +16,7 @@ def parse_args():
16 parse = argparse.ArgumentParser()
17 parse.add_argument("--host", type=str, default="localhost")
18 parse.add_argument("--port", type=int, default=8000)
19 - parse.add_argument("--model",
20 - type=str,
21 - default="jason9693/Qwen2.5-1.5B-apeach")
19 + parse.add_argument("--model", type=str, default="jason9693/Qwen2.5-1.5B-apeach")
20 return parse.parse_args()
21
22
examples/online_serving/openai_completion_client.py
+5 -4
@@ -11,9 +11,9 @@ openai_api_base = "http://localhost:8000/v1"
11
12 def parse_args():
13 parser = argparse.ArgumentParser(description="Client for vLLM API server")
14 - parser.add_argument("--stream",
15 - action="store_true",
16 - help="Enable streaming response")
14 + parser.add_argument(
15 + "--stream", action="store_true", help="Enable streaming response"
16 + )
17 return parser.parse_args()
18
19
@@ -34,7 +34,8 @@ def main(args):
34 echo=False,
35 n=2,
36 stream=args.stream,
37 - logprobs=3)
37 + logprobs=3,
38 + )
39
40 print("-" * 50)
41 print("Completion results:")
examples/online_serving/openai_cross_encoder_score.py
+4 -9
@@ -4,6 +4,7 @@ Example online usage of Score API.
4
5 Run `vllm serve <model> --task score` to start up the server in vLLM.
6 """
7 +
8 import argparse
9 import pprint
10
@@ -38,9 +39,7 @@ def main(args):
39 pprint.pprint(score_response.json())
40
41 text_1 = "What is the capital of France?"
41 - text_2 = [
42 - "The capital of Brazil is Brasilia.", "The capital of France is Paris."
43 - ]
42 + text_2 = ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]
43 prompt = {"model": model_name, "text_1": text_1, "text_2": text_2}
44 score_response = post_http_request(prompt=prompt, api_url=api_url)
45 print("\nPrompt when text_1 is string and text_2 is a list:")
@@ -48,12 +47,8 @@ def main(args):
47 print("\nScore Response:")
48 pprint.pprint(score_response.json())
49
51 - text_1 = [
52 - "What is the capital of Brazil?", "What is the capital of France?"
53 - ]
54 - text_2 = [
55 - "The capital of Brazil is Brasilia.", "The capital of France is Paris."
56 - ]
50 + text_1 = ["What is the capital of Brazil?", "What is the capital of France?"]
51 + text_2 = ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]
52 prompt = {"model": model_name, "text_1": text_1, "text_2": text_2}
53 score_response = post_http_request(prompt=prompt, api_url=api_url)
54 print("\nPrompt when text_1 and text_2 are both lists:")
examples/online_serving/openai_embedding_client.py
+1 -1
@@ -21,7 +21,7 @@ def main():
21 # ruff: noqa: E501
22 input=[
23 "Hello my name is",
24 - "The best thing about vLLM is that it supports many different models"
24 + "The best thing about vLLM is that it supports many different models",
25 ],
26 model=model,
27 )
examples/online_serving/openai_pooling_client.py
+9 -12
@@ -5,6 +5,7 @@ Example online usage of Pooling API.
5 Run `vllm serve <model> --task <embed|classify|reward|score>`
6 to start up the server in vLLM.
7 """
8 +
9 import argparse
10 import pprint
11
@@ -21,9 +22,7 @@ def parse_args():
22 parser = argparse.ArgumentParser()
23 parser.add_argument("--host", type=str, default="localhost")
24 parser.add_argument("--port", type=int, default=8000)
24 - parser.add_argument("--model",
25 - type=str,
26 - default="jason9693/Qwen2.5-1.5B-apeach")
25 + parser.add_argument("--model", type=str, default="jason9693/Qwen2.5-1.5B-apeach")
26
27 return parser.parse_args()
28
@@ -42,15 +41,13 @@ def main(args):
41
42 # Input like Chat API
43 prompt = {
45 - "model":
46 - model_name,
47 - "messages": [{
48 - "role": "user",
49 - "content": [{
50 - "type": "text",
51 - "text": "vLLM is great!"
52 - }],
53 - }]
44 + "model": model_name,
45 + "messages": [
46 + {
47 + "role": "user",
48 + "content": [{"type": "text", "text": "vLLM is great!"}],
49 + }
50 + ],
51 }
52 pooling_response = post_http_request(prompt=prompt, api_url=api_url)
53 print("Pooling Response:")
examples/online_serving/openai_transcription_client.py
+14 -16
@@ -7,8 +7,8 @@ from openai import OpenAI
7
8 from vllm.assets.audio import AudioAsset
9
10 -mary_had_lamb = AudioAsset('mary_had_lamb').get_local_path()
11 -winning_call = AudioAsset('winning_call').get_local_path()
10 +mary_had_lamb = AudioAsset("mary_had_lamb").get_local_path()
11 +winning_call = AudioAsset("winning_call").get_local_path()
12
13 # Modify OpenAI's API key and API base to use vLLM's API server.
14 openai_api_key = "EMPTY"
@@ -31,7 +31,8 @@ def sync_openai():
31 extra_body=dict(
32 seed=4419,
33 repetition_penalty=1.3,
34 - ))
34 + ),
35 + )
36 print("transcription result:", transcription.text)
37
38
@@ -42,33 +43,30 @@ sync_openai()
43 async def stream_openai_response():
44 data = {
45 "language": "en",
45 - 'stream': True,
46 + "stream": True,
47 "model": "openai/whisper-large-v3",
48 }
49 url = openai_api_base + "/audio/transcriptions"
50 headers = {"Authorization": f"Bearer {openai_api_key}"}
50 - print("transcription result:", end=' ')
51 + print("transcription result:", end=" ")
52 async with httpx.AsyncClient() as client:
53 with open(str(winning_call), "rb") as f:
53 - async with client.stream('POST',
54 - url,
55 - files={'file': f},
56 - data=data,
57 - headers=headers) as response:
54 + async with client.stream(
55 + "POST", url, files={"file": f}, data=data, headers=headers
56 + ) as response:
57 async for line in response.aiter_lines():
58 # Each line is a JSON object prefixed with 'data: '
59 if line:
61 - if line.startswith('data: '):
62 - line = line[len('data: '):]
60 + if line.startswith("data: "):
61 + line = line[len("data: ") :]
62 # Last chunk, stream ends
64 - if line.strip() == '[DONE]':
63 + if line.strip() == "[DONE]":
64 break
65 # Parse the JSON response
66 chunk = json.loads(line)
67 # Extract and print the content
69 - content = chunk['choices'][0].get('delta',
70 - {}).get('content')
71 - print(content, end='')
68 + content = chunk["choices"][0].get("delta", {}).get("content")
69 + print(content, end="")
70
71
72 # Run the asynchronous function
examples/online_serving/opentelemetry/dummy_client.py
+3 -6
@@ -1,14 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 import requests
4 -from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
5 - OTLPSpanExporter)
4 +from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
5 from opentelemetry.sdk.trace import TracerProvider
7 -from opentelemetry.sdk.trace.export import (BatchSpanProcessor,
8 - ConsoleSpanExporter)
6 +from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
7 from opentelemetry.trace import SpanKind, set_tracer_provider
10 -from opentelemetry.trace.propagation.tracecontext import (
11 - TraceContextTextMapPropagator)
8 +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
9
10 trace_provider = TracerProvider()
11 set_tracer_provider(trace_provider)
examples/online_serving/prompt_embed_inference_with_openai_client.py
+9 -11
@@ -26,6 +26,7 @@ Dependencies:
26 - torch
27 - openai
28 """
29 +
30 import base64
31 import io
32
@@ -44,17 +45,13 @@ def main():
45
46 # Transformers
47 tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
47 - transformers_model = transformers.AutoModelForCausalLM.from_pretrained(
48 - model_name)
48 + transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
49
50 # Refer to the HuggingFace repo for the correct format to use
51 - chat = [{
52 - "role": "user",
53 - "content": "Please tell me about the capital of France."
54 - }]
55 - token_ids = tokenizer.apply_chat_template(chat,
56 - add_generation_prompt=True,
57 - return_tensors='pt')
51 + chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
52 + token_ids = tokenizer.apply_chat_template(
53 + chat, add_generation_prompt=True, return_tensors="pt"
54 + )
55
56 embedding_layer = transformers_model.get_input_embeddings()
57 prompt_embeds = embedding_layer(token_ids).squeeze(0)
@@ -64,7 +61,7 @@ def main():
61 torch.save(prompt_embeds, buffer)
62 buffer.seek(0)
63 binary_data = buffer.read()
67 - encoded_embeds = base64.b64encode(binary_data).decode('utf-8')
64 + encoded_embeds = base64.b64encode(binary_data).decode("utf-8")
65
66 completion = client.completions.create(
67 model=model_name,
@@ -75,7 +72,8 @@ def main():
72 temperature=0.0,
73 # NOTE: The OpenAI client allows passing in extra JSON body via the
74 # `extra_body` argument.
78 - extra_body={"prompt_embeds": encoded_embeds})
75 + extra_body={"prompt_embeds": encoded_embeds},
76 + )
77
78 print("-" * 30)
79 print(completion.choices[0].text)
examples/online_serving/ray_serve_deepseek.py
+1 -3
@@ -28,9 +28,7 @@ llm_config = LLMConfig(
28 },
29 # Change to the accelerator type of the node
30 accelerator_type="H100",
31 - runtime_env={"env_vars": {
32 - "VLLM_USE_V1": "1"
33 - }},
31 + runtime_env={"env_vars": {"VLLM_USE_V1": "1"}},
32 # Customize engine arguments as needed (e.g. vLLM engine kwargs)
33 engine_kwargs={
34 "tensor_parallel_size": 8,
examples/online_serving/retrieval_augmented_generation_with_langchain.py
+60 -53
@@ -55,7 +55,7 @@ def load_and_split_documents(config: dict[str, Any]):
55 Load and split documents from web URL
56 """
57 try:
58 - loader = WebBaseLoader(web_paths=(config["url"], ))
58 + loader = WebBaseLoader(web_paths=(config["url"],))
59 docs = loader.load()
60
61 text_splitter = RecursiveCharacterTextSplitter(
@@ -121,64 +121,71 @@ def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate):
121 """
122 Set up question answering chain
123 """
124 - return ({
125 - "context": retriever | format_docs,
126 - "question": RunnablePassthrough(),
127 - }
128 - | prompt
129 - | llm
130 - | StrOutputParser())
124 + return (
125 + {
126 + "context": retriever | format_docs,
127 + "question": RunnablePassthrough(),
128 + }
129 + | prompt
130 + | llm
131 + | StrOutputParser()
132 + )
133
134
135 def get_parser() -> argparse.ArgumentParser:
136 """
137 Parse command line arguments
138 """
137 - parser = argparse.ArgumentParser(description='RAG with vLLM and langchain')
139 + parser = argparse.ArgumentParser(description="RAG with vLLM and langchain")
140
141 # Add command line arguments
140 - parser.add_argument('--vllm-api-key',
141 - default="EMPTY",
142 - help='API key for vLLM compatible services')
143 - parser.add_argument('--vllm-embedding-endpoint',
144 - default="http://localhost:8000/v1",
145 - help='Base URL for embedding service')
146 - parser.add_argument('--vllm-chat-endpoint',
147 - default="http://localhost:8001/v1",
148 - help='Base URL for chat service')
149 - parser.add_argument('--uri',
150 - default="./milvus.db",
151 - help='URI for Milvus database')
142 parser.add_argument(
153 - '--url',
154 - default=("https://docs.vllm.ai/en/latest/getting_started/"
155 - "quickstart.html"),
156 - help='URL of the document to process')
157 - parser.add_argument('--embedding-model',
158 - default="ssmits/Qwen2-7B-Instruct-embed-base",
159 - help='Model name for embeddings')
160 - parser.add_argument('--chat-model',
161 - default="qwen/Qwen1.5-0.5B-Chat",
162 - help='Model name for chat')
163 - parser.add_argument('-i',
164 - '--interactive',
165 - action='store_true',
166 - help='Enable interactive Q&A mode')
167 - parser.add_argument('-k',
168 - '--top-k',
169 - type=int,
170 - default=3,
171 - help='Number of top results to retrieve')
172 - parser.add_argument('-c',
173 - '--chunk-size',
174 - type=int,
175 - default=1000,
176 - help='Chunk size for document splitting')
177 - parser.add_argument('-o',
178 - '--chunk-overlap',
179 - type=int,
180 - default=200,
181 - help='Chunk overlap for document splitting')
143 + "--vllm-api-key", default="EMPTY", help="API key for vLLM compatible services"
144 + )
145 + parser.add_argument(
146 + "--vllm-embedding-endpoint",
147 + default="http://localhost:8000/v1",
148 + help="Base URL for embedding service",
149 + )
150 + parser.add_argument(
151 + "--vllm-chat-endpoint",
152 + default="http://localhost:8001/v1",
153 + help="Base URL for chat service",
154 + )
155 + parser.add_argument("--uri", default="./milvus.db", help="URI for Milvus database")
156 + parser.add_argument(
157 + "--url",
158 + default=("https://docs.vllm.ai/en/latest/getting_started/quickstart.html"),
159 + help="URL of the document to process",
160 + )
161 + parser.add_argument(
162 + "--embedding-model",
163 + default="ssmits/Qwen2-7B-Instruct-embed-base",
164 + help="Model name for embeddings",
165 + )
166 + parser.add_argument(
167 + "--chat-model", default="qwen/Qwen1.5-0.5B-Chat", help="Model name for chat"
168 + )
169 + parser.add_argument(
170 + "-i", "--interactive", action="store_true", help="Enable interactive Q&A mode"
171 + )
172 + parser.add_argument(
173 + "-k", "--top-k", type=int, default=3, help="Number of top results to retrieve"
174 + )
175 + parser.add_argument(
176 + "-c",
177 + "--chunk-size",
178 + type=int,
179 + default=1000,
180 + help="Chunk size for document splitting",
181 + )
182 + parser.add_argument(
183 + "-o",
184 + "--chunk-overlap",
185 + type=int,
186 + default=200,
187 + help="Chunk overlap for document splitting",
188 + )
189
190 return parser
191
@@ -198,7 +205,7 @@ def init_config(args: Namespace):
205 "url": args.url,
206 "chunk_size": args.chunk_size,
207 "chunk_overlap": args.chunk_overlap,
201 - "top_k": args.top_k
208 + "top_k": args.top_k,
209 }
210
211
@@ -230,7 +237,7 @@ def main():
237
238 while True:
239 question = input("\nPlease enter your question: ")
233 - if question.lower() in ['q', 'quit']:
240 + if question.lower() in ["q", "quit"]:
241 print("\nThank you for using! Goodbye!")
242 break
243
@@ -238,7 +245,7 @@ def main():
245 print(output)
246 else:
247 # Default single question mode
241 - question = ("How to install vLLM?")
248 + question = "How to install vLLM?"
249 output = qa_chain.invoke(question)
250 print("-" * 50)
251 print(output)
examples/online_serving/retrieval_augmented_generation_with_llamaindex.py
+52 -45
@@ -35,6 +35,7 @@ Notes:
35 - Default ports: 8000 (embedding), 8001 (chat)
36 - First run may take time to download models
37 """
38 +
39 import argparse
40 from argparse import Namespace
41 from typing import Any
@@ -59,7 +60,7 @@ def init_config(args: Namespace):
60 "db_path": args.db_path,
61 "chunk_size": args.chunk_size,
62 "chunk_overlap": args.chunk_overlap,
62 - "top_k": args.top_k
63 + "top_k": args.top_k,
64 }
65
66
@@ -117,52 +118,58 @@ def query_document(index: VectorStoreIndex, question: str, top_k: int):
118
119 def get_parser() -> argparse.ArgumentParser:
120 """Parse command line arguments"""
120 - parser = argparse.ArgumentParser(
121 - description='RAG with vLLM and LlamaIndex')
121 + parser = argparse.ArgumentParser(description="RAG with vLLM and LlamaIndex")
122
123 # Add command line arguments
124 parser.add_argument(
125 - '--url',
126 - default=("https://docs.vllm.ai/en/latest/getting_started/"
127 - "quickstart.html"),
128 - help='URL of the document to process')
129 - parser.add_argument('--embedding-model',
130 - default="ssmits/Qwen2-7B-Instruct-embed-base",
131 - help='Model name for embeddings')
132 - parser.add_argument('--chat-model',
133 - default="qwen/Qwen1.5-0.5B-Chat",
134 - help='Model name for chat')
135 - parser.add_argument('--vllm-api-key',
136 - default="EMPTY",
137 - help='API key for vLLM compatible services')
138 - parser.add_argument('--embedding-endpoint',
139 - default="http://localhost:8000/v1",
140 - help='Base URL for embedding service')
141 - parser.add_argument('--chat-endpoint',
142 - default="http://localhost:8001/v1",
143 - help='Base URL for chat service')
144 - parser.add_argument('--db-path',
145 - default="./milvus_demo.db",
146 - help='Path to Milvus database')
147 - parser.add_argument('-i',
148 - '--interactive',
149 - action='store_true',
150 - help='Enable interactive Q&A mode')
151 - parser.add_argument('-c',
152 - '--chunk-size',
153 - type=int,
154 - default=1000,
155 - help='Chunk size for document splitting')
156 - parser.add_argument('-o',
157 - '--chunk-overlap',
158 - type=int,
159 - default=200,
160 - help='Chunk overlap for document splitting')
161 - parser.add_argument('-k',
162 - '--top-k',
163 - type=int,
164 - default=3,
165 - help='Number of top results to retrieve')
125 + "--url",
126 + default=("https://docs.vllm.ai/en/latest/getting_started/quickstart.html"),
127 + help="URL of the document to process",
128 + )
129 + parser.add_argument(
130 + "--embedding-model",
131 + default="ssmits/Qwen2-7B-Instruct-embed-base",
132 + help="Model name for embeddings",
133 + )
134 + parser.add_argument(
135 + "--chat-model", default="qwen/Qwen1.5-0.5B-Chat", help="Model name for chat"
136 + )
137 + parser.add_argument(
138 + "--vllm-api-key", default="EMPTY", help="API key for vLLM compatible services"
139 + )
140 + parser.add_argument(
141 + "--embedding-endpoint",
142 + default="http://localhost:8000/v1",
143 + help="Base URL for embedding service",
144 + )
145 + parser.add_argument(
146 + "--chat-endpoint",
147 + default="http://localhost:8001/v1",
148 + help="Base URL for chat service",
149 + )
150 + parser.add_argument(
151 + "--db-path", default="./milvus_demo.db", help="Path to Milvus database"
152 + )
153 + parser.add_argument(
154 + "-i", "--interactive", action="store_true", help="Enable interactive Q&A mode"
155 + )
156 + parser.add_argument(
157 + "-c",
158 + "--chunk-size",
159 + type=int,
160 + default=1000,
161 + help="Chunk size for document splitting",
162 + )
163 + parser.add_argument(
164 + "-o",
165 + "--chunk-overlap",
166 + type=int,
167 + default=200,
168 + help="Chunk overlap for document splitting",
169 + )
170 + parser.add_argument(
171 + "-k", "--top-k", type=int, default=3, help="Number of top results to retrieve"
172 + )
173
174 return parser
175
@@ -193,7 +200,7 @@ def main():
200 question = input("\nEnter your question: ")
201
202 # Check for exit command
196 - if question.lower() in ['quit', 'exit', 'q']:
203 + if question.lower() in ["quit", "exit", "q"]:
204 print("Exiting interactive mode...")
205 break
206
examples/online_serving/streamlit_openai_chatbot_webserver.py
+29 -26
@@ -26,6 +26,7 @@ Usage:
26 streamlit run streamlit_openai_chatbot_webserver.py \
27 --logger.level=debug
28 """
29 +
30 import os
31 from datetime import datetime
32
@@ -33,8 +34,8 @@ import streamlit as st
34 from openai import OpenAI
35
36 # Get command line arguments from environment variables
36 -openai_api_key = os.getenv('VLLM_API_KEY', "EMPTY")
37 -openai_api_base = os.getenv('VLLM_API_BASE', "http://localhost:8000/v1")
37 +openai_api_key = os.getenv("VLLM_API_KEY", "EMPTY")
38 +openai_api_base = os.getenv("VLLM_API_BASE", "http://localhost:8000/v1")
39
40 # Initialize session states for managing chat sessions
41 if "sessions" not in st.session_state:
@@ -81,9 +82,9 @@ def get_llm_response(messages, model):
82 Streaming response object or error message string
83 """
84 try:
84 - response = client.chat.completions.create(model=model,
85 - messages=messages,
86 - stream=True)
85 + response = client.chat.completions.create(
86 + model=model, messages=messages, stream=True
87 + )
88 return response
89 except Exception as e:
90 st.error(f"Error details: {str(e)}")
@@ -92,8 +93,9 @@ def get_llm_response(messages, model):
93
94 # Sidebar - API Settings first
95 st.sidebar.title("API Settings")
95 -new_api_base = st.sidebar.text_input("API Base URL:",
96 - value=st.session_state.api_base_url)
96 +new_api_base = st.sidebar.text_input(
97 + "API Base URL:", value=st.session_state.api_base_url
98 +)
99 if new_api_base != st.session_state.api_base_url:
100 st.session_state.api_base_url = new_api_base
101 st.rerun()
@@ -109,16 +111,20 @@ if st.sidebar.button("New Session"):
111 for session_id in sorted(st.session_state.sessions.keys(), reverse=True):
112 # Mark the active session with a pinned button
113 if session_id == st.session_state.active_session:
112 - st.sidebar.button(f"📍 {session_id}",
113 - key=session_id,
114 - type="primary",
115 - on_click=switch_to_chat_session,
116 - args=(session_id, ))
114 + st.sidebar.button(
115 + f"📍 {session_id}",
116 + key=session_id,
117 + type="primary",
118 + on_click=switch_to_chat_session,
119 + args=(session_id,),
120 + )
121 else:
118 - st.sidebar.button(f"Session {session_id}",
119 - key=session_id,
120 - on_click=switch_to_chat_session,
121 - args=(session_id, ))
122 + st.sidebar.button(
123 + f"Session {session_id}",
124 + key=session_id,
125 + on_click=switch_to_chat_session,
126 + args=(session_id,),
127 + )
128
129 # Main interface
130 st.title("vLLM Chat Assistant")
@@ -145,18 +151,18 @@ for message in st.session_state.messages:
151 if prompt := st.chat_input("Type your message here..."):
152 # Save user message to session
153 st.session_state.messages.append({"role": "user", "content": prompt})
148 - st.session_state.sessions[
149 - st.session_state.current_session] = st.session_state.messages
154 + st.session_state.sessions[st.session_state.current_session] = (
155 + st.session_state.messages
156 + )
157
158 # Display user message
159 with st.chat_message("user"):
160 st.write(prompt)
161
162 # Prepare messages for llm
156 - messages_for_llm = [{
157 - "role": m["role"],
158 - "content": m["content"]
159 - } for m in st.session_state.messages]
163 + messages_for_llm = [
164 + {"role": m["role"], "content": m["content"]} for m in st.session_state.messages
165 + ]
166
167 # Generate and display llm response
168 with st.chat_message("assistant"):
@@ -179,7 +185,4 @@ if prompt := st.chat_input("Type your message here..."):
185 message_placeholder.markdown(full_response)
186
187 # Save llm response to session history
182 - st.session_state.messages.append({
183 - "role": "assistant",
184 - "content": full_response
185 - })
188 + st.session_state.messages.append({"role": "assistant", "content": full_response})
examples/online_serving/utils.py
+3 -3
@@ -16,10 +16,10 @@ def get_first_model(client: OpenAI) -> str:
16 f"{client.base_url} with API key {client.api_key}. Check\n"
17 "1. the server is running\n"
18 "2. the server URL is correct\n"
19 - "3. the API key is correct") from e
19 + "3. the API key is correct"
20 + ) from e
21
22 if len(models.data) == 0:
22 - raise RuntimeError(
23 - f"No models found on the vLLM server at {client.base_url}")
23 + raise RuntimeError(f"No models found on the vLLM server at {client.base_url}")
24
25 return models.data[0].id
examples/others/lmcache/cpu_offload_lmcache.py
+11 -13
@@ -20,6 +20,7 @@ Requirements: Linux, Python: 3.10 or higher, CUDA: 12.1
20 Learn more about LMCache environment setup, please refer to:
21 https://docs.lmcache.ai/getting_started/installation.html
22 """
23 +
24 import argparse
25 import contextlib
26 import os
@@ -49,8 +50,7 @@ def setup_environment_variables(vllm_version: str):
50
51
52 @contextlib.contextmanager
52 -def build_llm_with_lmcache(lmcache_connector: str, model: str,
53 - vllm_version: str):
53 +def build_llm_with_lmcache(lmcache_connector: str, model: str, vllm_version: str):
54 ktc = KVTransferConfig(
55 kv_connector=lmcache_connector,
56 kv_role="kv_both",
@@ -97,18 +97,19 @@ def print_output(
97 for output in outputs:
98 generated_text = output.outputs[0].text
99 print(f"Generated text: {generated_text!r}")
100 - print(f"Generation took {time.time() - start:.2f} seconds, "
101 - f"{req_str} request done.")
100 + print(f"Generation took {time.time() - start:.2f} seconds, {req_str} request done.")
101 print("-" * 50)
102
103
104 def parse_args():
105 parser = argparse.ArgumentParser()
107 - parser.add_argument("-v",
108 - "--version",
109 - choices=["v0", "v1"],
110 - default="v1",
111 - help="Specify vLLM version (default: v1)")
106 + parser.add_argument(
107 + "-v",
108 + "--version",
109 + choices=["v0", "v1"],
110 + default="v1",
111 + help="Specify vLLM version (default: v1)",
112 + )
113 return parser.parse_args()
114
115
@@ -125,7 +126,6 @@ def main():
126 setup_environment_variables(args.version)
127
128 with build_llm_with_lmcache(lmcache_connector, model, args.version) as llm:
128 -
129 # This example script runs two requests with a shared prefix.
130 # Define the shared prompt and specific prompts
131 shared_prompt = "Hello, how are you?" * 1000
@@ -136,9 +136,7 @@ def main():
136 shared_prompt + "Tell me a very long story",
137 ]
138
139 - sampling_params = SamplingParams(temperature=0,
140 - top_p=0.95,
141 - max_tokens=10)
139 + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
140
141 # Print the first output
142 print_output(llm, first_prompt, sampling_params, "first")
examples/others/lmcache/disagg_prefill_lmcache_v0.py
+33 -25
@@ -4,12 +4,13 @@ This file demonstrates the example usage of disaggregated prefilling
4 with LMCache.
5 We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
6 and launch an additional LMCache server.
7 -KV cache is transferred in the following manner:
7 +KV cache is transferred in the following manner:
8 vLLM prefill node -> LMCache server -> vLLM decode node.
9
10 Note that `pip install lmcache` is needed to run this example.
11 Learn more about LMCache in https://github.com/LMCache/LMCache.
12 """
13 +
14 import os
15 import subprocess
16 import time
@@ -49,19 +50,23 @@ def run_prefill(prefill_done, prompts):
50
51 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
52
52 - ktc = KVTransferConfig(kv_connector="LMCacheConnector",
53 - kv_role="kv_producer",
54 - kv_rank=0,
55 - kv_parallel_size=2)
53 + ktc = KVTransferConfig(
54 + kv_connector="LMCacheConnector",
55 + kv_role="kv_producer",
56 + kv_rank=0,
57 + kv_parallel_size=2,
58 + )
59 # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
60 # memory. Reduce the value if your GPU has less memory.
58 - llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
59 - kv_transfer_config=ktc,
60 - max_model_len=8000,
61 - gpu_memory_utilization=0.8,
62 - enforce_eager=True)
63 -
64 - #llm.generate(prompts, sampling_params)
61 + llm = LLM(
62 + model="mistralai/Mistral-7B-Instruct-v0.2",
63 + kv_transfer_config=ktc,
64 + max_model_len=8000,
65 + gpu_memory_utilization=0.8,
66 + enforce_eager=True,
67 + )
68 +
69 + # llm.generate(prompts, sampling_params)
70 outputs = llm.generate(prompts, sampling_params)
71 for output in outputs:
72 generated_text = output.outputs[0].text
@@ -79,17 +84,21 @@ def run_decode(prefill_done, prompts, timeout=1):
84
85 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
86
82 - ktc = KVTransferConfig(kv_connector="LMCacheConnector",
83 - kv_role="kv_consumer",
84 - kv_rank=1,
85 - kv_parallel_size=2)
87 + ktc = KVTransferConfig(
88 + kv_connector="LMCacheConnector",
89 + kv_role="kv_consumer",
90 + kv_rank=1,
91 + kv_parallel_size=2,
92 + )
93 # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
94 # of memory. Reduce the value if your GPU has less memory.
88 - llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
89 - kv_transfer_config=ktc,
90 - max_model_len=8000,
91 - gpu_memory_utilization=0.8,
92 - enforce_eager=True)
95 + llm = LLM(
96 + model="mistralai/Mistral-7B-Instruct-v0.2",
97 + kv_transfer_config=ktc,
98 + max_model_len=8000,
99 + gpu_memory_utilization=0.8,
100 + enforce_eager=True,
101 + )
102
103 print("Waiting for prefill node to finish...")
104 prefill_done.wait()
@@ -105,10 +114,9 @@ def run_decode(prefill_done, prompts, timeout=1):
114
115
116 def run_lmcache_server(port):
108 - server_proc = subprocess.Popen([
109 - "python", "-m", "lmcache.experimental.server", "localhost",
110 - str(port)
111 - ])
117 + server_proc = subprocess.Popen(
118 + ["python", "-m", "lmcache.experimental.server", "localhost", str(port)]
119 + )
120 return server_proc
121
122
examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py
+57 -43
@@ -17,13 +17,17 @@ async def lifespan(app: FastAPI):
17 Lifespan context manager to handle startup and shutdown events.
18 """
19 # Startup: Initialize clients
20 - prefiller_base_url = f'http://{global_args.prefiller_host}:{global_args.prefiller_port}/v1'
21 - decoder_base_url = f'http://{global_args.decoder_host}:{global_args.decoder_port}/v1'
22 -
23 - app.state.prefill_client = httpx.AsyncClient(timeout=None,
24 - base_url=prefiller_base_url)
25 - app.state.decode_client = httpx.AsyncClient(timeout=None,
26 - base_url=decoder_base_url)
20 + prefiller_base_url = (
21 + f"http://{global_args.prefiller_host}:{global_args.prefiller_port}/v1"
22 + )
23 + decoder_base_url = (
24 + f"http://{global_args.decoder_host}:{global_args.decoder_port}/v1"
25 + )
26 +
27 + app.state.prefill_client = httpx.AsyncClient(
28 + timeout=None, base_url=prefiller_base_url
29 + )
30 + app.state.decode_client = httpx.AsyncClient(timeout=None, base_url=decoder_base_url)
31
32 yield
33
@@ -37,7 +41,6 @@ app = FastAPI(lifespan=lifespan)
41
42
43 class StatsCalculator:
40 -
44 def __init__(self):
45 self._stats = []
46 self._last_log_time = time.time()
@@ -51,13 +54,18 @@ class StatsCalculator:
54 def _log_stats(self):
55 # Print average, median, and 99th percentile
56 np_arr = np.array(self._stats)
54 - output_str = f"\nNum requests: {len(self._stats)}" + \
55 - "\nPrefill node TTFT stats:" + \
56 - f"\n - Average (ms): {np.mean(np_arr)}" + \
57 - f"\n - Median (ms): {np.median(np_arr)}" + \
58 - f"\n - 99th Percentile (ms): {np.percentile(np_arr, 99)}\n"
59 - print("===============================", output_str,
60 - "===============================")
57 + output_str = (
58 + f"\nNum requests: {len(self._stats)}"
59 + + "\nPrefill node TTFT stats:"
60 + + f"\n - Average (ms): {np.mean(np_arr)}"
61 + + f"\n - Median (ms): {np.median(np_arr)}"
62 + + f"\n - 99th Percentile (ms): {np.percentile(np_arr, 99)}\n"
63 + )
64 + print(
65 + "===============================",
66 + output_str,
67 + "===============================",
68 + )
69
70
71 stats_calculator = StatsCalculator()
@@ -82,15 +90,16 @@ app.state.prefill_client = None
90 app.state.decode_client = None
91
92
85 -async def send_request_to_service(client: httpx.AsyncClient, endpoint: str,
86 - req_data: dict):
93 +async def send_request_to_service(
94 + client: httpx.AsyncClient, endpoint: str, req_data: dict
95 +):
96 """
97 Send a request to a service using a persistent client.
98 """
99 req_data = req_data.copy()
91 - req_data['max_tokens'] = 1
92 - if 'max_completion_tokens' in req_data:
93 - req_data['max_completion_tokens'] = 1
100 + req_data["max_tokens"] = 1
101 + if "max_completion_tokens" in req_data:
102 + req_data["max_completion_tokens"] = 1
103
104 headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
105 response = await client.post(endpoint, json=req_data, headers=headers)
@@ -98,14 +107,16 @@ async def send_request_to_service(client: httpx.AsyncClient, endpoint: str,
107 return response
108
109
101 -async def stream_service_response(client: httpx.AsyncClient, endpoint: str,
102 - req_data: dict):
110 +async def stream_service_response(
111 + client: httpx.AsyncClient, endpoint: str, req_data: dict
112 +):
113 """
114 Asynchronously stream the response from a service using a persistent client.
115 """
116 headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
107 - async with client.stream("POST", endpoint, json=req_data,
108 - headers=headers) as response:
117 + async with client.stream(
118 + "POST", endpoint, json=req_data, headers=headers
119 + ) as response:
120 response.raise_for_status()
121 async for chunk in response.aiter_bytes():
122 yield chunk
@@ -121,28 +132,28 @@ async def handle_completions(request: Request):
132 req_data = await request.json()
133
134 # Send request to prefill service, ignore the response
124 - await send_request_to_service(app.state.prefill_client, "/completions",
125 - req_data)
135 + await send_request_to_service(
136 + app.state.prefill_client, "/completions", req_data
137 + )
138
139 et = time.time()
140 stats_calculator.add(et - st)
141
142 # Stream response from decode service
143 async def generate_stream():
132 - async for chunk in stream_service_response(app.state.decode_client,
133 - "/completions",
134 - req_data):
144 + async for chunk in stream_service_response(
145 + app.state.decode_client, "/completions", req_data
146 + ):
147 yield chunk
148
137 - return StreamingResponse(generate_stream(),
138 - media_type="text/event-stream")
149 + return StreamingResponse(generate_stream(), media_type="text/event-stream")
150
151 except Exception as e:
152 import sys
153 import traceback
154 +
155 exc_info = sys.exc_info()
144 - print("Error occurred in disagg prefill proxy server"
145 - " - completions endpoint")
156 + print("Error occurred in disagg prefill proxy server - completions endpoint")
157 print(e)
158 print("".join(traceback.format_exception(*exc_info)))
159 raise
@@ -158,36 +169,39 @@ async def handle_chat_completions(request: Request):
169 req_data = await request.json()
170
171 # Send request to prefill service, ignore the response
161 - await send_request_to_service(app.state.prefill_client,
162 - "/chat/completions", req_data)
172 + await send_request_to_service(
173 + app.state.prefill_client, "/chat/completions", req_data
174 + )
175
176 et = time.time()
177 stats_calculator.add(et - st)
178
179 # Stream response from decode service
180 async def generate_stream():
169 - async for chunk in stream_service_response(app.state.decode_client,
170 - "/chat/completions",
171 - req_data):
181 + async for chunk in stream_service_response(
182 + app.state.decode_client, "/chat/completions", req_data
183 + ):
184 yield chunk
185
174 - return StreamingResponse(generate_stream(),
175 - media_type="text/event-stream")
186 + return StreamingResponse(generate_stream(), media_type="text/event-stream")
187
188 except Exception as e:
189 import sys
190 import traceback
191 +
192 exc_info = sys.exc_info()
181 - print("Error occurred in disagg prefill proxy server "
182 - " - chat completions endpoint")
193 + print(
194 + "Error occurred in disagg prefill proxy server - chat completions endpoint"
195 + )
196 print(e)
197 print("".join(traceback.format_exception(*exc_info)))
198 raise
199
200
188 -if __name__ == '__main__':
201 +if __name__ == "__main__":
202 global global_args
203 global_args = parse_args()
204
205 import uvicorn
206 +
207 uvicorn.run(app, host=global_args.host, port=global_args.port)
examples/others/lmcache/kv_cache_sharing_lmcache_v1.py
+21 -19
@@ -3,13 +3,14 @@
3 This file demonstrates the example usage of remote KV cache sharing
4 with LMCache.
5 We will launch 2 vllm instances, and launch an additional LMCache server.
6 -KV cache is transferred in the following manner:
6 +KV cache is transferred in the following manner:
7 (1) vLLM instance 1 -> LMCache server (KV cache store).
8 (2) LMCache server -> vLLM instance 2 (KV cache reuse/retrieve).
9
10 Note that lmcache needs to be installed to run this example.
11 Learn more about LMCache in https://github.com/LMCache/LMCache.
12 """
13 +
14 import os
15 import subprocess
16 import time
@@ -49,15 +50,16 @@ def run_store(store_done, prompts):
50
51 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
52
52 - ktc = KVTransferConfig(kv_connector="LMCacheConnectorV1",
53 - kv_role="kv_both")
53 + ktc = KVTransferConfig(kv_connector="LMCacheConnectorV1", kv_role="kv_both")
54 # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
55 # memory. Reduce the value if your GPU has less memory.
56 - llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
57 - kv_transfer_config=ktc,
58 - max_model_len=8000,
59 - gpu_memory_utilization=0.8,
60 - enforce_eager=True)
56 + llm = LLM(
57 + model="mistralai/Mistral-7B-Instruct-v0.2",
58 + kv_transfer_config=ktc,
59 + max_model_len=8000,
60 + gpu_memory_utilization=0.8,
61 + enforce_eager=True,
62 + )
63
64 outputs = llm.generate(prompts, sampling_params)
65 for output in outputs:
@@ -76,15 +78,16 @@ def run_retrieve(store_done, prompts, timeout=1):
78
79 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
80
79 - ktc = KVTransferConfig(kv_connector="LMCacheConnectorV1",
80 - kv_role="kv_both")
81 + ktc = KVTransferConfig(kv_connector="LMCacheConnectorV1", kv_role="kv_both")
82 # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
83 # of memory. Reduce the value if your GPU has less memory.
83 - llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
84 - kv_transfer_config=ktc,
85 - max_model_len=8000,
86 - gpu_memory_utilization=0.8,
87 - enforce_eager=True)
84 + llm = LLM(
85 + model="mistralai/Mistral-7B-Instruct-v0.2",
86 + kv_transfer_config=ktc,
87 + max_model_len=8000,
88 + gpu_memory_utilization=0.8,
89 + enforce_eager=True,
90 + )
91
92 print("Waiting for KV cache store to finish...")
93 store_done.wait()
@@ -100,10 +103,9 @@ def run_retrieve(store_done, prompts, timeout=1):
103
104
105 def run_lmcache_server(port):
103 - server_proc = subprocess.Popen([
104 - "python", "-m", "lmcache.experimental.server", "localhost",
105 - str(port)
106 - ])
106 + server_proc = subprocess.Popen(
107 + ["python", "-m", "lmcache.experimental.server", "localhost", str(port)]
108 + )
109 return server_proc
110
111
examples/others/tensorize_vllm_model.py
+5 -2
@@ -10,8 +10,11 @@ from vllm import LLM, SamplingParams
10 from vllm.engine.arg_utils import EngineArgs
11 from vllm.lora.request import LoRARequest
12 from vllm.model_executor.model_loader.tensorizer import (
13 - TensorizerArgs, TensorizerConfig, tensorize_lora_adapter,
14 - tensorize_vllm_model)
13 + TensorizerArgs,
14 + TensorizerConfig,
15 + tensorize_lora_adapter,
16 + tensorize_vllm_model,
17 +)
18 from vllm.utils import FlexibleArgumentParser
19
20 # yapf conflicts with isort for this docstring
examples/pyproject.toml new
+54
@@ -0,0 +1,54 @@
1 +# This local pyproject file is part of the migration from yapf to ruff format.
2 +# It uses the same core rules as the main pyproject.toml file, but with the
3 +# following differences:
4 +# - ruff line length is overridden to 88
5 +# - deprecated typing ignores (UP006, UP035) have been removed
6 +
7 +[tool.ruff]
8 +line-length = 88
9 +exclude = [
10 + # External file, leaving license intact
11 + "examples/other/fp8/quantizer/quantize.py",
12 + "vllm/vllm_flash_attn/flash_attn_interface.pyi"
13 +]
14 +
15 +[tool.ruff.lint.per-file-ignores]
16 +"vllm/third_party/**" = ["ALL"]
17 +"vllm/version.py" = ["F401"]
18 +"vllm/_version.py" = ["ALL"]
19 +
20 +[tool.ruff.lint]
21 +select = [
22 + # pycodestyle
23 + "E",
24 + # Pyflakes
25 + "F",
26 + # pyupgrade
27 + "UP",
28 + # flake8-bugbear
29 + "B",
30 + # flake8-simplify
31 + "SIM",
32 + # isort
33 + "I",
34 + # flake8-logging-format
35 + "G",
36 +]
37 +ignore = [
38 + # star imports
39 + "F405", "F403",
40 + # lambda expression assignment
41 + "E731",
42 + # Loop control variable not used within loop body
43 + "B007",
44 + # f-string format
45 + "UP032",
46 + # Can remove once 3.10+ is the minimum Python version
47 + "UP007",
48 +]
49 +
50 +[tool.ruff.lint.isort]
51 +known-first-party = ["vllm"]
52 +
53 +[tool.ruff.format]
54 +docstring-code-format = true
\ No newline at end of file
pyproject.toml
+2
@@ -57,6 +57,7 @@ ignore_patterns = [
57 ".buildkite/**",
58 "benchmarks/**",
59 "build/**",
60 + "examples/**",
61 ]
62
63 [tool.ruff]
@@ -144,6 +145,7 @@ skip = "tests/models/fixtures/*,tests/prompts/*,benchmarks/sonnet.txt,tests/lora
145 skip_glob = [
146 ".buildkite/*",
147 "benchmarks/*",
148 + "examples/*",
149 ]
150 use_parentheses = true
151 skip_gitignore = true