| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import zlib |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.pooling_params import LateInteractionParams, PoolingParams |
| 10 | |
| 11 | LATE_INTERACTION_MODE_CACHE_QUERY = "cache_query" |
| 12 | LATE_INTERACTION_MODE_SCORE_DOC = "score_doc" |
| 13 | |
| 14 | |
| 15 | def get_late_interaction_engine_index( |
| 16 | pooling_params: PoolingParams | None, |
| 17 | num_engines: int, |
| 18 | ) -> int | None: |
| 19 | if pooling_params is None or pooling_params.late_interaction_params is None: |
| 20 | return None |
| 21 | |
| 22 | late_interaction_params = pooling_params.late_interaction_params |
| 23 | mode = late_interaction_params.mode |
| 24 | if mode not in ( |
| 25 | LATE_INTERACTION_MODE_CACHE_QUERY, |
| 26 | LATE_INTERACTION_MODE_SCORE_DOC, |
| 27 | ): |
| 28 | return None |
| 29 | |
| 30 | query_key = late_interaction_params.query_key |
| 31 | if not isinstance(query_key, str) or not query_key: |
| 32 | return None |
| 33 | |
| 34 | # query embeddings are cached in process-local worker memory, |
| 35 | # pin requests sharing the same query key to the same engine. |
| 36 | return zlib.crc32(query_key.encode("utf-8")) % num_engines |
| 37 | |
| 38 | |
| 39 | def build_late_interaction_query_params( |
| 40 | query_key: str, |
| 41 | query_uses: int, |
| 42 | ) -> LateInteractionParams: |
| 43 | return LateInteractionParams( |
| 44 | mode=LATE_INTERACTION_MODE_CACHE_QUERY, |
| 45 | query_key=query_key, |
| 46 | query_uses=max(1, int(query_uses)), |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | def build_late_interaction_doc_params( |
| 51 | query_key: str, |
| 52 | ) -> LateInteractionParams: |
| 53 | return LateInteractionParams( |
| 54 | mode=LATE_INTERACTION_MODE_SCORE_DOC, |
| 55 | query_key=query_key, |
| 56 | ) |
| 57 | |
| 58 | |
| 59 | def compute_maxsim_score_batched( |
| 60 | q_embs: Sequence[torch.Tensor], |
| 61 | d_embs: Sequence[torch.Tensor], |
| 62 | max_batch_size: int = 64, |
| 63 | max_score_matrix_elements: int = 64_000_000, |
| 64 | ) -> list[torch.Tensor]: |
| 65 | """Compute MaxSim for multiple query/doc pairs in mini-batches.""" |
| 66 | if len(q_embs) != len(d_embs): |
| 67 | raise ValueError("q_embs and d_embs must have the same length") |
| 68 | |
| 69 | num_pairs = len(q_embs) |
| 70 | if num_pairs == 0: |
| 71 | return [] |
| 72 | |
| 73 | if max_batch_size <= 0: |
| 74 | raise ValueError("max_batch_size must be greater than 0") |
| 75 | if max_score_matrix_elements <= 0: |
| 76 | raise ValueError("max_score_matrix_elements must be greater than 0") |
| 77 | |
| 78 | for q_emb, d_emb in zip(q_embs, d_embs): |
| 79 | if q_emb.ndim != 2 or d_emb.ndim != 2: |
| 80 | raise ValueError("Each embedding tensor must be 2-D") |
| 81 | if q_emb.shape[1] != d_emb.shape[1]: |
| 82 | raise ValueError("Query and document embeddings must have same dim") |
| 83 | if q_emb.device != d_emb.device: |
| 84 | raise ValueError("Query and document embeddings must be on same device") |
| 85 | |
| 86 | scores: list[torch.Tensor] = [] |
| 87 | start = 0 |
| 88 | while start < num_pairs: |
| 89 | end = min(start + max_batch_size, num_pairs) |
| 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 | # keep score matrix bounded to avoid oversized allocations. |
| 94 | while ( |
| 95 | end - start > 1 |
| 96 | and (end - start) * max_q * max_d > max_score_matrix_elements |
| 97 | ): |
| 98 | end -= 1 |
| 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 | batch_q = q_embs[start:end] |
| 103 | batch_d = d_embs[start:end] |
| 104 | batch_size = end - start |
| 105 | device = batch_q[0].device |
| 106 | dim = int(batch_q[0].shape[1]) |
| 107 | |
| 108 | q_batch = torch.zeros( |
| 109 | (batch_size, max_q, dim), dtype=torch.float32, device=device |
| 110 | ) |
| 111 | d_batch = torch.zeros( |
| 112 | (batch_size, max_d, dim), dtype=torch.float32, device=device |
| 113 | ) |
| 114 | q_mask = torch.zeros((batch_size, max_q), dtype=torch.bool, device=device) |
| 115 | d_mask = torch.zeros((batch_size, max_d), dtype=torch.bool, device=device) |
| 116 | |
| 117 | # copy to padded tensors |
| 118 | for i, (q_emb, d_emb) in enumerate(zip(batch_q, batch_d)): |
| 119 | q_len = int(q_emb.shape[0]) |
| 120 | d_len = int(d_emb.shape[0]) |
| 121 | q_batch[i, :q_len] = q_emb.to(device=device, dtype=torch.float32) |
| 122 | d_batch[i, :d_len] = d_emb.to(device=device, dtype=torch.float32) |
| 123 | q_mask[i, :q_len] = True |
| 124 | d_mask[i, :d_len] = True |
| 125 | |
| 126 | token_scores = torch.bmm(q_batch, d_batch.transpose(1, 2)) |
| 127 | token_scores.masked_fill_(~d_mask.unsqueeze(1), float("-inf")) |
| 128 | max_per_query = token_scores.amax(dim=-1) |
| 129 | max_per_query.masked_fill_(~q_mask, 0.0) |
| 130 | batch_scores = max_per_query.sum(dim=-1) |
| 131 | scores.extend(batch_scores.unbind(0)) |
| 132 | start = end |
| 133 | |
| 134 | return scores |