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
[Feature][P/D]: Optimize NIXL Connector xfer Launch (#23887)
Signed-off-by: ycyaw66 <497410282@qq.com> Co-authored-by: ycyaw66 <497410282@qq.com>
WeiQing Chen committed
Sep 4, 2025 at 03:14 UTC
6adaed42f49ff683f80521b73daa3a3bde413baa
1 file changed
+18
-13
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+18
-13
@@ -14,6 +14,7 @@ from dataclasses import dataclass
14
from typing import TYPE_CHECKING, Any, Optional
15
16
import msgspec
17
+import numpy as np
18
import torch
19
import zmq
20
@@ -1191,8 +1192,8 @@ class NixlConnectorWorker:
1192
# workers will issue xfers to parts of the P worker remote kv caches.
1193
1194
# Get descs ids.
1194
- local_block_descs_ids: list[int] = []
1195
- remote_block_descs_ids: list[int] = []
1195
+ local_block_descs_ids: np.ndarray
1196
+ remote_block_descs_ids: np.ndarray
1197
if not self.block_window_per_layer:
1198
# Default case: assume global attention
1199
remote_block_descs_ids = self._get_block_descs_ids(
@@ -1202,6 +1203,8 @@ class NixlConnectorWorker:
1203
else:
1204
# TODO(mgoin): remove this once we have hybrid memory allocator
1205
# Optimization for models with local attention (Llama 4)
1206
+ local_descs_list = []
1207
+ remote_descs_list = []
1208
for layer_idx, block_window in enumerate(
1209
self.block_window_per_layer):
1210
# For each layer:
@@ -1221,8 +1224,11 @@ class NixlConnectorWorker:
1224
layer_remote_desc_ids = self._get_block_descs_ids(
1225
dst_engine_id, layer_remote_block_ids, layer_idx)
1226
1224
- local_block_descs_ids.extend(layer_local_desc_ids)
1225
- remote_block_descs_ids.extend(layer_remote_desc_ids)
1227
+ local_descs_list.append(layer_local_desc_ids)
1228
+ remote_descs_list.append(layer_remote_desc_ids)
1229
+
1230
+ local_block_descs_ids = np.concatenate(local_descs_list)
1231
+ remote_block_descs_ids = np.concatenate(remote_descs_list)
1232
1233
assert len(local_block_descs_ids) == len(remote_block_descs_ids)
1234
@@ -1247,14 +1253,14 @@ class NixlConnectorWorker:
1253
def _get_block_descs_ids(self,
1254
engine_id: str,
1255
block_ids: list[int],
1250
- layer_idx: Optional[int] = None) -> list[int]:
1256
+ layer_idx: Optional[int] = None) -> np.ndarray:
1257
"""
1258
Get the descs ids for a set of block ids.
1259
If layer_idx is provided, we use the region_ids for the given layer.
1260
Otherwise, we use all regions.
1261
"""
1262
if layer_idx is None:
1257
- region_ids = range(self.num_regions)
1263
+ region_ids = np.arange(self.num_regions)
1264
else:
1265
assert layer_idx < self.num_layers
1266
if self.num_layers < self.num_regions:
@@ -1262,20 +1268,19 @@ class NixlConnectorWorker:
1268
# the regions are organized as [K0, V0, K1, V1, ...]
1269
# and we select K_i and V_i
1270
assert 2 * self.num_layers == self.num_regions
1265
- region_ids = range(2 * layer_idx, 2 * layer_idx + 2)
1271
+ region_ids = np.arange(2 * layer_idx, 2 * layer_idx + 2)
1272
else:
1273
# Otherwise, we assume we have MLA and select i-th layer
1274
assert self.num_layers == self.num_regions
1269
- region_ids = range(layer_idx, layer_idx + 1)
1275
+ region_ids = np.arange(layer_idx, layer_idx + 1)
1276
1277
num_blocks = self.dst_num_blocks[engine_id]
1278
1279
# Compute the desc ids for each block.
1274
- descs_ids: list[int] = []
1275
- for reg_id in region_ids:
1276
- for block_id in block_ids:
1277
- descs_ids.append(reg_id * num_blocks + block_id)
1278
- return descs_ids
1280
+ region_ids = region_ids[:, None]
1281
+ block_ids = np.array(block_ids)[None, :]
1282
+ descs_ids = region_ids * num_blocks + block_ids
1283
+ return descs_ids.flatten()
1284
1285
def get_backend_aware_kv_block_len(self):
1286
"""