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 12, 01:00 UTC
[KVConnector] Remove v0-related kv connector components such as kv pipe and kv lookup buffer (#29705)
Signed-off-by: KuntaiDu <kuntai@uchicago.edu>
Kuntai Du committed
Dec 5, 2025 at 02:20 UTC
ece2825a29e6b54ce6b114c27ec7ea498c66b416
13 files changed
-1624
tests/kv_transfer/test_lookup_buffer.py
deleted
-160
@@ -1,160 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import os
5
-import random
6
-
7
-import torch
8
-from tqdm import tqdm
9
-
10
-from vllm.config import KVTransferConfig
11
-from vllm.distributed.kv_transfer.kv_lookup_buffer.simple_buffer import SimpleBuffer
12
-from vllm.distributed.kv_transfer.kv_pipe.pynccl_pipe import PyNcclPipe
13
-
14
-# TODO: the test depends on a lot of fields in the current implementation.
15
-# We should have standard interface instead direct field access
16
-
17
-
18
-def test_run(my_rank, buffer, device):
19
- # buffer should be empty in the beginning
20
- if my_rank == 0:
21
- assert buffer.buffer_size == 0
22
- assert len(buffer.buffer) == 0
23
-
24
- print(f"My rank: {my_rank}, device: {device}")
25
-
26
- # insert
27
- tokens = torch.tensor([1, 2, 3]).to(device)
28
- roi = tokens > 0
29
- if my_rank == 0:
30
- key = 2.0 * torch.ones([5, 6]).to(device)
31
- value = 3.0 * torch.ones([5, 6]).to(device)
32
-
33
- placeholder = torch.tensor([1]).to(device)
34
-
35
- buffer.insert(tokens, roi, key, value, placeholder)
36
-
37
- torch.distributed.barrier()
38
-
39
- # drop_select
40
- if my_rank == 1:
41
- tok, roi_, key, value, hidden = buffer.drop_select(tokens, roi)
42
- assert torch.allclose(tokens, tok)
43
- assert torch.allclose(roi, roi_)
44
- assert torch.allclose(key, 2.0 * torch.ones([5, 6], device=device))
45
- assert torch.allclose(value, 3.0 * torch.ones([5, 6], device=device))
46
- torch.distributed.barrier()
47
-
48
- if my_rank == 0:
49
- assert buffer.buffer_size == 0
50
- assert len(buffer.buffer) == 0
51
-
52
- print(f"My rank: {my_rank}, Test run passed!")
53
-
54
-
55
-def stress_test(my_rank, buf, device):
56
- torch.distributed.barrier()
57
- torch.manual_seed(100)
58
-
59
- reqs = [
60
- (
61
- torch.rand(100).to(device), # tokens
62
- torch.ones(100).bool().to(device), # roi
63
- torch.rand(100).to(device), # key
64
- torch.rand(100).to(device), # value
65
- torch.rand(100).to(device), # hidden
66
- )
67
- for i in tqdm(range(200))
68
- ]
69
-
70
- random.seed(my_rank)
71
- random.shuffle(reqs)
72
-
73
- torch.distributed.barrier()
74
-
75
- n = 0
76
-
77
- # the buffer size can only store 100 reqs
78
- # so the sender will occasionally block to wait for the receiver.
79
- for req in tqdm(reqs):
80
- if my_rank == 0:
81
- buf.insert(*req)
82
- else:
83
- tok, roi, k, v, h = req
84
- tok_, roi_, k_, v_, h_ = buf.drop_select(tok, roi)
85
-
86
- if tok_ is None:
87
- assert roi_ is None
88
- assert k_ is None
89
- assert v_ is None
90
- assert h_ is None
91
- n += 1
92
- else:
93
- assert torch.allclose(tok, tok_)
94
- assert torch.allclose(roi, roi_)
95
- assert torch.allclose(k, k_)
96
- assert torch.allclose(v, v_)
97
- assert torch.allclose(h, h_)
98
- print(f"Rank {my_rank} done")
99
- torch.distributed.barrier()
100
-
101
- if my_rank == 0:
102
- x = torch.tensor([0])
103
- torch.distributed.recv(x, 1)
104
- # the # of None received is the kv that are not selected
105
- assert x.item() == len(buf.buffer)
106
- # and the size of the buffer should be 2000 * buffer len
107
- print(buf.buffer_size)
108
- assert buf.buffer_size == 1700 * len(buf.buffer)
109
- else:
110
- torch.distributed.send(torch.tensor([n]), 0)
111
-
112
- print(f"My rank: {my_rank}, Passed stress test!")
113
-
114
-
115
-if __name__ == "__main__":
116
- my_rank = int(os.environ["RANK"])
117
-
118
- torch.distributed.init_process_group(
119
- backend="gloo",
120
- init_method="tcp://localhost:12398",
121
- world_size=2,
122
- rank=my_rank,
123
- )
124
-
125
- print(f"initialized! My rank is {my_rank}")
126
-
127
- config = KVTransferConfig(
128
- kv_connector="P2pNcclConnector",
129
- kv_buffer_device="cuda",
130
- kv_buffer_size=1e9,
131
- kv_rank=my_rank,
132
- kv_role="kv_both", # this arg doesn't matter in this test
133
- kv_parallel_size=2,
134
- kv_ip="127.0.0.1",
135
- kv_port=12345,
136
- )
137
-
138
- data_pipe = PyNcclPipe(
139
- local_rank=my_rank,
140
- config=config,
141
- device="cuda",
142
- port_offset=0,
143
- )
144
- cpu_pipe = PyNcclPipe(
145
- local_rank=my_rank,
146
- config=config,
147
- device="cpu",
148
- port_offset=1,
149
- )
150
-
151
- buffer = SimpleBuffer(cpu_pipe, data_pipe, 170000)
152
-
153
- test_run(my_rank, buffer, data_pipe.device)
154
-
155
- stress_test(my_rank, buffer, data_pipe.device)
156
-
157
- buffer.close()
158
- data_pipe.close()
159
- cpu_pipe.close()
160
- print("Done")
tests/kv_transfer/test_lookup_buffer.sh
deleted
-8
@@ -1,8 +0,0 @@
1
-#!/bin/bash
2
-RANK=0 python3 test_lookup_buffer.py &
3
-PID0=$!
4
-RANK=1 python3 test_lookup_buffer.py &
5
-PID1=$!
6
-
7
-wait $PID0
8
-wait $PID1
tests/kv_transfer/test_module.py
deleted
-62
@@ -1,62 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import subprocess
5
-import sys
6
-
7
-import pytest
8
-import torch
9
-
10
-
11
-def run_python_script(script_name, timeout):
12
- script_name = f"kv_transfer/{script_name}"
13
- try:
14
- # Start both processes asynchronously using Popen
15
- process0 = subprocess.Popen(
16
- [sys.executable, script_name],
17
- env={"RANK": "0"}, # Set the RANK environment variable for process 0
18
- stdout=sys.stdout, # Pipe stdout to current stdout
19
- stderr=sys.stderr, # Pipe stderr to current stderr
20
- )
21
-
22
- process1 = subprocess.Popen(
23
- [sys.executable, script_name],
24
- env={"RANK": "1"}, # Set the RANK environment variable for process 1
25
- stdout=sys.stdout, # Pipe stdout to current stdout
26
- stderr=sys.stderr, # Pipe stderr to current stderr
27
- )
28
-
29
- # Wait for both processes to complete, with a timeout
30
- process0.wait(timeout=timeout)
31
- process1.wait(timeout=timeout)
32
-
33
- # Check the return status of both processes
34
- if process0.returncode != 0:
35
- pytest.fail(f"Test {script_name} failed for RANK=0, {process0.returncode}")
36
- if process1.returncode != 0:
37
- pytest.fail(f"Test {script_name} failed for RANK=1, {process1.returncode}")
38
-
39
- except subprocess.TimeoutExpired:
40
- # If either process times out, terminate both and fail the test
41
- process0.terminate()
42
- process1.terminate()
43
- pytest.fail(f"Test {script_name} timed out")
44
- except Exception as e:
45
- pytest.fail(f"Test {script_name} failed with error: {str(e)}")
46
-
47
-
48
-# Define the test cases using pytest's parametrize
49
-@pytest.mark.parametrize(
50
- "script_name,timeout",
51
- [
52
- ("test_lookup_buffer.py", 60), # Second test case with a 60-second timeout
53
- ("test_send_recv.py", 120), # First test case with a 120-second timeout
54
- ],
55
-)
56
-def test_run_python_script(script_name, timeout):
57
- # Check the number of GPUs
58
- if torch.cuda.device_count() < 2:
59
- pytest.skip(f"Skipping test {script_name} because <2 GPUs are available")
60
-
61
- # Run the test if there are at least 2 GPUs
62
- run_python_script(script_name, timeout)
tests/kv_transfer/test_send_recv.py
deleted
-154
@@ -1,154 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import os
5
-import time
6
-
7
-import torch
8
-from tqdm import tqdm
9
-
10
-from vllm.config import KVTransferConfig
11
-from vllm.distributed.kv_transfer.kv_pipe.pynccl_pipe import PyNcclPipe
12
-
13
-
14
-def test_run(my_rank, pipe):
15
- print(f"rank {my_rank} test_run starts....")
16
- # test run
17
- x = torch.tensor([1]).to(pipe.device)
18
- y = torch.tensor([[2.0, 3.0, 4.0, 8.0]]).to(pipe.device)
19
- if my_rank == 0:
20
- pipe.send_tensor(x)
21
- print(f"rank {my_rank} sent tensor x")
22
- pipe.send_tensor(y)
23
- print(f"rank {my_rank} sent tensor y")
24
- x2 = pipe.recv_tensor()
25
- print(f"rank {my_rank} received x2 = ", x2)
26
- y2 = pipe.recv_tensor()
27
- print(f"rank {my_rank} received y2 = ", y2)
28
-
29
- else:
30
- x2 = pipe.recv_tensor()
31
- print(f"rank {my_rank} received x2 = ", x2)
32
- y2 = pipe.recv_tensor()
33
- print(f"rank {my_rank} received y2 = ", y2)
34
- pipe.send_tensor(x)
35
- print(f"rank {my_rank} sent tensor x")
36
- pipe.send_tensor(y)
37
- print(f"rank {my_rank} sent tensor y")
38
-
39
- assert torch.allclose(x, x2)
40
- assert torch.allclose(y, y2)
41
-
42
- print(f"rank {my_rank} test_run passed!")
43
-
44
-
45
-def stress_test(my_rank, pipe):
46
- print(f"rank {my_rank} stress_test starts....")
47
-
48
- tensors: list[torch.Tensor] = []
49
-
50
- torch.distributed.barrier()
51
- torch.manual_seed(0)
52
-
53
- for i in tqdm(range(500)):
54
- mean = torch.rand(1).item() * 100
55
- std = torch.rand(1).item() * 100
56
- size = torch.randint(900, 1000, (2,))
57
- x = torch.normal(mean * 1.0, std * 1.0, size=size.tolist()).to(pipe.device)
58
-
59
- # 5% probability of sending a None
60
- if torch.rand(1).item() < 0.05:
61
- tensors.append(None)
62
- tensors.append(None)
63
- tensors.append(None)
64
- else:
65
- tensors.append(x)
66
- tensors.append(x.mean().unsqueeze(0))
67
- tensors.append(x.std().unsqueeze(0))
68
-
69
- torch.distributed.barrier()
70
-
71
- for i in tqdm(range(500)):
72
- if my_rank == int((i % 10) > 3):
73
- pipe.send_tensor(tensors[3 * i])
74
- pipe.send_tensor(tensors[3 * i + 1])
75
- pipe.send_tensor(tensors[3 * i + 2])
76
- else:
77
- x = pipe.recv_tensor()
78
- mean = pipe.recv_tensor()
79
- std = pipe.recv_tensor()
80
-
81
- if x is None:
82
- assert mean is None
83
- assert std is None
84
- else:
85
- assert torch.allclose(x, tensors[3 * i])
86
- assert x.mean() == mean[0]
87
- assert x.std() == std[0]
88
-
89
- torch.distributed.barrier()
90
-
91
-
92
-def latency_test(my_rank, pipe, nelement, ntensor):
93
- latencies = []
94
-
95
- torch.distributed.barrier()
96
-
97
- for i in tqdm(range(500)):
98
- tensors = []
99
-
100
- if my_rank == 0:
101
- # create tensor
102
- tensors = [torch.rand(nelement).to(pipe.device) for _ in range(ntensor)]
103
-
104
- torch.distributed.barrier()
105
-
106
- if my_rank == 0:
107
- t = torch.tensor([time.time()], dtype=torch.float64).to(pipe.device)
108
- for tensor in tensors:
109
- pipe.send_tensor(tensor)
110
- pipe.send_tensor(t)
111
- else:
112
- for _ in range(ntensor):
113
- pipe.recv_tensor()
114
- t = pipe.recv_tensor()
115
- latencies.append(time.time() - t.item())
116
-
117
- torch.distributed.barrier()
118
-
119
- print("Latency test passed.")
120
- print("Latency:", torch.tensor(latencies).mean().item() * 1000, "ms")
121
-
122
-
123
-if __name__ == "__main__":
124
- my_rank = int(os.environ["RANK"])
125
-
126
- torch.distributed.init_process_group(
127
- backend="gloo",
128
- init_method="tcp://localhost:12398",
129
- world_size=2,
130
- rank=my_rank,
131
- )
132
-
133
- config = KVTransferConfig(
134
- kv_connector="P2pNcclConnector",
135
- kv_buffer_device="cuda",
136
- kv_buffer_size=1e9,
137
- kv_rank=my_rank,
138
- kv_role="kv_both", # this arg doesn't matter in this test
139
- kv_parallel_size=2,
140
- kv_ip="127.0.0.1",
141
- kv_port=12345,
142
- )
143
-
144
- pipe = PyNcclPipe(
145
- local_rank=my_rank,
146
- config=config,
147
- )
148
-
149
- test_run(my_rank, pipe)
150
-
151
- stress_test(my_rank, pipe)
152
-
153
- # Use this function if you want to test the latency of pipe impl.
154
- # latency_test(my_rank, pipe, 1024 * 8 * 128, 80)
tests/kv_transfer/test_send_recv.sh
deleted
-9
@@ -1,9 +0,0 @@
1
-#!/bin/bash
2
-
3
-RANK=0 python3 test_send_recv.py &
4
-PID0=$!
5
-RANK=1 python3 test_send_recv.py &
6
-PID1=$!
7
-
8
-wait $PID0
9
-wait $PID1
vllm/distributed/kv_transfer/kv_lookup_buffer/__init__.py
vllm/distributed/kv_transfer/kv_lookup_buffer/base.py
deleted
-179
@@ -1,179 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-This file contains a new class `KVLookupBufferBase` that allows developers to
5
-think of KV cache operations as inserting new KV cache entries (`insert`)
6
-into the lookup buffer and querying existing KV caches (`drop_select`)
7
-from the lookup buffer.
8
-
9
-This file also contains a new class `KVStoreBufferBase` that allows developers
10
-to manage the KVCache buffer as a simple key-value storage buffer with basic
11
-put/get operations.
12
-
13
-These classes above are abstracted behind class `KVCacheBufferBase`.
14
-"""
15
-
16
-from abc import ABC, abstractmethod
17
-
18
-import torch
19
-
20
-
21
-class KVCacheBufferBase(ABC):
22
- """
23
- Abstract base class for a KVCache buffer.
24
- """
25
-
26
- @abstractmethod
27
- def close(self) -> None:
28
- """Close the buffer and release resources.
29
-
30
- This method is responsible for cleaning up resources related to the
31
- KVCache buffer when it is no longer needed.
32
-
33
- Raises:
34
- NotImplementedError: This method must be implemented in subclasses.
35
- """
36
- raise NotImplementedError
37
-
38
-
39
-class KVLookupBufferBase(KVCacheBufferBase):
40
- """
41
- Abstract base class for a KVCache lookup buffer.
42
-
43
- This class provides an abstraction for a key-value (KV) cache lookup buffer.
44
-
45
- The key of the lookup buffer:
46
- - input_tokens: token IDs of the request
47
- - roi: a binary mask on top of input_tokens.
48
- - Purpose of roi: Since KV cache may only be available for a subset of
49
- tokens in the input (for example, when vLLM is connected to an external
50
- KV cache service), roi specifies the subset of tokens that the KV cache
51
- is associated with.
52
- - NOTE: roi can be further extended to describe which part of KV the
53
- current process is holding (each process may only hold a part of KV
54
- due to TP and PP). This is not implemented for now.
55
-
56
- The value of the lookup buffer:
57
- - key: the key tensor in the KV cache
58
- - value: the value tensor in the KV cache
59
- - hidden: the final hidden state generated by model forwarding. This allows
60
- vLLM to bypass further model forwarding by transmitting the hidden state.
61
- """
62
-
63
- @abstractmethod
64
- def insert(
65
- self,
66
- input_tokens: torch.Tensor,
67
- roi: torch.Tensor,
68
- key: torch.Tensor,
69
- value: torch.Tensor,
70
- hidden: torch.Tensor,
71
- ) -> None:
72
- """Insert into the lookup buffer.
73
-
74
- The functionality is similar to the following python statement
75
- ```
76
- buffer[input_tokens, roi] = [key, value, hidden]
77
- ```
78
-
79
- FIXME: in the future, we should only have two arguments, key and value,
80
- where key is a tensor dict and value is a tensor dict.
81
-
82
- FIXME: we should transmit both sampler outputs and the hidden states.
83
-
84
- Args:
85
- input_tokens (torch.Tensor): token IDs.
86
- roi (torch.Tensor): A binary mask on top of the input tokens
87
- key (torch.Tensor): The key tensor in the KV cache.
88
- value (torch.Tensor): The value tensor in the KV cache.
89
- hidden (torch.Tensor): The final hidden state tensor generated
90
- during model forwarding to bypass model
91
- forwarding.
92
-
93
- Raises:
94
- NotImplementedError: This method must be implemented in subclasses.
95
- """
96
- raise NotImplementedError
97
-
98
- @abstractmethod
99
- def drop_select(
100
- self, input_tokens: torch.Tensor | None, roi: torch.Tensor | None
101
- ) -> list[torch.Tensor | None]:
102
- """Select and *drop* KV cache entries from the lookup buffer.
103
-
104
- The functionality is similar to the following python statements
105
- ```
106
- ret = buffer.pop(input_tokens, roi)
107
- return ret
108
- ```
109
-
110
- If `input_tokens` and `roi` is `None`, it means selecting any of the
111
- KV caches in the buffer, return, and remove it from the buffer, useful
112
- when offloading KV cache to KV cache storage service.
113
-
114
- Args:
115
- input_tokens (torch.Tensor): token IDs.
116
- roi (torch.Tensor): A binary mask on top of the input tokens
117
-
118
- Returns:
119
- list[Optional[torch.Tensor]]: A list of tensors. Can be None.
120
-
121
- Raises:
122
- NotImplementedError: This method must be implemented in subclasses.
123
- """
124
- raise NotImplementedError
125
-
126
-
127
-class KVStoreBufferBase(KVCacheBufferBase):
128
- """
129
- Abstract base class for a KVCache storage buffer with key-value semantics.
130
- This class provides a simple key-value storage buffer abstract with basic
131
- put/get operations, which enables flexible KVCache transfer granular
132
- control.
133
-
134
- The functionality is similar to a distributed key-value store, where:
135
- - Key: A unique string identifier for the cached entry
136
- - Value:
137
- - Tensor to be stored and retrieved
138
- - None (indicating deletion or empty value)
139
- """
140
-
141
- @abstractmethod
142
- def put(
143
- self,
144
- key: str,
145
- value: torch.Tensor | None,
146
- ) -> None:
147
- """Store a key-value pair in the buffer.
148
-
149
- Args:
150
- key (str): Unique identifier for a tensor, this tensor could be the
151
- key cache tensor, value cache tensor, or hidden state tensor
152
- generated during model forwarding.
153
-
154
- value (Optional[torch.Tensor]): Tensor to be stored.
155
-
156
- Raises:
157
- NotImplementedError: This method must be implemented in subclasses.
158
- """
159
- raise NotImplementedError
160
-
161
- @abstractmethod
162
- def get(
163
- self,
164
- key: str,
165
- ) -> torch.Tensor | None:
166
- """Retrieve a value from the buffer by key.
167
-
168
- Args:
169
- key (str): Unique identifier for a tensor, this tensor could be the
170
- key cache tensor, value cache tensor, or hidden state tensor
171
- generated during model forwarding.
172
-
173
- Returns:
174
- Optional[torch.Tensor]: Stored tensor if exists, None otherwise.
175
-
176
- Raises:
177
- NotImplementedError: This method must be implemented in subclasses.
178
- """
179
- raise NotImplementedError
vllm/distributed/kv_transfer/kv_lookup_buffer/mooncake_store.py
deleted
-164
@@ -1,164 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-This file contains a new class `MooncakeStore` that allows developers to
5
-think of KV cache transfer operations as putting new KV cache entries
6
-into a remote KVStore-based lookup buffer and getting existing KV caches
7
-from this remote lookup buffer.
8
-"""
9
-
10
-import json
11
-import os
12
-from dataclasses import dataclass
13
-
14
-import torch
15
-from safetensors.torch import load as safetensors_load
16
-from safetensors.torch import save as safetensors_save
17
-
18
-from vllm.config import VllmConfig
19
-from vllm.distributed.kv_transfer.kv_lookup_buffer.base import KVStoreBufferBase
20
-from vllm.logger import init_logger
21
-
22
-DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 # 3.125 GiB
23
-DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 # 1.0 GiB
24
-
25
-logger = init_logger(__name__)
26
-
27
-
28
-@dataclass
29
-class MooncakeStoreConfig:
30
- local_hostname: str
31
- metadata_server: str
32
- global_segment_size: int
33
- local_buffer_size: int
34
- protocol: str
35
- device_name: str
36
- master_server_address: str
37
-
38
- @staticmethod
39
- def from_file(file_path: str) -> "MooncakeStoreConfig":
40
- """Load the config from a JSON file."""
41
- with open(file_path) as fin:
42
- config = json.load(fin)
43
- return MooncakeStoreConfig(
44
- local_hostname=config.get("local_hostname"),
45
- metadata_server=config.get("metadata_server"),
46
- global_segment_size=config.get(
47
- "global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE
48
- ),
49
- local_buffer_size=config.get(
50
- "local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE
51
- ),
52
- protocol=config.get("protocol", "tcp"),
53
- device_name=config.get("device_name", ""),
54
- master_server_address=config.get("master_server_address"),
55
- )
56
-
57
- @staticmethod
58
- def load_from_env() -> "MooncakeStoreConfig":
59
- """Load config from a file specified in the environment variable."""
60
- config_file_path = os.getenv("MOONCAKE_CONFIG_PATH")
61
- if config_file_path is None:
62
- raise ValueError(
63
- "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
64
- )
65
- return MooncakeStoreConfig.from_file(config_file_path)
66
-
67
-
68
-class MooncakeStore(KVStoreBufferBase):
69
- def __init__(
70
- self,
71
- config: VllmConfig,
72
- ):
73
- try:
74
- from mooncake.store import MooncakeDistributedStore
75
- except ImportError as e:
76
- raise ImportError(
77
- "Please install mooncake by following the instructions at "
78
- "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501
79
- "to run vLLM with MooncakeConnector."
80
- ) from e
81
-
82
- try:
83
- self.store = MooncakeDistributedStore()
84
- self.config = MooncakeStoreConfig.load_from_env()
85
- logger.info("Mooncake Configuration loaded successfully.")
86
-
87
- self.store.setup(
88
- self.config.local_hostname,
89
- self.config.metadata_server,
90
- self.config.global_segment_size,
91
- self.config.local_buffer_size,
92
- self.config.protocol,
93
- self.config.device_name,
94
- self.config.master_server_address,
95
- )
96
-
97
- except ValueError as e:
98
- logger.error("Configuration loading failed: %s", e)
99
- raise
100
- except Exception as exc:
101
- logger.error("An error occurred while loading the configuration: %s", exc)
102
- raise
103
-
104
- def close(self):
105
- # MooncakeDistributedStore will automatically call the destructor, so
106
- # it is unnecessary to close it manually.
107
- pass
108
-
109
- def put(
110
- self,
111
- key: str,
112
- value: torch.Tensor | None,
113
- ) -> None:
114
- # A message queue needs to be introduced before making it asynchronous.
115
- if value is not None:
116
- self._put_impl(key, value)
117
-
118
- def get(
119
- self,
120
- key: str,
121
- ) -> torch.Tensor | None:
122
- # A message queue needs to be introduced before making it asynchronous.
123
- value = self._get_impl(key)
124
- return value
125
-
126
- def _put_impl(
127
- self,
128
- key: str,
129
- value: torch.Tensor,
130
- ) -> None:
131
- """Put KVCache to Mooncake Store"""
132
- device_id = value.device.index if value.device.type == "cuda" else -1
133
- device_tensor = torch.tensor(device_id, dtype=torch.int32)
134
- value_bytes = safetensors_save({"tensor": value, "device_id": device_tensor})
135
- try:
136
- self.store.put(key, value_bytes)
137
- except TypeError as err:
138
- logger.error("Failed to put value into Mooncake Store: %s", err)
139
- raise TypeError("Mooncake Store Put Type Error.") from err
140
-
141
- def _get_impl(
142
- self,
143
- key: str,
144
- ) -> torch.Tensor | None:
145
- """Get KVCache from Mooncake Store"""
146
- try:
147
- data = self.store.get(key)
148
- except TypeError as err:
149
- logger.error("Failed to get value from Mooncake Store: %s", err)
150
- raise TypeError("Mooncake Store Get Type Error.") from err
151
-
152
- if data:
153
- loaded_tensors = safetensors_load(data)
154
- tensor = loaded_tensors["tensor"]
155
- device_id_tensor = loaded_tensors["device_id"]
156
- device_id = int(device_id_tensor.item())
157
- device = (
158
- torch.device("cuda", device_id)
159
- if device_id >= 0
160
- else torch.device("cpu")
161
- )
162
- return tensor.to(device)
163
-
164
- return None
vllm/distributed/kv_transfer/kv_lookup_buffer/simple_buffer.py
deleted
-242
@@ -1,242 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Implements a distributed key-value (KV) cache transfer mechanism.
5
-
6
-Key Features:
7
-- Distributed KV cache transmission using PyNccl pipes.
8
-- Non-blocking `insert`, blocking `drop_select`.
9
-- Use CPU signal pipe to avoid racing condition
10
-- Handles buffer size constraints and provide backpressure mechanism to
11
- stop the prefill instance when the decode instance is slow.
12
-"""
13
-
14
-import threading
15
-from collections import deque
16
-
17
-import torch
18
-
19
-from vllm.distributed.kv_transfer.kv_lookup_buffer.base import KVLookupBufferBase
20
-from vllm.distributed.kv_transfer.kv_pipe.base import KVPipeBase
21
-from vllm.logger import init_logger
22
-
23
-logger = init_logger(__name__)
24
-
25
-
26
-class SimpleBuffer(KVLookupBufferBase):
27
- def __init__(
28
- self, signal_pipe: KVPipeBase, data_pipe: KVPipeBase, buffer_size_thresh: float
29
- ):
30
- """
31
- signal_pipe: on CPU
32
-
33
- NOTE: on-device recv will block all threads in the process, making the
34
- KV cache producer unable to listen to new request while transmitting
35
- KV cache. Luckily CPU recv only blocks the current thread so we use
36
- CPU recv to listen to new request.
37
-
38
- data_pipe: on device (e.g. GPU)
39
- """
40
-
41
- self.buffer: deque[list[torch.Tensor]] = deque()
42
-
43
- self.buffer_size = 0
44
- self.buffer_size_threshold = buffer_size_thresh
45
- self.buffer_cv = threading.Condition()
46
- self.signal_pipe = signal_pipe
47
- self.data_pipe = data_pipe
48
- self.request_handling_thread: threading.Thread | None = None
49
-
50
- self.normal_signal = torch.tensor([0], device="cpu")
51
- self.end_signal = None
52
-
53
- def _matches(
54
- self,
55
- tokens_roi_sender: list[torch.Tensor],
56
- tokens_roi_recver: list[torch.Tensor],
57
- ):
58
- # tokens_roi_sender: tokens and roi of the producer (in the buffer)
59
- # tokens_roi_recver: tokens and roi of the consumer (query)
60
-
61
- tokens_sender = tokens_roi_sender[0]
62
- tokens_recver = tokens_roi_recver[0]
63
- roi_sender = tokens_roi_sender[1]
64
- roi_recver = tokens_roi_recver[1]
65
-
66
- if tokens_recver is None:
67
- # consumer sends an empty request
68
- # semantics: DROP SELECT * LIMIT 1
69
- # so any of the data in the buffer can be drop-selected
70
- return True
71
-
72
- # Assuming that roi is a binary mask on tokens
73
- tokens_sender = tokens_sender[roi_sender]
74
- tokens_recver = tokens_recver[roi_recver]
75
-
76
- # simple common prefix matching
77
- min_length = min(len(tokens_sender), len(tokens_recver))
78
- if torch.allclose(tokens_sender[:min_length], tokens_recver[:min_length]):
79
- return min_length
80
-
81
- return 0
82
-
83
- def _send_tensor_and_dec_size(self, tensor: torch.Tensor | None) -> None:
84
- assert tensor is not None, "Use self.data_pipe.send(None) instead"
85
- self.buffer_size -= tensor.element_size() * tensor.numel()
86
- if tensor.dtype == torch.bool:
87
- tensor = tensor.float()
88
- self.data_pipe.send_tensor(tensor)
89
-
90
- def _get_element_size(self, data: list | torch.Tensor | None):
91
- if isinstance(data, torch.Tensor):
92
- return data.element_size() * data.numel()
93
- if not data:
94
- # cannot perform `not data` on a tensor
95
- # so this check needs to go after the check above
96
- return 0
97
-
98
- raise AssertionError(f"Unknown data type {type(data)}")
99
-
100
- def _add_to_buffer(
101
- self,
102
- input_tokens: torch.Tensor,
103
- roi: torch.Tensor,
104
- key: torch.Tensor,
105
- value: torch.Tensor,
106
- hidden: torch.Tensor,
107
- ):
108
- if isinstance(input_tokens, torch.Tensor):
109
- input_tokens = input_tokens.clone()
110
- if isinstance(roi, torch.Tensor):
111
- roi = roi.clone()
112
- if isinstance(key, torch.Tensor):
113
- key = key.clone()
114
- if isinstance(value, torch.Tensor):
115
- value = value.clone()
116
- if isinstance(hidden, torch.Tensor):
117
- hidden = hidden.clone()
118
-
119
- buffer_item = [input_tokens, roi, key, value, hidden]
120
- data_size = sum([self._get_element_size(data) for data in buffer_item])
121
-
122
- with self.buffer_cv:
123
- if self.buffer_size + data_size > self.buffer_size_threshold:
124
- # log outside the while loop to avoid this message being logged
125
- # repeatedly.
126
- logger.debug("KV transfer buffer is full. Handling...")
127
- while self.buffer_size + data_size > self.buffer_size_threshold:
128
- self.buffer_cv.wait()
129
-
130
- self.buffer_size += data_size
131
- self.buffer.append(buffer_item)
132
- self.buffer_cv.notify()
133
-
134
- def _is_end_signal(self, signal):
135
- return signal is None
136
-
137
- def drop_select_handler(self):
138
- try:
139
- while True:
140
- signal = self.signal_pipe.recv_tensor()
141
- if self._is_end_signal(signal):
142
- logger.info("Received end signal!")
143
- break
144
-
145
- input_tokens = self.data_pipe.recv_tensor()
146
-
147
- roi = self.data_pipe.recv_tensor()
148
- assert roi is not None, (
149
- "Please provide the roi when sending drop-select request"
150
- )
151
- roi = roi > 0.5
152
- tokens_roi_recver = [input_tokens, roi]
153
-
154
- def is_buffer_available(
155
- tokens_roi_recver: list[torch.Tensor],
156
- ) -> bool:
157
- # perform input tokens and roi matching
158
- # FIXME: this matching is O(n), ideally it should be O(1)
159
- # but this buffer size won't (and shouldn't) be too large so
160
- # the fix is not urgent.
161
- for _ in range(len(self.buffer)):
162
- if self._matches(self.buffer[0], tokens_roi_recver) > 0:
163
- return True
164
- # rotate the element we just accessed to the end
165
- self.buffer.rotate(-1)
166
- return False
167
-
168
- with self.buffer_cv:
169
- while not is_buffer_available(tokens_roi_recver):
170
- logger.debug("KV transfer buffer is not available. Waiting...")
171
- self.buffer_cv.wait()
172
- # need to clone the tensor
173
- # in case the tensor is freed before sending finishes
174
- matched_item = self.buffer.popleft()
175
- for tensor in matched_item:
176
- self._send_tensor_and_dec_size(tensor)
177
- self.buffer_cv.notify()
178
-
179
- except RuntimeError as e:
180
- if "Connection closed by peer" not in str(e):
181
- raise e
182
-
183
- logger.debug("Closing drop_select_handler")
184
-
185
- def drop_select(
186
- self, input_tokens: torch.Tensor | None, roi: torch.Tensor | None
187
- ) -> list[torch.Tensor | None]:
188
- assert self.request_handling_thread is None, (
189
- "drop_select should be called by the KV cache consumer "
190
- "(e.g. the decode vLLM instance)"
191
- )
192
-
193
- if isinstance(input_tokens, torch.Tensor):
194
- input_tokens = input_tokens.clone()
195
- if isinstance(roi, torch.Tensor):
196
- roi = roi.clone().float()
197
-
198
- self.signal_pipe.send_tensor(self.normal_signal)
199
- self.data_pipe.send_tensor(input_tokens)
200
- self.data_pipe.send_tensor(roi)
201
-
202
- input_tokens = self.data_pipe.recv_tensor()
203
- roi = self.data_pipe.recv_tensor()
204
- if roi is not None:
205
- # convert from float tensor to bool tensor
206
- # as PyNccl does not support sending bool tensor
207
- roi = roi > 0.5
208
- key = self.data_pipe.recv_tensor()
209
- value = self.data_pipe.recv_tensor()
210
- hidden = self.data_pipe.recv_tensor()
211
-
212
- return [input_tokens, roi, key, value, hidden]
213
-
214
- def insert(
215
- self,
216
- input_tokens: torch.Tensor,
217
- roi: torch.Tensor,
218
- key: torch.Tensor,
219
- value: torch.Tensor,
220
- hidden: torch.Tensor,
221
- ) -> None:
222
- self._add_to_buffer(input_tokens, roi, key, value, hidden)
223
-
224
- # when calling the insert, the current process is a sender
225
- # need to launch the request handler and start listening to request.
226
- if self.request_handling_thread is None:
227
- self.request_handling_thread = threading.Thread(
228
- target=self.drop_select_handler
229
- )
230
- self.request_handling_thread.start()
231
-
232
- def close(self):
233
- if (
234
- hasattr(self, "request_handling_thread")
235
- and self.request_handling_thread is not None
236
- ):
237
- self.request_handling_thread.join()
238
-
239
- else:
240
- # TODO: have a explicit close signal and have a explicit way to
241
- # check if it's requester
242
- self.signal_pipe.send_tensor(self.end_signal)
vllm/distributed/kv_transfer/kv_pipe/__init__.py
vllm/distributed/kv_transfer/kv_pipe/base.py
deleted
-66
@@ -1,66 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-This file defines an interface `KVPipeBase`
5
-that provides an abstraction for sending and receiving tensors, or None, via
6
-distributed communications.
7
-
8
-All classes instantiated from this interface are assumed to be a FIFO pipe.
9
-
10
-If your distributed communication platform already supports key-value lookup,
11
-you can bypass this interface and directly start from `kv_lookup_buffer`.
12
-"""
13
-
14
-from abc import ABC, abstractmethod
15
-
16
-import torch
17
-
18
-
19
-class KVPipeBase(ABC):
20
- """
21
- This class provides an interface for sending and receiving tensors, or
22
- None, by distributed communications.
23
- """
24
-
25
- @abstractmethod
26
- def send_tensor(self, tensor: torch.Tensor | None) -> None:
27
- """Send a tensor, or None, via the pipe.
28
-
29
- Need to support sending None -- important for error handling.
30
-
31
- TODO: add a `key` argument so that we can use traditional
32
- key-value database as the distributed communication mechanism behind
33
- the pipe.
34
-
35
- Args:
36
- tensor (Optional[torch.Tensor]): The tensor to be sent. Can be None.
37
-
38
- Raises:
39
- NotImplementedError: This method must be implemented in subclasses.
40
- """
41
- raise NotImplementedError
42
-
43
- @abstractmethod
44
- def recv_tensor(self) -> torch.Tensor | None:
45
- """Receive a tensor (can be None) from the pipeline.
46
-
47
- Returns:
48
- Optional[torch.Tensor]: The tensor received from the pipeline. Can
49
- be None.
50
-
51
- Raises:
52
- NotImplementedError: This method must be implemented in subclasses.
53
- """
54
- raise NotImplementedError
55
-
56
- @abstractmethod
57
- def close(self) -> None:
58
- """Close the pipeline and release resources.
59
-
60
- This method is responsible for closing the communication pipeline
61
- and releasing any resources associated with it.
62
-
63
- Raises:
64
- NotImplementedError: This method must be implemented in subclasses.
65
- """
66
- raise NotImplementedError
vllm/distributed/kv_transfer/kv_pipe/mooncake_pipe.py
deleted
-295
@@ -1,295 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import json
5
-import os
6
-import struct
7
-from concurrent.futures import ThreadPoolExecutor
8
-from dataclasses import dataclass
9
-
10
-import torch
11
-import zmq
12
-from safetensors.torch import load as safetensors_load
13
-from safetensors.torch import save as safetensors_save
14
-
15
-from vllm.config.kv_transfer import KVTransferConfig
16
-from vllm.distributed.kv_transfer.kv_pipe.base import KVPipeBase
17
-from vllm.logger import init_logger
18
-from vllm.utils.network_utils import join_host_port, make_zmq_path, split_host_port
19
-
20
-logger = init_logger(__name__)
21
-NONE_INT = -150886311
22
-
23
-
24
-@dataclass
25
-class MooncakeTransferEngineConfig:
26
- prefill_url: str
27
- decode_url: str
28
- metadata_backend: str | None
29
- metadata_server: str
30
- protocol: str
31
- device_name: str
32
-
33
- @staticmethod
34
- def from_file(file_path: str) -> "MooncakeTransferEngineConfig":
35
- """Load the config from a JSON file."""
36
- with open(file_path) as fin:
37
- config = json.load(fin)
38
- return MooncakeTransferEngineConfig(
39
- prefill_url=config.get("prefill_url"),
40
- decode_url=config.get("decode_url"),
41
- metadata_backend=config.get("metadata_backend", None),
42
- metadata_server=config.get("metadata_server"),
43
- protocol=config.get("protocol", "tcp"),
44
- device_name=config.get("device_name", ""),
45
- )
46
-
47
- @staticmethod
48
- def load_from_env() -> "MooncakeTransferEngineConfig":
49
- """Load config from a file specified in the environment variable."""
50
- config_file_path = os.getenv("MOONCAKE_CONFIG_PATH")
51
- if config_file_path is None:
52
- raise ValueError(
53
- "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
54
- )
55
- return MooncakeTransferEngineConfig.from_file(config_file_path)
56
-
57
-
58
-class MooncakeTransferEngine:
59
- """Handles the transfer of data using mooncake_vllm_adaptor and ZeroMQ."""
60
-
61
- def __init__(self, kv_rank: int, local_rank: int):
62
- try:
63
- from mooncake.engine import TransferEngine
64
- except ImportError as e:
65
- raise ImportError(
66
- "Please install mooncake by following the instructions at "
67
- "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501
68
- "to run vLLM with MooncakeConnector."
69
- ) from e
70
-
71
- self.engine = TransferEngine()
72
- self.local_rank = local_rank
73
-
74
- try:
75
- self.config = MooncakeTransferEngineConfig.load_from_env()
76
- logger.info("Mooncake Configuration loaded successfully.")
77
- except ValueError as e:
78
- logger.error(e)
79
- raise
80
- except Exception as exc:
81
- logger.error("An error occurred while loading the configuration: %s", exc)
82
- raise
83
- prefill_host, base_prefill_port = split_host_port(self.config.prefill_url)
84
- decode_host, base_decode_port = split_host_port(self.config.decode_url)
85
-
86
- # Avoid ports conflict when running prefill and decode on the same node
87
- if prefill_host == decode_host and base_prefill_port == base_decode_port:
88
- base_decode_port = base_decode_port + 100
89
-
90
- prefill_port = base_prefill_port + self.local_rank
91
- decode_port = base_decode_port + self.local_rank
92
- self.prefill_url = join_host_port(prefill_host, prefill_port)
93
- self.decode_url = join_host_port(decode_host, decode_port)
94
-
95
- self.initialize(
96
- self.prefill_url if kv_rank == 0 else self.decode_url,
97
- self.config.metadata_server,
98
- self.config.protocol,
99
- self.config.device_name,
100
- self.config.metadata_backend,
101
- )
102
-
103
- self.remote_url = self.decode_url if kv_rank == 0 else self.prefill_url
104
-
105
- # Initialize ZeroMQ context and sockets
106
- self.context = zmq.Context() # type: ignore[attr-defined]
107
- self.sender_socket = self.context.socket(zmq.constants.PUSH)
108
- self.receiver_socket = self.context.socket(zmq.constants.PULL)
109
- self.sender_ack = self.context.socket(zmq.constants.PULL)
110
- self.receiver_ack = self.context.socket(zmq.constants.PUSH)
111
-
112
- self.buffer_cleaner = ThreadPoolExecutor(max_workers=1)
113
- self._setup_metadata_sockets(
114
- kv_rank, prefill_host, base_prefill_port, decode_host, base_decode_port
115
- )
116
-
117
- def _setup_metadata_sockets(
118
- self, kv_rank: int, p_host: str, p_port: int, d_host: str, d_port: int
119
- ) -> None:
120
- """Set up ZeroMQ sockets for sending and receiving data."""
121
- # Offsets < 8 are left for initialization in case tp and pp are enabled
122
- p_rank_offset = p_port + 8 + self.local_rank * 2
123
- d_rank_offset = d_port + 8 + self.local_rank * 2
124
- if kv_rank == 0:
125
- self.sender_socket.bind(make_zmq_path("tcp", p_host, p_rank_offset + 1))
126
- self.receiver_socket.connect(
127
- make_zmq_path("tcp", d_host, d_rank_offset + 1)
128
- )
129
- self.sender_ack.connect(make_zmq_path("tcp", d_host, d_rank_offset + 2))
130
- self.receiver_ack.bind(make_zmq_path("tcp", p_host, p_rank_offset + 2))
131
- else:
132
- self.receiver_socket.connect(
133
- make_zmq_path("tcp", p_host, p_rank_offset + 1)
134
- )
135
- self.sender_socket.bind(make_zmq_path("tcp", d_host, d_rank_offset + 1))
136
- self.receiver_ack.bind(make_zmq_path("tcp", d_host, d_rank_offset + 2))
137
- self.sender_ack.connect(make_zmq_path("tcp", p_host, p_rank_offset + 2))
138
-
139
- def initialize(
140
- self,
141
- local_hostname: str,
142
- metadata_server: str,
143
- protocol: str,
144
- device_name: str,
145
- metadata_backend: str | None,
146
- ) -> None:
147
- """Initialize the mooncake instance."""
148
- if metadata_backend is None:
149
- self.engine.initialize(
150
- local_hostname, metadata_server, protocol, device_name
151
- )
152
- else:
153
- supported_backend = ["etcd", "redis"]
154
- metadata_backend = metadata_backend.lower()
155
- if metadata_backend not in supported_backend:
156
- raise ValueError(
157
- "Mooncake Configuration error. `metadata_backend`"
158
- f" should be one of {supported_backend}."
159
- )
160
-
161
- self.engine.initialize_ext(
162
- local_hostname, metadata_server, protocol, device_name, metadata_backend
163
- )
164
-
165
- def allocate_managed_buffer(self, length: int) -> int:
166
- """Allocate a managed buffer of the specified length."""
167
- ret = self.engine.allocate_managed_buffer(length)
168
- if ret <= 0:
169
- logger.error("Allocation Return Error")
170
- raise Exception("Allocation Return Error")
171
- return ret
172
-
173
- def free_managed_buffer(self, buffer: int, length: int) -> int:
174
- """Free a previously allocated managed buffer."""
175
- return self.engine.free_managed_buffer(buffer, length)
176
-
177
- def transfer_sync(self, buffer: int, peer_buffer_address: int, length: int) -> int:
178
- """Synchronously transfer data to the specified address."""
179
- ret = self.engine.transfer_sync_read(
180
- self.remote_url, buffer, peer_buffer_address, length
181
- )
182
- if ret < 0:
183
- logger.error("Transfer Return Error")
184
- raise Exception("Transfer Return Error")
185
- return ret
186
-
187
- def write_bytes_to_buffer(self, buffer: int, user_data: bytes, length: int) -> int:
188
- """Write bytes to the allocated buffer."""
189
- return self.engine.write_bytes_to_buffer(buffer, user_data, length)
190
-
191
- def read_bytes_from_buffer(self, buffer: int, length: int) -> bytes:
192
- """Read bytes from the allocated buffer."""
193
- return self.engine.read_bytes_from_buffer(buffer, length)
194
-
195
- def wait_for_ack(self, src_ptr: int, length: int) -> None:
196
- """Asynchronously wait for ACK from the receiver."""
197
- ack = self.sender_ack.recv()
198
- if ack != b"ACK":
199
- logger.error("Failed to receive ACK from the receiver")
200
-
201
- self.free_managed_buffer(src_ptr, length)
202
-
203
- def send_bytes(self, user_data: bytes) -> None:
204
- """Send bytes to the remote process."""
205
- length = len(user_data)
206
- src_ptr = self.allocate_managed_buffer(length)
207
- self.write_bytes_to_buffer(src_ptr, user_data, length)
208
- self.sender_socket.send_multipart(
209
- [struct.pack("!Q", src_ptr), struct.pack("!Q", length)]
210
- )
211
- self.buffer_cleaner.submit(self.wait_for_ack, src_ptr, length)
212
-
213
- def recv_bytes(self) -> bytes:
214
- """Receive bytes from the remote process."""
215
- data = self.receiver_socket.recv_multipart()
216
- src_ptr = struct.unpack("!Q", data[0])[0]
217
- length = struct.unpack("!Q", data[1])[0]
218
- dst_ptr = self.allocate_managed_buffer(length)
219
- self.transfer_sync(dst_ptr, src_ptr, length)
220
- ret = self.read_bytes_from_buffer(dst_ptr, length)
221
-
222
- # Buffer cleanup
223
- self.receiver_ack.send(b"ACK")
224
- self.free_managed_buffer(dst_ptr, length)
225
-
226
- return ret
227
-
228
-
229
-class MooncakePipe(KVPipeBase):
230
- """MooncakeTransferEngine based Pipe implementation."""
231
-
232
- def __init__(
233
- self, local_rank: int, config: KVTransferConfig, device: str | None = None
234
- ):
235
- """Initialize the mooncake pipe and set related parameters."""
236
- self.config = config
237
- self.local_rank = local_rank
238
- self.kv_rank = self.config.kv_rank
239
- assert self.kv_rank is not None
240
- if device is None:
241
- self.device = self._select_device(self.config.kv_buffer_device)
242
- else:
243
- self.device = self._select_device(device)
244
-
245
- self.transfer_engine = MooncakeTransferEngine(self.kv_rank, self.local_rank)
246
- self.transport_thread: ThreadPoolExecutor | None = None
247
- self.none_tensor = torch.tensor([NONE_INT], device=self.device)
248
-
249
- def _select_device(self, device: str) -> torch.device:
250
- """Select available device (CUDA or CPU)."""
251
- logger.info("Selecting device: %s", device)
252
- if device == "cuda":
253
- return torch.device(f"cuda:{self.local_rank}")
254
- else:
255
- return torch.device("cpu")
256
-
257
- def tensor_hash(self, tensor: torch.Tensor) -> int:
258
- """Calculate the hash value of the tensor."""
259
- return hash(tensor.data_ptr())
260
-
261
- def _send_impl(self, tensor: torch.Tensor) -> None:
262
- """Implement the tensor sending logic using safetensors."""
263
- self.transfer_engine.send_bytes(safetensors_save({"tensor": tensor}))
264
-
265
- def _recv_impl(self) -> torch.Tensor:
266
- """Implement the tensor receiving logic using safetensors."""
267
- data = self.transfer_engine.recv_bytes()
268
- return safetensors_load(data)["tensor"].to(self.device)
269
-
270
- def send_tensor(self, tensor: torch.Tensor | None) -> None:
271
- """Send tensor to the target process."""
272
- if self.transport_thread is None:
273
- self.transport_thread = ThreadPoolExecutor(max_workers=1)
274
- tensor = tensor if tensor is not None else self.none_tensor
275
- assert len(tensor.shape) > 0
276
- self.transport_thread.submit(self._send_impl, tensor)
277
-
278
- def recv_tensor(self) -> torch.Tensor | None:
279
- """Receive tensor from other processes."""
280
- if self.transport_thread is None:
281
- self.transport_thread = ThreadPoolExecutor(max_workers=1)
282
- tensor = self.transport_thread.submit(self._recv_impl).result()
283
- if tensor.numel() == 1 and tensor.item() == NONE_INT:
284
- return None
285
- else:
286
- return tensor
287
-
288
- def close(self) -> None:
289
- """Cleanup logic when closing the pipe."""
290
- self.transfer_engine.sender_socket.close()
291
- self.transfer_engine.receiver_socket.close()
292
- self.transfer_engine.sender_ack.close()
293
- self.transfer_engine.receiver_ack.close()
294
- self.transfer_engine.context.term() # Terminate the ZMQ context
295
- logger.info("Closed the transfer engine and cleaned up resources.")
vllm/distributed/kv_transfer/kv_pipe/pynccl_pipe.py
deleted
-285
@@ -1,285 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-This module implements a PyNccl pipe for sending and receiving
5
-Optional[torch.Tensor] between distributed ranks with advanced
6
-communication features.
7
-
8
-Key Features:
9
-- Supports sending and receiving tensors with metadata
10
-- Handles both CUDA and CPU device communications
11
-- Implements a non-blocking tensor transfer mechanism
12
-- Manages buffer size and provides backpressure control
13
-- Supports distributed process groups with configurable parameters
14
-"""
15
-
16
-import threading
17
-import time
18
-from collections.abc import Callable
19
-from concurrent.futures import ThreadPoolExecutor
20
-
21
-import torch
22
-
23
-from vllm.config.kv_transfer import KVTransferConfig
24
-from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
25
-from vllm.distributed.kv_transfer.kv_pipe.base import KVPipeBase
26
-from vllm.distributed.utils import StatelessProcessGroup
27
-from vllm.logger import init_logger
28
-
29
-logger = init_logger(__name__)
30
-
31
-
32
-class BrokenPipeException(Exception):
33
- def __init__(self, message):
34
- self.message = message
35
- super().__init__(self.message)
36
-
37
-
38
-Metadata = dict[str, torch.Tensor | None]
39
-
40
-
41
-class PyNcclPipe(KVPipeBase):
42
- METADATA_LENGTH = 16
43
- MAX_TENSOR_DIMENSIONS = 14
44
- METADATA_DTYPE = torch.int64
45
-
46
- def __init__(
47
- self,
48
- local_rank: int,
49
- config: KVTransferConfig,
50
- device: str | None = None,
51
- port_offset: int = 0,
52
- ):
53
- self.config = config
54
- self.local_rank = local_rank
55
- self.kv_rank = self.config.kv_rank
56
- assert self.kv_rank is not None
57
- self.kv_parallel_size = self.config.kv_parallel_size
58
- if device is None:
59
- self.device = self._select_device(self.config.kv_buffer_device)
60
- else:
61
- self.device = self._select_device(device)
62
-
63
- # build distributed connection and send/recv implementation
64
- store_timeout = self.config.get_from_extra_config("store_timeout", 300)
65
- self.group = StatelessProcessGroup.create(
66
- host=self.config.kv_ip,
67
- port=self.config.kv_port + port_offset,
68
- rank=self.kv_rank,
69
- world_size=self.kv_parallel_size,
70
- store_timeout=store_timeout,
71
- )
72
- # add a barrier to make sure the connection is initiated properly
73
- self.group.barrier()
74
- impl = self._get_device_send_recv_impl(self.group)
75
- self.device_send_func, self.device_recv_func = impl
76
- # set target rank
77
- self.target_rank_for_send = (self.kv_rank + 1) % self.kv_parallel_size
78
- self.target_rank_for_recv = (self.kv_rank - 1) % self.kv_parallel_size
79
-
80
- # transportation-related variables
81
- self.transport_thread: ThreadPoolExecutor | None = None
82
- self.buffer_size = 0
83
- self.buffer_size_lock = threading.Lock()
84
- self.buffer_size_thresh = self.config.kv_buffer_size
85
-
86
- def _get_device_send_recv_impl(
87
- self, group: StatelessProcessGroup
88
- ) -> tuple[
89
- Callable[[torch.Tensor, int], None], Callable[[torch.Tensor, int], None]
90
- ]:
91
- send: Callable[[torch.Tensor, int], None]
92
- recv: Callable[[torch.Tensor, int], None]
93
- if self.device.type == "cuda":
94
- # use PyNCCL for send / recv
95
- comm = PyNcclCommunicator(group, device=self.local_rank)
96
- comm.disabled = False
97
- send, recv = comm.send, comm.recv # type: ignore
98
- else:
99
- # This send / recv implementation here is NOT intended to transfer
100
- # KV caches (and should NOT be repurposed to transfer KV caches).
101
- # Currently it is only used to transmit control-plane messages
102
- # for PyNcclBuffer.
103
- send = group.send_obj
104
-
105
- def my_recv(x, src):
106
- x[...] = group.recv_obj(src)
107
-
108
- recv = my_recv
109
-
110
- return send, recv
111
-
112
- def _select_device(self, device: str):
113
- logger.info("Selecting device: %s", device)
114
- if device == "cuda":
115
- return torch.device(f"cuda:{self.local_rank}")
116
- else:
117
- return torch.device("cpu")
118
-
119
- def _make_metadata(self, tensor: torch.Tensor | None) -> Metadata:
120
- """
121
- Create the metadata as a dictionary based on the input tensor.
122
-
123
- Args:
124
- tensor: The input tensor or None if no tensor is provided.
125
-
126
- Returns:
127
- metadata: A dictionary with the following keys:
128
- - "dtype": The data type of the tensor or None.
129
- - "shape": The shape of the tensor or None.
130
- """
131
- if tensor is None:
132
- return {"dtype": None, "shape": None}
133
- else:
134
- return {"dtype": tensor.dtype, "shape": tensor.shape}
135
-
136
- def _prepare_recv_buffer(self, metadata: Metadata) -> torch.Tensor:
137
- """
138
- Create a buffer to receive the tensor based on the provided metadata.
139
-
140
- Args:
141
- metadata: A dictionary with keys "dtype" and "shape",
142
- describing the tensor's data type and shape.
143
-
144
- Returns:
145
- buffer: A tensor of the specified type and shape,
146
- allocated on `self.device`.
147
- """
148
- return torch.empty(
149
- metadata["shape"], dtype=metadata["dtype"], device=self.device
150
- )
151
-
152
- def _send_metadata(self, metadata: Metadata):
153
- """
154
- Send the metadata dictionary to the target rank.
155
-
156
- Args:
157
- metadata: A dictionary with keys "dtype" and "shape".
158
- """
159
- self.group.send_obj(metadata, self.target_rank_for_send)
160
-
161
- def _recv_metadata(self) -> Metadata:
162
- """
163
- Receive the metadata dictionary from the target rank.
164
-
165
- Returns:
166
- metadata: A dictionary with keys "dtype" and "shape"
167
- describing the tensor.
168
- """
169
- return self.group.recv_obj(self.target_rank_for_recv)
170
-
171
- def _send_impl(self, tensor: torch.Tensor | None) -> None:
172
- """
173
- The actual implementation of sending the tensor and its metadata to the
174
- target rank.
175
-
176
- Args:
177
- tensor: The input tensor to be sent, or `None` if no tensor is
178
- being sent.
179
- """
180
- metadata = self._make_metadata(tensor)
181
- self._send_metadata(metadata)
182
- if tensor is not None:
183
- self.device_send_func(tensor.to(self.device), self.target_rank_for_send)
184
-
185
- def _recv_impl(self) -> torch.Tensor | None:
186
- """
187
- The actual implementation of receiving a tensor and its metadata from
188
- the target rank.
189
-
190
- Returns:
191
- buffer: The received tensor, or `None` if no tensor is received.
192
- """
193
- metadata = self._recv_metadata()
194
- if metadata["dtype"] is None:
195
- return None
196
- buffer = self._prepare_recv_buffer(metadata)
197
- self.device_recv_func(buffer, self.target_rank_for_recv)
198
-
199
- return buffer
200
-
201
- def send_tensor_wrapper(
202
- self, tensor: torch.Tensor | None, tensor_size: int
203
- ) -> None:
204
- """
205
- Wrapper for _send_impl to handle exceptions and update buffer size.
206
- """
207
- try:
208
- self._send_impl(tensor)
209
-
210
- with self.buffer_size_lock:
211
- self.buffer_size -= tensor_size
212
- except Exception as e:
213
- logger.error(
214
- "[rank%d]: Exception when trying to send %s, msg: %s",
215
- torch.distributed.get_rank(),
216
- str(tensor),
217
- str(e),
218
- )
219
- import traceback
220
-
221
- traceback.print_exc()
222
-
223
- def block_if_full(self):
224
- """
225
- Block the current thread if the buffer size is larger than the
226
- threshold.
227
- """
228
- while self.buffer_size > self.buffer_size_thresh:
229
- logger.debug("KV cache transfer pipe is full. Waiting...")
230
- time.sleep(0.05)
231
-
232
- def send_tensor(self, tensor: torch.Tensor | None) -> None:
233
- """
234
- Sends a tensor and its metadata to the destination rank in a
235
- non-blocking way.
236
-
237
- Args:
238
- tensor: The tensor to send, or `None` if no tensor is being sent.
239
- """
240
- if self.transport_thread is None:
241
- self.transport_thread = ThreadPoolExecutor(max_workers=1)
242
-
243
- if tensor is not None:
244
- tensor_size = tensor.element_size() * tensor.numel()
245
- else:
246
- tensor_size = 0
247
-
248
- self.block_if_full()
249
-
250
- with self.buffer_size_lock:
251
- self.buffer_size += tensor_size
252
-
253
- self.transport_thread.submit(self.send_tensor_wrapper, tensor, tensor_size)
254
-
255
- def recv_tensor(self) -> torch.Tensor | None:
256
- """
257
- Receives a tensor and its metadata from the source rank. Blocking call.
258
-
259
- Returns:
260
- The received tensor, or `None` if no tensor is received.
261
- """
262
- if self.transport_thread is None:
263
- self.transport_thread = ThreadPoolExecutor(max_workers=1)
264
-
265
- future = self.transport_thread.submit(self._recv_impl)
266
-
267
- try:
268
- tensor = future.result()
269
- except Exception as e:
270
- logger.error("Encountering exception in KV receiving thread")
271
- logger.error("%s", e)
272
- logger.error("My device: %s", self.device)
273
- import traceback
274
-
275
- traceback.print_exc()
276
- raise e
277
-
278
- return tensor
279
-
280
- def close(self):
281
- """
282
- Close the pipe and release associated resources.
283
- """
284
- if hasattr(self, "transport_thread") and self.transport_thread is not None:
285
- self.transport_thread.shutdown()