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 14, 05:30 UTC
[Model] Improve Pooling Model (#25149)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Sep 18, 2025 at 19:04 UTC
37970105fed95d58677f0a4635cb253a71e8817c
2 files changed
+7
-6
vllm/model_executor/layers/pooler.py
+6
-6
@@ -12,8 +12,9 @@ import torch.nn as nn
12
import torch.nn.functional as F
13
from transformers import PretrainedConfig
14
15
-from vllm.config import ModelConfig, PoolerConfig
15
+from vllm.config import ModelConfig, PoolerConfig, get_current_vllm_config
16
from vllm.logger import init_logger
17
+from vllm.model_executor.models.adapters import _load_st_projector
18
from vllm.pooling_params import PoolingParams
19
from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
20
from vllm.tasks import PoolingTask
@@ -377,7 +378,6 @@ class PoolerClassify(PoolerActivation):
378
super().__init__()
379
380
if static_num_labels:
380
- from vllm.config import get_current_vllm_config
381
vllm_config = get_current_vllm_config()
382
self.num_labels = getattr(vllm_config.model_config.hf_config,
383
"num_labels", 0)
@@ -427,8 +427,6 @@ class EmbeddingPoolerHead(PoolerHead):
427
super().__init__(activation=PoolerNormalize())
428
429
# Load ST projector if available
430
- from vllm.config import get_current_vllm_config
431
- from vllm.model_executor.models.adapters import _load_st_projector
430
431
vllm_config = get_current_vllm_config()
432
self.projector: Optional[nn.Module] = _load_st_projector(
@@ -489,7 +487,6 @@ class RewardPoolerHead(PoolerHead):
487
def __init__(self) -> None:
488
super().__init__(activation=PoolerClassify(static_num_labels=False))
489
492
- from vllm.config import get_current_vllm_config
490
vllm_config = get_current_vllm_config()
491
self.head_dtype = vllm_config.model_config.head_dtype
492
@@ -638,7 +635,6 @@ class ClassifierPooler(Pooler):
635
) -> None:
636
super().__init__()
637
641
- from vllm.config import get_current_vllm_config
638
vllm_config = get_current_vllm_config()
639
640
self.pooling = pooling
@@ -730,3 +726,7 @@ class DispatchPooler(Pooler):
726
offset += num_items
727
728
return PoolerOutput(outputs)
729
+
730
+ def extra_repr(self) -> str:
731
+ s = f"supported_task={self.get_supported_tasks()}"
732
+ return s
vllm/v1/worker/gpu_model_runner.py
+1
@@ -3151,6 +3151,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3151
3152
model = cast(VllmModelForPooling, self.get_model())
3153
dummy_pooling_params = PoolingParams(task=task)
3154
+ dummy_pooling_params.verify(task=task, model_config=self.model_config)
3155
to_update = model.pooler.get_pooling_updates(task)
3156
to_update.apply(dummy_pooling_params)
3157