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
[MM] Optimize memory profiling for scattered multimodal embeddings (#25810)
Signed-off-by: Roger Wang <hey@rogerw.io>
Roger Wang committed
Sep 27, 2025 at 19:17 UTC
69311446baa7c447708b591156c57d6996dda0a7
1 file changed
+17
vllm/v1/worker/gpu_model_runner.py
+17
@@ -3429,6 +3429,23 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3429
expected_num_items=max_mm_items_per_batch,
3430
)
3431
3432
+ # NOTE: This happens when encoder cache needs to store
3433
+ # the embeddings that encoder outputs are scattered onto.
3434
+ # In this case we create dummy embeddings of size
3435
+ # (encode_budget, hidden_size) and scatter encoder
3436
+ # output into it.
3437
+ encoder_output_shape = dummy_encoder_outputs[0].shape
3438
+ if encoder_output_shape[0] < encoder_budget:
3439
+ expanded_outputs = []
3440
+ for output in dummy_encoder_outputs:
3441
+ expanded = output.new_zeros(
3442
+ (encoder_budget, encoder_output_shape[-1]))
3443
+ num_tokens = output.shape[0]
3444
+ expanded[:num_tokens].copy_(output)
3445
+ expanded_outputs.append(expanded)
3446
+
3447
+ dummy_encoder_outputs = expanded_outputs
3448
+
3449
# Cache the dummy encoder outputs.
3450
self.encoder_cache["tmp"] = dict(
3451
enumerate(dummy_encoder_outputs))