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
refactor: Turn GPUModelRunner.inputs_embeds to a CpuGpuBuffer (#24345)
Signed-off-by: Andrew Sansom <andrew@protopia.ai>
Andrew Sansom committed
Sep 6, 2025 at 01:01 UTC
305a1cc0d27870568c9c0b2ee6f2124479f0cb3a
2 files changed
+37
-14
vllm/v1/utils.py
+17
-4
@@ -19,6 +19,8 @@ from vllm.utils import (get_open_port, get_open_zmq_ipc_path, get_tcp_uri,
19
kill_process_tree)
20
21
if TYPE_CHECKING:
22
+ import numpy as np
23
+
24
from vllm.v1.engine.coordinator import DPCoordinator
25
from vllm.v1.engine.utils import (CoreEngineActorManager,
26
CoreEngineProcManager)
@@ -97,20 +99,31 @@ class ConstantList(Generic[T], Sequence):
99
100
101
class CpuGpuBuffer:
102
+ """Buffer to easily copy tensors between CPU and GPU."""
103
104
def __init__(
105
self,
103
- *args,
106
+ *size: Union[int, torch.SymInt],
107
dtype: torch.dtype,
108
device: torch.device,
109
pin_memory: bool,
107
- ):
108
- self.cpu = torch.zeros(*args,
110
+ with_numpy: bool = True,
111
+ ) -> None:
112
+ self.cpu = torch.zeros(*size,
113
dtype=dtype,
114
device="cpu",
115
pin_memory=pin_memory)
112
- self.np = self.cpu.numpy()
116
self.gpu = self.cpu.to(device)
117
+ self.np: np.ndarray
118
+ # To keep type hints simple (avoiding generics and subclasses), we
119
+ # only conditionally create the numpy array attribute. This can cause
120
+ # AttributeError if `self.np` is accessed when `with_numpy=False`.
121
+ if with_numpy:
122
+ if dtype == torch.bfloat16:
123
+ raise ValueError(
124
+ "Bfloat16 torch tensors cannot be directly cast to a "
125
+ "numpy array, so call CpuGpuBuffer with with_numpy=False")
126
+ self.np = self.cpu.numpy()
127
128
def copy_to_gpu(self, n: Optional[int] = None) -> torch.Tensor:
129
if n is None:
vllm/v1/worker/gpu_model_runner.py
+20
-10
@@ -303,10 +303,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
303
self.query_start_loc = self._make_buffer(self.max_num_reqs + 1,
304
dtype=torch.int32)
305
self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
306
- self.inputs_embeds = torch.zeros(
307
- (self.max_num_tokens, self.hidden_size),
308
- dtype=self.dtype,
309
- device=self.device)
306
+ # Because inputs_embeds may be bfloat16 and we don't need a numpy
307
+ # version of this tensor, avoid a RuntimeError by not creating a
308
+ # numpy buffer.
309
+ self.inputs_embeds = self._make_buffer(self.max_num_tokens,
310
+ self.hidden_size,
311
+ dtype=self.dtype,
312
+ numpy=False)
313
314
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
315
if self.uses_mrope:
@@ -374,11 +377,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
377
device="cpu",
378
pin_memory=self.pin_memory)
379
377
- def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
378
- return CpuGpuBuffer(*args,
380
+ def _make_buffer(self,
381
+ *size: Union[int, torch.SymInt],
382
+ dtype: torch.dtype,
383
+ numpy: bool = True) -> CpuGpuBuffer:
384
+ # Bfloat16 torch tensors cannot be directly cast to a numpy array, so
385
+ # if a bfloat16 buffer is needed without a corresponding numpy array,
386
+ # don't bother instantiating the numpy array.
387
+ return CpuGpuBuffer(*size,
388
dtype=dtype,
389
device=self.device,
381
- pin_memory=self.pin_memory)
390
+ pin_memory=self.pin_memory,
391
+ with_numpy=numpy)
392
393
def _init_model_kwargs(self, num_tokens: int):
394
model_kwargs = dict[str, Any]()
@@ -1645,11 +1655,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1655
)
1656
1657
# TODO(woosuk): Avoid the copy. Optimize.
1648
- self.inputs_embeds[:num_scheduled_tokens].copy_(
1658
+ self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(
1659
inputs_embeds_scheduled)
1660
1661
input_ids = None
1652
- inputs_embeds = self.inputs_embeds[:num_input_tokens]
1662
+ inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
1663
model_kwargs = {
1664
**self._init_model_kwargs(num_scheduled_tokens),
1665
**self._extract_mm_kwargs(scheduler_output),
@@ -2484,7 +2494,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2494
num_scheduled_tokens, remove_lora):
2495
if self.supports_mm_inputs:
2496
input_ids = None
2487
- inputs_embeds = self.inputs_embeds[:num_tokens]
2497
+ inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
2498
model_kwargs = {
2499
**self._init_model_kwargs(num_tokens),
2500
**self._dummy_mm_kwargs(num_reqs),