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 9, 23:00 UTC
[Model Runner V2] Initialized communication buffer for DP (#32624)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 19, 2026 at 17:27 UTC
05dc4bfab6978f458dfeee2388efc83882d7bfaf
1 file changed
+7
vllm/v1/worker/gpu/model_runner.py
+7
@@ -12,6 +12,7 @@ import torch.nn as nn
12
13
from vllm.config import VllmConfig
14
from vllm.config.compilation import CUDAGraphMode
15
+from vllm.distributed.parallel_state import prepare_communication_buffer_for_model
16
from vllm.forward_context import set_forward_context
17
from vllm.logger import init_logger
18
from vllm.model_executor.model_loader import get_model_loader
@@ -206,6 +207,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
207
time_after_load - time_before_load,
208
)
209
210
+ prepare_communication_buffer_for_model(self.model)
211
+ if self.do_spec_decode:
212
+ speculator_model = getattr(self.speculator, "model", None)
213
+ if speculator_model is not None:
214
+ prepare_communication_buffer_for_model(speculator_model)
215
+
216
def get_model(self) -> nn.Module:
217
return self.model
218