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 7, 10:15 UTC
[Refactor] Fix maxsim cuda platform and add cli to control it (#35427)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Wentao Ye committed
Mar 2, 2026 at 23:48 UTC
c21d0039ecc85b06617034ff2166a1fb79309d53
5 files changed
+25
-1
vllm/entrypoints/cli/serve.py
+6
@@ -220,6 +220,12 @@ def run_multi_api_server(args: argparse.Namespace):
220
num_api_servers: int = args.api_server_count
221
assert num_api_servers > 0
222
223
+ if num_api_servers > 1 and getattr(args, "use_gpu_for_pooling_score", False):
224
+ # TODO(wentao): remove this once well tested
225
+ raise ValueError(
226
+ "--use-gpu-for-pooling-score cannot be used with api_server_count > 1 now"
227
+ )
228
+
229
if num_api_servers > 1:
230
setup_multiprocess_prometheus()
231
vllm/entrypoints/openai/cli_args.py
+4
@@ -278,6 +278,10 @@ class FrontendArgs(BaseFrontendArgs):
278
Enable offline FastAPI documentation for air-gapped environments.
279
Uses vendored static assets bundled with vLLM.
280
"""
281
+ use_gpu_for_pooling_score: bool = False
282
+ """If set, run pooling score MaxSim on GPU in the API server process.
283
+ Can significantly improve late-interaction scoring performance.
284
+ https://github.com/vllm-project/vllm/pull/35330"""
285
286
@classmethod
287
def _customize_cli_kwargs(
vllm/entrypoints/pooling/__init__.py
+1
@@ -115,6 +115,7 @@ def init_pooling_state(
115
request_logger=request_logger,
116
score_template=resolved_chat_template,
117
log_error_stack=args.log_error_stack,
118
+ use_gpu_for_pooling_score=getattr(args, "use_gpu_for_pooling_score", False),
119
)
120
if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
121
else None
vllm/entrypoints/pooling/score/serving.py
+3
@@ -56,6 +56,7 @@ class ServingScores(OpenAIServing):
56
request_logger: RequestLogger | None,
57
score_template: str | None = None,
58
log_error_stack: bool = False,
59
+ use_gpu_for_pooling_score: bool = False,
60
) -> None:
61
super().__init__(
62
engine_client=engine_client,
@@ -64,6 +65,7 @@ class ServingScores(OpenAIServing):
65
log_error_stack=log_error_stack,
66
)
67
self.score_template = score_template
68
+ self.use_gpu_for_pooling_score = use_gpu_for_pooling_score
69
70
self._tokenizer_executor = ThreadPoolExecutor(max_workers=1)
71
@@ -314,6 +316,7 @@ class ServingScores(OpenAIServing):
316
maxsim_scores = compute_maxsim_scores(
317
[emb.outputs.data for emb in emb_data_1],
318
[emb.outputs.data for emb in emb_data_2],
319
+ use_gpu_for_pooling_score=self.use_gpu_for_pooling_score,
320
)
321
322
scores: list[PoolingRequestOutput] = []
vllm/entrypoints/pooling/score/utils.py
+11
-1
@@ -25,6 +25,7 @@ from vllm.inputs.data import PromptType, TextPrompt
25
from vllm.model_executor.models.interfaces import supports_score_template
26
from vllm.multimodal.inputs import MultiModalDataDict, MultiModalUUIDDict
27
from vllm.outputs import PoolingRequestOutput
28
+from vllm.platforms import current_platform
29
from vllm.renderers.hf import safe_apply_chat_template
30
from vllm.tokenizers import TokenizerLike
31
@@ -53,11 +54,16 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens
54
return token_scores.amax(dim=-1).sum()
55
56
57
+def _should_use_gpu_for_maxsim(use_gpu_for_pooling_score: bool) -> bool:
58
+ return use_gpu_for_pooling_score and not current_platform.is_cpu()
59
+
60
+
61
def compute_maxsim_scores(
62
q_embs: Sequence[torch.Tensor],
63
d_embs: Sequence[torch.Tensor],
64
max_batch_size: int = 16,
65
max_score_matrix_elements: int = 16_000_000,
66
+ use_gpu_for_pooling_score: bool = False,
67
) -> list[torch.Tensor]:
68
"""Compute ColBERT MaxSim scores in padded mini-batches."""
69
if len(q_embs) != len(d_embs):
@@ -73,7 +79,11 @@ def compute_maxsim_scores(
79
if q_emb.shape[1] != d_emb.shape[1]:
80
raise ValueError("Query and document embeddings must have same dim")
81
76
- compute_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
82
+ compute_device = torch.device(
83
+ current_platform.device_type
84
+ if _should_use_gpu_for_maxsim(use_gpu_for_pooling_score)
85
+ else "cpu"
86
+ )
87
scores: list[torch.Tensor] = []
88
start = 0
89
while start < num_pairs: