serving_embedding.py 26.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import json
4
from collections.abc import AsyncGenerator, Mapping
5
from typing import Any, Final, cast
6

7
import torch
8
from fastapi import Request
9
10
from fastapi.responses import Response
from typing_extensions import assert_never, override
11

12
from vllm.engine.protocol import EngineClient
13
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
14
from vllm.entrypoints.logger import RequestLogger
15
from vllm.entrypoints.openai.protocol import (
16
    EmbeddingBytesResponse,
17
18
19
20
21
22
23
24
25
26
27
28
29
30
    EmbeddingChatRequest,
    EmbeddingCompletionRequest,
    EmbeddingRequest,
    EmbeddingResponse,
    EmbeddingResponseData,
    ErrorResponse,
    UsageInfo,
)
from vllm.entrypoints.openai.serving_engine import (
    EmbeddingServeContext,
    OpenAIServing,
    ServeContext,
    TextTokensPrompt,
)
31
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
32
from vllm.entrypoints.renderer import RenderConfig
33
from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
34
from vllm.logger import init_logger
35
36
37
38
39
40
from vllm.outputs import (
    EmbeddingRequestOutput,
    PoolingOutput,
    PoolingRequestOutput,
    RequestOutput,
)
41
from vllm.pooling_params import PoolingParams
42
43
from vllm.utils.async_utils import merge_async_iterators
from vllm.utils.collection_utils import chunk_list
44
45
46
47
48
49
50
from vllm.utils.serial_utils import (
    EmbedDType,
    EncodingFormat,
    Endianness,
    encode_pooling_bytes,
    encode_pooling_output,
)
51
52
53
54

logger = init_logger(__name__)


55
class EmbeddingMixin(OpenAIServing):
56
57
58
59
60
61
62
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        pooler_config = self.model_config.pooler_config

        # Avoid repeated attribute lookups
        self.supports_chunked_processing = bool(
63
64
65
66
67
68
69
            pooler_config and pooler_config.enable_chunked_processing
        )
        self.max_embed_len = (
            pooler_config.max_embed_len
            if pooler_config and pooler_config.max_embed_len
            else None
        )
70

71
    @override
72
    async def _preprocess(
73
        self,
74
        ctx: ServeContext,
75
    ) -> ErrorResponse | None:
76
        ctx = cast(EmbeddingServeContext, ctx)
77
        try:
78
            ctx.lora_request = self._maybe_get_adapters(ctx.request)
79

80
            tokenizer = await self.engine_client.get_tokenizer()
81
            renderer = self._get_renderer(tokenizer)
82

83
            if isinstance(ctx.request, EmbeddingChatRequest):
84
85
                (
                    _,
86
                    _,
87
                    ctx.engine_prompts,
88
                ) = await self._preprocess_chat(
89
                    ctx.request,
90
                    tokenizer,
91
                    ctx.request.messages,
92
93
                    chat_template=ctx.request.chat_template or ctx.chat_template,
                    chat_template_content_format=ctx.chat_template_content_format,
94
                    add_generation_prompt=ctx.request.add_generation_prompt,
95
                    continue_final_message=False,
96
                    add_special_tokens=ctx.request.add_special_tokens,
97
98
                )
            else:
99
100
                ctx.engine_prompts = await renderer.render_prompt(
                    prompt_or_prompts=ctx.request.input,
101
                    config=self._build_render_config(ctx.request),
102
                )
103
            return None
104
        except (ValueError, TypeError) as e:
105
106
            logger.exception("Error in preprocessing prompt inputs")
            return self.create_error_response(str(e))
107

108
    def _build_render_config(self, request: EmbeddingCompletionRequest) -> RenderConfig:
109
110
111
112
113
114
115
116
117
        # Set max_length based on chunked processing capability
        if self._should_use_chunked_processing(request):
            max_length = None
        else:
            max_length = self.max_embed_len or self.max_model_len

        return RenderConfig(
            max_length=max_length,
            truncate_prompt_tokens=request.truncate_prompt_tokens,
118
119
            add_special_tokens=request.add_special_tokens,
        )
120

121
    @override
122
    def _build_response(
123
        self,
124
        ctx: ServeContext,
125
    ) -> EmbeddingResponse | Response | ErrorResponse:
126
        final_res_batch_checked = cast(list[PoolingRequestOutput], ctx.final_res_batch)
127

128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
        encoding_format: EncodingFormat = ctx.request.encoding_format
        embed_dtype: EmbedDType = ctx.request.embed_dtype
        endianness: Endianness = ctx.request.endianness

        def encode_float_base64():
            items: list[EmbeddingResponseData] = []
            num_prompt_tokens = 0

            for idx, final_res in enumerate(final_res_batch_checked):
                item = EmbeddingResponseData(
                    index=idx,
                    embedding=encode_pooling_output(
                        final_res,
                        encoding_format=encoding_format,
                        embed_dtype=embed_dtype,
                        endianness=endianness,
                    ),
                )
                prompt_token_ids = final_res.prompt_token_ids

                items.append(item)
                num_prompt_tokens += len(prompt_token_ids)

            usage = UsageInfo(
                prompt_tokens=num_prompt_tokens,
                total_tokens=num_prompt_tokens,
154
155
            )

156
157
158
159
160
161
162
            return EmbeddingResponse(
                id=ctx.request_id,
                created=ctx.created_time,
                model=ctx.model_name,
                data=items,
                usage=usage,
            )
163

164
165
166
167
168
169
        def encode_bytes():
            body, items, usage = encode_pooling_bytes(
                pooling_outputs=final_res_batch_checked,
                embed_dtype=embed_dtype,
                endianness=endianness,
            )
170

171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
            metadata = {
                "id": ctx.request_id,
                "created": ctx.created_time,
                "model": ctx.model_name,
                "data": items,
                "usage": usage,
            }
            return EmbeddingBytesResponse(
                body=body,
                metadata=json.dumps(metadata),
            )

        if encoding_format == "float" or encoding_format == "base64":
            return encode_float_base64()
        elif encoding_format == "bytes":
            return encode_bytes()
        else:
            assert_never(encoding_format)
189

190
191
192
193
194
195
    def _get_max_position_embeddings(self) -> int:
        """Get the model's effective maximum sequence length for chunking."""
        return self.model_config.max_model_len

    def _should_use_chunked_processing(self, request) -> bool:
        """Check if chunked processing should be used for this request."""
196
197
198
199
        return (
            isinstance(request, (EmbeddingCompletionRequest, EmbeddingChatRequest))
            and self.supports_chunked_processing
        )
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216

    async def _process_chunked_request(
        self,
        ctx: EmbeddingServeContext,
        original_prompt: TextTokensPrompt,
        pooling_params,
        trace_headers,
        prompt_idx: int,
    ) -> list[AsyncGenerator[PoolingRequestOutput, None]]:
        """Process a single prompt using chunked processing."""
        generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
        token_ids = original_prompt["prompt_token_ids"]

        # Split into chunks using max_position_embeddings
        max_pos_embeddings = self._get_max_position_embeddings()
        # Process all chunks for MEAN aggregation
        for chunk_idx, chunk_tokens in enumerate(
217
218
            chunk_list(token_ids, max_pos_embeddings)
        ):
219
            # Create a request ID for this chunk
220
            chunk_request_id = f"{ctx.request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}"
221
222

            # Create engine prompt for this chunk
223
            chunk_engine_prompt = EngineTokensPrompt(prompt_token_ids=chunk_tokens)
224
225
226
227

            # Create chunk request prompt for logging
            chunk_text = ""
            chunk_request_prompt = TextTokensPrompt(
228
229
                prompt=chunk_text, prompt_token_ids=chunk_tokens
            )
230
231

            # Log the chunk
232
233
234
235
236
237
            self._log_inputs(
                chunk_request_id,
                chunk_request_prompt,
                params=pooling_params,
                lora_request=ctx.lora_request,
            )
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

            # Create generator for this chunk and wrap it to return indices
            original_generator = self.engine_client.encode(
                chunk_engine_prompt,
                pooling_params,
                chunk_request_id,
                lora_request=ctx.lora_request,
                trace_headers=trace_headers,
                priority=getattr(ctx.request, "priority", 0),
            )

            generators.append(original_generator)

        return generators

    def _validate_input(
        self,
        request,
        input_ids: list[int],
        input_text: str,
    ) -> TextTokensPrompt:
        """Override to support chunked processing for embedding requests."""
        token_num = len(input_ids)

        # Note: EmbeddingRequest doesn't have max_tokens
263
        if isinstance(request, (EmbeddingCompletionRequest, EmbeddingChatRequest)):
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
            # Check if chunked processing is enabled for pooling models
            enable_chunked = self._should_use_chunked_processing(request)

            # Use max_position_embeddings for chunked processing decisions
            max_pos_embeddings = self._get_max_position_embeddings()

            # Determine the effective max length for validation
            if self.max_embed_len is not None:
                # Use max_embed_len for validation instead of max_model_len
                length_type = "maximum embedding input length"
                max_length_value = self.max_embed_len
            else:
                # Fall back to max_model_len validation (original behavior)
                length_type = "maximum context length"
                max_length_value = self.max_model_len

            validation_error_msg = (
                "This model's {length_type} is {max_length_value} tokens. "
                "However, you requested {token_num} tokens in the input for "
283
284
                "embedding generation. Please reduce the length of the input."
            )
285
286
287
288
289

            chunked_processing_error_msg = (
                "This model's {length_type} is {max_length_value} tokens. "
                "However, you requested {token_num} tokens in the input for "
                "embedding generation. Please reduce the length of the input "
290
291
                "or enable chunked processing."
            )
292
293
294
295
296
297
298

            # Check if input exceeds max length
            if token_num > max_length_value:
                raise ValueError(
                    validation_error_msg.format(
                        length_type=length_type,
                        max_length_value=max_length_value,
299
300
301
                        token_num=token_num,
                    )
                )
302
303
304
305
306
307
308
309

            # Check for chunked processing
            # when exceeding max_position_embeddings
            if token_num > max_pos_embeddings:
                if enable_chunked:
                    # Allow long inputs when chunked processing is enabled
                    logger.info(
                        "Input length %s exceeds max_position_embeddings "
310
311
312
313
                        "%s, will use chunked processing",
                        token_num,
                        max_pos_embeddings,
                    )
314
315
316
317
318
                else:
                    raise ValueError(
                        chunked_processing_error_msg.format(
                            length_type="maximum position embeddings length",
                            max_length_value=max_pos_embeddings,
319
320
321
                            token_num=token_num,
                        )
                    )
322

323
            return TextTokensPrompt(prompt=input_text, prompt_token_ids=input_ids)
324
325
326
327
328
329

        # For other request types, use the parent's implementation
        return super()._validate_input(request, input_ids, input_text)

    def _is_text_tokens_prompt(self, prompt) -> bool:
        """Check if a prompt is a TextTokensPrompt (has prompt_token_ids)."""
330
331
332
333
334
        return (
            isinstance(prompt, dict)
            and "prompt_token_ids" in prompt
            and "prompt_embeds" not in prompt
        )
335
336
337
338

    async def _create_single_prompt_generator(
        self,
        ctx: EmbeddingServeContext,
339
        engine_prompt: EngineTokensPrompt,
340
        pooling_params: PoolingParams,
341
        trace_headers: Mapping[str, str] | None,
342
        prompt_index: int,
343
    ) -> AsyncGenerator[RequestOutput | PoolingRequestOutput, None]:
344
345
346
        """Create a generator for a single prompt using standard processing."""
        request_id_item = f"{ctx.request_id}-{prompt_index}"

347
348
349
350
351
352
        self._log_inputs(
            request_id_item,
            engine_prompt,
            params=pooling_params,
            lora_request=ctx.lora_request,
        )
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367

        # Return the original generator without wrapping
        return self.engine_client.encode(
            engine_prompt,
            pooling_params,
            request_id_item,
            lora_request=ctx.lora_request,
            trace_headers=trace_headers,
            priority=getattr(ctx.request, "priority", 0),
        )

    @override
    async def _prepare_generators(
        self,
        ctx: ServeContext,
368
    ) -> ErrorResponse | None:
369
370
371
372
373
374
375
376
377
378
379
        """Override to support chunked processing."""
        ctx = cast(EmbeddingServeContext, ctx)

        # Check if we should use chunked processing
        use_chunked = self._should_use_chunked_processing(ctx.request)

        # If no chunked processing needed, delegate to parent class
        if not use_chunked:
            return await super()._prepare_generators(ctx)

        # Custom logic for chunked processing
380
        generators: list[
381
            AsyncGenerator[RequestOutput | PoolingRequestOutput, None]
382
        ] = []
383
384

        try:
385
386
387
388
389
            trace_headers = (
                None
                if ctx.raw_request is None
                else await self._get_trace_headers(ctx.raw_request.headers)
            )
390
391
392
393
394
395
396
397
398
399
400
401

            pooling_params = self._create_pooling_params(ctx)
            if isinstance(pooling_params, ErrorResponse):
                return pooling_params

            # Verify and set the task for pooling params
            try:
                pooling_params.verify("embed", self.model_config)
            except ValueError as e:
                return self.create_error_response(str(e))

            if ctx.engine_prompts is None:
402
                return self.create_error_response("Engine prompts not available")
403
404
405
406
407

            max_pos_embeddings = self._get_max_position_embeddings()

            for i, engine_prompt in enumerate(ctx.engine_prompts):
                # Check if this specific prompt needs chunked processing
408
                if self._is_text_tokens_prompt(engine_prompt):
409
410
                    # Cast to TextTokensPrompt since we've verified
                    # prompt_token_ids
411
                    text_tokens_prompt = cast(TextTokensPrompt, engine_prompt)
412
                    if len(text_tokens_prompt["prompt_token_ids"]) > max_pos_embeddings:
413
414
                        # Use chunked processing for this prompt
                        chunk_generators = await self._process_chunked_request(
415
416
                            ctx, text_tokens_prompt, pooling_params, trace_headers, i
                        )
417
418
419
420
421
                        generators.extend(chunk_generators)
                        continue

                # Normal processing for short prompts or non-token prompts
                generator = await self._create_single_prompt_generator(
422
423
                    ctx, engine_prompt, pooling_params, trace_headers, i
                )
424
425
426
427
428
429
430
431
432
433
434
435
436
437
                generators.append(generator)

            ctx.result_generator = merge_async_iterators(*generators)

            return None

        except Exception as e:
            # TODO: Use a vllm-specific Validation Error
            return self.create_error_response(str(e))

    @override
    async def _collect_batch(
        self,
        ctx: ServeContext,
438
    ) -> ErrorResponse | None:
439
440
        """Collect and aggregate batch results
        with support for chunked processing.
441
442

        For chunked requests, performs online aggregation to
443
444
445
446
447
448
        minimize memory usage.
        For regular requests, collects results normally.
        """
        ctx = cast(EmbeddingServeContext, ctx)
        try:
            if ctx.engine_prompts is None:
449
                return self.create_error_response("Engine prompts not available")
450
451
452
453
454
455
456
457

            # Check if we used chunked processing
            use_chunked = self._should_use_chunked_processing(ctx.request)

            if not use_chunked:
                return await super()._collect_batch(ctx=ctx)

            if ctx.result_generator is None:
