[Pooling] Fix normalization of chunked long-text embeddings (#57498)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Sep 21, 2026 at 09:46 UTC
887cf91e30545373f90eb967c981b456c9cbbf8d
1 file changed
+22
-1
vllm/entrypoints/pooling/embed/io_processor.py
+22
-1
@@ -57,6 +57,8 @@ logger = init_logger(__name__)
57
class _ChunkedPromptAggregator:
58
weighted_sum: torch.Tensor | None = None
59
total_weight: int = 0
60
+ num_chunks: int = 0
61
+ all_chunks_l2_normalized: bool = True
62
63
64
class EmbedIOProcessor(PoolingIOProcessor):
@@ -204,7 +206,17 @@ class EmbedIOProcessor(PoolingIOProcessor):
206
207
weight = len(result.prompt_token_ids)
208
embedding_data = result.outputs.data
207
- weighted_embedding = embedding_data.to(dtype=torch.float32) * weight
209
+ float_embedding = embedding_data.to(dtype=torch.float32)
210
+ if ctx.pooling_params.use_activation:
211
+ # Custom activations need not L2-normalize; allow BF16 rounding.
212
+ embedding_norm = torch.linalg.vector_norm(float_embedding)
213
+ aggregator.all_chunks_l2_normalized &= torch.allclose(
214
+ embedding_norm,
215
+ torch.ones_like(embedding_norm),
216
+ rtol=1e-2,
217
+ atol=1e-5,
218
+ )
219
+ weighted_embedding = float_embedding * weight
220
221
if aggregator.weighted_sum is None:
222
# First chunk
@@ -214,6 +226,7 @@ class EmbedIOProcessor(PoolingIOProcessor):
226
aggregator.weighted_sum += weighted_embedding
227
228
aggregator.total_weight += weight
229
+ aggregator.num_chunks += 1
230
231
if ctx.original_engine_inputs is None:
232
raise ValueError("Original engine inputs not available")
@@ -239,6 +252,14 @@ class EmbedIOProcessor(PoolingIOProcessor):
252
):
253
# Compute final mean embedding
254
final_embedding = weighted_sum / total_weight
255
+ if (
256
+ ctx.pooling_params.use_activation
257
+ and aggregator.num_chunks > 1
258
+ and aggregator.all_chunks_l2_normalized
259
+ ):
260
+ final_embedding = torch.nn.functional.normalize(
261
+ final_embedding, dim=-1
262
+ )
263
264
# Create a PoolingRequestOutput
265
# for the aggregated result