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
[Bugfix] Fix cuda event usage with CPU model runner (#23643)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Aug 27, 2025 at 01:10 UTC
9b0187003e62bdb7311b23b5b5026ea8e4e207d3
2 files changed
+26
-4
vllm/v1/worker/cpu_model_runner.py
+25
-3
@@ -11,6 +11,7 @@ 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.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
@@ -21,7 +22,8 @@ logger = init_logger(__name__)
22
class CPUModelRunner(GPUModelRunner):
23
24
def __init__(self, vllm_config: VllmConfig, device: torch.device):
24
- super().__init__(vllm_config, device)
25
+ with _torch_cuda_wrapper():
26
+ super().__init__(vllm_config, device)
27
28
assert device == torch.device("cpu")
29
assert self.speculative_config is None, "spec decode is not supported."
@@ -71,8 +73,8 @@ class CPUModelRunner(GPUModelRunner):
73
setattr(obj, device_attr_name, cpu_tensor)
74
75
for k, v in vars(self).items():
74
- if k.endswith("_cpu") and isinstance(v, torch.Tensor):
75
- replace_tensor(self, k, k[:-4])
76
+ if isinstance(v, CpuGpuBuffer):
77
+ v.gpu = v.cpu
78
79
for k, v in vars(self.input_batch).items():
80
if k.endswith("_cpu_tensor") and isinstance(v, torch.Tensor):
@@ -108,6 +110,26 @@ class CPUModelRunner(GPUModelRunner):
110
def _sync_device(self) -> None:
111
pass
112
113
+ def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
114
+ return sampled_token_ids.tolist()
115
+
116
+
117
+@contextmanager
118
+def _torch_cuda_wrapper():
119
+
120
+ class _EventPlaceholder:
121
+
122
+ def __init__(self, *args, **kwargs) -> None:
123
+ self.record = lambda: None
124
+ self.synchronize = lambda: None
125
+
126
+ try:
127
+ cuda_event = torch.cuda.Event
128
+ torch.cuda.Event = _EventPlaceholder
129
+ yield
130
+ finally:
131
+ torch.cuda.Event = cuda_event
132
+
133
134
@contextmanager
135
def _set_global_compilation_settings(config: VllmConfig):
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -321,7 +321,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
321
(self.max_model_len, 1),
322
dtype=torch.int64,
323
device="cpu",
324
- pin_memory=True)
324
+ pin_memory=self.pin_memory)
325
326
def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
327
return CpuGpuBuffer(*args,