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, 16:30 UTC
[Perf] Optimize maxsim scores computation for pooling models, 13.9% E2E throughput improvement (#35330)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Feb 26, 2026 at 12:14 UTC
99c7892c5bf20afc90e2ef0e1ad0a89637ae67a9
3 files changed
+123
-11
tests/entrypoints/pooling/score/test_utils.py
+39
-1
@@ -4,10 +4,15 @@
4
from unittest.mock import patch
5
6
import pytest
7
+import torch
8
9
from vllm.config import ModelConfig
10
from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
10
-from vllm.entrypoints.pooling.score.utils import get_score_prompt
11
+from vllm.entrypoints.pooling.score.utils import (
12
+ compute_maxsim_score,
13
+ compute_maxsim_scores,
14
+ get_score_prompt,
15
+)
16
from vllm.inputs import TokensPrompt
17
from vllm.tokenizers import get_tokenizer
18
@@ -349,3 +354,36 @@ class TestGetScorePrompt:
354
assert_prompt_tokenization_consistent(
355
cross_encoder_tokenizer, full_prompt, engine_prompt
356
)
357
+
358
+
359
+def test_compute_maxsim_scores_matches_reference_per_pair() -> None:
360
+ generator = torch.Generator()
361
+ generator.manual_seed(7)
362
+
363
+ shared_query = torch.randn(5, 8, generator=generator)
364
+ q_embs = [
365
+ shared_query, # 1:N style shared query
366
+ shared_query,
367
+ torch.randn(2, 8, generator=generator),
368
+ torch.randn(4, 8, generator=generator),
369
+ ]
370
+ d_embs = [
371
+ torch.randn(6, 8, generator=generator),
372
+ torch.randn(3, 8, generator=generator),
373
+ torch.randn(5, 8, generator=generator),
374
+ torch.randn(7, 8, generator=generator),
375
+ ]
376
+
377
+ batched_scores = compute_maxsim_scores(
378
+ q_embs,
379
+ d_embs,
380
+ max_batch_size=4,
381
+ max_score_matrix_elements=40, # batch shrinking path.
382
+ )
383
+ reference_scores = [
384
+ compute_maxsim_score(q, d).to("cpu") for q, d in zip(q_embs, d_embs)
385
+ ]
386
+
387
+ assert len(batched_scores) == len(reference_scores)
388
+ for batched, reference in zip(batched_scores, reference_scores):
389
+ torch.testing.assert_close(batched, reference, rtol=1e-4, atol=1e-4)
vllm/entrypoints/pooling/score/serving.py
+7
-9
@@ -31,7 +31,7 @@ from vllm.entrypoints.pooling.score.utils import (
31
ScoreInputs,
32
_cosine_similarity,
33
compress_token_type_ids,
34
- compute_maxsim_score,
34
+ compute_maxsim_scores,
35
get_score_prompt,
36
parse_score_data_single,
37
validate_score_input,
@@ -311,19 +311,17 @@ class ServingScores(OpenAIServing):
311
# Compute MaxSim scores
312
from vllm.outputs import PoolingOutput
313
314
+ maxsim_scores = compute_maxsim_scores(
315
+ [emb.outputs.data for emb in emb_data_1],
316
+ [emb.outputs.data for emb in emb_data_2],
317
+ )
318
+
319
scores: list[PoolingRequestOutput] = []
320
padding: list[int] = []
321
if (pad_token_id := tokenizer.pad_token_id) is not None:
322
padding = [pad_token_id]
323
319
- for emb_1, emb_2 in zip(emb_data_1, emb_data_2):
320
- # emb_1.outputs.data: [query_len, dim]
321
- # emb_2.outputs.data: [doc_len, dim]
322
- q_emb = emb_1.outputs.data
323
- d_emb = emb_2.outputs.data
324
-
325
- maxsim_score = compute_maxsim_score(q_emb, d_emb)
326
-
324
+ for emb_1, emb_2, maxsim_score in zip(emb_data_1, emb_data_2, maxsim_scores):
325
tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
326
327
scores.append(
vllm/entrypoints/pooling/score/utils.py
+77
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Iterable
3
+from collections.abc import Iterable, Sequence
4
from typing import Any, TypeAlias, cast
5
6
import torch
@@ -53,6 +53,82 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens
53
return token_scores.amax(dim=-1).sum()
54
55
56
+def compute_maxsim_scores(
57
+ q_embs: Sequence[torch.Tensor],
58
+ d_embs: Sequence[torch.Tensor],
59
+ max_batch_size: int = 16,
60
+ max_score_matrix_elements: int = 16_000_000,
61
+) -> list[torch.Tensor]:
62
+ """Compute ColBERT MaxSim scores in padded mini-batches."""
63
+ if len(q_embs) != len(d_embs):
64
+ raise ValueError("q_embs and d_embs must have the same length")
65
+
66
+ num_pairs = len(q_embs)
67
+ if num_pairs == 0:
68
+ return []
69
+
70
+ for q_emb, d_emb in zip(q_embs, d_embs):
71
+ if q_emb.ndim != 2 or d_emb.ndim != 2:
72
+ raise ValueError("Each embedding tensor must be 2-D")
73
+ if q_emb.shape[1] != d_emb.shape[1]:
74
+ raise ValueError("Query and document embeddings must have same dim")
75
+
76
+ compute_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
77
+ scores: list[torch.Tensor] = []
78
+ start = 0
79
+ while start < num_pairs:
80
+ end = min(start + max_batch_size, num_pairs)
81
+ max_q = max(int(x.shape[0]) for x in q_embs[start:end])
82
+ max_d = max(int(x.shape[0]) for x in d_embs[start:end])
83
+
84
+ # keep score matrix bounded to avoid oversized allocations.
85
+ while (
86
+ end - start > 1
87
+ and (end - start) * max_q * max_d > max_score_matrix_elements
88
+ ):
89
+ end -= 1
90
+ max_q = max(int(x.shape[0]) for x in q_embs[start:end])
91
+ max_d = max(int(x.shape[0]) for x in d_embs[start:end])
92
+
93
+ batch_q = q_embs[start:end]
94
+ batch_d = d_embs[start:end]
95
+ batch_size = end - start
96
+ dim = int(batch_q[0].shape[1])
97
+ dtype = batch_q[0].dtype
98
+
99
+ q_batch = torch.zeros(
100
+ (batch_size, max_q, dim), dtype=dtype, device=compute_device
101
+ )
102
+ d_batch = torch.zeros(
103
+ (batch_size, max_d, dim), dtype=dtype, device=compute_device
104
+ )
105
+ q_mask = torch.zeros(
106
+ (batch_size, max_q), dtype=torch.bool, device=compute_device
107
+ )
108
+ d_mask = torch.zeros(
109
+ (batch_size, max_d), dtype=torch.bool, device=compute_device
110
+ )
111
+
112
+ # copy to padded tensors
113
+ for i, (q_emb, d_emb) in enumerate(zip(batch_q, batch_d)):
114
+ q_len = int(q_emb.shape[0])
115
+ d_len = int(d_emb.shape[0])
116
+ q_batch[i, :q_len] = q_emb.to(device=compute_device, dtype=dtype)
117
+ d_batch[i, :d_len] = d_emb.to(device=compute_device, dtype=dtype)
118
+ q_mask[i, :q_len] = True
119
+ d_mask[i, :d_len] = True
120
+
121
+ token_scores = torch.bmm(q_batch, d_batch.transpose(1, 2))
122
+ token_scores.masked_fill_(~d_mask.unsqueeze(1), float("-inf"))
123
+ max_per_query = token_scores.amax(dim=-1)
124
+ max_per_query.masked_fill_(~q_mask, 0)
125
+ batch_scores = max_per_query.sum(dim=-1).to("cpu")
126
+ scores.extend(batch_scores.unbind(0))
127
+ start = end
128
+
129
+ return [cast(torch.Tensor, score) for score in scores]
130
+
131
+
132
class ScoreMultiModalParam(TypedDict, total=False):
133
"""
134
A specialized parameter type for scoring multimodal content