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, 04:45 UTC
[Kernel] correct cpu worker function parameter type (#19745)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Ning Xie committed
Jun 20, 2025 at 18:50 UTC
71d1219545b5139bf8f00fc72bdd3682cce62775
2 files changed
+5
-5
vllm/attention/ops/ipex_attn.py
+1
-1
@@ -29,7 +29,7 @@ class _PagedAttention:
29
head_size: int,
30
*args,
31
) -> Tuple[int, ...]:
32
- return (2, num_blocks, block_size * num_kv_heads * head_size)
32
+ return 2, num_blocks, block_size * num_kv_heads * head_size
33
34
@staticmethod
35
def split_kv_cache(
vllm/worker/cpu_worker.py
+4
-4
@@ -3,7 +3,7 @@
3
"""A CPU worker class."""
4
import os
5
from importlib import util
6
-from typing import Dict, List, Optional, Set, Tuple, Type
6
+from typing import List, Optional, Set, Tuple, Type
7
8
import torch
9
import torch.distributed
@@ -88,13 +88,13 @@ class CPUCacheEngine:
88
torch.empty(kv_cache_shape, dtype=self.dtype, device="cpu"))
89
return kv_cache
90
91
- def swap_in(self, src_to_dst: Dict[int, int]) -> None:
91
+ def swap_in(self, src_to_dst: torch.Tensor) -> None:
92
raise NotImplementedError("Swap is not supported in CPUCacheEngine.")
93
94
- def swap_out(self, src_to_dst: Dict[int, int]) -> None:
94
+ def swap_out(self, src_to_dst: torch.Tensor) -> None:
95
raise NotImplementedError("Swap is not supported in CPUCacheEngine.")
96
97
- def copy(self, src_to_dsts: Dict[int, List[int]]) -> None:
97
+ def copy(self, src_to_dsts: torch.Tensor) -> None:
98
self.attn_backend.copy_blocks(self.cpu_cache, src_to_dsts)
99
100
@staticmethod