| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from fastapi.responses import JSONResponse, Response |
| 5 | |
| 6 | from vllm import PoolingParams |
| 7 | from vllm.engine.protocol import EngineClient |
| 8 | from vllm.entrypoints.serve.engine.protocol import UsageInfo |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput |
| 11 | from vllm.tasks import SCORE_TYPE_MAP, SupportedTask |
| 12 | from vllm.utils import random_uuid |
| 13 | from vllm.v1.pool.late_interaction import ( |
| 14 | build_late_interaction_doc_params, |
| 15 | build_late_interaction_query_params, |
| 16 | ) |
| 17 | |
| 18 | from ..base.io_processor import PoolingIOProcessor |
| 19 | from ..base.serving import PoolingServing |
| 20 | from .io_processor import ScoringIOProcessors, ScoringServeContext |
| 21 | from .protocol import ( |
| 22 | RerankDocument, |
| 23 | RerankRequest, |
| 24 | RerankResponse, |
| 25 | RerankResult, |
| 26 | RerankUsage, |
| 27 | ScoreRequest, |
| 28 | ScoreResponse, |
| 29 | ScoreResponseData, |
| 30 | ) |
| 31 | from .typing import ScoreInput |
| 32 | |
| 33 | logger = init_logger(__name__) |
| 34 | |
| 35 | |
| 36 | class ServingScores(PoolingServing): |
| 37 | request_id_prefix = "score" |
| 38 | |
| 39 | def __init__( |
| 40 | self, |
| 41 | engine_client: EngineClient, |
| 42 | *args, |
| 43 | supported_tasks: tuple[SupportedTask, ...], |
| 44 | enable_flash_late_interaction: bool = True, |
| 45 | **kwargs, |
| 46 | ): |
| 47 | pooling_task = engine_client.model_config.get_pooling_task(supported_tasks) |
| 48 | score_type = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type] |
| 49 | assert score_type is not None |
| 50 | |
| 51 | self.io_processor_name: str = score_type |
| 52 | self.enable_flash_late_interaction = ( |
| 53 | self.io_processor_name == "late-interaction" |
| 54 | and enable_flash_late_interaction |
| 55 | ) |
| 56 | |
| 57 | if self.enable_flash_late_interaction: |
| 58 | self.io_processor_name = "flash-late-interaction" |
| 59 | |
| 60 | if engine_client.model_config.architecture == "JinaForRanking": |
| 61 | self.io_processor_name = "jina-reranking-scoring" |
| 62 | self.enable_flash_late_interaction = False |
| 63 | |
| 64 | super().__init__(engine_client, *args, **kwargs) |
| 65 | |
| 66 | def init_io_processor(self, *args, **kwargs) -> PoolingIOProcessor: |
| 67 | return ScoringIOProcessors[self.io_processor_name](*args, **kwargs) |
| 68 | |
| 69 | async def __call__(self, *args, **kwargs) -> Response: |
| 70 | if not self.enable_flash_late_interaction: |
| 71 | return await super().__call__(*args, **kwargs) |
| 72 | |
| 73 | return await self.flash_late_interaction(*args, **kwargs) |
| 74 | |
| 75 | def _build_response( |
| 76 | self, |
| 77 | ctx: ScoringServeContext, |
| 78 | ) -> JSONResponse: |
| 79 | final_res_batch = ctx.final_res_batch |
| 80 | request_id = ctx.request_id |
| 81 | created_time = ctx.created_time |
| 82 | model_name = ctx.model_name |
| 83 | assert ctx.engine_inputs is not None |
| 84 | num_prompt_tokens = sum( |
| 85 | self._extract_prompt_len(engine_input["prompts"]) |
| 86 | for engine_input in ctx.engine_inputs |
| 87 | ) |
| 88 | |
| 89 | if isinstance(ctx.request, ScoreRequest): |
| 90 | return self._request_output_to_score_response( |
| 91 | final_res_batch, |
| 92 | request_id, |
| 93 | created_time, |
| 94 | model_name, |
| 95 | num_prompt_tokens, |
| 96 | ) |
| 97 | elif isinstance(ctx.request, RerankRequest): |
| 98 | return self._request_output_to_rerank_response( |
| 99 | final_res_batch, |
| 100 | request_id, |
| 101 | model_name, |
| 102 | ctx.request.documents, |
| 103 | ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch), |
| 104 | num_prompt_tokens, |
| 105 | ) |
| 106 | else: |
| 107 | raise ValueError(f"Invalid {self.request_id_prefix} request type") |
| 108 | |
| 109 | def _request_output_to_score_response( |
| 110 | self, |
| 111 | final_res_batch: list[PoolingRequestOutput], |
| 112 | request_id: str, |
| 113 | created_time: int, |
| 114 | model_name: str, |
| 115 | num_prompt_tokens: int, |
| 116 | ) -> JSONResponse: |
| 117 | items: list[ScoreResponseData] = [] |
| 118 | |
| 119 | for idx, final_res in enumerate(final_res_batch): |
| 120 | classify_res = ScoringRequestOutput.from_base(final_res) |
| 121 | |
| 122 | item = ScoreResponseData( |
| 123 | index=idx, |
| 124 | score=classify_res.outputs.score, |
| 125 | ) |
| 126 | items.append(item) |
| 127 | |
| 128 | usage = UsageInfo( |
| 129 | prompt_tokens=num_prompt_tokens, |
| 130 | total_tokens=num_prompt_tokens, |
| 131 | ) |
| 132 | |
| 133 | response = ScoreResponse( |
| 134 | id=request_id, |
| 135 | created=created_time, |
| 136 | model=model_name, |
| 137 | data=items, |
| 138 | usage=usage, |
| 139 | ) |
| 140 | |
| 141 | return JSONResponse(content=response.model_dump()) |
| 142 | |
| 143 | def _request_output_to_rerank_response( |
| 144 | self, |
| 145 | final_res_batch: list[PoolingRequestOutput], |
| 146 | request_id: str, |
| 147 | model_name: str, |
| 148 | documents: ScoreInput | list[ScoreInput], |
| 149 | top_n: int, |
| 150 | num_prompt_tokens: int, |
| 151 | ) -> JSONResponse: |
| 152 | if not isinstance(documents, list): |
| 153 | documents = [documents] |
| 154 | |
| 155 | results: list[RerankResult] = [] |
| 156 | for idx, final_res in enumerate(final_res_batch): |
| 157 | classify_res = ScoringRequestOutput.from_base(final_res) |
| 158 | |
| 159 | document = documents[idx] |
| 160 | if isinstance(document, str): |
| 161 | rerank_document = RerankDocument(text=document) |
| 162 | else: |
| 163 | rerank_document = RerankDocument( |
| 164 | multi_modal=document.get("content", []) |
| 165 | ) |
| 166 | |
| 167 | result = RerankResult( |
| 168 | index=idx, |
| 169 | document=rerank_document, |
| 170 | relevance_score=classify_res.outputs.score, |
| 171 | ) |
| 172 | results.append(result) |
| 173 | |
| 174 | # sort by relevance, then return the top n if set |
| 175 | results.sort(key=lambda x: x.relevance_score, reverse=True) |
| 176 | if top_n < len(documents): |
| 177 | results = results[:top_n] |
| 178 | |
| 179 | response = RerankResponse( |
| 180 | id=request_id, |
| 181 | model=model_name, |
| 182 | results=results, |
| 183 | usage=RerankUsage( |
| 184 | total_tokens=num_prompt_tokens, prompt_tokens=num_prompt_tokens |
| 185 | ), |
| 186 | ) |
| 187 | |
| 188 | return JSONResponse(content=response.model_dump()) |
| 189 | |
| 190 | ################################################################################### |
| 191 | ### Run pooling score MaxSim on worker side (GPU) in the API server process |
| 192 | ### Can significantly improve late-interaction scoring performance. |
| 193 | |
| 194 | async def flash_late_interaction(self, *args, **kwargs) -> Response: |
| 195 | ctx = await self._init_ctx(self.io_processor, *args, **kwargs) |
| 196 | await self._preprocessing(self.io_processor, ctx) |
| 197 | |
| 198 | try: |
| 199 | # stage 1: encode queries and cache token embeddings on workers. |
| 200 | await self._flash_late_interaction_encode_queries(ctx) |
| 201 | # stage 2: encode docs and return scalar scores from workers. |
| 202 | await self._flash_late_interaction_encode_docs(ctx) |
| 203 | except BaseException: |
| 204 | try: |
| 205 | await self._cleanup_flash_late_interaction(ctx) |
| 206 | except Exception: |
| 207 | logger.exception("Failed to clean up late-interaction query cache.") |
| 208 | raise |
| 209 | |
| 210 | return await self._postprocessing_async(self.io_processor, ctx) |
| 211 | |
| 212 | async def _cleanup_flash_late_interaction(self, ctx: ScoringServeContext) -> None: |
| 213 | query_keys = ctx.late_interaction_query_keys |
| 214 | if not query_keys: |
| 215 | return |
| 216 | |
| 217 | doc_keys = ctx.late_interaction_doc_keys or [] |
| 218 | |
| 219 | try: |
| 220 | # Stop documents before removing the query tensors they reference. |
| 221 | await self.engine_client.abort([*query_keys, *doc_keys]) |
| 222 | finally: |
| 223 | await self.engine_client.collective_rpc( |
| 224 | "release_late_interaction_query_cache", |
| 225 | args=(query_keys,), |
| 226 | ) |
| 227 | |
| 228 | async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext): |
| 229 | assert ctx.n_queries is not None |
| 230 | assert ctx.engine_inputs is not None |
| 231 | assert isinstance(ctx.pooling_params, PoolingParams) |
| 232 | |
| 233 | n_queries = ctx.n_queries |
| 234 | n_docs = len(ctx.engine_inputs) - n_queries |
| 235 | query_engine_inputs = ctx.engine_inputs[:n_queries] |
| 236 | |
| 237 | query_namespace = random_uuid() |
| 238 | query_keys = [ |
| 239 | f"late-interaction-{query_namespace}-query-{i}" for i in range(n_queries) |
| 240 | ] |
| 241 | ctx.late_interaction_query_keys = query_keys |
| 242 | query_uses = [n_docs if n_queries == 1 else 1] * n_queries |
| 243 | |
| 244 | for i in range(n_queries): |
| 245 | pooling_params = ctx.pooling_params.clone() |
| 246 | pooling_params.late_interaction_params = ( |
| 247 | build_late_interaction_query_params( |
| 248 | query_key=query_keys[i], |
| 249 | query_uses=query_uses[i], |
| 250 | ) |
| 251 | ) |
| 252 | query_engine_inputs[i]["params"] = pooling_params |
| 253 | |
| 254 | assert n_queries == len(query_engine_inputs) == len(query_keys) |
| 255 | |
| 256 | query_ctx = ScoringServeContext( |
| 257 | request=ctx.request, |
| 258 | raw_request=ctx.raw_request, |
| 259 | model_name=ctx.model_name, |
| 260 | request_id=ctx.request_id, |
| 261 | pooling_params=ctx.pooling_params, |
| 262 | prompt_request_ids=query_keys, |
| 263 | engine_inputs=query_engine_inputs, |
| 264 | lora_request=ctx.lora_request, |
| 265 | priorities=ctx.priorities, |
| 266 | prompt_extras=ctx.prompt_extras, |
| 267 | ) |
| 268 | |
| 269 | await self._prepare_generators(query_ctx) |
| 270 | await self._collect_batch(query_ctx) |
| 271 | ctx.query_final_res_batch = query_ctx.final_res_batch |
| 272 | |
| 273 | async def _flash_late_interaction_encode_docs(self, ctx: ScoringServeContext): |
| 274 | assert ctx.n_queries is not None |
| 275 | assert ctx.engine_inputs is not None |
| 276 | assert isinstance(ctx.pooling_params, PoolingParams) |
| 277 | |
| 278 | n_queries = ctx.n_queries |
| 279 | n_docs = len(ctx.engine_inputs) - n_queries |
| 280 | doc_engine_inputs = ctx.engine_inputs[n_queries:] |
| 281 | |
| 282 | query_keys = ctx.late_interaction_query_keys |
| 283 | if query_keys is None: |
| 284 | raise RuntimeError("Late-interaction query keys were not initialized.") |
| 285 | doc_namespace = random_uuid() |
| 286 | doc_keys = [f"late-interaction-{doc_namespace}-doc-{i}" for i in range(n_docs)] |
| 287 | ctx.late_interaction_doc_keys = doc_keys |
| 288 | |
| 289 | for i in range(n_docs): |
| 290 | query_idx = 0 if n_queries == 1 else i |
| 291 | pooling_params = ctx.pooling_params.clone() |
| 292 | pooling_params.late_interaction_params = build_late_interaction_doc_params( |
| 293 | query_key=query_keys[query_idx] |
| 294 | ) |
| 295 | doc_engine_inputs[i]["params"] = pooling_params |
| 296 | |
| 297 | assert n_docs == len(doc_engine_inputs) == len(doc_keys) |
| 298 | |
| 299 | doc_ctx = ScoringServeContext( |
| 300 | request=ctx.request, |
| 301 | raw_request=ctx.raw_request, |
| 302 | model_name=ctx.model_name, |
| 303 | request_id=ctx.request_id, |
| 304 | pooling_params=ctx.pooling_params, |
| 305 | prompt_request_ids=doc_keys, |
| 306 | engine_inputs=doc_engine_inputs, |
| 307 | lora_request=ctx.lora_request, |
| 308 | priorities=ctx.priorities, |
| 309 | prompt_extras=ctx.prompt_extras, |
| 310 | ) |
| 311 | |
| 312 | await self._prepare_generators(doc_ctx) |
| 313 | await self._collect_batch(doc_ctx) |
| 314 | |
| 315 | ctx.final_res_batch = doc_ctx.final_res_batch |