@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9f8cbfd8eb

Vectorize prep xfer list creation (#48209)

Signed-off-by: Ilia Yastrebov <iyastrebov@nvidia.com>

Ilia Yastrebov committed Jul 16, 2026 at 11:39 UTC 9f8cbfd8eb01de1310caa0b657840cbaf6c3b015

3 files changed +85 -74

tests/v1/kv_connector/unit/test_nixl_connector.py

+14 -7

@@ -13,6 +13,7 @@ from typing import Any, cast

13 from unittest.mock import MagicMock, patch

15 import msgspec

16 +import numpy as np

17 import pytest

18 import ray

19 import torch

@@ -750,7 +751,10 @@ class TestNixlHandshake:

751 worker.block_len_per_layer = [4096 * worker.block_size]

752 worker.num_blocks = 1

753 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks

753 - worker.src_blocks_data = [(0, worker.block_len_per_layer[0], worker.tp_rank)]

754 + worker.src_blocks_data = np.array(

755 + [(0, worker.block_len_per_layer[0], worker.tp_rank)],

756 + dtype=np.uint64,

757 + )

758 worker.num_descs = len(worker.src_blocks_data)

759

760 def check_handshake(remote_tp_size: int):

@@ -1124,8 +1128,8 @@ class TestNixlHandshake:

1128 == worker._mamba_ssm_size[1]

1129 )

1130

1127 - assert worker._build_fa_remote(plan, meta, block_size_ratio=1) == [

1128 - (0x1000 + local_block_len, local_block_len, 0)

1131 + assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [

1132 + [0x1000 + local_block_len, local_block_len, 0]

1133 ]

1134

1135 @patch(

@@ -1158,10 +1162,13 @@ class TestNixlHandshake:

1162 worker._region_is_mla = [False, True]

1163 worker.num_blocks = 1

1164 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks

1161 - worker.src_blocks_data = [

1162 - (0, fa_len, worker.tp_rank),

1163 - (0, idx_len, worker.tp_rank),

1164 - ]

1165 + worker.src_blocks_data = np.array(

1166 + [

1167 + (0, fa_len, worker.tp_rank),

1168 + (0, idx_len, worker.tp_rank),

1169 + ],

1170 + dtype=np.uint64,

1171 + )

1172 worker.num_descs = len(worker.src_blocks_data)

1173

1174 # D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio;

tests/v1/kv_connector/unit/test_tp_mapping.py

+13 -6

@@ -11,6 +11,7 @@ from __future__ import annotations

12 from types import SimpleNamespace

14 +import numpy as np

15 import pytest

17 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (

@@ -102,7 +103,10 @@ class TestBuildSrcSplitHandles:

103 )

104

105 worker = _make_mock_worker_for_splits((FullAttentionSpec,))

105 - src_blocks_data = [(0x2000 + i * 1024, 1024, 0) for i in range(8)]

106 + src_blocks_data = np.array(

107 + [(0x2000 + i * 1024, 1024, 0) for i in range(8)],

108 + dtype=np.uint64,

109 + )

110 num_descs = len(src_blocks_data)

111 splits = list(

112 worker._build_local_splits_from_plan(

@@ -135,11 +139,14 @@ class TestMambaPlanSplitHandles:

139

140 worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))

141 # 2 FA descs + 1 SSM desc

138 - src_blocks_data = [

139 - (1000, 200, 0), # FA desc 0

140 - (2000, 200, 0), # FA desc 1

141 - (3000, 400, 0), # SSM desc 0

142 - ]

142 + src_blocks_data = np.array(

143 + [

144 + (1000, 200, 0), # FA desc 0

145 + (2000, 200, 0), # FA desc 1

146 + (3000, 400, 0), # SSM desc 0

147 + ],

148 + dtype=np.uint64,

149 + )

150

151 splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2))

152

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+58 -61

@@ -160,7 +160,7 @@ class NixlBaseConnectorWorker:

160 def _build_local_splits_from_plan(

161 self,

162 plan: TPMapping,

163 - src_blocks_data: list[tuple[int, int, int]],

163 + src_blocks_data: np.ndarray,

164 num_fa_descs: int,

165 ) -> Iterator[list[tuple[int, int, int]]]:

166 """Build split handle data for P_TP > D_TP scenario.

@@ -188,12 +188,13 @@ class NixlBaseConnectorWorker:

188

189 # Per-FA-descriptor replicate flag, in _build_fa_local emission order.

190 fa_desc_replicated = self._fa_desc_replicated(num_fa_descs)

191 + src_blocks_list = src_blocks_data.tolist()

192

193 for p_idx, p_rank in enumerate(plan.all_source_ranks):

194 fa_slot = plan.rank_to_attention_slot.get(p_rank, 0)

195

196 handle: list[tuple[int, int, int]] = []

196 - for j, (addr, local_len, dev) in enumerate(src_blocks_data):

197 + for j, (addr, local_len, dev) in enumerate(src_blocks_list):

198 if j < num_fa_descs:

199 if fa_desc_replicated[j]:

200 # REPLICATE (MLA): whole block written on every rank.

@@ -1277,51 +1278,48 @@ class NixlBaseConnectorWorker:

1278 self,

1279 base_addresses: list[int],

1280 block_size_ratio: int,

1280 - ) -> list[tuple[int, int, int]]:

1281 + ) -> np.ndarray:

1282 """Build desc regions (conv sub-projections + ssm) per layer for

1282 - local mamba blocks with DS conv layout."""

1283 + local mamba blocks with DS conv layout, as an Nx3 uint64 array."""

1284 assert block_size_ratio == 1, (

1285 "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "

1286 f"Got block_size_ratio={block_size_ratio}."

1287 )

1288 + assert base_addresses, "Local KV cache base addresses must not be empty."

1289 assert self._conv_decomp is not None

1290 conv_offsets = self._conv_decomp.local_conv_offsets

1291 conv_size, ssm_size = self._mamba_ssm_size

1292 num_blocks = self._logical_num_blocks * block_size_ratio

1293 physical_per_logical = self._physical_blocks_per_logical_kv_block

1294 + device_id = self.device_id

1295 + block_arange = np.arange(num_blocks, dtype=np.uint64)

1296

1293 - result: list[tuple[int, int, int]] = []

1297 + parts: list[np.ndarray] = []

1298 for i, base_addr in enumerate(base_addresses):

1299 # Jump one page_size, but ssm page_size may be bigger when kernel

1300 # locks block size to a specific value (physical_per_logical scale).

1301 page_stride = (

1302 self.block_len_per_layer[i] // block_size_ratio * physical_per_logical

1303 )

1304 + blk_addrs = base_addr + block_arange * page_stride

1305 for off, sz in conv_offsets:

1301 - for blk in range(num_blocks):

1302 - result.append(

1303 - (base_addr + blk * page_stride + off, sz, self.device_id)

1304 - )

1306 + parts.append(self._stack_descs(blk_addrs + off, sz, device_id))

1307 # SSM temporal state follows the conv state.

1306 - for blk in range(num_blocks):

1307 - result.append(

1308 - (

1309 - base_addr + blk * page_stride + conv_size,

1310 - ssm_size,

1311 - self.device_id,

1312 - )

1313 - )

1314 - return result

1308 + parts.append(self._stack_descs(blk_addrs + conv_size, ssm_size, device_id))

1309 + return np.concatenate(parts)

1310

1311 def _build_mamba_remote(

1312 self,

1313 nixl_agent_meta: NixlAgentMetadata,

1314 tp_ratio: int,

1315 transfer_info: EngineTransferInfo,

1321 - ) -> list[tuple[int, int, int]]:

1316 + ) -> np.ndarray:

1317 """Build remote desc regions (conv sub-projections + ssm) per layer.

1318 For hetero-TP, each D rank reads only its sub-projection slice from

1324 - the P rank."""

