| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from contextlib import AbstractContextManager, contextmanager |
| 4 | from typing import Any |
| 5 | |
| 6 | import torch |
| 7 | import torch.nn as nn |
| 8 | |
| 9 | from vllm.config import ( |
| 10 | CompilationMode, |
| 11 | VllmConfig, |
| 12 | ) |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.model_loader import get_model |
| 15 | from vllm.tracing import instrument |
| 16 | from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig |
| 17 | from vllm.v1.utils import CpuGpuBuffer |
| 18 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner |
| 19 | |
| 20 | logger = init_logger(__name__) |
| 21 | |
| 22 | |
| 23 | class CPUModelRunner(GPUModelRunner): |
| 24 | def __init__(self, vllm_config: VllmConfig, device: torch.device): |
| 25 | # avoid calling accelerator APIs for methods inherited from super class |
| 26 | _set_torch_accelerator_to_noop() |
| 27 | |
| 28 | with _torch_cuda_wrapper(): |
| 29 | super().__init__(vllm_config, device) |
| 30 | |
| 31 | assert device == torch.device("cpu") |
| 32 | # Note: speculative decoding is now supported on CPU with C++ native impls |
| 33 | |
| 34 | self.use_cuda_graph = False |
| 35 | self.cascade_attn_enabled = False |
| 36 | |
| 37 | self._postprocess_tensors() |
| 38 | |
| 39 | def _postprocess_tensors(self) -> None: |
| 40 | # Note: replace device tensors with cpu tensors |
| 41 | def replace_tensor(obj: Any, cpu_attr_name: str, device_attr_name) -> None: |
| 42 | cpu_tensor = getattr(obj, cpu_attr_name, None) |
| 43 | device_tensor = getattr(obj, device_attr_name, None) |
| 44 | if isinstance(cpu_tensor, torch.Tensor) and isinstance( |
| 45 | device_tensor, torch.Tensor |
| 46 | ): |
| 47 | setattr(obj, device_attr_name, cpu_tensor) |
| 48 | |
| 49 | for v in vars(self).values(): |
| 50 | if isinstance(v, CpuGpuBuffer): |
| 51 | v.gpu = v.cpu |
| 52 | |
| 53 | for k, v in vars(self.input_batch).items(): |
| 54 | if k.endswith("_cpu_tensor") and isinstance(v, torch.Tensor): |
| 55 | replace_tensor(self.input_batch, k, k[:-11]) |
| 56 | |
| 57 | for block_table in self.input_batch.block_table.block_tables: |
| 58 | for v in vars(block_table).values(): |
| 59 | if isinstance(v, CpuGpuBuffer): |
| 60 | v.gpu = v.cpu |
| 61 | |
| 62 | @instrument(span_name="Loading (CPU)") |
| 63 | def load_model(self, load_dummy_weights: bool = False) -> None: |
| 64 | if load_dummy_weights: |
| 65 | raise ValueError( |
| 66 | "Loading dummy weights (needed for elastic EP scale-up) " |
| 67 | "Is not supported by the CPU Model Runner." |
| 68 | ) |
| 69 | logger.info("Starting to load model %s...", self.model_config.model) |
| 70 | self.model = get_model(vllm_config=self.vllm_config) |
| 71 | |
| 72 | if self.lora_config: |
| 73 | self.model = self.load_lora_model(self.model, self.vllm_config, self.device) |
| 74 | |
| 75 | if hasattr(self, "drafter"): |
| 76 | logger.info_once("Loading drafter model...") |
| 77 | self.drafter.load_model(self.model) |
| 78 | |
| 79 | self._setup_eagle3_aux_hidden_state_outputs() |
| 80 | |
| 81 | def get_model(self) -> nn.Module: |
| 82 | return self.model |
| 83 | |
| 84 | @instrument(span_name="Warmup (CPU)") |
| 85 | def warming_up_model(self) -> None: |
| 86 | if self.vllm_config.compilation_config.mode == CompilationMode.NONE: |
| 87 | return |
| 88 | logger.info("Warming up model for the compilation...") |
| 89 | # Only generate graph for the generic shape |
| 90 | with _set_global_compilation_settings(self.vllm_config): |
| 91 | self.profile_run() |
| 92 | logger.info("Warming up done.") |
| 93 | |
| 94 | def initialize_kv_cache( |
| 95 | self, |
| 96 | kv_cache_config: KVCacheConfig, |
| 97 | is_profiling: bool = False, |
| 98 | kv_cache_allocation_context: AbstractContextManager | None = None, |
| 99 | ) -> None: |
| 100 | super().initialize_kv_cache( |
| 101 | kv_cache_config, |
| 102 | is_profiling, |
| 103 | kv_cache_allocation_context=kv_cache_allocation_context, |
| 104 | ) |
| 105 | |
| 106 | if self.speculative_config: |
| 107 | if self.speculative_config.use_eagle(): |
| 108 | logger.info("EAGLE drafter KV cache initialized for CPU backend") |
| 109 | elif self.speculative_config.uses_draft_model(): |
| 110 | logger.info("Draft model KV cache initialized for CPU backend") |
| 111 | |
| 112 | def _init_device_properties(self) -> None: |
| 113 | pass |
| 114 | |
| 115 | def _sync_device(self) -> None: |
| 116 | pass |
| 117 | |
| 118 | def _zero_block_ids(self, block_ids: list[int]) -> None: |
| 119 | # Zero full-attention blocks to prevent stale data corruption on partial writes. |
| 120 | # Encoder-only (runner-only) layers are not FullAttentionSpec, so the |
| 121 | # spec filter below already excludes them; no runner-only skip needed. |
| 122 | seen_ptrs: set[int] = set() |
| 123 | for group in self.kv_cache_config.kv_cache_groups: |
| 124 | if not isinstance(group.kv_cache_spec, FullAttentionSpec): |
| 125 | continue |
| 126 | for layer_name in group.layer_names: |
| 127 | ctx = self.compilation_config.static_forward_context.get(layer_name) |
| 128 | if ctx is None: |
| 129 | continue |
| 130 | kv = ctx.kv_cache |
| 131 | if not isinstance(kv, torch.Tensor): |
| 132 | continue |
| 133 | if kv.data_ptr() in seen_ptrs: |
| 134 | continue |
| 135 | seen_ptrs.add(kv.data_ptr()) |
| 136 | for block_id in block_ids: |
| 137 | kv[block_id].zero_() |
| 138 | |
| 139 | def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]: |
| 140 | """CPU-safe version: direct tolist() without CUDA events.""" |
| 141 | return sampled_token_ids.tolist() |
| 142 | |
| 143 | |
| 144 | @contextmanager |
| 145 | def _torch_cuda_wrapper(): |
| 146 | class _EventPlaceholder: |
| 147 | def __init__(self, *args, **kwargs) -> None: |
| 148 | self.record = lambda: None |
| 149 | self.synchronize = lambda: None |
| 150 | |
| 151 | class _StreamPlaceholder: |
| 152 | def __init__(self, *args, **kwargs) -> None: |
| 153 | self.wait_stream = lambda *a, **kw: None |
| 154 | self.device = torch.device("cpu") |
| 155 | |
| 156 | cuda_event = torch.Event |
| 157 | cuda_stream = torch.cuda.Stream |
| 158 | try: |
| 159 | torch.Event = _EventPlaceholder |
| 160 | torch.cuda.Stream = _StreamPlaceholder |
| 161 | yield |
| 162 | finally: |
| 163 | torch.Event = cuda_event |
| 164 | torch.cuda.Stream = cuda_stream |
| 165 | |
| 166 | |
| 167 | @contextmanager |
| 168 | def _set_global_compilation_settings(config: VllmConfig): |
| 169 | import torch._inductor.config as torch_inductor_config |
| 170 | |
| 171 | inductor_config = config.compilation_config.inductor_compile_config |
| 172 | # Note: The MKLDNN and CPPGEMM backend requires freezing parameters. |
| 173 | freezing_value = torch_inductor_config.freezing |
| 174 | try: |
| 175 | if inductor_config.get("max_autotune", False): |
| 176 | torch_inductor_config.freezing = True |
| 177 | yield |
| 178 | finally: |
| 179 | torch_inductor_config.freezing = freezing_value |
| 180 | |
| 181 | |
| 182 | def _set_torch_accelerator_to_noop() -> None: |
| 183 | def noop(*args: Any, **kwargs: Any) -> None: |
| 184 | pass |
| 185 | |
| 186 | # Distinct no-op so empty_cache does not alias synchronize: Dynamo's |
| 187 | # handle_synchronize is keyed on that object and asserts on CPU-only hosts. |
| 188 | def empty_cache_noop(*args: Any, **kwargs: Any) -> None: |
| 189 | pass |
| 190 | |
| 191 | torch.accelerator.synchronize = noop |
| 192 | torch.accelerator.empty_cache = empty_cache_noop |