"""An OpenVINO worker class.""" from typing import Any, Dict, List, Optional, Tuple import openvino as ov import torch import torch.distributed import vllm.envs as envs from vllm.attention import get_attn_backend from vllm.config import (CacheConfig, DeviceConfig, ModelConfig, ParallelConfig, VllmConfig) from vllm.distributed import (broadcast_tensor_dict, ensure_model_parallel_initialized, init_distributed_environment) from vllm.inputs import INPUT_REGISTRY from vllm.logger import init_logger from vllm.model_executor import set_random_seed from vllm.model_executor.layers.sampler import SamplerOutput from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams from vllm.sequence import ExecuteModelRequest, SequenceGroupMetadata from vllm.worker.openvino_model_runner import OpenVINOModelRunner from vllm.worker.worker_base import LoraNotSupportedWorkerBase, WorkerBase logger = init_logger(__name__) class OpenVINOCacheEngine: """Manages the KV cache for OpenVINO backend. This class is responsible for initializing and managing CPU KV caches. It also provides methods for performing KV cache operations, such as copying. """ def __init__( self, cache_config: CacheConfig, model_config: ModelConfig, parallel_config: ParallelConfig, device_config: DeviceConfig, ov_core: ov.Core, ov_device: str, ) -> None: assert device_config.device_type == "openvino" self.cache_config = cache_config self.model_config = model_config self.parallel_config = parallel_config self.head_size = model_config.get_head_size() if device_config.device.type == "cpu" and \ cache_config.cache_dtype == ov.Type.u8: # Scale, zero point and quantized data will be stored together. # The layout for per token per head: # |scale(f32)|zeropoint(f32)|quantized data(u8,idx_1)|quantized data(u8,idx_2)|...|quantized data(u8,idx_head_size)| # noqa: E501 # so, we have to extend head_size by 8, which is sizeof(float) # for scale and sizeof(float) for zeropoint self.head_size += 8 self.num_layers = model_config.get_num_layers(parallel_config) self.num_kv_heads = model_config.get_num_kv_heads(parallel_config) self.block_size = cache_config.block_size # Note: In CacheConfig, num_gpu_blocks actual is num_cpu_blocks # for OpenVINO backend with a CPU target device, because we want # to reuse KV cache management in the scheduler. self.num_device_blocks = cache_config.num_gpu_blocks self.num_swap_blocks = cache_config.num_cpu_blocks # Get attention backend. self.attn_backend = get_attn_backend( self.head_size, self.model_config.dtype, self.cache_config.cache_dtype, self.block_size, self.model_config.is_attention_free, ) # Initialize the cache. self.kv_cache: List[Tuple[ov.Tensor, ov.Tensor]] = self._allocate_kv_cache( self.num_device_blocks, ov_core, ov_device) # Initialize the swap. self.swap_cache: List[Tuple[ov.Tensor, ov.Tensor]] = self._allocate_swap_cache( self.num_swap_blocks, ov_device) def _allocate_kv_cache( self, num_blocks: int, ov_core: ov.Core, ov_device: str, ) -> List[Tuple[ov.Tensor, ov.Tensor]]: """Allocates KV cache.""" k_block_shape = v_block_shape = self.attn_backend.get_kv_cache_shape( num_blocks, self.block_size, self.num_kv_heads, self.head_size)[1:] kv_cache: List[Tuple[ov.Tensor, ov.Tensor]] = [] if current_platform.is_openvino_cpu(): for _ in range(self.num_layers): key_blocks = ov.Tensor(self.cache_config.cache_dtype, k_block_shape) value_blocks = ov.Tensor(self.cache_config.cache_dtype, v_block_shape) kv_cache.append((key_blocks, value_blocks)) else: # Update key_cache shape: k_block_shape = (v_block_shape[0], v_block_shape[1], v_block_shape[3], v_block_shape[2]) remote_context = ov_core.get_default_context(ov_device) for _ in range(self.num_layers): key_blocks = \ remote_context.create_tensor(self.cache_config.cache_dtype, ov.Shape(k_block_shape), {}) value_blocks = \ remote_context.create_tensor(self.cache_config.cache_dtype, ov.Shape(v_block_shape), {}) kv_cache.append((key_blocks, value_blocks)) return kv_cache def _allocate_swap_cache( self, num_blocks: int, ov_device: str, ) -> List[Tuple[ov.Tensor, ov.Tensor]]: """Allocates swap cache.""" k_block_shape = v_block_shape = self.attn_backend.get_kv_cache_shape( num_blocks, self.block_size, self.num_kv_heads, self.head_size)[1:] swap_cache: List[Tuple[ov.Tensor, ov.Tensor]] = [] if num_blocks == 0: return swap_cache assert not current_platform.is_openvino_cpu(), \ "CPU device isn't supposed to have swap cache" # Update key_cache shape: k_block_shape = (v_block_shape[0], v_block_shape[1], v_block_shape[3], v_block_shape[2]) for _ in range(self.num_layers): key_blocks = ov.Tensor(self.cache_config.cache_dtype, k_block_shape) value_blocks = ov.Tensor(self.cache_config.cache_dtype, v_block_shape) swap_cache.append((key_blocks, value_blocks)) return swap_cache def swap_in(self, src_to_dst: List[Tuple[int, int]]) -> None: for i in range(self.num_layers): for swap_tensor, kv_tensor in zip(self.swap_cache[i], self.kv_cache[i]): self.attn_backend.swap_blocks(swap_tensor, kv_tensor, src_to_dst) def swap_out(self, src_to_dst: List[Tuple[int, int]]) -> None: for i in range(self.num_layers): for swap_tensor, kv_tensor in zip(self.swap_cache[i], self.kv_cache[i]): self.attn_backend.swap_blocks(kv_tensor, swap_tensor, src_to_dst) def copy(self, src_to_dsts: List[Tuple[int, int]]) -> None: if (len(src_to_dsts) > 0): self.attn_backend.copy_blocks(self.kv_cache, src_to_dsts) @staticmethod def get_cache_block_size( block_size: int, cache_dtype: ov.Type, model_config: ModelConfig, parallel_config: ParallelConfig, ) -> int: head_size = model_config.get_head_size() num_kv_heads = model_config.get_num_kv_heads(parallel_config) num_layers = model_config.get_num_layers(parallel_config) if cache_dtype == ov.Type.u8: # Scale, zero point and quantized data will be stored together. # The layout for per token per head: # |scale(f32)|zeropoint(f32)|quantized data(u8,idx_1)|quantized data(u8,idx_2)|...|quantized data(u8,idx_head_size)| # noqa: E501 # so, we have to extend head_size by 8, which is sizeof(float) # for scale and sizeof(float) for zeropoint head_size += 8 key_cache_block = block_size * num_kv_heads * head_size value_cache_block = key_cache_block total = num_layers * (key_cache_block + value_cache_block) dtype_size = cache_dtype.size return dtype_size * total class OpenVINOWorker(LoraNotSupportedWorkerBase): """A worker class that executes the model on OpenVINO backend. Each worker is associated with a single OpenVINO device. The worker is responsible for maintaining the KV cache and executing the model on the OpenVINO backend. """ def __init__( self, ov_core: ov.Core, vllm_config: VllmConfig, local_rank: int, rank: int, distributed_init_method: str, kv_cache_dtype: Optional[ov.Type] = ov.Type.undefined, is_driver_worker: bool = False, ) -> None: self.ov_core = ov_core WorkerBase.__init__(self, vllm_config) self.parallel_config.rank = rank self.local_rank = local_rank self.rank = rank self.distributed_init_method = distributed_init_method self.is_driver_worker = is_driver_worker if self.is_driver_worker: assert self.rank == 0, "The driver worker must have rank 0." if self.model_config.trust_remote_code: # note: lazy import to avoid importing torch before initializing from vllm.utils import init_cached_hf_modules init_cached_hf_modules() self.model_runner = OpenVINOModelRunner( self.ov_core, vllm_config=self.vllm_config, kv_cache_dtype=kv_cache_dtype, is_driver_worker=is_driver_worker, ) # Uninitialized cache engine. Will be initialized by # initialize_cache. self.cache_engine: OpenVINOCacheEngine self.kv_cache: List[Tuple[ov.Tensor, ov.Tensor]] def init_device(self) -> None: self.init_distributed_environment() # Set random seed. set_random_seed(self.model_config.seed) def load_model(self): self.model_runner.load_model() def determine_num_available_blocks(self) -> Tuple[int, int]: """Determine the number of blocks available for the KV cache. This determines how many KV blocks can fit into the configured KV cache space. """ # For OpenVINO backend, in case of CPU device, the block number will be # calculated based on the openvino_kvcache_space_bytes. cache_block_size = self.get_cache_block_size_bytes() kvcache_space_bytes = self.cache_config.openvino_kvcache_space_bytes if current_platform.is_openvino_cpu(): num_device_blocks = int(kvcache_space_bytes // cache_block_size) num_swap_blocks = 0 else: if kvcache_space_bytes > 0: logger.info("KV_CACHE size was explicitly configured via " "VLLM_OPENVINO_KVCACHE_SPACE environment " "variable, ignoring profiling run.") kv_cache_size = kvcache_space_bytes else: try: kv_cache_size = self.profile_run() except Exception as err: raise RuntimeError( "The error occurred during profile run. This might be " "due to insufficient GPU memory. Consider decreasing " "`max_model_len` to limit the maximum simultaneously " "processed tokens.") from err num_device_blocks = int(kv_cache_size // cache_block_size) num_swap_blocks = int(self.cache_config.swap_space_bytes // cache_block_size) return num_device_blocks, num_swap_blocks def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: """Initialize the KV cache. Swappable CPU memory is only supported on GPU. For CPU, we use the num_gpu_blocks to determine how many non-swappable CPU blocks to allocate. """ num_device_blocks = num_gpu_blocks num_swap_blocks = num_cpu_blocks if current_platform.is_openvino_cpu(): assert (num_swap_blocks == 0 ), f"{type(self)} does not support swappable cache for CPU" self._validate_num_blocks(num_device_blocks) self.cache_config.num_gpu_blocks = num_device_blocks self.cache_config.num_cpu_blocks = num_swap_blocks # Initialize the cache. self._init_cache_engine() def _validate_num_blocks(self, num_blocks: int) -> None: """Raise errors if the num_blocks is invalid.""" if num_blocks <= 0: raise ValueError( "No available memory for the cache blocks. " "Try increasing `VLLM_OPENVINO_KVCACHE_SPACE` when " "initializing the engine.") max_seq_len = self.cache_config.block_size * num_blocks if self.model_config.max_model_len > max_seq_len: raise ValueError( f"The model's max seq len ({self.model_config.max_model_len}) " "is larger than the maximum number of tokens that can be " f"stored in KV cache ({max_seq_len}). Try increasing " "`VLLM_OPENVINO_KVCACHE_SPACE` or decreasing `max_model_len` " "when initializing the engine.") def _init_cache_engine(self) -> None: ov_device = envs.VLLM_OPENVINO_DEVICE self.cache_engine = OpenVINOCacheEngine( self.cache_config, self.model_config, self.parallel_config, self.device_config, self.ov_core, ov_device, ) self.kv_cache = self.cache_engine.kv_cache self.model_runner.block_size = self.cache_engine.block_size assert self.kv_cache is not None # Populate the cache to warmup the memory if current_platform.is_openvino_cpu(): for key_cache, value_cache in self.kv_cache: key_cache.data[:] = 0 value_cache.data[:] = 0 def cache_swap_in(self, src_to_dst: List[Tuple[int, int]]) -> None: self.cache_engine.swap_in(src_to_dst) def cache_swap_out(self, src_to_dst: List[Tuple[int, int]]) -> None: self.cache_engine.swap_out(src_to_dst) def cache_copy( self, blocks_to_copy: List[Tuple[int, int]], ) -> None: self.cache_engine.copy(blocks_to_copy) # type: ignore @torch.inference_mode() def execute_model( self, execute_model_req: Optional[ExecuteModelRequest] = None, ) -> List[SamplerOutput]: if execute_model_req is None: seq_group_metadata_list = None else: seq_group_metadata_list = execute_model_req.seq_group_metadata_list if self.is_driver_worker: assert seq_group_metadata_list is not None num_seq_groups: int = len(seq_group_metadata_list) assert execute_model_req is not None blocks_to_copy = execute_model_req.blocks_to_copy blocks_to_swap_in = execute_model_req.blocks_to_swap_in blocks_to_swap_out = execute_model_req.blocks_to_swap_out data: Dict[str, Any] = { "num_seq_groups": num_seq_groups, "blocks_to_copy": execute_model_req.blocks_to_copy, "blocks_to_swap_in": execute_model_req.blocks_to_swap_in, "blocks_to_swap_out": execute_model_req.blocks_to_swap_out, } broadcast_tensor_dict(data, src=0) else: data = broadcast_tensor_dict(src=0) num_seq_groups = data["num_seq_groups"] blocks_to_copy = data["blocks_to_copy"] blocks_to_swap_in = data["blocks_to_swap_in"] blocks_to_swap_out = data["blocks_to_swap_out"] if current_platform.is_openvino_cpu(): assert len(execute_model_req.blocks_to_swap_in) == 0 assert len(execute_model_req.blocks_to_swap_out) == 0 else: self.cache_swap_in(blocks_to_swap_in) self.cache_swap_out(blocks_to_swap_out) self.cache_copy(blocks_to_copy) # If there is no input, we don't need to execute the model. if num_seq_groups == 0: return [] output = self.model_runner.execute_model(seq_group_metadata_list, self.kv_cache) # OpenVINO worker only supports single-step execution. return [output] def init_distributed_environment(self) -> None: """Initialize the distributed environment.""" parallel_config = self.parallel_config rank = self.rank distributed_init_method = self.distributed_init_method init_distributed_environment( world_size=parallel_config.world_size, rank=rank, distributed_init_method=distributed_init_method, backend="gloo", ) # A small all_reduce for warmup. torch.distributed.all_reduce(torch.zeros(1).cpu()) ensure_model_parallel_initialized( parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size, ) def get_cache_block_size_bytes(self) -> int: """Return the size in bytes of a single KV cache block.""" return OpenVINOCacheEngine.get_cache_block_size( self.cache_config.block_size, self.cache_config.cache_dtype, self.model_config, self.parallel_config, ) def profile_run(self) -> int: ov_device = envs.VLLM_OPENVINO_DEVICE assert not current_platform.is_openvino_cpu(), \ "CPU device isn't supposed to use profile run." import openvino.properties.device as device import openvino.properties.intel_gpu as intel_gpu ov_core = self.ov_core cache_config = self.cache_config model_config = self.model_config parallel_config = self.parallel_config device_config = self.device_config input_registry = INPUT_REGISTRY mm_registry = MULTIMODAL_REGISTRY mm_registry.init_mm_limits_per_prompt(model_config) # Execute a forward pass with dummy inputs to profile the memory usage # of the model. def model_profile_run(): top_k = model_config.get_vocab_size() - 1 sampling_params = SamplingParams(top_p=0.99, top_k=top_k) max_num_batched_tokens = \ self.scheduler_config.max_num_batched_tokens max_num_seqs = self.scheduler_config.max_num_seqs tmp_cache_config = CacheConfig(cache_config.block_size, cache_config.gpu_memory_utilization, cache_config.swap_space_bytes, "auto") tmp_cache_config.num_gpu_blocks = 1 tmp_cache_config.num_cpu_blocks = 0 tmp_cache_config.cache_dtype = cache_config.cache_dtype profiling_cache_engine = OpenVINOCacheEngine( tmp_cache_config, model_config, parallel_config, device_config, ov_core, ov_device) # Profile memory usage with max_num_sequences sequences and the # total # number of tokens equal to max_num_batched_tokens. seqs: List[SequenceGroupMetadata] = [] for group_id in range(max_num_seqs): seq_len = (max_num_batched_tokens // max_num_seqs + (group_id < max_num_batched_tokens % max_num_seqs)) block_size = cache_config.block_size seq_num_blocks = (seq_len + block_size - 1) // block_size seq_data, dummy_multi_modal_data = input_registry \ .dummy_data_for_profiling(model_config, seq_len, mm_registry) block_tables = [[0] * seq_num_blocks] * max_num_seqs seq = SequenceGroupMetadata( request_id=str(group_id), is_prompt=True, seq_data={group_id: seq_data}, sampling_params=sampling_params, block_tables=block_tables, lora_request=None, multi_modal_data=dummy_multi_modal_data) seqs.append(seq) self.model_runner.block_size = tmp_cache_config.block_size # Run the model with the dummy inputs. self.model_runner.execute_model(seqs, profiling_cache_engine.kv_cache) # explicitly delete temporary KV cache manager to free KV cache # when real inputs will be passed to OV del profiling_cache_engine logger.info( "Start profiling run with dummy inputs to evaluate " "memory usage for %s. It might take a while.", ov_device) model_profile_run() gpu_device_type = ov_core.get_property(ov_device, device.type) memory_statistics = \ ov_core.get_property(ov_device, intel_gpu.memory_statistics) memory_utilization = cache_config.gpu_memory_utilization if gpu_device_type == device.Type.INTEGRATED and \ memory_utilization >= 0.9: logger.warning( "iGPU is used with high gpu_memory_utilization=%f " "value. This may cause low performance due to " "occupying the majority of available system " "memory. Please consider decreasing " "gpu_memory_utilization or explicitly setting" "`VLLM_OPENVINO_KVCACHE_SPACE` (GB) environment " "variable.", memory_utilization) # sum up all used device memory device_memory_types = ["cl_mem", "usm_device"] used_device_mem = \ sum(memory_statistics.get(key, 0) for key in device_memory_types) if gpu_device_type == device.Type.INTEGRATED: used_device_mem += memory_statistics.get("usm_host", 0) # there could be unaccounted extra memory reserved by kernels, kept # in memory pools, etc # therefore, add a threshold to account for this used_memory_threshold = 1.1 used_device_mem *= used_memory_threshold total_device_memory = \ ov_core.get_property(ov_device, intel_gpu.device_total_mem_size) def format_memory_size(size) -> str: units = ["B", "KB", "MB", "GB"] unit_index = 0 while size > 1024 and unit_index < len(units) - 1: size /= 1024 unit_index += 1 return f"{size:.2f} {units[unit_index]}" total_device_memory_str = \ format(format_memory_size(total_device_memory)) used_device_memory_str = \ format(format_memory_size(used_device_mem)) logger.info( "Total %s memory: %s. " "Amount of memory required to run the model with " "max_num_batched_tokens=%d: %s.", ov_device, total_device_memory_str, self.scheduler_config.max_num_batched_tokens, used_device_memory_str) if used_device_mem >= total_device_memory: raise RuntimeError( f"The required memory size {used_device_memory_str} for model " "is higher than the total available device " "memory {total_device_memory_str}. Please consider to " "decrease `max_num_batched_tokens` or increase " "`gpu_memory_utilization`") return total_device_memory * memory_utilization - used_device_mem