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, 01:00 UTC
[CI/Build] Fix model executor tests (#21387)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Jul 23, 2025 at 11:25 UTC
c401c64b4ccdf7c16a985ffa7d2c5eaf0422a04f
2 files changed
+9
-5
.buildkite/test-pipeline.yaml
-1
@@ -434,7 +434,6 @@ steps:
434
435
- label: Model Executor Test
436
mirror_hardwares: [amdexperimental, amdproduction]
437
- soft_fail: true
437
source_file_dependencies:
438
- vllm/model_executor
439
- tests/model_executor
tests/model_executor/test_model_load_with_params.py
+9
-4
@@ -5,7 +5,8 @@ import os
5
6
import pytest
7
8
-from vllm.model_executor.layers.pooler import CLSPool, MeanPool, PoolingType
8
+from vllm.model_executor.layers.pooler import (CLSPool, DispatchPooler,
9
+ MeanPool, PoolingType)
10
from vllm.model_executor.models.bert import BertEmbeddingModel
11
from vllm.model_executor.models.roberta import RobertaEmbeddingModel
12
from vllm.platforms import current_platform
@@ -49,7 +50,8 @@ def test_model_loading_with_params(vllm_runner):
50
51
def check_model(model):
52
assert isinstance(model, BertEmbeddingModel)
52
- assert isinstance(model.pooler.pooling, CLSPool)
53
+ assert isinstance(pooler := model.pooler, DispatchPooler)
54
+ assert isinstance(pooler.poolers_by_task["embed"].pooling, CLSPool)
55
56
vllm_model.apply_model(check_model)
57
@@ -87,7 +89,9 @@ def test_roberta_model_loading_with_params(vllm_runner):
89
90
def check_model(model):
91
assert isinstance(model, RobertaEmbeddingModel)
90
- assert isinstance(model.pooler.pooling, MeanPool)
92
+ assert isinstance(pooler := model.pooler, DispatchPooler)
93
+ assert isinstance(pooler.poolers_by_task["embed"].pooling,
94
+ MeanPool)
95
96
vllm_model.apply_model(check_model)
97
@@ -114,7 +118,8 @@ def test_facebook_roberta_model_loading_with_params(vllm_runner):
118
def check_model(model):
119
assert isinstance(model, RobertaEmbeddingModel)
120
assert not hasattr(model, "lm_head")
117
- assert isinstance(model.pooler.pooling, CLSPool)
121
+ assert isinstance(pooler := model.pooler, DispatchPooler)
122
+ assert isinstance(pooler.poolers_by_task["embed"].pooling, CLSPool)
123
124
vllm_model.apply_model(check_model)
125