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
[Misc] Move CpuGpuBuffer to vllm/v1/utils.py (#23728)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 27, 2025 at 03:25 UTC
04ff1e43fb6e2e675170d0c90399290f8925abb7
4 files changed
+33
-33
vllm/v1/utils.py
+29
@@ -96,6 +96,35 @@ class ConstantList(Generic[T], Sequence):
96
return f"ConstantList({self._x})"
97
98
99
+class CpuGpuBuffer:
100
+
101
+ def __init__(
102
+ self,
103
+ *args,
104
+ dtype: torch.dtype,
105
+ device: torch.device,
106
+ pin_memory: bool,
107
+ ):
108
+ self.cpu = torch.zeros(*args,
109
+ dtype=dtype,
110
+ device="cpu",
111
+ pin_memory=pin_memory)
112
+ self.np = self.cpu.numpy()
113
+ self.gpu = self.cpu.to(device)
114
+
115
+ def copy_to_gpu(self, n: Optional[int] = None) -> torch.Tensor:
116
+ if n is None:
117
+ return self.gpu.copy_(self.cpu, non_blocking=True)
118
+ return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)
119
+
120
+ def copy_to_cpu(self, n: Optional[int] = None) -> torch.Tensor:
121
+ """NOTE: Because this method is non-blocking, explicit synchronization
122
+ is needed to ensure the data is copied to CPU."""
123
+ if n is None:
124
+ return self.cpu.copy_(self.gpu, non_blocking=True)
125
+ return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True)
126
+
127
+
128
def get_engine_client_zmq_addr(local_only: bool,
129
host: str,
130
port: int = 0) -> str:
vllm/v1/worker/cpu_model_runner.py
+1
-1
@@ -10,8 +10,8 @@ from vllm.config import VllmConfig
10
from vllm.logger import init_logger
11
from vllm.model_executor.model_loader import get_model
12
from vllm.v1.attention.backends.cpu_attn import TorchSDPAMetadataBuilderV1
13
+from vllm.v1.utils import CpuGpuBuffer
14
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
14
-from vllm.v1.worker.utils import CpuGpuBuffer
15
16
if TYPE_CHECKING:
17
from vllm.v1.core.sched.output import SchedulerOutput
vllm/v1/worker/gpu_model_runner.py
+3
-3
@@ -78,14 +78,14 @@ from vllm.v1.spec_decode.eagle import EagleProposer
78
from vllm.v1.spec_decode.medusa import MedusaProposer
79
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
80
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
81
+from vllm.v1.utils import CpuGpuBuffer
82
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
83
from vllm.v1.worker.kv_connector_model_runner_mixin import (
84
KVConnectorModelRunnerMixin, KVConnectorOutput)
85
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
86
86
-from .utils import (AttentionGroup, CpuGpuBuffer, MultiModalBudget,
87
- bind_kv_cache, gather_mm_placeholders,
88
- initialize_kv_cache_for_kv_sharing,
87
+from .utils import (AttentionGroup, MultiModalBudget, bind_kv_cache,
88
+ gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,
89
sanity_check_mm_encoder_outputs, scatter_mm_placeholders)
90
91
if TYPE_CHECKING:
vllm/v1/worker/utils.py
-29
@@ -303,32 +303,3 @@ def bind_kv_cache(
303
for layer_name, kv_cache in kv_caches.items():
304
# NOTE: Use list because of v0 PP virtual engine.
305
forward_context[layer_name].kv_cache = [kv_cache]
306
-
307
-
308
-class CpuGpuBuffer:
309
-
310
- def __init__(
311
- self,
312
- *args,
313
- dtype: torch.dtype,
314
- device: torch.device,
315
- pin_memory: bool,
316
- ):
317
- self.cpu = torch.zeros(*args,
318
- dtype=dtype,
319
- device="cpu",
320
- pin_memory=pin_memory)
321
- self.np = self.cpu.numpy()
322
- self.gpu = self.cpu.to(device)
323
-
324
- def copy_to_gpu(self, n: Optional[int] = None) -> None:
325
- if n is None:
326
- return self.gpu.copy_(self.cpu, non_blocking=True)
327
- return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)
328
-
329
- def copy_to_cpu(self, n: Optional[int] = None) -> None:
330
- """NOTE: Because this method is non-blocking, explicit synchronization
331
- is needed to ensure the data is copied to CPU."""
332
- if n is None:
333
- return self.cpu.copy_(self.gpu, non_blocking=True)
334
- return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True)