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 15, 03:30 UTC
[Kernel] refactor cpu worker v0 cache dtype (#20080)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Ning Xie committed
Jul 3, 2025 at 16:39 UTC
fb14d53cf6e68866cc13e348db18f361da8fb7ec
1 file changed
+20
-18
vllm/worker/cpu_worker.py
+20
-18
@@ -18,7 +18,7 @@ from vllm.logger import init_logger
18
from vllm.lora.request import LoRARequest
19
from vllm.model_executor import set_random_seed
20
from vllm.sequence import ExecuteModelRequest
21
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, bind_kv_cache
21
+from vllm.utils import bind_kv_cache
22
from vllm.worker.cpu_enc_dec_model_runner import CPUEncoderDecoderModelRunner
23
from vllm.worker.cpu_model_runner import CPUModelRunner, CPUModelRunnerBase
24
from vllm.worker.cpu_pooling_model_runner import CPUPoolingModelRunner
@@ -54,13 +54,8 @@ class CPUCacheEngine:
54
# in the scheduler.
55
self.num_cpu_blocks = cache_config.num_gpu_blocks
56
57
- if cache_config.cache_dtype == "auto":
58
- self.dtype = model_config.dtype
59
- elif cache_config.cache_dtype in ["fp8", "fp8_e5m2"]:
60
- self.dtype = torch.float8_e5m2
61
- else:
62
- raise NotImplementedError(f"Unsupported KV cache type "
63
- f"{cache_config.cache_dtype}.")
57
+ self.dtype = CPUCacheEngine.get_kv_cache_dtype(cache_config,
58
+ model_config)
59
60
# Get attention backend.
61
self.attn_backend = get_attn_backend(
@@ -97,10 +92,20 @@ class CPUCacheEngine:
92
def copy(self, src_to_dsts: torch.Tensor) -> None:
93
self.attn_backend.copy_blocks(self.cpu_cache, src_to_dsts)
94
95
+ @staticmethod
96
+ def get_kv_cache_dtype(cache_config: CacheConfig,
97
+ model_config: ModelConfig):
98
+ if cache_config.cache_dtype == "auto":
99
+ return model_config.dtype
100
+ elif cache_config.cache_dtype in ["fp8", "fp8_e5m2"]:
101
+ return torch.float8_e5m2
102
+ else:
103
+ raise NotImplementedError(f"Unsupported KV cache type "
104
+ f"{cache_config.cache_dtype}.")
105
+
106
@staticmethod
107
def get_cache_block_size(
102
- block_size: int,
103
- cache_dtype: str,
108
+ cache_config: CacheConfig,
109
model_config: ModelConfig,
110
parallel_config: ParallelConfig,
111
) -> int:
@@ -108,13 +113,10 @@ class CPUCacheEngine:
113
num_heads = model_config.get_num_kv_heads(parallel_config)
114
num_layers = model_config.get_num_layers(parallel_config)
115
111
- key_cache_block = block_size * num_heads * head_size
116
+ key_cache_block = cache_config.block_size * num_heads * head_size
117
value_cache_block = key_cache_block if not model_config.use_mla else 0
118
total = num_layers * (key_cache_block + value_cache_block)
114
- if cache_dtype == "auto":
115
- dtype = model_config.dtype
116
- else:
117
- dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
119
+ dtype = CPUCacheEngine.get_kv_cache_dtype(cache_config, model_config)
120
dtype_size = torch.tensor([], dtype=dtype).element_size()
121
return dtype_size * total
122
@@ -399,9 +401,9 @@ class CPUWorker(LocalOrDistributedWorkerBase):
401
def get_cache_block_size_bytes(self) -> int:
402
"""Return the size in bytes of a single KV cache block.
403
"""
402
- return CPUCacheEngine.get_cache_block_size(
403
- self.cache_config.block_size, self.cache_config.cache_dtype,
404
- self.model_config, self.parallel_config)
404
+ return CPUCacheEngine.get_cache_block_size(self.cache_config,
405
+ self.model_config,
406
+ self.parallel_config)
407
408
def get_cpus_id_binding_based_on_numa_nodes(self) -> str:
409
"""Return CPUs id binding based on NUMA nodes.