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, 04:45 UTC
[Bugfix][CPU] Fix InputBatch for pooling models in the CPU v1 (#20014)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Jun 24, 2025 at 21:20 UTC
53da4cd397821195e1b8bc481ada9f83a2007d05
2 files changed
+5
-1
tests/models/language/pooling/test_reward.py
+1
-1
@@ -101,4 +101,4 @@ def test_prm_models(
101
hf_output = torch.tensor(hf_output)
102
vllm_output = torch.tensor(vllm_output)
103
104
- assert torch.allclose(hf_output, vllm_output, 1e-2)
104
+ assert torch.allclose(hf_output, vllm_output, 1.5e-2)
vllm/v1/worker/cpu_model_runner.py
+4
@@ -7,6 +7,7 @@ import torch
7
from vllm.config import VllmConfig
8
from vllm.logger import init_logger
9
from vllm.model_executor.model_loader import get_model
10
+from vllm.model_executor.models.interfaces import has_step_pooler
11
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
12
13
logger = init_logger(__name__)
@@ -52,6 +53,9 @@ class CPUModelRunner(GPUModelRunner):
53
logger.info("Starting to load model %s...", self.model_config.model)
54
self.model = get_model(vllm_config=self.vllm_config)
55
56
+ if has_step_pooler(self.model):
57
+ self.input_batch.logits_processing_needs_token_ids = True
58
+
59
if self.lora_config:
60
self.model = self.load_lora_model(self.model, self.model_config,
61
self.scheduler_config,