We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
CUDA error: out of memory
I successfully installed vLLM in WSL2, when I was trying to run the sample code, I got error info like this:

```
from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
llm = LLM(model="/mnt/d/github/text-generation-webui/models/facebook_opt-125m")

outputs = llm.generate(prompts, sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
```

INFO 06-21 21:40:02 llm_engine.py:59] Initializing an LLM engine with config: model='/mnt/d/github/text-generation-webui/models/facebook_opt-125m', dtype=torch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)
INFO 06-21 21:40:12 llm_engine.py:128] # GPU blocks: 37375, # CPU blocks: 7281
Traceback (most recent call last):
  File "/mnt/d/01Projects/vllm/prac_1.py", line 11, in <module>
    llm = LLM(model="/mnt/d/github/text-generation-webui/models/facebook_opt-125m")
  File "/mnt/d/github/vllm/vllm/entrypoints/llm.py", line 55, in __init__
    self.llm_engine = LLMEngine.from_engine_args(engine_args)
  File "/mnt/d/github/vllm/vllm/engine/llm_engine.py", line 145, in from_engine_args
    engine = cls(*engine_configs, distributed_init_method, devices,
  File "/mnt/d/github/vllm/vllm/engine/llm_engine.py", line 102, in __init__
    self._init_cache()
  File "/mnt/d/github/vllm/vllm/engine/llm_engine.py", line 134, in _init_cache
    self._run_workers("init_cache_engine", cache_config=self.cache_config)
  File "/mnt/d/github/vllm/vllm/engine/llm_engine.py", line 307, in _run_workers
    output = executor(*args, **kwargs)
  File "/mnt/d/github/vllm/vllm/worker/worker.py", line 126, in init_cache_engine
    self.cache_engine = CacheEngine(
  File "/mnt/d/github/vllm/vllm/worker/cache_engine.py", line 41, in __init__
    self.cpu_cache = self.allocate_cpu_cache()
  File "/mnt/d/github/vllm/vllm/worker/cache_engine.py", line 89, in allocate_cpu_cache
    key_blocks = torch.empty(
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.


Python: 3.10.11
GPU: RTX 3090 24G
Linux: WSL2, Ubuntu 20.04.6 LTS
Can anyone help to answer this?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `vllm/utils.py`
Content:
```
1 import enum
2 import uuid
3 
4 import psutil
5 import torch
6 
7 
8 class Device(enum.Enum):
9     GPU = enum.auto()
10     CPU = enum.auto()
11 
12 
13 class Counter:
14 
15     def __init__(self, start: int = 0) -> None:
16         self.counter = start
17 
18     def __next__(self) -> int:
19         id = self.counter
20         self.counter += 1
21         return id
22 
23     def reset(self) -> None:
24         self.counter = 0
25 
26 
27 def get_gpu_memory(gpu: int = 0) -> int:
28     """Returns the total memory of the GPU in bytes."""
29     return torch.cuda.get_device_properties(gpu).total_memory
30 
31 
32 def get_cpu_memory() -> int:
33     """Returns the total CPU memory of the node in bytes."""
34     return psutil.virtual_memory().total
35 
36 
37 def random_uuid() -> str:
38     return str(uuid.uuid4().hex)
39
```

Path: `vllm/worker/cache_engine.py`
Content:
```
1 """CacheEngine class for managing the KV cache."""
2 from typing import Dict, List, Tuple
3 
4 import torch
5 
6 from vllm import cache_ops
7 from vllm.config import CacheConfig, ModelConfig, ParallelConfig
8 
9 KVCache = Tuple[torch.Tensor, torch.Tensor]
10 
11 
12 class CacheEngine:
13     """Manages the KV cache.
14 
15     This class is responsible for initializing and managing the GPU and CPU KV
16     caches. It also provides methods for performing KV cache operations, such
17     as swapping and copying.
18     """
19 
20     def __init__(
21         self,
22         cache_config: CacheConfig,
23         model_config: ModelConfig,
24         parallel_config: ParallelConfig,
25     ) -> None:
26         self.cache_config = cache_config
27         self.model_config = model_config
28         self.parallel_config = parallel_config
29 
30         self.head_size = model_config.get_head_size()
31         self.num_layers = model_config.get_num_layers(parallel_config)
32         self.num_heads = model_config.get_num_heads(parallel_config)
33         self.dtype = model_config.dtype
34 
35         self.block_size = cache_config.block_size
36         self.num_gpu_blocks = cache_config.num_gpu_blocks
37         self.num_cpu_blocks = cache_config.num_cpu_blocks
38 
39         # Initialize the cache.
40         self.gpu_cache = self.allocate_gpu_cache()
41         self.cpu_cache = self.allocate_cpu_cache()
42 
43         # Initialize the stream for caching operations.
44         self.cache_stream = torch.cuda.Stream()
45         assert self.cache_stream != torch.cuda.current_stream()
46         # Initialize the events for stream synchronization.
47         self.events = [torch.cuda.Event() for _ in range(self.num_layers)]
48 
49     def get_key_block_shape(self) -> Tuple[int, int, int, int]:
50         element_size = torch.tensor([], dtype=self.dtype).element_size()
51         x = 16 // element_size
52         return (
53             self.num_heads,
54             self.head_size // x,
55             self.block_size,
56             x,
57         )
58 
59     def get_value_block_shape(self) -> Tuple[int, int, int]:
60         return (
61             self.num_heads,
62             self.head_size,
63             self.block_size,
64         )
65 
66     def allocate_gpu_cache(self) -> List[KVCache]:
67         gpu_cache: List[KVCache] = []
68         key_block_shape = self.get_key_block_shape()
69         value_block_shape = self.get_value_block_shape()
70         for _ in range(self.num_layers):
71             key_blocks = torch.empty(
72                 size=(self.num_gpu_blocks, *key_block_shape),
73                 dtype=self.dtype,
74                 device="cuda",
75             )
76             value_blocks = torch.empty(
77                 size=(self.num_gpu_blocks, *value_block_shape),
78                 dtype=self.dtype,
79                 device="cuda",
80             )
81             gpu_cache.append((key_blocks, value_blocks))
82         return gpu_cache
83 
84     def allocate_cpu_cache(self) -> List[KVCache]:
85         cpu_cache: List[KVCache] = []
86         key_block_shape = self.get_key_block_shape()
87         value_block_shape = self.get_value_block_shape()
88         for _ in range(self.num_layers):
89             key_blocks = torch.empty(
90                 size=(self.num_cpu_blocks, *key_block_shape),
91                 dtype=self.dtype,
92                 pin_memory=True,
93             )
94             value_blocks = torch.empty(
95                 size=(self.num_cpu_blocks, *value_block_shape),
96                 dtype=self.dtype,
97                 pin_memory=True,
98             )
99             cpu_cache.append((key_blocks, value_blocks))
100         return cpu_cache
101 
102     def _swap(
103         self,
104         src: List[KVCache],
105         dst: List[KVCache],
106         src_to_dst: Dict[int, int],
107     ) -> None:
108         with torch.cuda.stream(self.cache_stream):
109             for i in range(self.num_layers):
110                 src_key_cache, src_value_cache = src[i]
111                 dst_key_cache, dst_value_cache = dst[i]
112                 # Copy the key blocks.
113                 cache_ops.swap_blocks(
114                     src_key_cache, dst_key_cache, src_to_dst)
115                 # Copy the value blocks.
116                 cache_ops.swap_blocks(
117                     src_value_cache, dst_value_cache, src_to_dst)
118                 event = self.events[i]
119                 event.record(stream=self.cache_stream)
120 
121     def swap_in(self, src_to_dst: Dict[int, int]) -> None:
122         self._swap(self.cpu_cache, self.gpu_cache, src_to_dst)
123 
124     def swap_out(self, src_to_dst: Dict[int, int]) -> None:
125         self._swap(self.gpu_cache, self.cpu_cache, src_to_dst)
126 
127     def copy(self, src_to_dsts: Dict[int, List[int]]) -> None:
128         key_caches = [key_cache for key_cache, _ in self.gpu_cache]
129         value_caches = [value_cache for _, value_cache in self.gpu_cache]
130         # NOTE(woosuk): This operation implicitly synchronizes the CPU and GPU.
131         cache_ops.copy_blocks(key_caches, value_caches, src_to_dsts)
132 
133     @staticmethod
134     def get_cache_block_size(
135         block_size: int,
136         model_config: ModelConfig,
137         parallel_config: ParallelConfig,
138     ) -> int:
139         head_size = model_config.get_head_size()
140         num_heads = model_config.get_num_heads(parallel_config)
141         num_layers = model_config.get_num_layers(parallel_config)
142 
143         key_cache_block = block_size * num_heads * head_size
144         value_cache_block = key_cache_block
145         total = num_layers * (key_cache_block + value_cache_block)
146         dtype_size = _get_dtype_size(model_config.dtype)
147         return dtype_size * total
148 
149 
150 def _get_dtype_size(dtype: torch.dtype) -> int:
151     return torch.tensor([], dtype=dtype).element_size()
152
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/vllm/utils.py b/vllm/utils.py
--- a/vllm/utils.py
+++ b/vllm/utils.py
@@ -1,4 +1,5 @@
 import enum
+from platform import uname
 import uuid
 
 import psutil
@@ -36,3 +37,7 @@
 
 def random_uuid() -> str:
     return str(uuid.uuid4().hex)
+
+def in_wsl() -> bool:
+    # Reference: https://github.com/microsoft/WSL/issues/4071
+    return "microsoft" in " ".join(uname()).lower()
diff --git a/vllm/worker/cache_engine.py b/vllm/worker/cache_engine.py
--- a/vllm/worker/cache_engine.py
+++ b/vllm/worker/cache_engine.py
@@ -5,6 +5,10 @@
 
 from vllm import cache_ops
 from vllm.config import CacheConfig, ModelConfig, ParallelConfig
+from vllm.logger import init_logger
+from vllm.utils import in_wsl
+
+logger = init_logger(__name__)
 
 KVCache = Tuple[torch.Tensor, torch.Tensor]
 
@@ -85,16 +89,22 @@
         cpu_cache: List[KVCache] = []
         key_block_shape = self.get_key_block_shape()
         value_block_shape = self.get_value_block_shape()
+        pin_memory = not in_wsl()
+        if not pin_memory:
+            # Pinning memory in WSL is not supported.
+            # https://docs.nvidia.com/cuda/wsl-user-guide/index.html#known-limitations-for-linux-cuda-applications
+            logger.warn("Using 'pin_memory=False' as WSL is detected. "
+                        "This may slow down the performance.")
         for _ in range(self.num_layers):
             key_blocks = torch.empty(
                 size=(self.num_cpu_blocks, *key_block_shape),
                 dtype=self.dtype,
-                pin_memory=True,
+                pin_memory=pin_memory,
             )
             value_blocks = torch.empty(
                 size=(self.num_cpu_blocks, *value_block_shape),
                 dtype=self.dtype,
-                pin_memory=True,
+                pin_memory=pin_memory,
             )
             cpu_cache.append((key_blocks, value_blocks))
         return cpu_cache

{"golden_diff": "diff --git a/vllm/utils.py b/vllm/utils.py\n--- a/vllm/utils.py\n+++ b/vllm/utils.py\n@@ -1,4 +1,5 @@\n import enum\n+from platform import uname\n import uuid\n \n import psutil\n@@ -36,3 +37,7 @@\n \n def random_uuid() -> str:\n     return str(uuid.uuid4().hex)\n+\n+def in_wsl() -> bool:\n+    # Reference: https://github.com/microsoft/WSL/issues/4071\n+    return \"microsoft\" in \" \".join(uname()).lower()\ndiff --git a/vllm/worker/cache_engine.py b/vllm/worker/cache_engine.py\n--- a/vllm/worker/cache_engine.py\n+++ b/vllm/worker/cache_engine.py\n@@ -5,6 +5,10 @@\n \n from vllm import cache_ops\n from vllm.config import CacheConfig, ModelConfig, ParallelConfig\n+from vllm.logger import init_logger\n+from vllm.utils import in_wsl\n+\n+logger = init_logger(__name__)\n \n KVCache = Tuple[torch.Tensor, torch.Tensor]\n \n@@ -85,16 +89,22 @@\n         cpu_cache: List[KVCache] = []\n         key_block_shape = self.get_key_block_shape()\n         value_block_shape = self.get_value_block_shape()\n+        pin_memory = not in_wsl()\n+        if not pin_memory:\n+            # Pinning memory in WSL is not supported.\n+            # https://docs.nvidia.com/cuda/wsl-user-guide/index.html#known-limitations-for-linux-cuda-applications\n+            logger.warn(\"Using 'pin_memory=False' as WSL is detected. \"\n+                        \"This may slow down the performance.\")\n         for _ in range(self.num_layers):\n             key_blocks = torch.empty(\n                 size=(self.num_cpu_blocks, *key_block_shape),\n                 dtype=self.dtype,\n-                pin_memory=True,\n+                pin_memory=pin_memory,\n             )\n             value_blocks = torch.empty(\n                 size=(self.num_cpu_blocks, *value_block_shape),\n                 dtype=self.dtype,\n-                pin_memory=True,\n+                pin_memory=pin_memory,\n             )\n             cpu_cache.append((key_blocks, value_blocks))\n         return cpu_cache\n", "issue": "CUDA error: out of memory\nI successfully installed vLLM in WSL2, when I was trying to run the sample code, I got error info like this:\r\n\r\n```\r\nfrom vllm import LLM, SamplingParams\r\n\r\nprompts = [\r\n    \"Hello, my name is\",\r\n    \"The president of the United States is\",\r\n    \"The capital of France is\",\r\n    \"The future of AI is\",\r\n]\r\nsampling_params = SamplingParams(temperature=0.8, top_p=0.95)\r\nllm = LLM(model=\"/mnt/d/github/text-generation-webui/models/facebook_opt-125m\")\r\n\r\noutputs = llm.generate(prompts, sampling_params)\r\n\r\n# Print the outputs.\r\nfor output in outputs:\r\n    prompt = output.prompt\r\n    generated_text = output.outputs[0].text\r\n    print(f\"Prompt: {prompt!r}, Generated text: {generated_text!r}\")\r\n```\r\n\r\nINFO 06-21 21:40:02 llm_engine.py:59] Initializing an LLM engine with config: model='/mnt/d/github/text-generation-webui/models/facebook_opt-125m', dtype=torch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)\r\nINFO 06-21 21:40:12 llm_engine.py:128] # GPU blocks: 37375, # CPU blocks: 7281\r\nTraceback (most recent call last):\r\n  File \"/mnt/d/01Projects/vllm/prac_1.py\", line 11, in <module>\r\n    llm = LLM(model=\"/mnt/d/github/text-generation-webui/models/facebook_opt-125m\")\r\n  File \"/mnt/d/github/vllm/vllm/entrypoints/llm.py\", line 55, in __init__\r\n    self.llm_engine = LLMEngine.from_engine_args(engine_args)\r\n  File \"/mnt/d/github/vllm/vllm/engine/llm_engine.py\", line 145, in from_engine_args\r\n    engine = cls(*engine_configs, distributed_init_method, devices,\r\n  File \"/mnt/d/github/vllm/vllm/engine/llm_engine.py\", line 102, in __init__\r\n    self._init_cache()\r\n  File \"/mnt/d/github/vllm/vllm/engine/llm_engine.py\", line 134, in _init_cache\r\n    self._run_workers(\"init_cache_engine\", cache_config=self.cache_config)\r\n  File \"/mnt/d/github/vllm/vllm/engine/llm_engine.py\", line 307, in _run_workers\r\n    output = executor(*args, **kwargs)\r\n  File \"/mnt/d/github/vllm/vllm/worker/worker.py\", line 126, in init_cache_engine\r\n    self.cache_engine = CacheEngine(\r\n  File \"/mnt/d/github/vllm/vllm/worker/cache_engine.py\", line 41, in __init__\r\n    self.cpu_cache = self.allocate_cpu_cache()\r\n  File \"/mnt/d/github/vllm/vllm/worker/cache_engine.py\", line 89, in allocate_cpu_cache\r\n    key_blocks = torch.empty(\r\nRuntimeError: CUDA error: out of memory\r\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\r\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1.\r\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\r\n\r\n\r\nPython: 3.10.11\r\nGPU: RTX 3090 24G\r\nLinux: WSL2, Ubuntu 20.04.6 LTS\r\nCan anyone help to answer this?\n", "before_files": [{"content": "import enum\nimport uuid\n\nimport psutil\nimport torch\n\n\nclass Device(enum.Enum):\n    GPU = enum.auto()\n    CPU = enum.auto()\n\n\nclass Counter:\n\n    def __init__(self, start: int = 0) -> None:\n        self.counter = start\n\n    def __next__(self) -> int:\n        id = self.counter\n        self.counter += 1\n        return id\n\n    def reset(self) -> None:\n        self.counter = 0\n\n\ndef get_gpu_memory(gpu: int = 0) -> int:\n    \"\"\"Returns the total memory of the GPU in bytes.\"\"\"\n    return torch.cuda.get_device_properties(gpu).total_memory\n\n\ndef get_cpu_memory() -> int:\n    \"\"\"Returns the total CPU memory of the node in bytes.\"\"\"\n    return psutil.virtual_memory().total\n\n\ndef random_uuid() -> str:\n    return str(uuid.uuid4().hex)\n", "path": "vllm/utils.py"}, {"content": "\"\"\"CacheEngine class for managing the KV cache.\"\"\"\nfrom typing import Dict, List, Tuple\n\nimport torch\n\nfrom vllm import cache_ops\nfrom vllm.config import CacheConfig, ModelConfig, ParallelConfig\n\nKVCache = Tuple[torch.Tensor, torch.Tensor]\n\n\nclass CacheEngine:\n    \"\"\"Manages the KV cache.\n\n    This class is responsible for initializing and managing the GPU and CPU KV\n    caches. It also provides methods for performing KV cache operations, such\n    as swapping and copying.\n    \"\"\"\n\n    def __init__(\n        self,\n        cache_config: CacheConfig,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n    ) -> None:\n        self.cache_config = cache_config\n        self.model_config = model_config\n        self.parallel_config = parallel_config\n\n        self.head_size = model_config.get_head_size()\n        self.num_layers = model_config.get_num_layers(parallel_config)\n        self.num_heads = model_config.get_num_heads(parallel_config)\n        self.dtype = model_config.dtype\n\n        self.block_size = cache_config.block_size\n        self.num_gpu_blocks = cache_config.num_gpu_blocks\n        self.num_cpu_blocks = cache_config.num_cpu_blocks\n\n        # Initialize the cache.\n        self.gpu_cache = self.allocate_gpu_cache()\n        self.cpu_cache = self.allocate_cpu_cache()\n\n        # Initialize the stream for caching operations.\n        self.cache_stream = torch.cuda.Stream()\n        assert self.cache_stream != torch.cuda.current_stream()\n        # Initialize the events for stream synchronization.\n        self.events = [torch.cuda.Event() for _ in range(self.num_layers)]\n\n    def get_key_block_shape(self) -> Tuple[int, int, int, int]:\n        element_size = torch.tensor([], dtype=self.dtype).element_size()\n        x = 16 // element_size\n        return (\n            self.num_heads,\n            self.head_size // x,\n            self.block_size,\n            x,\n        )\n\n    def get_value_block_shape(self) -> Tuple[int, int, int]:\n        return (\n            self.num_heads,\n            self.head_size,\n            self.block_size,\n        )\n\n    def allocate_gpu_cache(self) -> List[KVCache]:\n        gpu_cache: List[KVCache] = []\n        key_block_shape = self.get_key_block_shape()\n        value_block_shape = self.get_value_block_shape()\n        for _ in range(self.num_layers):\n            key_blocks = torch.empty(\n                size=(self.num_gpu_blocks, *key_block_shape),\n                dtype=self.dtype,\n                device=\"cuda\",\n            )\n            value_blocks = torch.empty(\n                size=(self.num_gpu_blocks, *value_block_shape),\n                dtype=self.dtype,\n                device=\"cuda\",\n            )\n            gpu_cache.append((key_blocks, value_blocks))\n        return gpu_cache\n\n    def allocate_cpu_cache(self) -> List[KVCache]:\n        cpu_cache: List[KVCache] = []\n        key_block_shape = self.get_key_block_shape()\n        value_block_shape = self.get_value_block_shape()\n        for _ in range(self.num_layers):\n            key_blocks = torch.empty(\n                size=(self.num_cpu_blocks, *key_block_shape),\n                dtype=self.dtype,\n                pin_memory=True,\n            )\n            value_blocks = torch.empty(\n                size=(self.num_cpu_blocks, *value_block_shape),\n                dtype=self.dtype,\n                pin_memory=True,\n            )\n            cpu_cache.append((key_blocks, value_blocks))\n        return cpu_cache\n\n    def _swap(\n        self,\n        src: List[KVCache],\n        dst: List[KVCache],\n        src_to_dst: Dict[int, int],\n    ) -> None:\n        with torch.cuda.stream(self.cache_stream):\n            for i in range(self.num_layers):\n                src_key_cache, src_value_cache = src[i]\n                dst_key_cache, dst_value_cache = dst[i]\n                # Copy the key blocks.\n                cache_ops.swap_blocks(\n                    src_key_cache, dst_key_cache, src_to_dst)\n                # Copy the value blocks.\n                cache_ops.swap_blocks(\n                    src_value_cache, dst_value_cache, src_to_dst)\n                event = self.events[i]\n                event.record(stream=self.cache_stream)\n\n    def swap_in(self, src_to_dst: Dict[int, int]) -> None:\n        self._swap(self.cpu_cache, self.gpu_cache, src_to_dst)\n\n    def swap_out(self, src_to_dst: Dict[int, int]) -> None:\n        self._swap(self.gpu_cache, self.cpu_cache, src_to_dst)\n\n    def copy(self, src_to_dsts: Dict[int, List[int]]) -> None:\n        key_caches = [key_cache for key_cache, _ in self.gpu_cache]\n        value_caches = [value_cache for _, value_cache in self.gpu_cache]\n        # NOTE(woosuk): This operation implicitly synchronizes the CPU and GPU.\n        cache_ops.copy_blocks(key_caches, value_caches, src_to_dsts)\n\n    @staticmethod\n    def get_cache_block_size(\n        block_size: int,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n    ) -> int:\n        head_size = model_config.get_head_size()\n        num_heads = model_config.get_num_heads(parallel_config)\n        num_layers = model_config.get_num_layers(parallel_config)\n\n        key_cache_block = block_size * num_heads * head_size\n        value_cache_block = key_cache_block\n        total = num_layers * (key_cache_block + value_cache_block)\n        dtype_size = _get_dtype_size(model_config.dtype)\n        return dtype_size * total\n\n\ndef _get_dtype_size(dtype: torch.dtype) -> int:\n    return torch.tensor([], dtype=dtype).element_size()\n", "path": "vllm/worker/cache_engine.py"}], "after_files": [{"content": "import enum\nfrom platform import uname\nimport uuid\n\nimport psutil\nimport torch\n\n\nclass Device(enum.Enum):\n    GPU = enum.auto()\n    CPU = enum.auto()\n\n\nclass Counter:\n\n    def __init__(self, start: int = 0) -> None:\n        self.counter = start\n\n    def __next__(self) -> int:\n        id = self.counter\n        self.counter += 1\n        return id\n\n    def reset(self) -> None:\n        self.counter = 0\n\n\ndef get_gpu_memory(gpu: int = 0) -> int:\n    \"\"\"Returns the total memory of the GPU in bytes.\"\"\"\n    return torch.cuda.get_device_properties(gpu).total_memory\n\n\ndef get_cpu_memory() -> int:\n    \"\"\"Returns the total CPU memory of the node in bytes.\"\"\"\n    return psutil.virtual_memory().total\n\n\ndef random_uuid() -> str:\n    return str(uuid.uuid4().hex)\n\ndef in_wsl() -> bool:\n    # Reference: https://github.com/microsoft/WSL/issues/4071\n    return \"microsoft\" in \" \".join(uname()).lower()\n", "path": "vllm/utils.py"}, {"content": "\"\"\"CacheEngine class for managing the KV cache.\"\"\"\nfrom typing import Dict, List, Tuple\n\nimport torch\n\nfrom vllm import cache_ops\nfrom vllm.config import CacheConfig, ModelConfig, ParallelConfig\nfrom vllm.logger import init_logger\nfrom vllm.utils import in_wsl\n\nlogger = init_logger(__name__)\n\nKVCache = Tuple[torch.Tensor, torch.Tensor]\n\n\nclass CacheEngine:\n    \"\"\"Manages the KV cache.\n\n    This class is responsible for initializing and managing the GPU and CPU KV\n    caches. It also provides methods for performing KV cache operations, such\n    as swapping and copying.\n    \"\"\"\n\n    def __init__(\n        self,\n        cache_config: CacheConfig,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n    ) -> None:\n        self.cache_config = cache_config\n        self.model_config = model_config\n        self.parallel_config = parallel_config\n\n        self.head_size = model_config.get_head_size()\n        self.num_layers = model_config.get_num_layers(parallel_config)\n        self.num_heads = model_config.get_num_heads(parallel_config)\n        self.dtype = model_config.dtype\n\n        self.block_size = cache_config.block_size\n        self.num_gpu_blocks = cache_config.num_gpu_blocks\n        self.num_cpu_blocks = cache_config.num_cpu_blocks\n\n        # Initialize the cache.\n        self.gpu_cache = self.allocate_gpu_cache()\n        self.cpu_cache = self.allocate_cpu_cache()\n\n        # Initialize the stream for caching operations.\n        self.cache_stream = torch.cuda.Stream()\n        assert self.cache_stream != torch.cuda.current_stream()\n        # Initialize the events for stream synchronization.\n        self.events = [torch.cuda.Event() for _ in range(self.num_layers)]\n\n    def get_key_block_shape(self) -> Tuple[int, int, int, int]:\n        element_size = torch.tensor([], dtype=self.dtype).element_size()\n        x = 16 // element_size\n        return (\n            self.num_heads,\n            self.head_size // x,\n            self.block_size,\n            x,\n        )\n\n    def get_value_block_shape(self) -> Tuple[int, int, int]:\n        return (\n            self.num_heads,\n            self.head_size,\n            self.block_size,\n        )\n\n    def allocate_gpu_cache(self) -> List[KVCache]:\n        gpu_cache: List[KVCache] = []\n        key_block_shape = self.get_key_block_shape()\n        value_block_shape = self.get_value_block_shape()\n        for _ in range(self.num_layers):\n            key_blocks = torch.empty(\n                size=(self.num_gpu_blocks, *key_block_shape),\n                dtype=self.dtype,\n                device=\"cuda\",\n            )\n            value_blocks = torch.empty(\n                size=(self.num_gpu_blocks, *value_block_shape),\n                dtype=self.dtype,\n                device=\"cuda\",\n            )\n            gpu_cache.append((key_blocks, value_blocks))\n        return gpu_cache\n\n    def allocate_cpu_cache(self) -> List[KVCache]:\n        cpu_cache: List[KVCache] = []\n        key_block_shape = self.get_key_block_shape()\n        value_block_shape = self.get_value_block_shape()\n        pin_memory = not in_wsl()\n        if not pin_memory:\n            # Pinning memory in WSL is not supported.\n            # https://docs.nvidia.com/cuda/wsl-user-guide/index.html#known-limitations-for-linux-cuda-applications\n            logger.warn(\"Using 'pin_memory=False' as WSL is detected. \"\n                        \"This may slow down the performance.\")\n        for _ in range(self.num_layers):\n            key_blocks = torch.empty(\n                size=(self.num_cpu_blocks, *key_block_shape),\n                dtype=self.dtype,\n                pin_memory=pin_memory,\n            )\n            value_blocks = torch.empty(\n                size=(self.num_cpu_blocks, *value_block_shape),\n                dtype=self.dtype,\n                pin_memory=pin_memory,\n            )\n            cpu_cache.append((key_blocks, value_blocks))\n        return cpu_cache\n\n    def _swap(\n        self,\n        src: List[KVCache],\n        dst: List[KVCache],\n        src_to_dst: Dict[int, int],\n    ) -> None:\n        with torch.cuda.stream(self.cache_stream):\n            for i in range(self.num_layers):\n                src_key_cache, src_value_cache = src[i]\n                dst_key_cache, dst_value_cache = dst[i]\n                # Copy the key blocks.\n                cache_ops.swap_blocks(\n                    src_key_cache, dst_key_cache, src_to_dst)\n                # Copy the value blocks.\n                cache_ops.swap_blocks(\n                    src_value_cache, dst_value_cache, src_to_dst)\n                event = self.events[i]\n                event.record(stream=self.cache_stream)\n\n    def swap_in(self, src_to_dst: Dict[int, int]) -> None:\n        self._swap(self.cpu_cache, self.gpu_cache, src_to_dst)\n\n    def swap_out(self, src_to_dst: Dict[int, int]) -> None:\n        self._swap(self.gpu_cache, self.cpu_cache, src_to_dst)\n\n    def copy(self, src_to_dsts: Dict[int, List[int]]) -> None:\n        key_caches = [key_cache for key_cache, _ in self.gpu_cache]\n        value_caches = [value_cache for _, value_cache in self.gpu_cache]\n        # NOTE(woosuk): This operation implicitly synchronizes the CPU and GPU.\n        cache_ops.copy_blocks(key_caches, value_caches, src_to_dsts)\n\n    @staticmethod\n    def get_cache_block_size(\n        block_size: int,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n    ) -> int:\n        head_size = model_config.get_head_size()\n        num_heads = model_config.get_num_heads(parallel_config)\n        num_layers = model_config.get_num_layers(parallel_config)\n\n        key_cache_block = block_size * num_heads * head_size\n        value_cache_block = key_cache_block\n        total = num_layers * (key_cache_block + value_cache_block)\n        dtype_size = _get_dtype_size(model_config.dtype)\n        return dtype_size * total\n\n\ndef _get_dtype_size(dtype: torch.dtype) -> int:\n    return torch.tensor([], dtype=dtype).element_size()\n", "path": "vllm/worker/cache_engine.py"}]}