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 3, 11:45 UTC
[ROCm][CI] Stabilize ROCm pooling and multimodal CI (#42909)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
May 17, 2026 at 22:57 UTC
b50646e5effd7cb5884cd96fdff4c53c18521198
4 files changed
+51
-17
tests/models/language/pooling/test_gritlm.py
+9
-3
@@ -7,7 +7,7 @@ from scipy.spatial.distance import cosine
7
from vllm import LLM, SamplingParams
8
from vllm.config import ModelConfig
9
10
-from ....utils import RemoteOpenAIServer
10
+from ....utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
11
from .embed_utils import run_client_embeddings
12
13
MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
@@ -126,9 +126,15 @@ def test_gritlm_offline_embedding(vllm_runner):
126
async def test_gritlm_api_server_embedding():
127
queries, q_instruction, documents, d_instruction = get_test_data()
128
129
- args = ["--runner", "pooling", "--max_model_len", str(MAX_MODEL_LEN)]
129
+ args = [
130
+ "--runner",
131
+ "pooling",
132
+ "--max_model_len",
133
+ str(MAX_MODEL_LEN),
134
+ *ROCM_EXTRA_ARGS,
135
+ ]
136
131
- with RemoteOpenAIServer(MODEL_NAME, args) as server:
137
+ with RemoteOpenAIServer(MODEL_NAME, args, env_dict=ROCM_ENV_OVERRIDES) as server:
138
client_embedding = server.get_async_client()
139
140
d_rep = await run_client_embeddings(
tests/models/language/pooling/test_max_tokens_per_doc.py
+31
-12
@@ -17,6 +17,7 @@ from vllm.entrypoints.pooling.scoring.protocol import RerankResponse
17
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
18
19
TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template")
20
+ExpectedPromptTokens = int | tuple[int, ...]
21
22
long_query = "What is the capital of France?" * 20
23
long_doc = "The capital of France is Paris. " * 20
@@ -26,10 +27,10 @@ long_doc = "The capital of France is Paris. " * 20
27
class TestConfig:
28
model: str
29
args: list[str]
29
- without_truncated_prompt_tokens: int
30
- with_max_tokens_per_query_prompt_tokens: int
31
- with_max_tokens_per_doc_prompt_tokens: int
32
- with_max_tokens_per_query_and_doc_prompt_tokens: int
30
+ without_truncated_prompt_tokens: ExpectedPromptTokens
31
+ with_max_tokens_per_query_prompt_tokens: ExpectedPromptTokens
32
+ with_max_tokens_per_doc_prompt_tokens: ExpectedPromptTokens
33
+ with_max_tokens_per_query_and_doc_prompt_tokens: ExpectedPromptTokens
34
35
36
RERANK_CONFIGS = [
@@ -79,8 +80,10 @@ RERANK_CONFIGS = [
80
"512",
81
"--trust-remote-code",
82
],
82
- without_truncated_prompt_tokens=286,
83
- with_max_tokens_per_query_prompt_tokens=156,
83
+ # This model has produced both prompt-token totals in CI/local cache;
84
+ # keep truncation checks exact while tolerating the boundary delta.
85
+ without_truncated_prompt_tokens=(285, 286),
86
+ with_max_tokens_per_query_prompt_tokens=(155, 156),
87
with_max_tokens_per_doc_prompt_tokens=155,
88
with_max_tokens_per_query_and_doc_prompt_tokens=25,
89
),
@@ -115,6 +118,13 @@ RERANK_CONFIGS = [
118
]
119
120
121
+def assert_prompt_tokens(actual: int, expected: ExpectedPromptTokens) -> None:
122
+ if isinstance(expected, int):
123
+ assert actual == expected
124
+ else:
125
+ assert actual in expected
126
+
127
+
128
@pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model)
129
def server(request):
130
config: TestConfig = request.param
@@ -136,7 +146,10 @@ def test_without_truncated(server):
146
assert rerank.id is not None
147
assert rerank.results is not None
148
assert len(rerank.results) == 1
139
- assert rerank.usage.prompt_tokens == config.without_truncated_prompt_tokens
149
+ assert_prompt_tokens(
150
+ rerank.usage.prompt_tokens,
151
+ config.without_truncated_prompt_tokens,
152
+ )
153
154
155
def test_max_tokens_per_query(server):
@@ -158,7 +171,10 @@ def test_max_tokens_per_query(server):
171
assert rerank.id is not None
172
assert rerank.results is not None
173
assert len(rerank.results) == 1
161
- assert rerank.usage.prompt_tokens == config.with_max_tokens_per_query_prompt_tokens
174
+ assert_prompt_tokens(
175
+ rerank.usage.prompt_tokens,
176
+ config.with_max_tokens_per_query_prompt_tokens,
177
+ )
178
179
180
def test_max_tokens_per_doc(server):
@@ -180,7 +196,10 @@ def test_max_tokens_per_doc(server):
196
assert rerank.id is not None
197
assert rerank.results is not None
198
assert len(rerank.results) == 1
183
- assert rerank.usage.prompt_tokens == config.with_max_tokens_per_doc_prompt_tokens
199
+ assert_prompt_tokens(
200
+ rerank.usage.prompt_tokens,
201
+ config.with_max_tokens_per_doc_prompt_tokens,
202
+ )
203
204
205
def test_max_tokens_per_query_and_doc(server):
@@ -203,7 +222,7 @@ def test_max_tokens_per_query_and_doc(server):
222
assert rerank.id is not None
223
assert rerank.results is not None
224
assert len(rerank.results) == 1
206
- assert (
207
- rerank.usage.prompt_tokens
208
- == config.with_max_tokens_per_query_and_doc_prompt_tokens
225
+ assert_prompt_tokens(
226
+ rerank.usage.prompt_tokens,
227
+ config.with_max_tokens_per_query_and_doc_prompt_tokens,
228
)
tests/models/multimodal/generation/test_qwen2_5_vl.py
+3
-2
@@ -34,6 +34,7 @@ WINDOW_ATTN_IMAGE_PROMPT = qwen2_5_vl_chat_template(
34
IMAGE_PLACEHOLDER,
35
"Describe the image.",
36
)
37
+IMAGE_ONLY_LIMIT_MM_PER_PROMPT = {"image": 1, "video": 0}
38
39
40
def _window_attention_regression_image():
@@ -193,7 +194,7 @@ def test_qwen2_5_vl_window_attention_image(
194
runner="generate",
195
max_model_len=4096,
196
dtype=dtype,
196
- limit_mm_per_prompt={"image": 1},
197
+ limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT,
198
compilation_config=_encoder_cudagraph_config(max_vision_items=1),
199
) as vllm_model:
200
outputs = vllm_model.generate_greedy(prompt, max_tokens, images=images)
@@ -231,7 +232,7 @@ def test_qwen2_5_vl_window_attention_image_batch(
232
max_model_len=4096,
233
max_num_seqs=2,
234
dtype=dtype,
234
- limit_mm_per_prompt={"image": 1},
235
+ limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT,
236
compilation_config=_encoder_cudagraph_config(max_vision_items=2),
237
) as vllm_model:
238
outputs = vllm_model.generate_greedy(prompts, max_tokens, images=images)
vllm/model_executor/models/transformers/base.py
+8
@@ -353,6 +353,14 @@ class Base(
353
for source, target in ccm.items():
354
orig_to_new_regex[re.compile(source)] = target
355
356
+ # Gemma3 checkpoints saved with older Transformers versions include an
357
+ # extra `vision_model` level that the current AutoModel no longer has.
358
+ vision_tower = getattr(self.model, "vision_tower", None)
359
+ if vision_tower is not None and not hasattr(vision_tower, "vision_model"):
360
+ orig_to_new_regex[
361
+ re.compile(r"^(?:model\.)?vision_tower\.vision_model\.(.+)")
362
+ ] = r"model.vision_tower.\1"
363
+
364
# Handle unexpected weights which should be ignored
365
if self.model._keys_to_ignore_on_load_unexpected is not None:
366
for key in self.model._keys_to_ignore_on_load_unexpected: