@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 653124bd46
Fork 0

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 14, 05:30 UTC

[Frontend] Add chunked processing to handle long inputs in embedding models (#22280)

Signed-off-by: x22x22 <wadeking@qq.com> Signed-off-by: Kdump <rootshellexp@gmail.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Maximilien de Bayser <maxdebayser@gmail.com> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>

Kdump committed Aug 13, 2025 at 19:14 UTC 653124bd46c57770b151eb58cc2a59170753daa5
6 files changed +1603 -3
examples/online_serving/openai_embedding_long_text/README.md new
+186
@@ -0,0 +1,186 @@
1 +# Long Text Embedding with Chunked Processing
2 +
3 +This directory contains examples for using vLLM's **chunked processing** feature to handle long text embedding that exceeds the model's maximum context length.
4 +
5 +## 🚀 Quick Start
6 +
7 +### Start the Server
8 +
9 +Use the provided script to start a vLLM server with chunked processing enabled:
10 +
11 +```bash
12 +# Basic usage (supports very long texts up to ~3M tokens)
13 +./service.sh
14 +
15 +# Custom configuration with different models
16 +MODEL_NAME="jinaai/jina-embeddings-v3" \
17 +MAX_EMBED_LEN=1048576 \
18 +./service.sh
19 +
20 +# For extremely long documents
21 +MODEL_NAME="intfloat/multilingual-e5-large" \
22 +MAX_EMBED_LEN=3072000 \
23 +./service.sh
24 +```
25 +
26 +### Test Long Text Embedding
27 +
28 +Run the comprehensive test client:
29 +
30 +```bash
31 +python client.py
32 +```
33 +
34 +## 📁 Files
35 +
36 +| File | Description |
37 +|------|-------------|
38 +| `service.sh` | Server startup script with chunked processing enabled |
39 +| `client.py` | Comprehensive test client for long text embedding |
40 +
41 +## ⚙️ Configuration
42 +
43 +### Server Configuration
44 +
45 +The key parameters for chunked processing are in the `--override-pooler-config`:
46 +
47 +```json
48 +{
49 + "pooling_type": "auto",
50 + "normalize": true,
51 + "enable_chunked_processing": true,
52 + "max_embed_len": 3072000
53 +}
54 +```
55 +
56 +!!! note
57 + `pooling_type` sets the model's own pooling strategy for processing within each chunk. The cross-chunk aggregation automatically uses MEAN strategy when input exceeds the model's native maximum length.
58 +
59 +#### Chunked Processing Behavior
60 +
61 +Chunked processing uses **MEAN aggregation** for cross-chunk combination when input exceeds the model's native maximum length:
62 +
63 +| Component | Behavior | Description |
64 +|-----------|----------|-------------|
65 +| **Within chunks** | Model's native pooling | Uses the model's configured pooling strategy |
66 +| **Cross-chunk aggregation** | Always MEAN | Weighted averaging based on chunk token counts |
67 +| **Performance** | Optimal | All chunks processed for complete semantic coverage |
68 +
69 +### Environment Variables
70 +
71 +| Variable | Default | Description |
72 +|----------|---------|-------------|
73 +| `MODEL_NAME` | `intfloat/multilingual-e5-large` | Embedding model to use (supports multiple models) |
74 +| `PORT` | `31090` | Server port |
75 +| `GPU_COUNT` | `1` | Number of GPUs to use |
76 +| `MAX_EMBED_LEN` | `3072000` | Maximum embedding input length (supports very long documents) |
77 +| `POOLING_TYPE` | `auto` | Model's native pooling type: `auto`, `MEAN`, `CLS`, `LAST` (only affects within-chunk pooling, not cross-chunk aggregation) |
78 +| `API_KEY` | `EMPTY` | API key for authentication |
79 +
80 +## 🔧 How It Works
81 +
82 +1. **Enhanced Input Validation**: `max_embed_len` allows accepting inputs longer than `max_model_len` without environment variables
83 +2. **Smart Chunking**: Text is split based on `max_position_embeddings` to maintain semantic integrity
84 +3. **Unified Processing**: All chunks processed separately through the model using its configured pooling strategy
85 +4. **MEAN Aggregation**: When input exceeds model's native length, results combined using token count-based weighted averaging across all chunks
86 +5. **Consistent Output**: Final embeddings maintain the same dimensionality as standard processing
87 +
88 +### Input Length Handling
89 +
90 +- **Within max_embed_len**: Input is accepted and processed (up to 3M+ tokens)
91 +- **Exceeds max_position_embeddings**: Chunked processing is automatically triggered
92 +- **Exceeds max_embed_len**: Input is rejected with clear error message
93 +- **No environment variables required**: Works without `VLLM_ALLOW_LONG_MAX_MODEL_LEN`
94 +
95 +### Extreme Long Text Support
96 +
97 +With `MAX_EMBED_LEN=3072000`, you can process:
98 +
99 +- **Academic papers**: Full research papers with references
100 +- **Legal documents**: Complete contracts and legal texts
101 +- **Books**: Entire chapters or small books
102 +- **Code repositories**: Large codebases and documentation
103 +
104 +## 📊 Performance Characteristics
105 +
106 +### Chunked Processing Performance
107 +
108 +| Aspect | Behavior | Performance |
109 +|--------|----------|-------------|
110 +| **Chunk Processing** | All chunks processed with native pooling | Consistent with input length |
111 +| **Cross-chunk Aggregation** | MEAN weighted averaging | Minimal overhead |
112 +| **Memory Usage** | Proportional to number of chunks | Moderate, scalable |
113 +| **Semantic Quality** | Complete text coverage | Optimal for long documents |
114 +
115 +## 🧪 Test Cases
116 +
117 +The test client demonstrates:
118 +
119 +- ✅ **Short text**: Normal processing (baseline)
120 +- ✅ **Medium text**: Single chunk processing
121 +- ✅ **Long text**: Multi-chunk processing with aggregation
122 +- ✅ **Very long text**: Many chunks processing
123 +- ✅ **Extreme long text**: Document-level processing (100K+ tokens)
124 +- ✅ **Batch processing**: Mixed-length inputs in one request
125 +- ✅ **Consistency**: Reproducible results across runs
126 +
127 +## 🐛 Troubleshooting
128 +
129 +### Common Issues
130 +
131 +1. **Chunked processing not enabled**:
132 +
133 + ```log
134 + ValueError: This model's maximum position embeddings length is 4096 tokens...
135 + ```
136 +
137 + **Solution**: Ensure `enable_chunked_processing: true` in pooler config
138 +
139 +2. **Input exceeds max_embed_len**:
140 +
141 + ```log
142 + ValueError: This model's maximum embedding input length is 3072000 tokens...
143 + ```
144 +
145 + **Solution**: Increase `max_embed_len` in pooler config or reduce input length
146 +
147 +3. **Memory errors**:
148 +
149 + ```log
150 + RuntimeError: CUDA out of memory
151 + ```
152 +
153 + **Solution**: Reduce chunk size by adjusting model's `max_position_embeddings` or use fewer GPUs
154 +
155 +4. **Slow processing**:
156 + **Expected**: Long text takes more time due to multiple inference calls
157 +
158 +### Debug Information
159 +
160 +Server logs show chunked processing activity:
161 +
162 +```log
163 +INFO: Input length 150000 exceeds max_position_embeddings 4096, will use chunked processing
164 +INFO: Split input of 150000 tokens into 37 chunks (max_chunk_size: 4096)
165 +```
166 +
167 +## 🤝 Contributing
168 +
169 +To extend chunked processing support to other embedding models:
170 +
171 +1. Check model compatibility with the pooling architecture
172 +2. Test with various text lengths
173 +3. Validate embedding quality compared to single-chunk processing
174 +4. Submit PR with test cases and documentation updates
175 +
176 +## 🆕 Enhanced Features
177 +
178 +### max_embed_len Parameter
179 +
180 +The new `max_embed_len` parameter provides:
181 +
182 +- **Simplified Configuration**: No need for `VLLM_ALLOW_LONG_MAX_MODEL_LEN` environment variable
183 +- **Flexible Input Validation**: Accept inputs longer than `max_model_len` up to `max_embed_len`
184 +- **Extreme Length Support**: Process documents with millions of tokens
185 +- **Clear Error Messages**: Better feedback when inputs exceed limits
186 +- **Backward Compatibility**: Existing configurations continue to work
examples/online_serving/openai_embedding_long_text/client.py new
+366
@@ -0,0 +1,366 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Example script demonstrating long text embedding with chunked processing in vLLM.
6 +
7 +This example shows how to use vLLM's chunked processing feature to handle text
8 +inputs that exceed the model's maximum token length. The feature automatically
9 +splits long text into chunks and handles different pooling types optimally.
10 +
11 +Prerequisites:
12 +1. Start vLLM server with chunked processing enabled:
13 +
14 + # MEAN pooling (processes all chunks, recommended for complete coverage)
15 + vllm serve intfloat/multilingual-e5-large \
16 + --override-pooler-config \
17 + '{"pooling_type": "MEAN", "normalize": true, ' \
18 + '"enable_chunked_processing": true, "max_embed_len": 3072000}' \
19 + --served-model-name multilingual-e5-large \
20 + --trust-remote-code \
21 + --port 31090 \
22 + --api-key your-api-key
23 +
24 + # OR CLS pooling (native CLS within chunks, MEAN aggregation across chunks)
25 + vllm serve BAAI/bge-large-en-v1.5 \
26 + --override-pooler-config \
27 + '{"pooling_type": "CLS", "normalize": true, ' \
28 + '"enable_chunked_processing": true, "max_embed_len": 1048576}' \
29 + --served-model-name bge-large-en-v1.5 \
30 + --trust-remote-code \
31 + --port 31090 \
32 + --api-key your-api-key
33 +
34 +2. Install required dependencies:
35 + pip install openai requests
36 +"""
37 +
38 +import time
39 +
40 +import numpy as np
41 +from openai import OpenAI
42 +
43 +# Configuration
44 +API_KEY = "your-api-key" # Replace with your actual API key
45 +BASE_URL = "http://localhost:31090/v1"
46 +MODEL_NAME = "multilingual-e5-large"
47 +
48 +
49 +def generate_long_text(base_text: str, repeat_count: int) -> str:
50 + """Generate long text by repeating base text."""
51 + return base_text * repeat_count
52 +
53 +
54 +def test_embedding_with_different_lengths():
55 + """Test embedding generation with different text lengths."""
56 + client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
57 +
58 + # Test cases with different text lengths
59 + test_cases = [
60 + {
61 + "name": "Short Text",
62 + "text": "Hello, this is a short text for embedding.",
63 + "expected_chunks": 1,
64 + },
65 + {
66 + "name": "Medium Text",
67 + "text": generate_long_text(
68 + "This is a medium-length text that should fit within the "
69 + "model's context window. " * 20,
70 + 2,
71 + ),
72 + "expected_chunks": 1,
73 + },
74 + {
75 + "name": "Long Text (2 chunks)",
76 + "text": generate_long_text(
77 + "This is a very long text that will exceed the model's "
78 + "maximum context length and trigger chunked processing. " * 50,
79 + 5,
80 + ),
81 + "expected_chunks": 2,
82 + },
83 + {
84 + "name": "Very Long Text (3+ chunks)",
85 + "text": generate_long_text(
86 + "This text is extremely long and will definitely "
87 + "require multiple chunks for processing. " * 100,
88 + 10,
89 + ),
90 + "expected_chunks": 3,
91 + },
92 + ]
93 +
94 + print("🧪 Testing vLLM Long Text Embedding with Chunked Processing")
95 + print("=" * 70)
96 +
97 + for i, test_case in enumerate(test_cases, 1):
98 + print(f"\n📝 Test {i}: {test_case['name']}")
99 + print(f"Text length: {len(test_case['text'])} characters")
100 +
101 + try:
102 + start_time = time.time()
103 +
104 + response = client.embeddings.create(
105 + input=test_case["text"], model=MODEL_NAME, encoding_format="float"
106 + )
107 +
108 + end_time = time.time()
109 + processing_time = end_time - start_time
110 +
111 + # Extract embedding data
112 + embedding = response.data[0].embedding
113 + embedding_dim = len(embedding)
114 +
115 + print("✅ Success!")
116 + print(f" - Embedding dimension: {embedding_dim}")
117 + print(f" - Processing time: {processing_time:.2f}s")
118 + print(f" - Expected chunks: ~{test_case['expected_chunks']}")
119 + print(f" - First 5 values: {embedding[:5]}")
120 +
121 + except Exception as e:
122 + print(f"❌ Failed: {str(e)}")
123 +
124 +
125 +def test_batch_embedding():
126 + """Test batch embedding with mixed-length inputs."""
127 + client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
128 +
129 + print("\n🔄 Testing Batch Embedding with Mixed Lengths")
130 + print("=" * 50)
131 +
132 + # Mix of short and long texts
133 + batch_inputs = [
134 + "Short text 1",
135 + generate_long_text("Medium length text that fits in one chunk. " * 20, 1),
136 + "Another short text",
137 + generate_long_text("Long text requiring chunked processing. " * 100, 5),
138 + ]
139 +
140 + try:
141 + start_time = time.time()
142 +
143 + response = client.embeddings.create(
144 + input=batch_inputs, model=MODEL_NAME, encoding_format="float"
145 + )
146 +
147 + end_time = time.time()
148 + processing_time = end_time - start_time
149 +
150 + print("✅ Batch processing successful!")
151 + print(f" - Number of inputs: {len(batch_inputs)}")
152 + print(f" - Number of embeddings: {len(response.data)}")
153 + print(f" - Total processing time: {processing_time:.2f}s")
154 + print(
155 + f" - Average time per input: {processing_time / len(batch_inputs):.2f}s"
156 + )
157 +
158 + for i, data in enumerate(response.data):
159 + input_length = len(batch_inputs[i])
160 + embedding_dim = len(data.embedding)
161 + print(
162 + f" - Input {i + 1}: {input_length} chars → {embedding_dim}D embedding"
163 + )
164 +
165 + except Exception as e:
166 + print(f"❌ Batch processing failed: {str(e)}")
167 +
168 +
169 +def test_multiple_long_texts_batch():
170 + """Test batch processing with multiple long texts to verify chunk ID uniqueness."""
171 + client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
172 +
173 + print("\n🔧 Testing Multiple Long Texts in Batch (Chunk ID Fix Verification)")
174 + print("=" * 70)
175 +
176 + # Create multiple distinct long texts that will all require chunking
177 + # Note: All pooling types now use MEAN aggregation across chunks:
178 + # - Native pooling (MEAN/CLS/LAST) is used within each chunk
179 + # - MEAN aggregation combines results across all chunks
180 + # - Full semantic coverage for all pooling types
181 + long_texts = [
182 + generate_long_text(
183 + "First long document about artificial intelligence and machine learning. "
184 + * 80,
185 + 6,
186 + ),
187 + generate_long_text(
188 + "Second long document about natural language processing and transformers. "
189 + * 80,
190 + 6,
191 + ),
192 + generate_long_text(
193 + "Third long document about computer vision and neural networks. " * 80, 6
194 + ),
195 + ]
196 +
197 + # Add some short texts to mix things up
198 + batch_inputs = [
199 + "Short text before long texts",
200 + long_texts[0],
201 + "Short text between long texts",
202 + long_texts[1],
203 + long_texts[2],
204 + "Short text after long texts",
205 + ]
206 +
207 + print("📊 Batch composition:")
208 + for i, text in enumerate(batch_inputs):
209 + length = len(text)
210 + text_type = "Long (will be chunked)" if length > 5000 else "Short"
211 + print(f" - Input {i + 1}: {length} chars ({text_type})")
212 +
213 + try:
214 + start_time = time.time()
215 +
216 + response = client.embeddings.create(
217 + input=batch_inputs, model=MODEL_NAME, encoding_format="float"
218 + )
219 +
220 + end_time = time.time()
221 + processing_time = end_time - start_time
222 +
223 + print("\n✅ Multiple long texts batch processing successful!")
224 + print(f" - Number of inputs: {len(batch_inputs)}")
225 + print(f" - Number of embeddings returned: {len(response.data)}")
226 + print(f" - Total processing time: {processing_time:.2f}s")
227 +
228 + # Verify each embedding is different (no incorrect aggregation)
229 + embeddings = [data.embedding for data in response.data]
230 +
231 + if len(embeddings) >= 3:
232 + import numpy as np
233 +
234 + # Compare embeddings of the long texts (indices 1, 3, 4)
235 + long_embeddings = [
236 + np.array(embeddings[1]), # First long text
237 + np.array(embeddings[3]), # Second long text
238 + np.array(embeddings[4]), # Third long text
239 + ]
240 +
241 + print("\n🔍 Verifying embedding uniqueness:")
242 + for i in range(len(long_embeddings)):
243 + for j in range(i + 1, len(long_embeddings)):
244 + cosine_sim = np.dot(long_embeddings[i], long_embeddings[j]) / (
245 + np.linalg.norm(long_embeddings[i])
246 + * np.linalg.norm(long_embeddings[j])
247 + )
248 + print(
249 + f" - Similarity between long text {i + 1} and {j + 1}: "
250 + f"{cosine_sim:.4f}"
251 + )
252 +
253 + if (
254 + cosine_sim < 0.9
255 + ): # Different content should have lower similarity
256 + print(" ✅ Good: Embeddings are appropriately different")
257 + else:
258 + print(
259 + " ⚠️ High similarity - may indicate chunk "
260 + "aggregation issue"
261 + )
262 +
263 + print("\n📋 Per-input results:")
264 + for i, data in enumerate(response.data):
265 + input_length = len(batch_inputs[i])
266 + embedding_dim = len(data.embedding)
267 + embedding_norm = np.linalg.norm(data.embedding)
268 + print(
269 + f" - Input {i + 1}: {input_length} chars → {embedding_dim}D "
270 + f"embedding (norm: {embedding_norm:.4f})"
271 + )
272 +
273 + print(
274 + "\n✅ This test verifies the fix for chunk ID collisions in "
275 + "batch processing"
276 + )
277 + print(" - Before fix: Multiple long texts would have conflicting chunk IDs")
278 + print(" - After fix: Each prompt's chunks have unique IDs with prompt index")
279 +
280 + except Exception as e:
281 + print(f"❌ Multiple long texts batch test failed: {str(e)}")
282 + print(" This might indicate the chunk ID collision bug is present!")
283 +
284 +
285 +def test_embedding_consistency():
286 + """Test that chunked processing produces consistent results."""
287 + client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
288 +
289 + print("\n🔍 Testing Embedding Consistency")
290 + print("=" * 40)
291 +
292 + # Use the same long text multiple times
293 + long_text = generate_long_text(
294 + "Consistency test text for chunked processing validation. " * 50, 3
295 + )
296 +
297 + embeddings = []
298 +
299 + try:
300 + for i in range(3):
301 + response = client.embeddings.create(
302 + input=long_text, model=MODEL_NAME, encoding_format="float"
303 + )
304 + embeddings.append(response.data[0].embedding)
305 + print(f" - Generated embedding {i + 1}")
306 +
307 + # Check consistency (embeddings should be identical)
308 + if len(embeddings) >= 2:
309 + # Calculate similarity between first two embeddings
310 +
311 + emb1 = np.array(embeddings[0])
312 + emb2 = np.array(embeddings[1])
313 +
314 + # Cosine similarity
315 + cosine_sim = np.dot(emb1, emb2) / (
316 + np.linalg.norm(emb1) * np.linalg.norm(emb2)
317 + )
318 +
319 + print("✅ Consistency test completed!")
320 + print(f" - Cosine similarity between runs: {cosine_sim:.6f}")
321 + print(" - Expected: ~1.0 (identical embeddings)")
322 +
323 + if cosine_sim > 0.999:
324 + print(" - ✅ High consistency achieved!")
325 + else:
326 + print(" - ⚠️ Consistency may vary due to numerical precision")
327 +
328 + except Exception as e:
329 + print(f"❌ Consistency test failed: {str(e)}")
330 +
331 +
332 +def main():
333 + """Main function to run all tests."""
334 + print("🚀 vLLM Long Text Embedding Client")
335 + print(f"📡 Connecting to: {BASE_URL}")
336 + print(f"🤖 Model: {MODEL_NAME}")
337 + masked_key = "*" * (len(API_KEY) - 4) + API_KEY[-4:] if len(API_KEY) > 4 else "****"
338 + print(f"🔑 API Key: {masked_key}")
339 +
340 + # Run all test cases
341 + test_embedding_with_different_lengths()
342 + test_batch_embedding()
343 + test_multiple_long_texts_batch()
344 + test_embedding_consistency()
345 +
346 + print("\n" + "=" * 70)
347 + print("🎉 All tests completed!")
348 + print("\n💡 Key Features Demonstrated:")
349 + print(" - ✅ Automatic chunked processing for long text")
350 + print(" - ✅ Seamless handling of mixed-length batches")
351 + print(" - ✅ Multiple long texts in single batch (chunk ID fix)")
352 + print(" - ✅ Unified chunked processing:")
353 + print(" • Native pooling used within each chunk")
354 + print(" • MEAN aggregation across all chunks")
355 + print(" • Complete semantic coverage for all pooling types")
356 + print(" - ✅ Consistent embedding generation")
357 + print(" - ✅ Backward compatibility with short text")
358 + print("\n📚 For more information, see:")
359 + print(
360 + " - Documentation: https://docs.vllm.ai/en/latest/models/pooling_models.html"
361 + )
362 + print(" - Chunked Processing Guide: openai_embedding_long_text.md")
363 +
364 +
365 +if __name__ == "__main__":
366 + main()
examples/online_serving/openai_embedding_long_text/service.sh new
+137
@@ -0,0 +1,137 @@
1 +#!/bin/bash
2 +
3 +# SPDX-License-Identifier: Apache-2.0
4 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
5 +
6 +# vLLM Embedding Server with Enhanced Chunked Processing
7 +# This script starts a vLLM server with chunked processing enabled for long text embedding.
8 +# Now supports proper pooling type validation and model-specific configurations.
9 +
10 +set -euo pipefail
11 +
12 +# Configuration
13 +MODEL_NAME=${MODEL_NAME:-"intfloat/multilingual-e5-large"}
14 +MODEL_CODE=${MODEL_CODE:-"multilingual-e5-large"}
15 +
16 +PORT=${PORT:-31090}
17 +GPU_COUNT=${GPU_COUNT:-1}
18 +MAX_EMBED_LEN=${MAX_EMBED_LEN:-3072000}
19 +API_KEY=${API_KEY:-"your-api-key"}
20 +
21 +# Enhanced pooling configuration with model-specific defaults
22 +POOLING_TYPE=${POOLING_TYPE:-"auto"} # auto, MEAN, CLS, LAST
23 +export VLLM_ENABLE_CHUNKED_PROCESSING=true
24 +export CUDA_VISIBLE_DEVICES=2,3,4,5
25 +# export VLLM_ATTENTION_BACKEND=XFORMERS
26 +
27 +echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing"
28 +echo "=================================================================="
29 +
30 +# Environment variables for optimization
31 +export VLLM_WORKER_MULTIPROC_METHOD=spawn
32 +
33 +# Function to determine optimal pooling type for known models
34 +get_optimal_pooling_type() {
35 + local model="$1"
36 + case "$model" in
37 + *"e5-"* | *"multilingual-e5"*)
38 + echo "MEAN" # E5 series native pooling
39 + ;;
40 + *"bge-"*)
41 + echo "CLS" # BGE series native pooling
42 + ;;
43 + *"gte-"*)
44 + echo "LAST" # GTE series native pooling
45 + ;;
46 + *"sentence-t5"* | *"st5"*)
47 + echo "MEAN" # Sentence-T5 native pooling
48 + ;;
49 + *"jina-embeddings"*)
50 + echo "MEAN" # Jina embeddings native pooling
51 + ;;
52 + *"Qwen"*"Embedding"*)
53 + echo "LAST" # Qwen embeddings native pooling
54 + ;;
55 + *)
56 + echo "MEAN" # Default native pooling for unknown models
57 + ;;
58 + esac
59 +}
60 +
61 +# Auto-detect pooling type if not explicitly set
62 +if [ "$POOLING_TYPE" = "auto" ]; then
63 + POOLING_TYPE=$(get_optimal_pooling_type "$MODEL_NAME")
64 + echo "🔍 Auto-detected pooling type: $POOLING_TYPE for model $MODEL_NAME"
65 +fi
66 +
67 +# Display configuration
68 +echo "📋 Configuration:"
69 +echo " - Model: $MODEL_NAME"
70 +echo " - Port: $PORT"
71 +echo " - GPU Count: $GPU_COUNT"
72 +echo " - Enhanced Chunked Processing: ${VLLM_ENABLE_CHUNKED_PROCESSING}"
73 +echo " - Max Embed Length: ${MAX_EMBED_LEN} tokens"
74 +echo " - Native Pooling Type: $POOLING_TYPE + Normalization"
75 +echo " - Cross-chunk Aggregation: MEAN (automatic)"
76 +echo ""
77 +
78 +# Validate GPU availability
79 +if command -v nvidia-smi &> /dev/null; then
80 + gpu_count=$(nvidia-smi --list-gpus | wc -l)
81 + echo "🖥️ Available GPUs: $gpu_count"
82 + if [ "$GPU_COUNT" -gt "$gpu_count" ]; then
83 + echo "⚠️ Warning: Requested $GPU_COUNT GPUs but only $gpu_count available"
84 + echo " Adjusting to use $gpu_count GPUs"
85 + GPU_COUNT=$gpu_count
86 + fi
87 +else
88 + echo "⚠️ Warning: nvidia-smi not found. GPU detection skipped."
89 +fi
90 +
91 +# Chunked processing uses unified MEAN aggregation
92 +echo "ℹ️ Chunked Processing: Using $POOLING_TYPE pooling within chunks, MEAN aggregation across chunks"
93 +echo " - All chunks processed for complete semantic coverage"
94 +echo " - Weighted averaging based on chunk token counts"
95 +
96 +echo ""
97 +echo "🔧 Starting server with enhanced chunked processing configuration..."
98 +
99 +# Build pooler config JSON
100 +POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"normalize\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}"
101 +
102 +# Start vLLM server with enhanced chunked processing
103 +vllm serve "$MODEL_NAME" \
104 + --tensor-parallel-size "$GPU_COUNT" \
105 + --enforce-eager \
106 + --override-pooler-config "$POOLER_CONFIG" \
107 + --served-model-name ${MODEL_CODE} \
108 + --api-key "$API_KEY" \
109 + --trust-remote-code \
110 + --port "$PORT" \
111 + --host 0.0.0.0
112 +
113 +echo ""
114 +echo "✅ vLLM Embedding Server started successfully!"
115 +echo ""
116 +echo "📡 Server Information:"
117 +echo " - Base URL: http://localhost:$PORT"
118 +echo " - Model Code: ${MODEL_CODE}"
119 +echo " - API Key: $API_KEY"
120 +echo " - Native Pooling: $POOLING_TYPE | Cross-chunk: MEAN"
121 +echo ""
122 +echo "🧪 Test the server with:"
123 +echo " python examples/online_serving/openai_embedding_long_text_client.py"
124 +echo ""
125 +echo "📚 Enhanced features enabled:"
126 +echo " ✅ Intelligent native pooling type detection"
127 +echo " ✅ Unified MEAN aggregation for chunked processing"
128 +echo " ✅ Model-specific native pooling optimization"
129 +echo " ✅ Enhanced max embedding length (${MAX_EMBED_LEN} tokens)"
130 +echo " ✅ Complete semantic coverage for all pooling types"
131 +echo " ✅ OpenAI-compatible API"
132 +echo " ✅ GPU acceleration"
133 +echo ""
134 +echo "🔧 Advanced usage:"
135 +echo " - Set POOLING_TYPE=MEAN|CLS|LAST to override auto-detection"
136 +echo " - Set MAX_EMBED_LEN to adjust maximum input length"
137 +echo " - All pooling types use MEAN aggregation across chunks"
tests/entrypoints/openai/test_embedding_long_text.py new
+441
@@ -0,0 +1,441 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Test cases for long text embedding with automatic chunking mechanism.
5 +
6 +This test suite validates vLLM's automatic chunking functionality for handling
7 +text inputs that exceed the model's maximum token length, specifically targeting
8 +the intfloat/multilingual-e5-small model (max token length: 512).
9 +"""
10 +
11 +import random
12 +
13 +import openai
14 +import pytest
15 +import pytest_asyncio
16 +
17 +from vllm.entrypoints.openai.protocol import EmbeddingResponse
18 +
19 +from ...utils import RemoteOpenAIServer
20 +
21 +
22 +def _generate_random_text(word_count: int) -> str:
23 + """Generate random text with approximately the specified word count."""
24 + # Common English words with focus on verbs and nouns for realistic text
25 + common_words = [
26 + # Essential articles and pronouns (minimal)
27 + "the",
28 + "and",
29 + "you",
30 + "they",
31 + "this",
32 + "that",
33 + "these",
34 + "those",
35 +
36 + # Action verbs
37 + "create",
38 + "build",
39 + "develop",
40 + "design",
41 + "implement",
42 + "execute",
43 + "analyze",
44 + "process",
45 + "generate",
46 + "calculate",
47 + "evaluate",
48 + "optimize",
49 + "transform",
50 + "integrate",
51 + "configure",
52 + "deploy",
53 + "monitor",
54 + "manage",
55 + "discover",
56 + "explore",
57 + "investigate",
58 + "research",
59 + "study",
60 + "examine",
61 + "improve",
62 + "enhance",
63 + "upgrade",
64 + "modify",
65 + "update",
66 + "maintain",
67 + "solve",
68 + "resolve",
69 + "handle",
70 + "address",
71 + "tackle",
72 + "overcome",
73 + "communicate",
74 + "collaborate",
75 + "coordinate",
76 + "organize",
77 + "plan",
78 + "achieve",
79 + "accomplish",
80 + "complete",
81 + "finish",
82 + "deliver",
83 + "provide",
84 +
85 + # Technology and science nouns
86 + "system",
87 + "application",
88 + "software",
89 + "hardware",
90 + "network",
91 + "database",
92 + "algorithm",
93 + "model",
94 + "framework",
95 + "platform",
96 + "interface",
97 + "protocol",
98 + "architecture",
99 + "infrastructure",
100 + "component",
101 + "module",
102 + "service",
103 + "technology",
104 + "innovation",
105 + "solution",
106 + "methodology",
107 + "approach",
108 + "artificial",
109 + "intelligence",
110 + "machine",
111 + "learning",
112 + "neural",
113 + "network",
114 + "computer",
115 + "processor",
116 + "memory",
117 + "storage",
118 + "computation",
119 + "data",
120 + "information",
121 + "knowledge",
122 + "insight",
123 + "pattern",
124 + "trend",
125 + "analysis",
126 + "research",
127 + "development",
128 + "engineering",
129 + "science",
130 + "mathematics",
131 + "statistics",
132 + "probability",
133 + "optimization",
134 + "performance",
135 + "efficiency",
136 +
137 + # General nouns
138 + "project",
139 + "team",
140 + "organization",
141 + "company",
142 + "business",
143 + "industry",
144 + "market",
145 + "customer",
146 + "user",
147 + "client",
148 + "product",
149 + "feature",
150 + "function",
151 + "requirement",
152 + "specification",
153 + "documentation",
154 + "report",
155 + "result",
156 + "outcome",
157 + "impact",
158 + "benefit",
159 + "advantage",
160 + "challenge",
161 + "problem",
162 + "opportunity",
163 + "strategy",
164 + "goal",
165 + "objective",
166 + "target",
167 + "milestone",
168 + "process",
169 + "procedure",
170 + "workflow",
171 + "pipeline",
172 + "operation",
173 + "task",
174 + "activity",
175 + "event",
176 + "session",
177 + "meeting",
178 + "discussion",
179 + "decision"
180 + ]
181 +
182 + words = []
183 + for _ in range(word_count):
184 + words.append(random.choice(common_words))
185 +
186 + # Add some punctuation for more realistic text
187 + text = " ".join(words)
188 + # Add periods every 10-20 words
189 + words_list = text.split()
190 + result = []
191 + for i, word in enumerate(words_list):
192 + result.append(word)
193 + if ((i + 1) % random.randint(10, 20) == 0 and i < len(words_list) - 1):
194 + result[-1] += "."
195 +
196 + return " ".join(result)
197 +
198 +
199 +MODEL_NAME = "intfloat/multilingual-e5-small"
200 +DTYPE = "bfloat16"
201 +
202 +# Test text: Generate text with approximately 1500 words to exceed 1024 tokens
203 +LONG_TEXT_1500_WORDS = _generate_random_text(1500)
204 +
205 +# Test text: Generate text with approximately 2500 words to exceed 2048 tokens
206 +LONG_TEXT_2500_WORDS = _generate_random_text(2500)
207 +
208 +
209 +@pytest.fixture(scope="module")
210 +def server_with_chunked_processing():
211 + """Start server with automatic chunking processing enabled."""
212 + args = [
213 + "--runner",
214 + "pooling",
215 + "--dtype",
216 + DTYPE,
217 + "--enforce-eager",
218 + "--max-model-len",
219 + "512", # Set smaller max_model_len to trigger chunking mechanism
220 + '--override-pooler-config',
221 + ('{"pooling_type": "MEAN", "normalize": true, '
222 + '"enable_chunked_processing": true, "max_embed_len": 10000}'),
223 + "--gpu-memory-utilization",
224 + "0.8",
225 + ]
226 +
227 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
228 + yield remote_server
229 +
230 +
231 +@pytest_asyncio.fixture
232 +async def client_with_chunked_processing(server_with_chunked_processing):
233 + """Create async client with chunking processing support."""
234 + async with server_with_chunked_processing.get_async_client(
235 + ) as async_client:
236 + yield async_client
237 +
238 +
239 +@pytest.mark.asyncio
240 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
241 +async def test_long_text_embedding_1500_chars(
242 + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str):
243 + """Test embedding processing for ~1500 character long text
244 + (~1028 tokens, exceeding 512 token limit)."""
245 +
246 + # Verify text length
247 + # Verify text has sufficient word count (approximately 1500 words)
248 + word_count = len(LONG_TEXT_1500_WORDS.split())
249 + assert word_count >= 1400, (
250 + f"Test text word count insufficient: {word_count} words")
251 +
252 + # Send embedding request
253 + embedding_response = await client_with_chunked_processing.embeddings.create(
254 + model=model_name,
255 + input=[LONG_TEXT_1500_WORDS],
256 + encoding_format="float",
257 + )
258 +
259 + # Verify response structure
260 + embeddings = EmbeddingResponse.model_validate(
261 + embedding_response.model_dump(mode="json"))
262 +
263 + assert embeddings.id is not None
264 + assert len(embeddings.data) == 1
265 + assert len(embeddings.data[0].embedding
266 + ) == 384 # multilingual-e5-small embedding dimension
267 + assert embeddings.usage.completion_tokens == 0
268 + # Due to chunked processing, token count should
269 + # reflect actual processed tokens
270 + # With ~1500 words, we expect roughly
271 + # 1024+ tokens (exceeding 512 token limit)
272 + # Should exceed single chunk limit of 512
273 + assert embeddings.usage.prompt_tokens > 800
274 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
275 +
276 + # Verify embedding vector validity
277 + embedding_vector = embeddings.data[0].embedding
278 + assert all(
279 + isinstance(x, float)
280 + for x in embedding_vector), "Embedding vector should contain floats"
281 + assert not all(
282 + x == 0
283 + for x in embedding_vector), "Embedding vector should not be all zeros"
284 +
285 +
286 +@pytest.mark.asyncio
287 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
288 +async def test_long_text_embedding_2500_chars(
289 + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str):
290 + """Test embedding processing for ~2500 character long text
291 + (~2048 tokens, requiring multiple chunks)."""
292 +
293 + # Verify text length
294 + # Verify text has sufficient word count (approximately 2500 words)
295 + word_count = len(LONG_TEXT_2500_WORDS.split())
296 + assert word_count >= 2300, (
297 + f"Test text word count insufficient: {word_count} words")
298 +
299 + # Send embedding request
300 + embedding_response = await client_with_chunked_processing.embeddings.create(
301 + model=model_name,
302 + input=[LONG_TEXT_2500_WORDS],
303 + encoding_format="float",
304 + )
305 +
306 + # Verify response structure
307 + embeddings = EmbeddingResponse.model_validate(
308 + embedding_response.model_dump(mode="json"))
309 +
310 + assert embeddings.id is not None
311 + assert len(embeddings.data) == 1
312 + assert len(embeddings.data[0].embedding
313 + ) == 384 # multilingual-e5-small embedding dimension
314 + assert embeddings.usage.completion_tokens == 0
315 + # Due to chunked processing, token count should
316 + # reflect actual processed tokens
317 + # With ~2500 words, we expect
318 + # roughly 2048+ tokens (requiring multiple chunks)
319 + # Should require multiple chunks for processing
320 + assert embeddings.usage.prompt_tokens > 1500
321 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
322 +
323 + # Verify embedding vector validity
324 + embedding_vector = embeddings.data[0].embedding
325 + assert all(
326 + isinstance(x, float)
327 + for x in embedding_vector), "Embedding vector should contain floats"
328 + assert not all(
329 + x == 0
330 + for x in embedding_vector), "Embedding vector should not be all zeros"
331 +
332 +
333 +@pytest.mark.asyncio
334 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
335 +async def test_batch_long_text_embedding(
336 + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str):
337 + """Test batch long text embedding processing."""
338 +
339 + input_texts = [
340 + LONG_TEXT_1500_WORDS,
341 + LONG_TEXT_2500_WORDS,
342 + "This is a short text test.", # Short text for comparison
343 + ]
344 +
345 + # Send batch embedding request
346 + embedding_response = await client_with_chunked_processing.embeddings.create(
347 + model=model_name,
348 + input=input_texts,
349 + encoding_format="float",
350 + )
351 +
352 + # Verify response structure
353 + embeddings = EmbeddingResponse.model_validate(
354 + embedding_response.model_dump(mode="json"))
355 +
356 + assert embeddings.id is not None
357 + assert len(embeddings.data) == 3 # Three input texts
358 +
359 + # Verify each embedding dimension
360 + for i, embedding_data in enumerate(embeddings.data):
361 + assert len(embedding_data.embedding) == 384
362 + assert embedding_data.index == i
363 +
364 + # Verify embedding vector validity
365 + embedding_vector = embedding_data.embedding
366 + assert all(isinstance(x, float) for x in embedding_vector)
367 + assert not all(x == 0 for x in embedding_vector)
368 +
369 + # Verify token usage
370 + assert embeddings.usage.completion_tokens == 0
371 + # Total token count should be very substantial
372 + assert embeddings.usage.prompt_tokens > 1000
373 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
374 +
375 +
376 +@pytest.mark.asyncio
377 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
378 +async def test_chunked_vs_normal_consistency(
379 + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str):
380 + """Test consistency between chunked and
381 + normal processing (using short text)."""
382 +
383 + # Use a short text within the 512 token limit
384 + short_text = ("Artificial intelligence technology is changing our world, "
385 + "bringing unprecedented opportunities and challenges.")
386 +
387 + # Send embedding request
388 + embedding_response = await client_with_chunked_processing.embeddings.create(
389 + model=model_name,
390 + input=[short_text],
391 + encoding_format="float",
392 + )
393 +
394 + # Verify response structure
395 + embeddings = EmbeddingResponse.model_validate(
396 + embedding_response.model_dump(mode="json"))
397 +
398 + assert embeddings.id is not None
399 + assert len(embeddings.data) == 1
400 + assert len(embeddings.data[0].embedding) == 384
401 + assert embeddings.usage.completion_tokens == 0
402 + # Short text should not require chunked processing
403 + assert embeddings.usage.prompt_tokens < 512
404 + assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
405 +
406 + # 验证embedding向量的有效性
407 + embedding_vector = embeddings.data[0].embedding
408 + assert all(isinstance(x, float) for x in embedding_vector)
409 + assert not all(x == 0 for x in embedding_vector)
410 +
411 +
412 +@pytest.mark.asyncio
413 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
414 +async def test_chunked_processing_response_format(
415 + client_with_chunked_processing: openai.AsyncOpenAI, model_name: str):
416 + """Test response format and structure during chunked processing."""
417 +
418 + # Test with long text to trigger chunking
419 + embedding_response = await client_with_chunked_processing.embeddings.create(
420 + model=model_name,
421 + input=[LONG_TEXT_1500_WORDS],
422 + encoding_format="float",
423 + )
424 +
425 + # Verify response structure
426 + embeddings = EmbeddingResponse.model_validate(
427 + embedding_response.model_dump(mode="json"))
428 +
429 + assert embeddings.id is not None
430 + assert len(embeddings.data) == 1
431 + assert embeddings.data[0].object == "embedding"
432 + assert embeddings.data[0].index == 0
433 +
434 + # Verify embedding vector properties
435 + embedding_vector = embeddings.data[0].embedding
436 + import math
437 + vector_norm = math.sqrt(sum(x * x for x in embedding_vector))
438 + # Check that the vector is normalized
439 + # (default behavior for most embedding models)
440 + assert 0.8 < vector_norm < 1.2, (
441 + f"Vector norm should be reasonable, actual: {vector_norm}")
vllm/config/__init__.py
+19
@@ -2598,6 +2598,25 @@ class PoolerConfig:
2598 ``math-shepherd-mistral-7b-prm`` model.
2599 """
2600
2601 + enable_chunked_processing: Optional[bool] = None
2602 + """
2603 + Whether to enable chunked processing for long inputs that exceed the model's
2604 + maximum position embeddings. When enabled, long inputs will be split into
2605 + chunks, processed separately, and then aggregated using weighted averaging.
2606 + This allows embedding models to handle arbitrarily long text without CUDA
2607 + errors. Defaults to False.
2608 + """
2609 +
2610 + max_embed_len: Optional[int] = None
2611 + """
2612 + Maximum input length allowed for embedding generation. When set, allows
2613 + inputs longer than max_embed_len to be accepted for embedding models.
2614 + This parameter enables accepting long inputs without requiring
2615 + VLLM_ALLOW_LONG_MAX_MODEL_LEN environment variable. When an input exceeds
2616 + max_embed_len, it will be handled according to the original max_model_len
2617 + validation logic. Defaults to None (i.e. set to max_model_len).
2618 + """
2619 +
2620 def compute_hash(self) -> str:
2621 """
2622 WARNING: Whenever a new field is added to this config,
vllm/entrypoints/openai/serving_embedding.py
+454 -3
@@ -2,9 +2,11 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import base64
5 -from typing import Final, Literal, Optional, Union, cast
5 +from collections.abc import AsyncGenerator, Mapping
6 +from typing import Any, Final, Literal, Optional, Union, cast
7
8 import numpy as np
9 +import torch
10 from fastapi import Request
11 from typing_extensions import assert_never, override
12
@@ -12,19 +14,28 @@ from vllm.config import ModelConfig
14 from vllm.engine.protocol import EngineClient
15 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
16 from vllm.entrypoints.logger import RequestLogger
17 +# yapf conflicts with isort for this docstring
18 +# yapf: disable
19 from vllm.entrypoints.openai.protocol import (EmbeddingChatRequest,
20 + EmbeddingCompletionRequest,
21 EmbeddingRequest,
22 EmbeddingResponse,
23 EmbeddingResponseData,
24 ErrorResponse, UsageInfo)
25 from vllm.entrypoints.openai.serving_engine import (EmbeddingServeContext,
26 OpenAIServing,
22 - ServeContext)
27 + RequestPrompt,
28 + ServeContext,
29 + TextTokensPrompt)
30 +# yapf: enable
31 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
32 +from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt
33 +from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
34 from vllm.logger import init_logger
35 from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput,
26 - PoolingRequestOutput)
36 + PoolingOutput, PoolingRequestOutput, RequestOutput)
37 from vllm.pooling_params import PoolingParams
38 +from vllm.utils import chunk_list
39
40 logger = init_logger(__name__)
41
@@ -46,6 +57,17 @@ def _get_embedding(
57
58 class EmbeddingMixin(OpenAIServing):
59
60 + def __init__(self, *args, **kwargs):
61 + super().__init__(*args, **kwargs)
62 +
63 + pooler_config = self.model_config.pooler_config
64 +
65 + # Avoid repeated attribute lookups
66 + self.supports_chunked_processing = bool(
67 + pooler_config and pooler_config.enable_chunked_processing)
68 + self.max_embed_len = (pooler_config.max_embed_len if pooler_config
69 + and pooler_config.max_embed_len else None)
70 +
71 @override
72 async def _preprocess(
73 self,
@@ -129,6 +151,435 @@ class EmbeddingMixin(OpenAIServing):
151 usage=usage,
152 )
153
154 + def _get_max_position_embeddings(self) -> int:
155 + """Get the model's effective maximum sequence length for chunking."""
156 + return self.model_config.max_model_len
157 +
158 + def _should_use_chunked_processing(self, request) -> bool:
159 + """Check if chunked processing should be used for this request."""
160 + return isinstance(
161 + request,
162 + (EmbeddingCompletionRequest,
163 + EmbeddingChatRequest)) and self.supports_chunked_processing
164 +
165 + async def _process_chunked_request(
166 + self,
167 + ctx: EmbeddingServeContext,
168 + original_prompt: TextTokensPrompt,
169 + pooling_params,
170 + trace_headers,
171 + prompt_idx: int,
172 + ) -> list[AsyncGenerator[PoolingRequestOutput, None]]:
173 + """Process a single prompt using chunked processing."""
174 + generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
175 + token_ids = original_prompt["prompt_token_ids"]
176 +
177 + # Split into chunks using max_position_embeddings
178 + max_pos_embeddings = self._get_max_position_embeddings()
179 + # Process all chunks for MEAN aggregation
180 + for chunk_idx, chunk_tokens in enumerate(
181 + chunk_list(token_ids, max_pos_embeddings)):
182 + # Create a request ID for this chunk
183 + chunk_request_id = (f"{ctx.request_id}-prompt-{prompt_idx}-"
184 + f"chunk-{chunk_idx}")
185 +
186 + # Create engine prompt for this chunk
187 + chunk_engine_prompt = EngineTokensPrompt(
188 + prompt_token_ids=chunk_tokens)
189 +
190 + # Create chunk request prompt for logging
191 + chunk_text = ""
192 + chunk_request_prompt = TextTokensPrompt(
193 + prompt=chunk_text, prompt_token_ids=chunk_tokens)
194 +
195 + # Log the chunk
196 + self._log_inputs(chunk_request_id,
197 + chunk_request_prompt,
198 + params=pooling_params,
199 + lora_request=ctx.lora_request)
200 +
201 + # Create generator for this chunk and wrap it to return indices
202 + original_generator = self.engine_client.encode(
203 + chunk_engine_prompt,
204 + pooling_params,
205 + chunk_request_id,
206 + lora_request=ctx.lora_request,
207 + trace_headers=trace_headers,
208 + priority=getattr(ctx.request, "priority", 0),
209 + )
210 +
211 + generators.append(original_generator)
212 +
213 + return generators
214 +
215 + def _validate_input(
216 + self,
217 + request,
218 + input_ids: list[int],
219 + input_text: str,
220 + ) -> TextTokensPrompt:
221 + """Override to support chunked processing for embedding requests."""
222 + token_num = len(input_ids)
223 +
224 + # Note: EmbeddingRequest doesn't have max_tokens
225 + if isinstance(request,
226 + (EmbeddingCompletionRequest, EmbeddingChatRequest)):
227 + # Check if chunked processing is enabled for pooling models
228 + enable_chunked = self._should_use_chunked_processing(request)
229 +
230 + # Use max_position_embeddings for chunked processing decisions
231 + max_pos_embeddings = self._get_max_position_embeddings()
232 +
233 + # Determine the effective max length for validation
234 + if self.max_embed_len is not None:
235 + # Use max_embed_len for validation instead of max_model_len
236 + length_type = "maximum embedding input length"
237 + max_length_value = self.max_embed_len
238 + else:
239 + # Fall back to max_model_len validation (original behavior)
240 + length_type = "maximum context length"
241 + max_length_value = self.max_model_len
242 +
243 + validation_error_msg = (
244 + "This model's {length_type} is {max_length_value} tokens. "
245 + "However, you requested {token_num} tokens in the input for "
246 + "embedding generation. Please reduce the length of the input.")
247 +
248 + chunked_processing_error_msg = (
249 + "This model's {length_type} is {max_length_value} tokens. "
250 + "However, you requested {token_num} tokens in the input for "
251 + "embedding generation. Please reduce the length of the input "
252 + "or enable chunked processing.")
253 +
254 + # Check if input exceeds max length
255 + if token_num > max_length_value:
256 + raise ValueError(
257 + validation_error_msg.format(
258 + length_type=length_type,
259 + max_length_value=max_length_value,
260 + token_num=token_num))
261 +
262 + # Check for chunked processing
263 + # when exceeding max_position_embeddings
264 + if token_num > max_pos_embeddings:
265 + if enable_chunked:
266 + # Allow long inputs when chunked processing is enabled
267 + logger.info(
268 + "Input length %s exceeds max_position_embeddings "
269 + "%s, will use chunked processing", token_num,
270 + max_pos_embeddings)
271 + else:
272 + raise ValueError(
273 + chunked_processing_error_msg.format(
274 + length_type="maximum position embeddings length",
275 + max_length_value=max_pos_embeddings,
276 + token_num=token_num))
277 +
278 + return TextTokensPrompt(prompt=input_text,
279 + prompt_token_ids=input_ids)
280 +
281 + # For other request types, use the parent's implementation
282 + return super()._validate_input(request, input_ids, input_text)
283 +
284 + def _is_text_tokens_prompt(self, prompt) -> bool:
285 + """Check if a prompt is a TextTokensPrompt (has prompt_token_ids)."""
286 + return (isinstance(prompt, dict) and "prompt_token_ids" in prompt
287 + and "prompt_embeds" not in prompt)
288 +
289 + async def _create_single_prompt_generator(
290 + self,
291 + ctx: EmbeddingServeContext,
292 + engine_prompt: Union[EngineTokensPrompt, EngineEmbedsPrompt],
293 + request_prompt: RequestPrompt,
294 + pooling_params: PoolingParams,
295 + trace_headers: Optional[Mapping[str, str]],
296 + prompt_index: int,
297 + ) -> AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]:
298 + """Create a generator for a single prompt using standard processing."""
299 + request_id_item = f"{ctx.request_id}-{prompt_index}"
300 +
301 + self._log_inputs(request_id_item,
302 + request_prompt,
303 + params=pooling_params,
304 + lora_request=ctx.lora_request)
305 +
306 + # Mypy has an existing bug related to inferring the variance
307 + # of TypedDicts with `builtins.enumerate`:
308 + # https://github.com/python/mypy/issues/8586#issuecomment-2867698435
309 + engine_prompt = cast(Union[EngineTokensPrompt, EngineEmbedsPrompt],
310 + engine_prompt)
311 +
312 + # Return the original generator without wrapping
313 + return self.engine_client.encode(
314 + engine_prompt,
315 + pooling_params,
316 + request_id_item,
317 + lora_request=ctx.lora_request,
318 + trace_headers=trace_headers,
319 + priority=getattr(ctx.request, "priority", 0),
320 + )
321 +
322 + @override
323 + async def _prepare_generators(
324 + self,
325 + ctx: ServeContext,
326 + ) -> Optional[ErrorResponse]:
327 + """Override to support chunked processing."""
328 + ctx = cast(EmbeddingServeContext, ctx)
329 +
330 + # Check if we should use chunked processing
331 + use_chunked = self._should_use_chunked_processing(ctx.request)
332 +
333 + # If no chunked processing needed, delegate to parent class
334 + if not use_chunked:
335 + return await super()._prepare_generators(ctx)
336 +
337 + # Custom logic for chunked processing
338 + generators: list[AsyncGenerator[Union[RequestOutput,
339 + PoolingRequestOutput],
340 + None]] = []
341 +
342 + try:
343 + trace_headers = (None if ctx.raw_request is None else await
344 + self._get_trace_headers(ctx.raw_request.headers))
345 +
346 + pooling_params = self._create_pooling_params(ctx)
347 + if isinstance(pooling_params, ErrorResponse):
348 + return pooling_params
349 +
350 + # Verify and set the task for pooling params
351 + try:
352 + pooling_params.verify("embed", self.model_config)
353 + except ValueError as e:
354 + return self.create_error_response(str(e))
355 +
356 + if ctx.engine_prompts is None:
357 + return self.create_error_response(
358 + "Engine prompts not available")
359 +
360 + if ctx.request_prompts is None:
361 + return self.create_error_response(
362 + "Request prompts not available")
363 +
364 + max_pos_embeddings = self._get_max_position_embeddings()
365 +
366 + for i, engine_prompt in enumerate(ctx.engine_prompts):
367 + request_prompt = ctx.request_prompts[i]
368 +
369 + # Check if this specific prompt needs chunked processing
370 + if self._is_text_tokens_prompt(request_prompt):
371 + # Cast to TextTokensPrompt since we've verified
372 + # prompt_token_ids
373 + text_tokens_prompt = cast(TextTokensPrompt, request_prompt)
374 + if (len(text_tokens_prompt["prompt_token_ids"])
375 + > max_pos_embeddings):
376 + # Use chunked processing for this prompt
377 + chunk_generators = await self._process_chunked_request(
378 + ctx, text_tokens_prompt, pooling_params,
379 + trace_headers, i)
380 + generators.extend(chunk_generators)
381 + continue
382 +
383 + # Normal processing for short prompts or non-token prompts
384 + # Cast engine_prompt to the expected type for mypy
385 + engine_prompt_typed = cast(
386 + Union[EngineTokensPrompt, EngineEmbedsPrompt],
387 + engine_prompt)
388 + generator = await self._create_single_prompt_generator(
389 + ctx, engine_prompt_typed, request_prompt, pooling_params,
390 + trace_headers, i)
391 + generators.append(generator)
392 +
393 + from vllm.utils import merge_async_iterators
394 + ctx.result_generator = merge_async_iterators(*generators)
395 +
396 + return None
397 +
398 + except Exception as e:
399 + # TODO: Use a vllm-specific Validation Error
400 + return self.create_error_response(str(e))
401 +
402 + @override
403 + async def _collect_batch(
404 + self,
405 + ctx: ServeContext,
406 + ) -> Optional[ErrorResponse]:
407 + """Collect and aggregate batch results
408 + with support for chunked processing.
409 +
410 + For chunked requests, performs online aggregation to
411 + minimize memory usage.
412 + For regular requests, collects results normally.
413 + """
414 + ctx = cast(EmbeddingServeContext, ctx)
415 + try:
416 + if ctx.engine_prompts is None:
417 + return self.create_error_response(
418 + "Engine prompts not available")
419 +
420 + # Check if we used chunked processing
421 + use_chunked = self._should_use_chunked_processing(ctx.request)
422 +
423 + if not use_chunked:
424 + return await super()._collect_batch(ctx=ctx)
425 +
426 + if ctx.request_prompts is None:
427 + return self.create_error_response(
428 + "Request prompts not available")
429 +
430 + if ctx.result_generator is None:
431 + return self.create_error_response(
432 + "Result generator not available")
433 +
434 + # Online aggregation for chunked requests to
435 + # minimize memory usage
436 + # Track aggregation state for each prompt
437 + prompt_aggregators: dict[int, dict[str, Any]] = {}
438 + short_prompts_results: dict[int, PoolingRequestOutput] = {}
439 +
440 + async for result_idx, result in ctx.result_generator:
441 + if "-chunk-" in result.request_id:
442 + # Extract prompt_idx from chunked request_id
443 + parts = result.request_id.split("-")
444 + try:
445 + prompt_idx = int(parts[parts.index("prompt") + 1])
446 + except (ValueError, IndexError):
447 + # Fallback: extract from result_idx if parsing fails
448 + prompt_idx = result_idx
449 +
450 + # Initialize aggregator for this prompt if needed
451 + if prompt_idx not in prompt_aggregators:
452 + prompt_aggregators[prompt_idx] = {
453 + 'weighted_sum': None,
454 + 'total_weight': 0,
455 + 'chunk_count': 0,
456 + 'request_id': result.request_id.split("-chunk-")[0]
457 + }
458 +
459 + aggregator = prompt_aggregators[prompt_idx]
460 +
461 + # MEAN pooling with online weighted averaging
462 + # Ensure result is PoolingRequestOutput
463 + # for embedding processing
464 + if not isinstance(result, PoolingRequestOutput):
465 + return self.create_error_response(
466 + f"Expected PoolingRequestOutput for "
467 + f"chunked embedding, got "
468 + f"{type(result).__name__}")
469 +
470 + # Handle both PoolingOutput and
471 + # EmbeddingOutput types
472 + if hasattr(result.outputs, 'data'):
473 + # PoolingOutput case
474 + embedding_data = result.outputs.data
475 + elif hasattr(result.outputs, 'embedding'):
476 + # EmbeddingOutput case -
477 + # convert embedding list to tensor
478 + embedding_data = result.outputs.embedding
479 + else:
480 + return self.create_error_response(
481 + f"Unsupported output type: "
482 + f"{type(result.outputs).__name__}")
483 +
484 + if not isinstance(embedding_data, torch.Tensor):
485 + embedding_data = torch.tensor(embedding_data,
486 + dtype=torch.float32)
487 +
488 + if result.prompt_token_ids is None:
489 + return self.create_error_response(
490 + "prompt_token_ids cannot be None for "
491 + "chunked processing")
492 + weight = len(result.prompt_token_ids)
493 +
494 + weighted_embedding = embedding_data.to(
495 + dtype=torch.float32) * weight
496 +
497 + if aggregator['weighted_sum'] is None:
498 + # First chunk
499 + aggregator['weighted_sum'] = weighted_embedding
500 + else:
501 + # Accumulate
502 + aggregator['weighted_sum'] += weighted_embedding
503 +
504 + aggregator['total_weight'] += weight
505 + aggregator['chunk_count'] += 1
506 + else:
507 + # Non-chunked result - extract prompt_idx from request_id
508 + parts = result.request_id.split("-")
509 + try:
510 + # Last part should be prompt index
511 + prompt_idx = int(parts[-1])
512 + except (ValueError, IndexError):
513 + prompt_idx = result_idx # Fallback to result_idx
514 +
515 + short_prompts_results[prompt_idx] = cast(
516 + PoolingRequestOutput, result)
517 +
518 + # Finalize aggregated results
519 + final_res_batch: list[Union[PoolingRequestOutput,
520 + EmbeddingRequestOutput]] = []
521 + num_prompts = len(ctx.engine_prompts)
522 +
523 + for prompt_idx in range(num_prompts):
524 + if prompt_idx in prompt_aggregators:
525 + # Finalize MEAN aggregation for this chunked prompt
526 + aggregator = prompt_aggregators[prompt_idx]
527 +
528 + weighted_sum = aggregator['weighted_sum']
529 + total_weight = aggregator['total_weight']
530 +
531 + if (weighted_sum is not None
532 + and isinstance(weighted_sum, torch.Tensor)
533 + and isinstance(total_weight,
534 + (int, float)) and total_weight > 0):
535 +
536 + # Compute final mean embedding
537 + final_embedding = weighted_sum / total_weight
538 +
539 + # Create a PoolingRequestOutput
540 + # for the aggregated result
541 + pooling_output_data = PoolingOutput(
542 + data=final_embedding)
543 +
544 + # Get original prompt token IDs for this prompt
545 + original_prompt = ctx.request_prompts[prompt_idx]
546 + if not self._is_text_tokens_prompt(original_prompt):
547 + return self.create_error_response(
548 + f"Chunked prompt {prompt_idx} is not a "
549 + f"TextTokensPrompt")
550 +
551 + original_token_ids = cast(
552 + TextTokensPrompt,
553 + original_prompt)["prompt_token_ids"]
554 +
555 + pooling_request_output = PoolingRequestOutput(
556 + request_id=aggregator['request_id'],
557 + prompt_token_ids=original_token_ids,
558 + outputs=pooling_output_data,
559 + finished=True)
560 +
561 + final_res_batch.append(pooling_request_output)
562 + else:
563 + return self.create_error_response(
564 + f"Failed to aggregate chunks "
565 + f"for prompt {prompt_idx}")
566 + elif prompt_idx in short_prompts_results:
567 + final_res_batch.append(
568 + cast(PoolingRequestOutput,
569 + short_prompts_results[prompt_idx]))
570 + else:
571 + return self.create_error_response(
572 + f"Result not found for prompt {prompt_idx}")
573 +
574 + ctx.final_res_batch = cast(
575 + list[Union[RequestOutput, PoolingRequestOutput]],
576 + final_res_batch)
577 +
578 + return None
579 +
580 + except Exception as e:
581 + return self.create_error_response(str(e))
582 +
583
584 class OpenAIServingEmbedding(EmbeddingMixin):
585 request_id_prefix = "embd"