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 6, 20:30 UTC
[Perf] Optimize compute maxsim using batched version, 3.2% E2E throughput improvement (#36710)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 11, 2026 at 20:37 UTC
c34ba6b9619f2398cfc4e87bf35555eff3590bf0
5 files changed
+141
-127
tests/entrypoints/pooling/score/test_utils.py
-36
@@ -4,13 +4,10 @@
4
from unittest.mock import patch
5
6
import pytest
7
-import torch
7
8
from vllm.config import ModelConfig
9
from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
10
from vllm.entrypoints.pooling.score.utils import (
12
- compute_maxsim_score,
13
- compute_maxsim_scores,
11
get_score_prompt,
12
)
13
from vllm.inputs import TokensPrompt
@@ -354,36 +351,3 @@ class TestGetScorePrompt:
351
assert_prompt_tokenization_consistent(
352
cross_encoder_tokenizer, full_prompt, engine_prompt
353
)
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)
tests/v1/worker/test_late_interaction_runner.py
+41
@@ -64,6 +64,47 @@ def test_postprocess_scores_and_releases_query_cache():
64
)
65
66
67
+def test_postprocess_scores_docs_in_batch():
68
+ runner = LateInteractionRunner()
69
+ query_key = "query-batch"
70
+ query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
71
+ doc_emb_1 = torch.tensor([[1.0, 0.0], [0.5, 0.5]], dtype=torch.float32)
72
+ doc_emb_2 = torch.tensor([[0.0, 1.0], [0.3, 0.7], [1.0, 0.0]], dtype=torch.float32)
73
+
74
+ query_params = _make_pooling_params(
75
+ build_late_interaction_query_params(query_key=query_key, query_uses=2)
76
+ )
77
+ runner.postprocess_pooler_output(
78
+ raw_pooler_output=[query_emb],
79
+ pooling_params=[query_params],
80
+ req_ids=["query-req"],
81
+ finished_mask=[True],
82
+ )
83
+
84
+ doc_params = _make_pooling_params(
85
+ build_late_interaction_doc_params(query_key=query_key)
86
+ )
87
+ doc_output = runner.postprocess_pooler_output(
88
+ raw_pooler_output=[doc_emb_1, doc_emb_2],
89
+ pooling_params=[doc_params, doc_params],
90
+ req_ids=["doc-req-1", "doc-req-2"],
91
+ finished_mask=[True, True],
92
+ )
93
+ assert isinstance(doc_output, list)
94
+ assert doc_output[0] is not None
95
+ assert doc_output[1] is not None
96
+ assert torch.allclose(doc_output[0], compute_maxsim_score(query_emb, doc_emb_1))
97
+ assert torch.allclose(doc_output[1], compute_maxsim_score(query_emb, doc_emb_2))
98
+
99
+ with pytest.raises(ValueError, match="query cache miss"):
100
+ runner.postprocess_pooler_output(
101
+ raw_pooler_output=[doc_emb_1],
102
+ pooling_params=[doc_params],
103
+ req_ids=["doc-req-3"],
104
+ finished_mask=[True],
105
+ )
106
+
107
+
108
def test_finished_request_releases_unscored_doc_use():
109
runner = LateInteractionRunner()
110
query_key = "query-cancel"
vllm/entrypoints/pooling/score/utils.py
+1
-87
@@ -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, Sequence
3
+from collections.abc import Iterable
4
from typing import Any, TypeAlias, cast
5
6
import torch
@@ -25,7 +25,6 @@ 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
28
from vllm.renderers.hf import safe_apply_chat_template
29
from vllm.tokenizers import TokenizerLike
30
@@ -54,91 +53,6 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens
53
return token_scores.amax(dim=-1).sum()
54
55
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):
70
- raise ValueError("q_embs and d_embs must have the same length")
71
-
72
- num_pairs = len(q_embs)
73
- if num_pairs == 0:
74
- return []
75
-
76
- for q_emb, d_emb in zip(q_embs, d_embs):
77
- if q_emb.ndim != 2 or d_emb.ndim != 2:
78
- raise ValueError("Each embedding tensor must be 2-D")
79
- if q_emb.shape[1] != d_emb.shape[1]:
80
- raise ValueError("Query and document embeddings must have same dim")
81
-
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:
90
- end = min(start + max_batch_size, num_pairs)
91
- max_q = max(int(x.shape[0]) for x in q_embs[start:end])
92
- max_d = max(int(x.shape[0]) for x in d_embs[start:end])
93
-
94
- # keep score matrix bounded to avoid oversized allocations.
95
- while (
96
- end - start > 1
97
- and (end - start) * max_q * max_d > max_score_matrix_elements
98
- ):
99
- end -= 1
100
- max_q = max(int(x.shape[0]) for x in q_embs[start:end])
101
- max_d = max(int(x.shape[0]) for x in d_embs[start:end])
102
-
103
- batch_q = q_embs[start:end]
104
- batch_d = d_embs[start:end]
105
- batch_size = end - start
106
- dim = int(batch_q[0].shape[1])
107
- dtype = batch_q[0].dtype
108
-
109
- q_batch = torch.zeros(
110
- (batch_size, max_q, dim), dtype=dtype, device=compute_device
111
- )
112
- d_batch = torch.zeros(
113
- (batch_size, max_d, dim), dtype=dtype, device=compute_device
114
- )
115
- q_mask = torch.zeros(
116
- (batch_size, max_q), dtype=torch.bool, device=compute_device
117
- )
118
- d_mask = torch.zeros(
119
- (batch_size, max_d), dtype=torch.bool, device=compute_device
120
- )
121
-
122
- # copy to padded tensors
123
- for i, (q_emb, d_emb) in enumerate(zip(batch_q, batch_d)):
124
- q_len = int(q_emb.shape[0])
125
- d_len = int(d_emb.shape[0])
126
- q_batch[i, :q_len] = q_emb.to(device=compute_device, dtype=dtype)
127
- d_batch[i, :d_len] = d_emb.to(device=compute_device, dtype=dtype)
128
- q_mask[i, :q_len] = True
129
- d_mask[i, :d_len] = True
130
-
131
- token_scores = torch.bmm(q_batch, d_batch.transpose(1, 2))
132
- token_scores.masked_fill_(~d_mask.unsqueeze(1), float("-inf"))
133
- max_per_query = token_scores.amax(dim=-1)
134
- max_per_query.masked_fill_(~q_mask, 0)
135
- batch_scores = max_per_query.sum(dim=-1).to("cpu")
136
- scores.extend(batch_scores.unbind(0))
137
- start = end
138
-
139
- return [cast(torch.Tensor, score) for score in scores]
140
-
141
-
56
class ScoreMultiModalParam(TypedDict, total=False):
57
"""
58
A specialized parameter type for scoring multimodal content
vllm/v1/pool/late_interaction.py
+79
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import zlib
5
+from collections.abc import Sequence
6
7
import torch
8
@@ -62,3 +63,81 @@ def compute_maxsim_score(
63
# compute in float32 for numerical stability
64
token_scores = torch.matmul(q_emb.float(), d_emb.float().T)
65
return token_scores.amax(dim=-1).sum()
66
+
67
+
68
+def compute_maxsim_scores(
69
+ q_embs: Sequence[torch.Tensor],
70
+ d_embs: Sequence[torch.Tensor],
71
+ max_batch_size: int = 64,
72
+ max_score_matrix_elements: int = 64_000_000,
73
+) -> list[torch.Tensor]:
74
+ """Compute MaxSim for multiple query/doc pairs in mini-batches."""
75
+ if len(q_embs) != len(d_embs):
76
+ raise ValueError("q_embs and d_embs must have the same length")
77
+
78
+ num_pairs = len(q_embs)
79
+ if num_pairs == 0:
80
+ return []
81
+
82
+ if max_batch_size <= 0:
83
+ raise ValueError("max_batch_size must be greater than 0")
84
+ if max_score_matrix_elements <= 0:
85
+ raise ValueError("max_score_matrix_elements must be greater than 0")
86
+
87
+ for q_emb, d_emb in zip(q_embs, d_embs):
88
+ if q_emb.ndim != 2 or d_emb.ndim != 2:
89
+ raise ValueError("Each embedding tensor must be 2-D")
90
+ if q_emb.shape[1] != d_emb.shape[1]:
91
+ raise ValueError("Query and document embeddings must have same dim")
92
+ if q_emb.device != d_emb.device:
93
+ raise ValueError("Query and document embeddings must be on same device")
94
+
95
+ scores: list[torch.Tensor] = []
96
+ start = 0
97
+ while start < num_pairs:
98
+ end = min(start + max_batch_size, num_pairs)
99
+ max_q = max(int(x.shape[0]) for x in q_embs[start:end])
100
+ max_d = max(int(x.shape[0]) for x in d_embs[start:end])
101
+
102
+ # keep score matrix bounded to avoid oversized allocations.
103
+ while (
104
+ end - start > 1
105
+ and (end - start) * max_q * max_d > max_score_matrix_elements
106
+ ):
107
+ end -= 1
108
+ max_q = max(int(x.shape[0]) for x in q_embs[start:end])
109
+ max_d = max(int(x.shape[0]) for x in d_embs[start:end])
110
+
111
+ batch_q = q_embs[start:end]
112
+ batch_d = d_embs[start:end]
113
+ batch_size = end - start
114
+ device = batch_q[0].device
115
+ dim = int(batch_q[0].shape[1])
116
+
117
+ q_batch = torch.zeros(
118
+ (batch_size, max_q, dim), dtype=torch.float32, device=device
119
+ )
120
+ d_batch = torch.zeros(
121
+ (batch_size, max_d, dim), dtype=torch.float32, device=device
122
+ )
123
+ q_mask = torch.zeros((batch_size, max_q), dtype=torch.bool, device=device)
124
+ d_mask = torch.zeros((batch_size, max_d), dtype=torch.bool, device=device)
125
+
126
+ # copy to padded tensors
127
+ for i, (q_emb, d_emb) in enumerate(zip(batch_q, batch_d)):
128
+ q_len = int(q_emb.shape[0])
129
+ d_len = int(d_emb.shape[0])
130
+ q_batch[i, :q_len] = q_emb.to(device=device, dtype=torch.float32)
131
+ d_batch[i, :d_len] = d_emb.to(device=device, dtype=torch.float32)
132
+ q_mask[i, :q_len] = True
133
+ d_mask[i, :d_len] = True
134
+
135
+ token_scores = torch.bmm(q_batch, d_batch.transpose(1, 2))
136
+ token_scores.masked_fill_(~d_mask.unsqueeze(1), float("-inf"))
137
+ max_per_query = token_scores.amax(dim=-1)
138
+ max_per_query.masked_fill_(~q_mask, 0.0)
139
+ batch_scores = max_per_query.sum(dim=-1)
140
+ scores.extend(batch_scores.unbind(0))
141
+ start = end
142
+
143
+ return scores
vllm/v1/worker/gpu/pool/late_interaction_runner.py
+20
-4
@@ -9,7 +9,7 @@ from vllm.v1.outputs import PoolerOutput
9
from vllm.v1.pool.late_interaction import (
10
LATE_INTERACTION_MODE_CACHE_QUERY,
11
LATE_INTERACTION_MODE_SCORE_DOC,
12
- compute_maxsim_score,
12
+ compute_maxsim_scores,
13
)
14
15
@@ -72,6 +72,11 @@ class LateInteractionRunner:
72
return raw_pooler_output
73
74
outputs: list[torch.Tensor | None] = list(raw_pooler_output)
75
+ score_indices: list[int] = []
76
+ score_req_ids: list[str] = []
77
+ score_query_keys: list[str] = []
78
+ score_queries: list[torch.Tensor] = []
79
+ score_docs: list[torch.Tensor] = []
80
for i, (req_id, output, params, finished) in enumerate(
81
zip(req_ids, outputs, pooling_params, finished_mask)
82
):
@@ -101,13 +106,24 @@ class LateInteractionRunner:
106
"before their paired document requests."
107
)
108
104
- outputs[i] = compute_maxsim_score(query_output, output)
105
- self._doc_query_keys.pop(req_id, None)
106
- self._release_query_use(query_key)
109
+ score_indices.append(i)
110
+ score_req_ids.append(req_id)
111
+ score_query_keys.append(query_key)
112
+ score_queries.append(query_output)
113
+ score_docs.append(output)
114
continue
115
116
raise ValueError(f"Unsupported late-interaction mode: {mode!r}")
117
118
+ if score_indices:
119
+ score_values = compute_maxsim_scores(score_queries, score_docs)
120
+ for i, req_id, query_key, score in zip(
121
+ score_indices, score_req_ids, score_query_keys, score_values
122
+ ):
123
+ outputs[i] = score
124
+ self._doc_query_keys.pop(req_id, None)
125
+ self._release_query_use(query_key)
126
+
127
return outputs
128
129
def _release_query_use(self, query_key: str) -> None: