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, 12:15 UTC
[Frontend] Support `chat_template_kwargs` in `LLM.chat` (#17356)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Apr 29, 2025 at 22:03 UTC
88ad9ec6b23b79c358ce279b02a67e7c96e2c8b9
2 files changed
+106
-24
tests/entrypoints/llm/test_chat.py
+93
-16
@@ -1,15 +1,31 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import weakref
3
4
import pytest
5
6
from vllm import LLM
7
+from vllm.distributed import cleanup_dist_env_and_memory
8
9
from ..openai.test_vision import TEST_IMAGE_URLS
10
11
10
-def test_chat():
11
- llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct")
12
+@pytest.fixture(scope="function")
13
+def text_llm():
14
+ # pytest caches the fixture so we use weakref.proxy to
15
+ # enable garbage collection
16
+ llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct",
17
+ enforce_eager=True,
18
+ seed=0)
19
20
+ with llm.deprecate_legacy_api():
21
+ yield weakref.proxy(llm)
22
+
23
+ del llm
24
+
25
+ cleanup_dist_env_and_memory()
26
+
27
+
28
+def test_chat(text_llm):
29
prompt1 = "Explain the concept of entropy."
30
messages = [
31
{
@@ -21,13 +37,11 @@ def test_chat():
37
"content": prompt1
38
},
39
]
24
- outputs = llm.chat(messages)
40
+ outputs = text_llm.chat(messages)
41
assert len(outputs) == 1
42
43
28
-def test_multi_chat():
29
- llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct")
30
-
44
+def test_multi_chat(text_llm):
45
prompt1 = "Explain the concept of entropy."
46
prompt2 = "Explain what among us is."
47
@@ -55,13 +69,14 @@ def test_multi_chat():
69
70
messages = [conversation1, conversation2]
71
58
- outputs = llm.chat(messages)
72
+ outputs = text_llm.chat(messages)
73
assert len(outputs) == 2
74
75
62
-@pytest.mark.parametrize("image_urls",
63
- [[TEST_IMAGE_URLS[0], TEST_IMAGE_URLS[1]]])
64
-def test_chat_multi_image(image_urls: list[str]):
76
+@pytest.fixture(scope="function")
77
+def vision_llm():
78
+ # pytest caches the fixture so we use weakref.proxy to
79
+ # enable garbage collection
80
llm = LLM(
81
model="microsoft/Phi-3.5-vision-instruct",
82
max_model_len=4096,
@@ -69,8 +84,20 @@ def test_chat_multi_image(image_urls: list[str]):
84
enforce_eager=True,
85
trust_remote_code=True,
86
limit_mm_per_prompt={"image": 2},
87
+ seed=0,
88
)
89
90
+ with llm.deprecate_legacy_api():
91
+ yield weakref.proxy(llm)
92
+
93
+ del llm
94
+
95
+ cleanup_dist_env_and_memory()
96
+
97
+
98
+@pytest.mark.parametrize("image_urls",
99
+ [[TEST_IMAGE_URLS[0], TEST_IMAGE_URLS[1]]])
100
+def test_chat_multi_image(vision_llm, image_urls: list[str]):
101
messages = [{
102
"role":
103
"user",
@@ -87,16 +114,15 @@ def test_chat_multi_image(image_urls: list[str]):
114
},
115
],
116
}]
90
- outputs = llm.chat(messages)
117
+ outputs = vision_llm.chat(messages)
118
assert len(outputs) >= 0
119
120
94
-def test_llm_chat_tokenization_no_double_bos():
121
+def test_llm_chat_tokenization_no_double_bos(text_llm):
122
"""
123
LLM.chat() should not add special tokens when using chat templates.
124
Check we get a single BOS token for llama chat.
125
"""
99
- llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True)
126
messages = [
127
{
128
"role": "system",
@@ -107,13 +133,64 @@ def test_llm_chat_tokenization_no_double_bos():
133
"content": "Hello!"
134
},
135
]
110
- outputs = llm.chat(messages)
136
+ outputs = text_llm.chat(messages)
137
assert len(outputs) == 1
112
- prompt_token_ids = getattr(outputs[0], "prompt_token_ids", None)
138
+
139
+ prompt_token_ids = outputs[0].prompt_token_ids
140
assert prompt_token_ids is not None
141
115
- bos_token = llm.get_tokenizer().bos_token_id
142
+ bos_token = text_llm.get_tokenizer().bos_token_id
143
144
# Ensure we have a single BOS
145
assert prompt_token_ids[0] == bos_token
146
assert prompt_token_ids[1] != bos_token, "Double BOS"
147
+
148
+
149
+@pytest.fixture(scope="function")
150
+def thinking_llm():
151
+ # pytest caches the fixture so we use weakref.proxy to
152
+ # enable garbage collection
153
+ llm = LLM(
154
+ model="Qwen/Qwen3-0.6B",
155
+ max_model_len=4096,
156
+ enforce_eager=True,
157
+ seed=0,
158
+ )
159
+
160
+ with llm.deprecate_legacy_api():
161
+ yield weakref.proxy(llm)
162
+
163
+ del llm
164
+
165
+ cleanup_dist_env_and_memory()
166
+
167
+
168
+@pytest.mark.parametrize("enable_thinking", [True, False])
169
+def test_chat_extra_kwargs(thinking_llm, enable_thinking):
170
+ messages = [
171
+ {
172
+ "role": "system",
173
+ "content": "You are a helpful assistant"
174
+ },
175
+ {
176
+ "role": "user",
177
+ "content": "What is 1+1?"
178
+ },
179
+ ]
180
+
181
+ outputs = thinking_llm.chat(
182
+ messages,
183
+ chat_template_kwargs={"enable_thinking": enable_thinking},
184
+ )
185
+ assert len(outputs) == 1
186
+
187
+ prompt_token_ids = outputs[0].prompt_token_ids
188
+ assert prompt_token_ids is not None
189
+
190
+ think_id = thinking_llm.get_tokenizer().get_vocab()["<think>"]
191
+
192
+ if enable_thinking:
193
+ assert think_id not in prompt_token_ids
194
+ else:
195
+ # The chat template includes dummy thinking process
196
+ assert think_id in prompt_token_ids
vllm/entrypoints/llm.py
+13
-8
@@ -656,6 +656,7 @@ class LLM:
656
add_generation_prompt: bool = True,
657
continue_final_message: bool = False,
658
tools: Optional[list[dict[str, Any]]] = None,
659
+ chat_template_kwargs: Optional[dict[str, Any]] = None,
660
mm_processor_kwargs: Optional[dict[str, Any]] = None,
661
) -> list[RequestOutput]:
662
"""
@@ -696,6 +697,8 @@ class LLM:
697
continue_final_message: If True, continues the final message in
698
the conversation instead of starting a new one. Cannot be
699
``True`` if ``add_generation_prompt`` is also ``True``.
700
+ chat_template_kwargs: Additional kwargs to pass to the chat
701
+ template.
702
mm_processor_kwargs: Multimodal processor kwarg overrides for this
703
chat request. Only used for offline requests.
704
@@ -726,6 +729,14 @@ class LLM:
729
trust_remote_code=model_config.trust_remote_code,
730
)
731
732
+ _chat_template_kwargs: dict[str, Any] = dict(
733
+ chat_template=chat_template,
734
+ add_generation_prompt=add_generation_prompt,
735
+ continue_final_message=continue_final_message,
736
+ tools=tools,
737
+ )
738
+ _chat_template_kwargs.update(chat_template_kwargs or {})
739
+
740
prompts: list[Union[TokensPrompt, TextPrompt]] = []
741
742
for msgs in list_of_messages:
@@ -743,20 +754,14 @@ class LLM:
754
prompt_token_ids = apply_mistral_chat_template(
755
tokenizer,
756
messages=msgs,
746
- chat_template=chat_template,
747
- tools=tools,
748
- add_generation_prompt=add_generation_prompt,
749
- continue_final_message=continue_final_message,
757
+ **_chat_template_kwargs,
758
)
759
else:
760
prompt_str = apply_hf_chat_template(
761
tokenizer,
762
trust_remote_code=model_config.trust_remote_code,
763
conversation=conversation,
756
- chat_template=chat_template,
757
- tools=tools,
758
- add_generation_prompt=add_generation_prompt,
759
- continue_final_message=continue_final_message,
764
+ **_chat_template_kwargs,
765
)
766
# Special tokens are already included in chat templates so
767
# should not be added by the tokenizer in this case.