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 14, 05:30 UTC
[Bugfix][V1][P/D]Fix the issue where repeated requests for the same input produce abnormal outputs for P2pNcclConnector (#23403)
Signed-off-by: Abatom <abzhonghua@gmail.com>
Zhonghua Deng committed
Aug 26, 2025 at 03:57 UTC
9188ae7cb5e78e6ecf95f41b587d3b279c231609
3 files changed
+26
-34
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+21
-4
@@ -245,16 +245,33 @@ class P2pNcclConnector(KVConnectorBase_V1):
245
246
assert self.p2p_nccl_engine is not None
247
248
+ def extract_kv_from_layer(
249
+ layer: torch.Tensor,
250
+ slot_mapping: torch.Tensor,
251
+ ) -> torch.Tensor:
252
+ """Extract the KV cache from the layer.
253
+
254
+ Assume the shape of the layer is (2, num_pages, page_size, xxx)
255
+ if MLA is not used, and (num_pages, page_size, xxx) otherwise.
256
+ """
257
+ if isinstance(attn_metadata, MLACommonMetadata):
258
+ num_pages, page_size = layer.shape[0], layer.shape[1]
259
+ return layer.reshape(num_pages * page_size, -1)[slot_mapping,
260
+ ...]
261
+ num_pages, page_size = layer.shape[1], layer.shape[2]
262
+ return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
263
+ ...]
264
+
265
connector_metadata = self._get_connector_metadata()
266
assert isinstance(connector_metadata, P2pNcclConnectorMetadata)
267
for request in connector_metadata.requests:
268
request_id = request.request_id
269
ip, port = self.parse_request_id(request_id, True)
270
remote_address = ip + ":" + str(port + self._rank)
254
- self.p2p_nccl_engine.send_tensor(
255
- request_id + "#" + layer_name, kv_layer, remote_address,
256
- request.slot_mapping,
257
- isinstance(attn_metadata, MLACommonMetadata))
271
+
272
+ kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping)
273
+ self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
274
+ kv_cache, remote_address)
275
276
def wait_for_save(self):
277
if self.is_producer:
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
+2
-28
@@ -62,8 +62,6 @@ class SendQueueItem:
62
tensor_id: str
63
remote_address: str
64
tensor: torch.Tensor
65
- slot_mapping: torch.Tensor
66
- is_mla: bool
65
66
67
class P2pNcclEngine:
@@ -202,8 +200,6 @@ class P2pNcclEngine:
200
tensor_id: str,
201
tensor: torch.Tensor,
202
remote_address: typing.Optional[str] = None,
205
- slot_mapping: torch.Tensor = None,
206
- is_mla: bool = False,
203
) -> bool:
204
if remote_address is None:
205
with self.recv_store_cv:
@@ -213,9 +209,7 @@ class P2pNcclEngine:
209
210
item = SendQueueItem(tensor_id=tensor_id,
211
remote_address=remote_address,
216
- tensor=tensor,
217
- slot_mapping=slot_mapping,
218
- is_mla=is_mla)
212
+ tensor=tensor)
213
214
if self.send_type == "PUT":
215
return self.send_sync(item)
@@ -433,9 +427,7 @@ class P2pNcclEngine:
427
if item.remote_address not in self.socks:
428
self.create_connect(item.remote_address)
429
436
- with self.send_stream:
437
- tensor = self.extract_kv_from_layer(item.is_mla, item.tensor,
438
- item.slot_mapping)
430
+ tensor = item.tensor
431
432
sock = self.socks[item.remote_address]
433
comm, rank = self.comms[item.remote_address]
@@ -548,21 +540,3 @@ class P2pNcclEngine:
540
self._send_thread.join()
541
if self._ping_thread is not None:
542
self._ping_thread.join()
551
-
552
- @staticmethod
553
- def extract_kv_from_layer(
554
- is_mla: bool,
555
- layer: torch.Tensor,
556
- slot_mapping: torch.Tensor,
557
- ) -> torch.Tensor:
558
- """Extract the KV cache from the layer.
559
- Assume the shape of the layer is (2, num_pages, page_size, xxx)
560
- if MLA is not used, and (num_pages, page_size, xxx) otherwise.
561
- """
562
- if is_mla:
563
- num_pages, page_size = layer.shape[0], layer.shape[1]
564
- return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]
565
-
566
- num_pages, page_size = layer.shape[1], layer.shape[2]
567
- return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
568
- ...]
vllm/distributed/kv_transfer/kv_connector/v1/p2p/tensor_memory_pool.py
+3
-2
@@ -99,8 +99,9 @@ class TensorMemoryPool:
99
addr=self.base_address)
100
self.free_lists[self.max_block_size][
101
initial_block.addr] = initial_block
102
- logger.debug("TensorMemoryPool, base_address:", self.base_address,
103
- self.base_address % self.max_block_size)
102
+
103
+ logger.debug("TensorMemoryPool, base_address:%d, max_block_size:%d",
104
+ self.base_address, self.max_block_size)
105
106
def allocate(self, size: int) -> int:
107
"""Allocates a memory block of at least the requested size.