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, 19:45 UTC
[Config][Disaggregated] Add timeout configuration for the torch.store and add KVTransferConfig.kv_connector_extra_config (#14367)
Signed-off-by: Mathis Felardos <mathis@mistral.ai>
Mathis Felardos committed
Mar 13, 2025 at 04:15 UTC
1bd32bc8dd685b1bcddc0e7408a46a7c637dae8a
4 files changed
+22
-11
vllm/config.py
+6
@@ -2837,6 +2837,9 @@ class KVTransferConfig(BaseModel):
2837
# The KV connector port, used to build distributed connection
2838
kv_port: int = 14579
2839
2840
+ # any extra config that the connector may need
2841
+ kv_connector_extra_config: dict[str, Any] = {}
2842
+
2843
def compute_hash(self) -> str:
2844
"""
2845
WARNING: Whenever a new field is added to this config,
@@ -2896,6 +2899,9 @@ class KVTransferConfig(BaseModel):
2899
return self.kv_connector is not None and \
2900
self.kv_role in ["kv_consumer", "kv_both"]
2901
2902
+ def get_from_extra_config(self, key, default) -> Any:
2903
+ return self.kv_connector_extra_config.get(key, default)
2904
+
2905
2906
class CompilationLevel:
2907
# constants for the levels of the compilation process
vllm/distributed/kv_transfer/kv_lookup_buffer/simple_buffer.py
+1
-1
@@ -6,7 +6,7 @@
6
- Distributed KV cache transmission using PyNccl pipes.
7
- Non-blocking `insert`, blocking `drop_select`.
8
- Use CPU signal pipe to avoid racing condition
9
- - Handles buffer size constraints and provide backpressure mechanism to
9
+ - Handles buffer size constraints and provide backpressure mechanism to
10
stop the prefill instance when the decode instance is slow.
11
"""
12
import threading
vllm/distributed/kv_transfer/kv_pipe/pynccl_pipe.py
+12
-10
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""
3
- This module implements a PyNccl pipe for sending and receiving
4
- Optional[torch.Tensor] between distributed ranks with advanced
3
+ This module implements a PyNccl pipe for sending and receiving
4
+ Optional[torch.Tensor] between distributed ranks with advanced
5
communication features.
6
7
Key Features:
@@ -59,11 +59,13 @@ class PyNcclPipe(KVPipeBase):
59
self.device = self._select_device(device)
60
61
# build distributed connection and send/recv implementation
62
+ store_timeout = self.config.get_from_extra_config("store_timeout", 300)
63
self.group = StatelessProcessGroup.create(
64
host=self.config.kv_ip,
65
port=self.config.kv_port + port_offset,
66
rank=self.kv_rank,
67
world_size=self.kv_parallel_size,
68
+ store_timeout=store_timeout,
69
)
70
# add a barrier to make sure the connection is initiated properly
71
self.group.barrier()
@@ -134,11 +136,11 @@ class PyNcclPipe(KVPipeBase):
136
Create a buffer to receive the tensor based on the provided metadata.
137
138
Parameters:
137
- - metadata: A dictionary with keys "dtype" and "shape", describing
139
+ - metadata: A dictionary with keys "dtype" and "shape", describing
140
the tensor's data type and shape.
141
142
Returns:
141
- - buffer: A tensor of the specified type and shape, allocated on
143
+ - buffer: A tensor of the specified type and shape, allocated on
144
self.device.
145
"""
146
return torch.empty(metadata["shape"],
@@ -159,18 +161,18 @@ class PyNcclPipe(KVPipeBase):
161
Receive the metadata dictionary from the target rank.
162
163
Returns:
162
- - metadata: A dictionary with keys "dtype" and "shape" describing
164
+ - metadata: A dictionary with keys "dtype" and "shape" describing
165
the tensor.
166
"""
167
return self.group.recv_obj(self.target_rank_for_recv)
168
169
def _send_impl(self, tensor: Optional[torch.Tensor]) -> None:
170
"""
169
- The actual implementation of sending the tensor and its metadata to the
171
+ The actual implementation of sending the tensor and its metadata to the
172
target rank.
173
174
Parameters:
173
- - tensor: The input tensor to be sent, or None if no tensor is
175
+ - tensor: The input tensor to be sent, or None if no tensor is
176
being sent.
177
"""
178
metadata = self._make_metadata(tensor)
@@ -181,7 +183,7 @@ class PyNcclPipe(KVPipeBase):
183
184
def _recv_impl(self) -> Optional[torch.Tensor]:
185
"""
184
- The actual implementation of receiving a tensor and its metadata from
186
+ The actual implementation of receiving a tensor and its metadata from
187
the target rank.
188
189
Returns:
@@ -213,7 +215,7 @@ class PyNcclPipe(KVPipeBase):
215
216
def block_if_full(self):
217
"""
216
- Block the current thread if the buffer size is larger than the
218
+ Block the current thread if the buffer size is larger than the
219
threshold.
220
"""
221
while self.buffer_size > self.buffer_size_thresh:
@@ -222,7 +224,7 @@ class PyNcclPipe(KVPipeBase):
224
225
def send_tensor(self, tensor: Optional[torch.Tensor]) -> None:
226
"""
225
- Sends a tensor and its metadata to the destination rank in a
227
+ Sends a tensor and its metadata to the destination rank in a
228
non-blocking way.
229
230
Parameters:
vllm/distributed/utils.py
+3
@@ -5,6 +5,7 @@
5
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py
6
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
7
import dataclasses
8
+import datetime
9
import pickle
10
import time
11
from collections import deque
@@ -217,6 +218,7 @@ class StatelessProcessGroup:
218
rank: int,
219
world_size: int,
220
data_expiration_seconds: int = 3600,
221
+ store_timeout: int = 300,
222
) -> "StatelessProcessGroup":
223
"""A replacement for `torch.distributed.init_process_group` that does not
224
pollute the global state.
@@ -238,6 +240,7 @@ class StatelessProcessGroup:
240
port=port,
241
world_size=world_size,
242
is_master=(rank == 0),
243
+ timeout=datetime.timedelta(seconds=store_timeout),
244
)
245
246
return StatelessProcessGroup(