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 8, 01:15 UTC
[Model Runner V2] Minor CPU optimizations (#34856)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 19, 2026 at 16:05 UTC
40b2f1c3d9c1dbcec185e8b6911fd273524f5b88
4 files changed
+28
-10
vllm/distributed/device_communicators/shm_broadcast.py
+2
-2
@@ -513,8 +513,8 @@ class MessageQueue:
513
assert self._is_local_reader, "Only readers can acquire read"
514
start_time = time.monotonic()
515
n_warning = 1
516
- while True:
517
- with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
516
+ with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
517
+ while True:
518
# Memory fence ensures we see the latest writes from the writer.
519
# Without this, we may read stale flags from our CPU cache
520
# and spin indefinitely even though writer has updated them.
vllm/v1/worker/gpu/async_utils.py
+16
-3
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import contextlib
4
5
import numpy as np
6
import torch
@@ -14,6 +15,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
15
model_runner_output: ModelRunnerOutput,
16
sampler_output: SamplerOutput,
17
num_sampled_tokens: torch.Tensor,
18
+ main_stream: torch.cuda.Stream,
19
copy_stream: torch.cuda.Stream,
20
copy_event: torch.cuda.Event,
21
):
@@ -25,9 +27,8 @@ class AsyncOutput(AsyncModelRunnerOutput):
27
self.num_sampled_tokens = num_sampled_tokens
28
self.copy_event = copy_event
29
28
- default_stream = torch.cuda.current_stream()
29
- with torch.cuda.stream(copy_stream):
30
- copy_stream.wait_stream(default_stream)
30
+ with stream(copy_stream, main_stream):
31
+ copy_stream.wait_stream(main_stream)
32
33
self.sampled_token_ids = async_copy_to_np(sampler_output.sampled_token_ids)
34
self.logprobs_tensors: LogprobsTensors | None = None
@@ -71,3 +72,15 @@ class AsyncOutput(AsyncModelRunnerOutput):
72
73
def async_copy_to_np(x: torch.Tensor) -> np.ndarray:
74
return x.to("cpu", non_blocking=True).numpy()
75
+
76
+
77
+@contextlib.contextmanager
78
+def stream(to_stream: torch.cuda.Stream, from_stream: torch.cuda.Stream):
79
+ """Lightweight version of torch.cuda.stream() context manager which
80
+ avoids current_stream and device lookups.
81
+ """
82
+ try:
83
+ torch.cuda.set_stream(to_stream)
84
+ yield
85
+ finally:
86
+ torch.cuda.set_stream(from_stream)
vllm/v1/worker/gpu/buffer_utils.py
+3
-5
@@ -22,7 +22,6 @@ def async_copy_to_gpu(
22
if isinstance(x, np.ndarray):
23
x = torch.from_numpy(x)
24
assert x.is_cpu
25
- assert not x.is_pinned()
25
26
if out is None:
27
assert device is not None
@@ -30,6 +29,8 @@ def async_copy_to_gpu(
29
30
# CPU-to-CPU copy
31
tmp = x.pin_memory()
32
+ assert tmp is not x
33
+
34
# CPU-to-GPU copy
35
return out.copy_(tmp, non_blocking=True)
36
@@ -75,11 +76,8 @@ class UvaBufferPool:
76
out: torch.Tensor | None = None,
77
) -> torch.Tensor:
78
uva = self.copy_to_uva(x)
78
- if out is None:
79
- # CPU-to-GPU copy
80
- return uva.clone()
79
# CPU-to-GPU copy
82
- return out.copy_(uva, non_blocking=True)
80
+ return uva.clone() if out is None else out.copy_(uva, non_blocking=True)
81
82
83
class UvaBackedTensor:
vllm/v1/worker/gpu/model_runner.py
+7
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
import gc
5
import time
6
from copy import deepcopy
@@ -239,6 +240,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
240
def get_model(self) -> nn.Module:
241
return self.model
242
243
+ @functools.cached_property
244
+ def main_stream(self) -> torch.cuda.Stream:
245
+ # Cache the default CUDA stream to avoid lookup overhead.
246
+ return torch.cuda.current_stream(self.device)
247
+
248
def get_kv_cache_spec(self):
249
return get_kv_cache_spec(self.vllm_config)
250
@@ -1065,6 +1071,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1071
model_runner_output=model_runner_output,
1072
sampler_output=sampler_output,
1073
num_sampled_tokens=num_sampled,
1074
+ main_stream=self.main_stream,
1075
copy_stream=self.output_copy_stream,
1076
copy_event=self.output_copy_event,
1077
)