458
                return self.create_error_response("Result generator not available")
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478

            # Online aggregation for chunked requests to
            # minimize memory usage
            # Track aggregation state for each prompt
            prompt_aggregators: dict[int, dict[str, Any]] = {}
            short_prompts_results: dict[int, PoolingRequestOutput] = {}

            async for result_idx, result in ctx.result_generator:
                if "-chunk-" in result.request_id:
                    # Extract prompt_idx from chunked request_id
                    parts = result.request_id.split("-")
                    try:
                        prompt_idx = int(parts[parts.index("prompt") + 1])
                    except (ValueError, IndexError):
                        # Fallback: extract from result_idx if parsing fails
                        prompt_idx = result_idx

                    # Initialize aggregator for this prompt if needed
                    if prompt_idx not in prompt_aggregators:
                        prompt_aggregators[prompt_idx] = {
479
480
481
482
                            "weighted_sum": None,
                            "total_weight": 0,
                            "chunk_count": 0,
                            "request_id": result.request_id.split("-chunk-")[0],
483
484
485
486
487
488
489
490
491
492
493
                        }

                    aggregator = prompt_aggregators[prompt_idx]

                    # MEAN pooling with online weighted averaging
                    # Ensure result is PoolingRequestOutput
                    # for embedding processing
                    if not isinstance(result, PoolingRequestOutput):
                        return self.create_error_response(
                            f"Expected PoolingRequestOutput for "
                            f"chunked embedding, got "
494
495
                            f"{type(result).__name__}"
                        )
496
497
498

                    # Handle both PoolingOutput and
                    # EmbeddingOutput types
499
                    if hasattr(result.outputs, "data"):
500
501
                        # PoolingOutput case
                        embedding_data = result.outputs.data
502
                    elif hasattr(result.outputs, "embedding"):
503
504
505
506
507
                        # EmbeddingOutput case -
                        # convert embedding list to tensor
                        embedding_data = result.outputs.embedding
                    else:
                        return self.create_error_response(
508
509
                            f"Unsupported output type: {type(result.outputs).__name__}"
                        )
510
511

                    if not isinstance(embedding_data, torch.Tensor):
512
513
514
                        embedding_data = torch.tensor(
                            embedding_data, dtype=torch.float32
                        )
515
516
517

                    if result.prompt_token_ids is None:
                        return self.create_error_response(
518
519
                            "prompt_token_ids cannot be None for chunked processing"
                        )
520
521
                    weight = len(result.prompt_token_ids)

522
                    weighted_embedding = embedding_data.to(dtype=torch.float32) * weight
523

524
                    if aggregator["weighted_sum"] is None:
525
                        # First chunk
526
                        aggregator["weighted_sum"] = weighted_embedding
527
528
                    else:
                        # Accumulate
529
                        aggregator["weighted_sum"] += weighted_embedding
530

531
532
                    aggregator["total_weight"] += weight
                    aggregator["chunk_count"] += 1
533
534
535
536
537
538
539
540
541
542
                else:
                    # Non-chunked result - extract prompt_idx from request_id
                    parts = result.request_id.split("-")
                    try:
                        # Last part should be prompt index
                        prompt_idx = int(parts[-1])
                    except (ValueError, IndexError):
                        prompt_idx = result_idx  # Fallback to result_idx

                    short_prompts_results[prompt_idx] = cast(
543
544
                        PoolingRequestOutput, result
                    )
545
546

            # Finalize aggregated results
547
            final_res_batch: list[PoolingRequestOutput | EmbeddingRequestOutput] = []
548
549
550
551
552
553
554
            num_prompts = len(ctx.engine_prompts)

            for prompt_idx in range(num_prompts):
                if prompt_idx in prompt_aggregators:
                    # Finalize MEAN aggregation for this chunked prompt
                    aggregator = prompt_aggregators[prompt_idx]

555
556
                    weighted_sum = aggregator["weighted_sum"]
                    total_weight = aggregator["total_weight"]
557

558
559
560
561
562
563
                    if (
                        weighted_sum is not None
                        and isinstance(weighted_sum, torch.Tensor)
                        and isinstance(total_weight, (int, float))
                        and total_weight > 0
                    ):
564
565
566
567
568
                        # Compute final mean embedding
                        final_embedding = weighted_sum / total_weight

                        # Create a PoolingRequestOutput
                        # for the aggregated result
569
                        pooling_output_data = PoolingOutput(data=final_embedding)
570
571

                        # Get original prompt token IDs for this prompt
572
                        original_prompt = ctx.engine_prompts[prompt_idx]
573
574
                        if not self._is_text_tokens_prompt(original_prompt):
                            return self.create_error_response(
575
576
                                f"Chunked prompt {prompt_idx} is not a TextTokensPrompt"
                            )
577