1319 + the P rank. Returns an Nx3 uint64 array."""

1320 + assert nixl_agent_meta.kv_caches_base_addr, (

1321 + "Remote KV cache base addresses must not be empty."

1322 + )

1323 assert self._conv_decomp is not None

1324 effective_ratio = max(tp_ratio, 1)

1325 # Mamba conv state is always TP-sharded, even when attention KV

@@ -1338,35 +1336,41 @@ class NixlBaseConnectorWorker:

1336 remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical

1337 num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical

1338 device_id = nixl_agent_meta.device_id

1339 + block_arange = np.arange(num_blocks, dtype=np.uint64)

1340

1342 - result: list[tuple[int, int, int]] = []

1341 + parts: list[np.ndarray] = []

1342 # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case

1343 # block lengths vary across layers (e.g. MLA).

1344 for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):

1345 page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical

1346 + blk_addrs = base_addr + block_arange * page_stride

1347 for off, sz in conv_offsets:

1348 - for blk in range(num_blocks):

1349 - result.append((base_addr + blk * page_stride + off, sz, device_id))

1348 + parts.append(self._stack_descs(blk_addrs + off, sz, device_id))

1349 # SSM temporal state is also TP-sharded on the heads dimension.

1351 - for blk in range(num_blocks):

1352 - ssm_addr = (

1353 - base_addr

1354 - + blk * page_stride

1355 - + conv_size_remote

1356 - + local_offset * ssm_read_size

1357 - )

1358 - result.append((ssm_addr, ssm_read_size, device_id))

1359 - return result

1350 + ssm_addrs = blk_addrs + conv_size_remote + local_offset * ssm_read_size

1351 + parts.append(self._stack_descs(ssm_addrs, ssm_read_size, device_id))

1352 + return np.concatenate(parts)

1353 +

1354 + @staticmethod

1355 + def _stack_descs(addrs: np.ndarray, length: int, device_id: int) -> np.ndarray:

1356 + out = np.empty((addrs.shape[0], 3), dtype=np.uint64)

1357 + out[:, 0] = addrs

1358 + out[:, 1] = length

1359 + out[:, 2] = device_id

1360 + return out

1361

1362 def _build_fa_local(

1363 self,

1364 base_addresses: list[int],

1365 block_size_ratio: int,

1365 - ) -> list[tuple[int, int, int]]:

1366 - """Build local FA descriptors for all layers."""

1366 + ) -> np.ndarray:

1367 + """Build local FA descriptors for all layers as an Nx3 uint64 array."""

1368 assert self.transfer_topo is not None

1369 + assert base_addresses, "Local KV cache base addresses must not be empty."

1370 num_blocks = self.num_blocks * block_size_ratio

1369 - result: list[tuple[int, int, int]] = []

1371 + device_id = self.device_id

1372 + block_arange = np.arange(num_blocks, dtype=np.uint64)

1373 + parts: list[np.ndarray] = []

1374 for i, base_addr in enumerate(base_addresses):

1375 kv_block_len = (

1376 self.get_backend_aware_kv_block_len(

@@ -1375,20 +1379,21 @@ class NixlBaseConnectorWorker:

1379 // block_size_ratio

1380 )

1381 page_stride = self.block_len_per_layer[i] // block_size_ratio

1378 - for block_id in range(num_blocks):

1379 - block_offset = block_id * page_stride

1380 - addr = base_addr + block_offset

1381 - result.append((addr, kv_block_len, self.device_id))

1382 - return result

1382 + addrs = base_addr + block_arange * page_stride

1383 + parts.append(self._stack_descs(addrs, kv_block_len, device_id))

1384 + return np.concatenate(parts)

1385

1386 def _build_fa_remote(

1387 self,

1388 plan: TPMapping,

1389 nixl_agent_meta: NixlAgentMetadata,

1390 block_size_ratio: int,

1389 - ) -> list[tuple[int, int, int]]:

1390 - """Build remote FA descriptors for all layers."""

1391 + ) -> np.ndarray:

1392 + """Build remote FA descriptors for all layers as an Nx3 uint64 array."""

1393 assert self.transfer_topo is not None

1394 + assert nixl_agent_meta.kv_caches_base_addr, (

1395 + "Remote KV cache base addresses must not be empty."

1396 + )

1397 fa_group_idx = next(

1398 i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)

1399 )

@@ -1396,7 +1401,9 @@ class NixlBaseConnectorWorker:

1401 # per-rank offset; REPLICATE regions read the whole block once.

1402 split_reads = len(plan.source_ranks_per_group[fa_group_idx])

1403 num_blocks = nixl_agent_meta.num_blocks

1399 - result: list[tuple[int, int, int]] = []

1404 + device_id = nixl_agent_meta.device_id

1405 + block_arange = np.arange(num_blocks, dtype=np.uint64)

1406 + parts: list[np.ndarray] = []

1407 for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):

1408 replicated = self._is_region_replicated(i)

1409 # Read our whole local region size from remote..

@@ -1417,18 +1424,14 @@ class NixlBaseConnectorWorker:

1424 local_block_len = local_block_len // num_reads

1425

1426 page_size = nixl_agent_meta.block_lens[i]

1420 - for block_id in range(num_blocks):

1421 - block_offset = block_id * page_size

1422 - # For each block, grab the kv heads chunk belonging to current local

1423 - # tp rank of size local_block_len.

1424 - addr = base_addr + block_offset + rank_offset

1425 - result.append((addr, local_block_len, nixl_agent_meta.device_id))

1426 - return result

1427 + addrs = base_addr + rank_offset + block_arange * page_size

1428 + parts.append(self._stack_descs(addrs, local_block_len, device_id))

1429 + return np.concatenate(parts)

1430

1431 def register_local_xfer_handler(

1432 self,

1433 block_size: int,

1431 - ) -> tuple[int, list[tuple[int, int, int]]]:

1434 + ) -> tuple[int, np.ndarray]:

1435 """

1436 Function used for register local xfer handler with local block_size or

1437 Remote block_size.

@@ -1461,9 +1464,8 @@ class NixlBaseConnectorWorker:

1464 # remote has been seen. Currently we always register 4 regions

1465 # because local descs are created before knowing the remote TP.

1466 logger.debug("Registering local Mamba descriptors (4 regions/layer)")

1464 - blocks_data.extend(

1465 - self._build_mamba_local(local_base_addresses, block_size_ratio)

1466 - )

1467 + mamba = self._build_mamba_local(local_base_addresses, block_size_ratio)

1468 + blocks_data = np.concatenate([blocks_data, mamba])

1469

1470 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

1471 # NIXL_INIT_AGENT to be used for preparations of local descs.

@@ -1650,13 +1652,8 @@ class NixlBaseConnectorWorker:

1652 engine_id,

1653 remote_tp_rank,

1654 )

1653 - blocks_data.extend(

1654 - self._build_mamba_remote(

1655 - nixl_agent_meta,

1656 - tp_ratio,

1657 - transfer_info,

1658 - )

1659 - )

1655 + mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)

1656 + blocks_data = np.concatenate([blocks_data, mamba])

1657

1658 # Register with NIXL.

1659 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)