[Frontend] Add chunked processing to handle long inputs in embedding models #20837
Conversation
x22x22
requested review from
hmellor,
simon-mo,
WoosukKwon,
youkaichao,
robertgshaw2-redhat,
mgoin,
tlrmchlsmth,
houseroad and
aarnphm
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @x22x22, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces a significant enhancement to vLLM's embedding capabilities by adding support for chunked processing of long text inputs. This feature directly addresses issues like CUDA crashes and memory constraints when processing text longer than a model's max_model_len, allowing vLLM to generate embeddings for arbitrarily long documents without requiring changes to the underlying models. The solution is implemented at the serving layer, providing a configurable and robust mechanism for handling large inputs while maintaining embedding quality through smart aggregation.
Highlights
- New Feature: Chunked Processing for Embedding Models: Introduced automatic chunked processing at the serving layer for embedding models. This enables vLLM to handle text inputs that exceed the model's maximum context length by splitting them into manageable chunks, processing each independently, and aggregating the results.
- Configuration and Activation: The chunked processing feature is configurable via the
PoolerConfigby settingenable_chunked_processing: true. It automatically detects when input exceedsmax_model_lenand triggers the chunking logic. - Intelligent Aggregation: Implemented a FastChat-style weighted averaging algorithm to combine embeddings from multiple chunks. This method uses token counts as weights, ensuring that longer chunks contribute proportionally more to the final aggregated embedding, preserving semantic quality.
- Backward Compatibility and Model Support: The implementation maintains backward compatibility for short text inputs and requires zero modifications to existing model code. Initially,
intfloat/multilingual-e5-largeis explicitly supported, with an extensible architecture for other embedding models. - Documentation and Examples: Added comprehensive documentation detailing the feature, its configuration, how it works, performance characteristics, and limitations. New example scripts (server and client) are provided to demonstrate how to configure and utilize chunking processing for long text embeddings.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
mergify
bot
added
documentation
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable feature for handling long text embeddings by implementing chunked processing. The code is generally well-written, and the documentation and examples are thorough. I've identified a critical bug in the batch processing logic that could lead to incorrect results when multiple long prompts are sent in a single request. I've also provided several suggestions to improve code quality, maintainability, and performance. Once the critical issue is addressed, this will be a great addition to the project.
x22x22
force-pushed
the
feat/support-long-text-embedding
branch
from
b5f245d to
5398bbd
Compare
x22x22
changed the title
|
In fact, embedding models are not very suitable for handling extremely long inputs, as too much content can lead to embeddings that are not able to effectively distinguish between similar content. Here's a simple way to confirm that automatic chunked processing is working effectively: Reference mteb_test_embed_models in vllm/tests/models/language/pooling Keeping only the very front part of long context, such as 2048 or even 512, is an extremely high baseline.
Do the following three comparative experiments
If automatic chunked processing using multilingual-e5-large on mteb/T2Reranking dataset(or any test with a context exceeding 8K), can achieve comparable results indicates that automatic chunked processing is effective |
@noooop I've manually tested using text chunks exceeding 1,000 tokens in vector databases, and confirmed that short user queries or task descriptions (~100 tokens) can successfully retrieve relevant text fragments. While this verification isn't scientifically rigorous, it demonstrates a viable practical solution. I'll allocate time later to run the benchmark tests you recommended - appreciate the suggestion. |
|
After some investigation, intfloat/multilingual-e5-large uses the classic BERT architecture with a context length of 512, which appears very weak in 2025. Please perform a comparative test using jina-embeddings-v3, which has a maximum context length of 8192 and uses mean pooling. Unless you use VLLM_ALLOW_LONG_MAX_MODEL_LEN or similar, you Should Not Allow set the context of intfloat/multilingual-e5-large beyond 512, as it will exceed position_embeddings and cause an out-of-bounds error. It is not a bug. Please weaken or remove the content related to CUDA crashes. |
@noooop The purpose is to enable models like While this approach may not deliver optimal embedding performance, it provides a practical low-cost solution for RAG scenarios requiring simultaneous processing of both short and long texts. Crucially, no performance penalty occurs when input stays within a model's native context limit (e.g. ≤512 for E5, ≤8192 for Jina), as no special chunking gets triggered. Would you be open to continuing this discussion more efficiently via https://slack.vllm.ai? I've requested access to the Slack workspace but haven't received approval yet - perhaps we could connect there once I'm onboarded. |
|
I looked through the code carefully. You can add a new parameter such as max_embed_len, but do not modify any code related to max_model_len, That will cause a huge number of bugs. And do not use VLLM_ALLOW_LONG_MAX_MODEL_LEN. I think we should remove VLLM_ALLOW_LONG_MAX_MODEL_LEN. I can’t think of any use case that would require this flag. |
@noooop Regarding communication, would you be open to continuing this discussion through a more efficient channel? I'd appreciate if we could connect either via:
Would either of these options work better for real-time collaboration? Thank you for your guidance on this implementation! |
I’m extremely socially anxious. |
I completely understand, I also have social anxiety. This way of communicating is pretty good too 😄 I'll modify the code according to your suggestions, expecting to have it done by tomorrow~ If there's anything else I need to pay attention to, please feel free to communicate anytime, thank you! |
|
@noooop
{
"pooling_type": "CLS",
"normalize": true,
"enable_chunked_processing": true,
"max_embed_len": 10240
} |
|
@x22x22 , please correct me if I'm wrong, but it seems that the aggregation is based on the assumptions that taking the mean of the embedding chunks would be correct: To work two requirements are necessary:
However, the BERT-type models don't satisfy the second requirement. As @noooop mentioned, there are newer models that are decoder models. For these models, if the pooling type is LAST, we already support chunked prefill.
@noooop , this var is useful for testing or bypassing restrictions of misconfigured models. |
Thank you for your feedback - you're absolutely right! I've updated the implementation to use MEAN pooling and made the following improvements: Correct pooling configuration: multilingual-e5-large now uses MEAN pooling by default Automatic detection: Added support for automatic configuration of various popular models Manual specification: Users can manually specify the pooling type with helpful prompts and guidance Safety warnings: Users are now alerted about potential impacts when using non-MEAN pooling methods Flexible configuration: All parameters can be customized through environment variables This ensures users can safely utilize the chunked processing functionality without worrying about pooling type mismatches! For usage reference, please check out the new example startup script: Your point about BERT-type models not satisfying the causal attention requirement is particularly important. The weighted aggregation approach works best with models that use mean pooling and have the appropriate attention patterns. The automatic detection and safety warnings should help users avoid potential issues with incompatible model architectures. |
Signed-off-by: Nick Hill <[email protected]> Signed-off-by: x22x22 <[email protected]>
…el installation (vllm-project#21635) Signed-off-by: Ming Yang <[email protected]> Signed-off-by: x22x22 <[email protected]>
) Signed-off-by: cascade812 <[email protected]> Signed-off-by: x22x22 <[email protected]>
… requests (vllm-project#20272) Signed-off-by: x22x22 <[email protected]>
…vllm-project#21627) Signed-off-by: linzebing <[email protected]> Signed-off-by: x22x22 <[email protected]>
… in python (vllm-project#21763) Signed-off-by: mgoin <[email protected]> Signed-off-by: x22x22 <[email protected]>
Signed-off-by: Sanchit Gandhi <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: x22x22 <[email protected]>
Signed-off-by: mgoin <[email protected]> Signed-off-by: x22x22 <[email protected]>
…llm-project#21818) Signed-off-by: Jee Jee Li <[email protected]> Signed-off-by: x22x22 <[email protected]>
…lm-project#21973) Signed-off-by: mgoin <[email protected]> Signed-off-by: x22x22 <[email protected]>
…roject#21830) Signed-off-by: Andy Xie <[email protected]> Signed-off-by: x22x22 <[email protected]>
Signed-off-by: DarkLight1337 <[email protected]> Signed-off-by: x22x22 <[email protected]>
…1599) Signed-off-by: Daniele Trifirò <[email protected]> Signed-off-by: x22x22 <[email protected]>
Signed-off-by: x22x22 <[email protected]>
…deleted, and the code and configurations that are no longer in use have been cleaned up. Signed-off-by: x22x22 <[email protected]>
…deleted, and the code and configurations that are no longer in use have been cleaned up. Signed-off-by: x22x22 <[email protected]>
…ized, with the use of mean aggregation enforced and support for other aggregation types removed. Relevant log information has been updated to reflect the new processing approach. Signed-off-by: x22x22 <[email protected]>
…n aggregation is uniformly adopted, and support for other aggregation types has been removed. Relevant documents and configurations have been updated to reflect the new processing approach. Configuration options that are no longer in use have been removed to ensure the code's cleanliness. Signed-off-by: x22x22 <[email protected]>
…wline Files should end with a single newline character 117 Error: docs/models/supported_models.md:777:265 MD047/single-trailing-newline Files should end with a single newline character 118 Error: examples/online_serving/openai_embedding_long_text.md:96 MD032/blanks-around-lists Lists should be surrounded by blank lines [Context: "- **Academic papers**: Full re..."] 119 Error: examples/online_serving/openai_embedding_long_text.md:130 MD040/fenced-code-language Fenced code blocks should have a language specified [Context: "```"] 120 Error: examples/online_serving/openai_embedding_long_text.md:138 MD040/fenced-code-language Fenced code blocks should have a language specified [Context: "```"] 121 Error: examples/online_serving/openai_embedding_long_text.md:146 MD040/fenced-code-language Fenced code blocks should have a language specified [Context: "```"] 122 Error: examples/online_serving/openai_embedding_long_text.md:159 MD040/fenced-code-language Fenced code blocks should have a language specified [Context: "```"] 123 Signed-off-by: x22x22 <[email protected]>
…vice scripts, incorporating chunk processing support. The README documentation has been revised to include a quick start guide and comprehensive configuration instructions. Server startup scripts have been enhanced with automatic detection of optimal pooling types, significantly improving performance and compatibility for long-text processing. Signed-off-by: x22x22 <[email protected]>
x22x22
force-pushed
the
feat/support-long-text-embedding
branch
from
ca252d7 to
b1562c5
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
x22x22
force-pushed
the
feat/support-long-text-embedding
branch
from
e10ef93 to
a835f52
Compare
…g, and update relevant documentation and examples. New example scripts and service startup scripts are added to demonstrate how to configure and utilize chunking processing. Update the model configuration to support long - text processing and implement the chunking processing logic in the code.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
Add chunked processing support for long text embeddings to resolve CUDA crashes when input text exceeds model's maximum context length.
Problem Solved
max_model_lenSolution
This PR implements automatic chunked processing at the serving layer that:
Key Features
enable_chunked_processing: truein pooler configSupported Models
intfloat/multilingual-e5-large(initially)This enables vLLM to handle embedding requests of any length without crashes, significantly expanding its utility for RAG applications and long document processing.
Test Plan
Long Text Embedding with Chunked Processing
Test Result
Before modification
After modification
(Optional) Documentation Update