578
579
580
                        original_token_ids = cast(TextTokensPrompt, original_prompt)[
                            "prompt_token_ids"
                        ]
581
582

                        pooling_request_output = PoolingRequestOutput(
583
                            request_id=aggregator["request_id"],
584
585
                            prompt_token_ids=original_token_ids,
                            outputs=pooling_output_data,
586
                            num_cached_tokens=0,
587
588
                            finished=True,
                        )
589
590
591
592

                        final_res_batch.append(pooling_request_output)
                    else:
                        return self.create_error_response(
593
594
                            f"Failed to aggregate chunks for prompt {prompt_idx}"
                        )
595
596
                elif prompt_idx in short_prompts_results:
                    final_res_batch.append(
597
598
                        cast(PoolingRequestOutput, short_prompts_results[prompt_idx])
                    )
599
600
                else:
                    return self.create_error_response(
601
602
                        f"Result not found for prompt {prompt_idx}"
                    )
603
604

            ctx.final_res_batch = cast(
605
                list[RequestOutput | PoolingRequestOutput], final_res_batch
606
            )
607
608
609
610
611
612

            return None

        except Exception as e:
            return self.create_error_response(str(e))

613
614
615
616
617
618
619
620
621

class OpenAIServingEmbedding(EmbeddingMixin):
    request_id_prefix = "embd"

    def __init__(
        self,
        engine_client: EngineClient,
        models: OpenAIServingModels,
        *,
622
623
        request_logger: RequestLogger | None,
        chat_template: str | None,
624
        chat_template_content_format: ChatTemplateContentFormatOption,
625
        trust_request_chat_template: bool = False,
626
        log_error_stack: bool = False,
627
    ) -> None:
628
629
630
631
632
633
        super().__init__(
            engine_client=engine_client,
            models=models,
            request_logger=request_logger,
            log_error_stack=log_error_stack,
        )
634
635
636

        self.chat_template = chat_template
        self.chat_template_content_format: Final = chat_template_content_format
637
        self.trust_request_chat_template = trust_request_chat_template
638
639
640
641

    async def create_embedding(
        self,
        request: EmbeddingRequest,
642
643
        raw_request: Request | None = None,
    ) -> EmbeddingResponse | ErrorResponse:
644
645
646
647
648
649
        """
        Embedding API similar to OpenAI's API.

        See https://platform.openai.com/docs/api-reference/embeddings/create
        for the API specification. This API mimics the OpenAI Embedding API.
        """
650
        model_name = self.models.model_name()
651
652
        request_id = (
            f"{self.request_id_prefix}-"
653
654
            f"{self._base_request_id(raw_request, request.request_id)}"
        )
655
656
657
658
659
660
661
662
663
664
665
666

        ctx = EmbeddingServeContext(
            request=request,
            raw_request=raw_request,
            model_name=model_name,
            request_id=request_id,
            chat_template=self.chat_template,
            chat_template_content_format=self.chat_template_content_format,
        )

        return await super().handle(ctx)  # type: ignore

667
668
669
670
    @override
    def _create_pooling_params(
        self,
        ctx: ServeContext[EmbeddingRequest],
671
    ) -> PoolingParams | ErrorResponse:
672
673
674
        pooling_params = super()._create_pooling_params(ctx)
        if isinstance(pooling_params, ErrorResponse):
            return pooling_params
675
676

        try:
677
            pooling_params.verify("embed", self.model_config)
678
679
680
        except ValueError as e:
            return self.create_error_response(str(e))

681
        return pooling_params
682
683
684
685

    async def _preprocess(
        self,
        ctx: ServeContext,
686
    ) -> ErrorResponse | None:
687
688
689
690
691
692
693
694
695
        if isinstance(ctx.request, EmbeddingChatRequest):
            error_check_ret = self._validate_chat_template(
                request_chat_template=ctx.request.chat_template,
                chat_template_kwargs=ctx.request.chat_template_kwargs,
                trust_request_chat_template=self.trust_request_chat_template,
            )
            if error_check_ret is not None:
                return error_check_ret
        return await super()._preprocess(ctx)