@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 225 lines 6.96 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import json
4 from dataclasses import asdict
5 from typing import TYPE_CHECKING, Any
6
7 import pytest
8 import torch
9
10 from vllm import SamplingParams
11 from vllm.logprobs import Logprob, SampleLogprobs
12 from vllm.model_executor.models.pixtral import _make_packed_sequence_metadata
13 from vllm.platforms import current_platform
14 from vllm.v1.attention.backends.registry import AttentionBackendEnum
15
16 from ....utils import VLLM_PATH, large_gpu_test
17 from ...utils import check_logprobs_close
18
19 if TYPE_CHECKING:
20 from _typeshed import StrPath
21
22 PIXTRAL_ID = "mistralai/Pixtral-12B-2409"
23 MISTRAL_SMALL_3_1_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
24 MINISTRAL_3B_ID = "mistralai/Ministral-3-3B-Instruct-2512"
25
26 MODELS = [PIXTRAL_ID, MISTRAL_SMALL_3_1_ID]
27
28 IMG_URLS = [
29 "237-400x300.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
30 "231-200x300.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
31 "27-500x500.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
32 "17-150x600.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
33 ]
34 PROMPT = "Describe each image in one short sentence."
35
36
37 def _create_msg_format(urls: list[str]) -> list[dict[str, Any]]:
38 return [
39 {
40 "role": "user",
41 "content": [
42 {
43 "type": "text",
44 "text": PROMPT,
45 }
46 ]
47 + [{"type": "image_url", "image_url": {"url": url}} for url in urls],
48 }
49 ]
50
51
52 SAMPLING_PARAMS = SamplingParams(max_tokens=512, temperature=0.0, logprobs=5)
53 LIMIT_MM_PER_PROMPT = dict(image=4)
54
55 MAX_MODEL_LEN = [8192, 65536]
56
57 FIXTURES_PATH = VLLM_PATH / "tests/models/fixtures"
58 assert FIXTURES_PATH.exists()
59
60 FIXTURE_LOGPROBS_CHAT = {
61 PIXTRAL_ID: FIXTURES_PATH / "pixtral_chat.json",
62 MISTRAL_SMALL_3_1_ID: FIXTURES_PATH / "mistral_small_3_chat.json",
63 MINISTRAL_3B_ID: FIXTURES_PATH / "ministral_3b_chat.json",
64 }
65
66 OutputsLogprobs = list[tuple[list[int], str, SampleLogprobs | None]]
67
68
69 @pytest.mark.parametrize(
70 "backend",
71 [
72 AttentionBackendEnum.FLASH_ATTN,
73 AttentionBackendEnum.FLASHINFER,
74 AttentionBackendEnum.TORCH_SDPA,
75 ],
76 )
77 def test_packed_sequence_metadata(backend: AttentionBackendEnum) -> None:
78 cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(
79 [4, 6],
80 backend,
81 hidden_size=64,
82 tp_size=1,
83 device=torch.device("cpu"),
84 )
85
86 assert cu_seqlens.dtype == torch.int32
87 assert max_seqlen.dtype == torch.int32
88 if backend == AttentionBackendEnum.FLASHINFER:
89 assert max_seqlen.item() >= 6
90 assert sequence_lengths is not None
91 assert sequence_lengths.dtype == torch.int32
92 assert len(cu_seqlens) % 2 == 0
93 else:
94 expected_max_seqlen = 6 if backend == AttentionBackendEnum.FLASH_ATTN else 0
95 assert max_seqlen.item() == expected_max_seqlen
96 assert cu_seqlens.tolist() == [0, 4, 10]
97 assert sequence_lengths is None
98
99
100 # For the test author to store golden output in JSON
101 def _dump_outputs_w_logprobs(
102 outputs: OutputsLogprobs,
103 filename: "StrPath",
104 ) -> None:
105 json_data = [
106 (
107 tokens,
108 text,
109 [
110 {k: asdict(v) for k, v in token_logprobs.items()}
111 for token_logprobs in (logprobs or [])
112 ],
113 )
114 for tokens, text, logprobs in outputs
115 ]
116
117 with open(filename, "w") as f:
118 json.dump(json_data, f)
119
120
121 def load_outputs_w_logprobs(filename: "StrPath") -> OutputsLogprobs:
122 with open(filename, "rb") as f:
123 json_data = json.load(f)
124
125 return [
126 (
127 tokens,
128 text,
129 [
130 {int(k): Logprob(**v) for k, v in token_logprobs.items()}
131 for token_logprobs in logprobs
132 ],
133 )
134 for tokens, text, logprobs in json_data
135 ]
136
137
138 @large_gpu_test(min_gb=80)
139 @pytest.mark.parametrize("model", MODELS)
140 @pytest.mark.parametrize("max_model_len", MAX_MODEL_LEN)
141 @pytest.mark.parametrize("dtype", ["bfloat16"])
142 def test_chat(
143 vllm_runner, max_model_len: int, model: str, dtype: str, local_asset_server
144 ) -> None:
145 if (
146 model == MISTRAL_SMALL_3_1_ID
147 and max_model_len == 65536
148 and current_platform.is_rocm()
149 ):
150 pytest.skip(
151 "OOM on ROCm: 24B model with 65536 context length exceeds GPU memory"
152 )
153
154 EXPECTED_CHAT_LOGPROBS = load_outputs_w_logprobs(FIXTURE_LOGPROBS_CHAT[model])
155 with vllm_runner(
156 model,
157 dtype=dtype,
158 tokenizer_mode="mistral",
159 load_format="mistral",
160 config_format="mistral",
161 max_model_len=max_model_len,
162 limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
163 ) as vllm_model:
164 outputs = []
165
166 urls_all = [local_asset_server.url_for(u) for u in IMG_URLS]
167 msgs = [
168 _create_msg_format(urls_all[:1]),
169 _create_msg_format(urls_all[:2]),
170 _create_msg_format(urls_all),
171 ]
172 for msg in msgs:
173 output = vllm_model.llm.chat(msg, sampling_params=SAMPLING_PARAMS)
174
175 outputs.extend(output)
176
177 logprobs = vllm_runner._final_steps_generate_w_logprobs(outputs)
178 # Remove last `None` prompt_logprobs to compare with fixture
179 for i in range(len(logprobs)):
180 assert logprobs[i][-1] is None
181 logprobs[i] = logprobs[i][:-1]
182 check_logprobs_close(
183 outputs_0_lst=EXPECTED_CHAT_LOGPROBS,
184 outputs_1_lst=logprobs,
185 name_0="h100_ref",
186 name_1="output",
187 )
188
189
190 @large_gpu_test(min_gb=16)
191 @pytest.mark.parametrize("dtype", ["bfloat16"])
192 def test_chat_consolidated(vllm_runner, dtype: str, local_asset_server) -> None:
193 EXPECTED_CHAT_LOGPROBS = load_outputs_w_logprobs(
194 FIXTURE_LOGPROBS_CHAT[MINISTRAL_3B_ID]
195 )
196 with vllm_runner(
197 MINISTRAL_3B_ID,
198 dtype=dtype,
199 tokenizer_mode="mistral",
200 load_format="mistral",
201 config_format="mistral",
202 max_model_len=8192,
203 limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
204 ) as vllm_model:
205 outputs = []
206 urls_all = [local_asset_server.url_for(u) for u in IMG_URLS]
207 msgs = [
208 _create_msg_format(urls_all[:1]),
209 _create_msg_format(urls_all[:2]),
210 _create_msg_format(urls_all),
211 ]
212 for msg in msgs:
213 output = vllm_model.llm.chat(msg, sampling_params=SAMPLING_PARAMS)
214 outputs.extend(output)
215
216 logprobs = vllm_runner._final_steps_generate_w_logprobs(outputs)
217 for i in range(len(logprobs)):
218 assert logprobs[i][-1] is None
219 logprobs[i] = logprobs[i][:-1]
220 check_logprobs_close(
221 outputs_0_lst=EXPECTED_CHAT_LOGPROBS,
222 outputs_1_lst=logprobs,
223 name_0="h100_ref",
224 name_1="output",
225 )