[CPU] Remove speculative decoding stream overrides from CPUModelRunner (#47162)
Signed-off-by: jmamou <jonathan.mamou@intel.com>
Jonathan Mamou committed
Jul 1, 2026 at 09:12 UTC
df802a87b77f21e0bf5f8d358633d5dd00ac5400
2 files changed
+3
-63
vllm/v1/worker/cpu/shm.py
+1
@@ -30,6 +30,7 @@ class _EventPlaceholder:
30
class _StreamPlaceholder:
31
def __init__(self, *args, **kwargs) -> None:
32
self.wait_stream = noop
33
+ self.device = torch.device("cpu")
34
35
def __enter__(self, *args, **kwargs):
36
return self
vllm/v1/worker/cpu_model_runner.py
+2
-63
@@ -12,7 +12,6 @@ from vllm.config import VllmConfig
12
from vllm.logger import init_logger
13
from vllm.model_executor.model_loader import get_model
14
from vllm.tracing import instrument
15
-from vllm.v1.core.sched.output import SchedulerOutput
15
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
16
from vllm.v1.utils import CpuGpuBuffer
17
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
@@ -189,67 +188,6 @@ class CPUModelRunner(GPUModelRunner):
188
for block_id in block_ids:
189
kv[block_id].zero_()
190
192
- # =========================================================================
193
- # CPU-safe overrides for speculative decoding methods
194
- # These methods override GPU-specific implementations that use CUDA streams
195
- # =========================================================================
196
-
197
- def _copy_draft_token_ids_to_cpu(
198
- self, scheduler_output: "SchedulerOutput", zeros_only: bool = False
199
- ) -> None:
200
- """CPU-safe version: no async copy needed, tensors already on CPU."""
201
- if self.use_async_scheduling and not (
202
- scheduler_output.has_structured_output_requests
203
- or self.input_batch.sampling_metadata.output_token_ids
204
- ):
205
- return
206
- self._draft_token_req_ids = self.input_batch.req_ids.copy()
207
-
208
- draft_token_ids: torch.Tensor = self._draft_token_ids
209
- if not torch.is_tensor(draft_token_ids):
210
- return
211
-
212
- num_reqs = draft_token_ids.shape[0]
213
- if self.draft_token_ids_cpu is not None:
214
- if not zeros_only:
215
- self.draft_token_ids_cpu[:num_reqs].copy_(draft_token_ids)
216
- else:
217
- self.draft_token_ids_cpu[:num_reqs] = 0
218
-
219
- def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]:
220
- """CPU-safe version: no event synchronization needed."""
221
- if isinstance(self._draft_token_ids, list):
222
- return self._draft_token_ids, self.input_batch.req_ids
223
- req_ids = self._draft_token_req_ids
224
- if req_ids is None:
225
- return [], []
226
- if self.draft_token_ids_cpu is not None:
227
- return self.draft_token_ids_cpu[: len(req_ids)].tolist(), req_ids
228
- return [], []
229
-
230
- def _copy_valid_sampled_token_count(
231
- self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor
232
- ) -> None:
233
- """CPU-safe version: direct copy without CUDA streams."""
234
- if self.valid_sampled_token_count_cpu is None:
235
- return
236
-
237
- counts = valid_sampled_tokens_count
238
- counts_cpu = self.valid_sampled_token_count_cpu
239
- counts_cpu[: counts.shape[0]].copy_(counts)
240
- self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1)
241
-
242
- def _get_valid_sampled_token_count(self) -> list[int]:
243
- """CPU-safe version: no event synchronization needed."""
244
- prev_sampled_token_ids = self.input_batch.prev_sampled_token_ids
245
- if prev_sampled_token_ids is None:
246
- return []
247
-
248
- counts_cpu = self.valid_sampled_token_count_cpu
249
- if counts_cpu is None:
250
- return []
251
- return counts_cpu[: prev_sampled_token_ids.shape[0]].tolist()
252
-
191
def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
192
"""CPU-safe version: direct tolist() without CUDA events."""
193
return sampled_token_ids.tolist()
@@ -264,7 +202,8 @@ def _torch_cuda_wrapper():
202
203
class _StreamPlaceholder:
204
def __init__(self, *args, **kwargs) -> None:
267
- pass
205
+ self.wait_stream = lambda *a, **kw: None
206
+ self.device = torch.device("cpu")
207
208
cuda_event = torch.Event
209
cuda_stream = torch.cuda.Stream