@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c9e665852a

Fork 0

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 13, 04:30 UTC

[NIXL] heterogeneous block_size support (#26759)

Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Chendi.Xue <chendi.xue@intel.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>

Chendi.Xue committed Nov 14, 2025 at 23:51 UTC c9e665852abbd42d7404a4f6dad7d47478ca95f8

3 files changed +257 -59

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -49,6 +49,8 @@ NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # Default to 1

49 PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}

50 DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}

51 GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}

52 +PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-16}

53 +DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-16}

55 # Find the git repository root directory

56 GIT_ROOT=$(git rev-parse --show-toplevel)

@@ -136,6 +138,7 @@ run_tests_for_model() {

138 vllm serve $model_name \

139 --port $PORT \

140 --enforce-eager \

141 + --block-size ${PREFILL_BLOCK_SIZE} \

142 --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \

143 --tensor-parallel-size $PREFILLER_TP_SIZE \

144 --kv-transfer-config '$KV_CONFIG'"

@@ -177,6 +180,7 @@ run_tests_for_model() {

180 vllm serve $model_name \

181 --port $PORT \

182 --enforce-eager \

183 + --block-size ${DECODE_BLOCK_SIZE} \

184 --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \

185 --kv-transfer-config '$KV_CONFIG'"

186

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -407,6 +407,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

407 # `self.kv_cache_layout` is only forced to HND when vllm engine

408 # is started. We mock HND here.

409 kv_cache_layout="HND",

410 + block_size=self.block_size,

411 ),

412 remote_tp_size=remote_tp_size,

413 )

@@ -652,6 +653,7 @@ class TestNixlHandshake:

653 block_lens=worker.block_len_per_layer,

654 attn_backend_name=worker.backend_name,

655 kv_cache_layout=mismatched_layout,

656 + block_size=worker.block_size,

657 )

658

659 with pytest.raises(RuntimeError):

@@ -706,6 +708,7 @@ class TestNixlHandshake:

708 block_lens=[i * 2 for i in worker.block_len_per_layer],

709 attn_backend_name=worker.backend_name,

710 kv_cache_layout="HND",

711 + block_size=worker.block_size,

712 )

713

714 # We don't check layout for homogeneous TP and MLA for now, as the

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+250 -59

@@ -108,6 +108,7 @@ class NixlAgentMetadata(KVConnectorHandshakeMetadata):

108 block_lens: list[int]

109 attn_backend_name: str

110 kv_cache_layout: str

111 + block_size: int

112

113

114 @dataclass

@@ -709,6 +710,9 @@ class NixlConnectorWorker:

710 self.is_mla or self._use_pallas or self.is_kv_layout_blocks_first

711 )

712

713 + block_size: int

714 + remote_block_size: dict[EngineId, int]

715 +

716 def tp_ratio(

717 self,

718 remote_tp_size: int,

@@ -725,6 +729,19 @@ class NixlConnectorWorker:

729 )

730 return self.tp_size // remote_tp_size

731

732 + def block_size_ratio(

733 + self,

734 + remote_block_size: int,

735 + ) -> float:

736 + """

737 + Calculate the block size ratio between local and remote TP.

738 + """

739 + assert self.block_size % remote_block_size == 0, (

740 + f"Local block size {self.block_size} is not divisible "

741 + f"by remote block size {remote_block_size} or vice versa."

742 + )

743 + return self.block_size // remote_block_size

744 +

745 def tp_ratio_from_engine_id(

746 self,

747 remote_engine_id: EngineId,

@@ -732,6 +749,13 @@ class NixlConnectorWorker:

749 remote_tp_size = self.remote_tp_size[remote_engine_id]

750 return self.tp_ratio(remote_tp_size)

751

752 + def block_size_ratio_from_engine_id(

753 + self,

754 + remote_engine_id: EngineId,

755 + ) -> float:

756 + remote_block_size = self.remote_block_size[remote_engine_id]

757 + return self.block_size_ratio(remote_block_size)

758 +

759 def is_kv_replicated(self, engine_id: EngineId) -> bool:

760 """

761 Whether the KV cache is replicated across TP workers due to the

@@ -866,6 +890,7 @@ class NixlConnectorWorker:

890

891 # nixl_prepped_dlist_handle.

892 self.src_xfer_side_handle: int = 0

893 + self.src_xfer_side_handles: dict[int, int] = {}

894 # Map of engine_id -> nixl_prepped_dlist_handle (int)].

895 self.dst_xfer_side_handles: dict[EngineId, int] = {}

896

@@ -925,6 +950,7 @@ class NixlConnectorWorker:

950 logger.debug("Detected kv cache layout %s", self.kv_cache_layout)

951

952 self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}

953 + self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}

954 # With heterogeneous TP, P must wait for all assigned D TP workers to

955 # finish reading before safely freeing the blocks.

956 self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)

@@ -936,6 +962,8 @@ class NixlConnectorWorker:

962 remote_tp_size=self._tp_size, # shared state

963 is_mla=self.use_mla,

964 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

965 + block_size=self.block_size,

966 + remote_block_size=self._block_size,

967 attn_backend=backend,

968 )

969 self._use_pallas = self.kv_topo._use_pallas

@@ -987,9 +1015,13 @@ class NixlConnectorWorker:

1015 )

1016

1017 # Register Remote agent.

1018 + assert metadata.block_size <= self.block_size, (

1019 + "nP > nD is not supported yet."

1020 + )

1021 remote_agent_name = self.add_remote_agent(

1022 metadata, p_remote_rank, remote_tp_size

1023 )

1024 +

1025 setup_agent_time = time.perf_counter()

1026 logger.debug(

1027 "NIXL handshake: add agent took: %s",

@@ -1217,43 +1249,10 @@ class NixlConnectorWorker:

1249 self.num_regions *= 2

1250

1251 # Register local/src descr for NIXL xfer.

1220 - blocks_data = []

1221 - for i, base_addr in enumerate(seen_base_addresses):

1222 - kv_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)

1223 - # NOTE With heter-TP, more blocks are prepared than what are

1224 - # needed as self.num_blocks >= nixl_agent_meta.num_blocks. We

1225 - # could create fewer, but then _get_block_descs_ids needs to

1226 - # select agent_meta.num_blocks instead of self.num_blocks for

1227 - # local descr, and that makes handling regular flow less clean.

1228 - for block_id in range(self.num_blocks):

1229 - block_offset = block_id * self.block_len_per_layer[i]

1230 - addr = base_addr + block_offset

1231 - # (addr, len, device id)

1232 - blocks_data.append((addr, kv_block_len, self.device_id))

1233 -

1234 - if self.kv_topo.is_kv_layout_blocks_first:

1235 - # Separate and interleave K/V regions to maintain the same

1236 - # descs ordering. This is needed for selecting contiguous heads

1237 - # when split across TP ranks.

1238 - for block_id in range(self.num_blocks):

1239 - block_offset = block_id * self.block_len_per_layer[i]

1240 - addr = base_addr + block_offset

1241 - # Register addresses for V cache (K registered first).

1242 - v_addr = addr + kv_block_len

1243 - blocks_data.append((v_addr, kv_block_len, self.device_id))

1244 - logger.debug(

1245 - "Created %s blocks for src engine %s and rank %s on device id %s",

1246 - len(blocks_data),

1247 - self.engine_id,

1248 - self.tp_rank,

1249 - self.device_id,

1250 - )

1252 + self.seen_base_addresses = seen_base_addresses

1253 + self.src_xfer_side_handle = self.register_local_xfer_handler(self.block_size)

1254

1252 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

1253 - # NIXL_INIT_AGENT to be used for preparations of local descs.

1254 - self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(

1255 - "NIXL_INIT_AGENT", descs

1256 - )

1255 + self.src_xfer_side_handles[self.block_size] = self.src_xfer_side_handle

1256

1257 # TODO(mgoin): Hybrid memory allocator is currently disabled for

1258 # models with local attention (Llama 4). Can remove this once enabled.

@@ -1289,8 +1288,62 @@ class NixlConnectorWorker:

1288 kv_cache_layout=self.kv_cache_layout

1289 if not self.use_host_buffer

1290 else self.host_buffer_kv_cache_layout,

1291 + block_size=self.block_size,

1292 + )

1293 +

1294 + def register_local_xfer_handler(

1295 + self,

1296 + block_size: int,

1297 + ) -> int:

1298 + """

1299 + Function used for register local xfer handler with local block_size or

1300 + Remote block_size.

1301 +

1302 + When local block_size is same as remote block_size, we use local block_size

1303 + to register local_xfer_handler during init.

1304 +

1305 + When remote block size is less than local block size, we need to use

1306 + register another local_xfer_handler using remote block len to ensure

1307 + data copy correctness.

1308 + """

1309 + block_size_ratio = self.block_size // block_size

1310 + blocks_data = []

1311 + for i, base_addr in enumerate(self.seen_base_addresses):

1312 + # The new block_len is using prefill block_len;

1313 + # and num_blocks is multiple with N

1314 + kv_block_len = (

1315 + self.get_backend_aware_kv_block_len(layer_idx=i) // block_size_ratio

1316 + )

1317 + block_len_per_layer = self.block_len_per_layer[i] // block_size_ratio

1318 + num_blocks = self.num_blocks * block_size_ratio

1319 + for block_id in range(num_blocks):

1320 + block_offset = block_id * block_len_per_layer

1321 + addr = base_addr + block_offset

1322 + # (addr, len, device id)

1323 + blocks_data.append((addr, kv_block_len, self.device_id))

1324 +

1325 + if self.kv_topo.is_kv_layout_blocks_first:

1326 + # Separate and interleave K/V regions to maintain the same

1327 + # descs ordering. This is needed for selecting contiguous heads

1328 + # when split across TP ranks.

1329 + for block_id in range(num_blocks):

1330 + block_offset = block_id * block_len_per_layer

1331 + addr = base_addr + block_offset

1332 + # Register addresses for V cache (K registered first).

1333 + v_addr = addr + kv_block_len

1334 + blocks_data.append((v_addr, kv_block_len, self.device_id))

1335 + logger.debug(

1336 + "Created %s blocks for src engine %s and rank %s on device id %s",

1337 + len(blocks_data),

1338 + self.engine_id,

1339 + self.tp_rank,

1340 + self.device_id,

1341 )

1342

1343 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

1344 + # NIXL_INIT_AGENT to be used for preparations of local descs.

1345 + return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)

1346 +

1347 def add_remote_agent(

1348 self,

1349 nixl_agent_meta: NixlAgentMetadata,

@@ -1349,6 +1402,8 @@ class NixlConnectorWorker:

1402 ### Register remote agent metadata

1403 if engine_id not in self._tp_size:

1404 self._tp_size[engine_id] = remote_tp_size

1405 + if engine_id not in self._block_size:

1406 + self._block_size[engine_id] = nixl_agent_meta.block_size

1407

1408 remote_agent_name = self.nixl_wrapper.add_remote_agent(

1409 nixl_agent_meta.agent_metadata

@@ -1359,6 +1414,13 @@ class NixlConnectorWorker:

1414

1415 # Create dst descs and xfer side handles. TP workers have same #blocks

1416 # so we only register once per engine_id.

1417 + # Example:

1418 + # block_size_ratio > 1:

1419 + # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|

1420 + # local origin:| 0| 1| 8| 12|

1421 + # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|

1422 + block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)

1423 +

1424 if engine_id not in self.dst_num_blocks:

1425 self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks

1426

@@ -1381,8 +1443,14 @@ class NixlConnectorWorker:

1443 # Register all remote blocks, but only the corresponding kv heads.

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

1445 kv_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)

1446 + remote_kv_block_len = kv_block_len // block_size_ratio

1447 + if block_size_ratio > 1:

1448 + # using remote kv_block_len as transfer unit

1449 + kv_block_len = remote_kv_block_len

1450 rank_offset = (

1385 - self.tp_rank % tp_ratio * kv_block_len if not replicates_kv_cache else 0

1451 + self.tp_rank % tp_ratio * remote_kv_block_len

1452 + if not replicates_kv_cache

1453 + else 0

1454 )

1455 for block_id in range(nixl_agent_meta.num_blocks):

1456 block_offset = block_id * nixl_agent_meta.block_lens[i]

@@ -1417,6 +1485,13 @@ class NixlConnectorWorker:

1485 remote_agent_name, descs

1486 )

1487

1488 + if block_size_ratio > 1:

1489 + # when prefill with smaller block_size, we need to init a

1490 + # new handler with same block_len to match

1491 + self.src_xfer_side_handles[nixl_agent_meta.block_size] = (

1492 + self.register_local_xfer_handler(nixl_agent_meta.block_size)

1493 + )

1494 +

1495 return remote_agent_name

1496

1497 def _validate_remote_agent_handshake(

@@ -1433,6 +1508,9 @@ class NixlConnectorWorker:

1508 assert nixl_agent_meta.attn_backend_name == self.backend_name

1509

1510 tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)

1511 + block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(

1512 + remote_engine_id

1513 + )

1514 assert tp_ratio > 0, "Decode TP cannot be smaller than prefill TP"

1515 assert not self._use_pallas or tp_ratio == 1, (

1516 "TPU (pallas_v1) DOES NOT support heterogeneous TP yet."

@@ -1463,33 +1541,26 @@ class NixlConnectorWorker:

1541 remote_block_len = nixl_agent_meta.block_lens[0]

1542 if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id):

1543 # With replicated KV cache, only the number of blocks can differ.

1466 - assert self.block_len_per_layer == nixl_agent_meta.block_lens, (

1467 - "KV cache sizes must match between P and D when replicated"

1468 - )

1469 - remote_block_size = remote_block_len // (self.slot_size_per_layer[0])

1544 + for i in range(len(self.block_len_per_layer)):

1545 + assert (

1546 + self.block_len_per_layer[i] // block_size_ratio

1547 + == nixl_agent_meta.block_lens[i]

1548 + ), "KV cache sizes must match between P and D when replicated"

1549 else:

1550 # When MLA is not used, this is a list of the same block length

1551 for block_len in nixl_agent_meta.block_lens:

1552 assert block_len == remote_block_len, (

1553 "All remote layers must have the same block size"

1554 )

1476 - remote_block_size = remote_block_len // (

1477 - self.slot_size_per_layer[0] * tp_ratio

1478 - )

1479 - if self.kv_topo.is_kv_layout_blocks_first:

1480 - # With flashinfer, KV are sent in the same message.

1481 - remote_block_size //= 2

1555

1483 - assert remote_block_len == self.block_len_per_layer[0] * tp_ratio, (

1556 + assert (

1557 + remote_block_len

1558 + == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio

1559 + ), (

1560 "Remote P worker KV layer cache must be of shape [2, N, "

1561 "local_kv_heads*tp_ratio, block_size, head_dim] and same dtype."

1562 )

1563

1488 - assert self.block_size == remote_block_size, (

1489 - "Remote P worker with different page/block size is not supported "

1490 - f"{self.block_size=}, {remote_block_size=}"

1491 - )

1492 -

1564 # TP workers have same #blocks.

1565 assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks

1566

@@ -1576,6 +1647,56 @@ class NixlConnectorWorker:

1647 )

1648 cache.index_copy_(0, indices, permuted_blocks)

1649

1650 + def blocksize_post_process(self, block_ids_per_ratio: dict[float, list[list[int]]]):

1651 + def _process_local_gt_remote(blocks_to_update, block_size_ratio):

1652 + n_kv_heads, block_size, head_size = blocks_to_update.shape[1:]

1653 + remote_block_size = block_size // block_size_ratio

1654 + n_blocks = block_size_ratio

1655 + # actual permute is to convert

1656 + # for local blocksize > remote blocksize

1657 + # ex: local blocksize = 16 tokens, remote blocksize = 4 tokens

1658 + # local block[0] = remote block[0, 1, 2, 3]

1659 + # remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|...

1660 + # local is |h0-b0..................|h1-b0..................|...

1661 + # permute is to:

1662 + # 1. view => view remote as n_blocks * remote_shape(H,remoteN,D)

1663 + # 2. permute => (H, nblocks, remoteN, D)

1664 + # 3. flatten => (H, localN, D)

1665 + permuted_blocks = (

1666 + blocks_to_update.reshape(

1667 + -1, n_blocks, n_kv_heads, remote_block_size, head_size

1668 + )

1669 + .permute(0, 2, 1, 3, 4)

1670 + .flatten(2, 3)

1671 + )

1672 + return permuted_blocks

1673 +

1674 + if len(self.device_kv_caches) == 0:

1675 + return

1676 + split_k_and_v = not (

1677 + self.use_mla or self._use_pallas or self.kv_topo.is_kv_layout_blocks_first

1678 + )

1679 + sample_cache = list(self.device_kv_caches.values())[0][0]

1680 + for block_size_ratio, block_ids_list in block_ids_per_ratio.items():

1681 + assert block_size_ratio > 1, "Only nP < nD supported currently."

1682 + block_ids_list = [[item for sublist in block_ids_list for item in sublist]]

1683 +

1684 + for block_ids in block_ids_list:

1685 + indices = torch.tensor(block_ids, device=sample_cache.device)

1686 +

1687 + for _, cache_or_caches in self.device_kv_caches.items():

1688 + cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]

1689 + for cache in cache_list:

1690 + blocks_to_update = cache.index_select(0, indices)

1691 + # because kv_cache is always using original layout NHD as

1692 + # virtual shape while stride can be either HND / NHD at

1693 + # initialization.

1694 + # we need to firstly get physical view of the tensor

1695 + permuted_blocks = _process_local_gt_remote(

1696 + blocks_to_update.permute(0, 2, 1, 3), block_size_ratio

1697 + ).permute(0, 2, 1, 3)

1698 + cache.index_copy_(0, indices, permuted_blocks)

1699 +

1700 def get_finished(self) -> tuple[set[str], set[str]]:

1701 """

1702 Get requests that are done sending or recving on this specific worker.

@@ -1599,6 +1720,7 @@ class NixlConnectorWorker:

1720 )

1721

1722 block_ids_to_permute = []

1723 + block_ids_for_blocksize_post_process = defaultdict(list)

1724 for req_id in done_recving:

1725 # clean up metadata for completed requests

1726 meta = self._recving_metadata.pop(req_id, None)

@@ -1607,6 +1729,20 @@ class NixlConnectorWorker:

1729 self.sync_recved_kv_to_device(req_id, meta)

1730 if self.enable_permute_local_kv:

1731 block_ids_to_permute += meta.local_physical_block_ids

1732 +

1733 + # post processing for heteroblocksize

1734 + block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(

1735 + meta.remote_engine_id

1736 + )

1737 + if (

1738 + not self.use_mla

1739 + and block_size_ratio > 1

1740 + and self.kv_cache_layout == "HND"

1741 + ):

1742 + block_ids_for_blocksize_post_process[block_size_ratio].append(

1743 + meta.local_block_ids

1744 + )

1745 + self.blocksize_post_process(block_ids_for_blocksize_post_process)

1746 if len(block_ids_to_permute) > 0:

1747 self.permute_device_kv(block_ids_to_permute)

1748

@@ -1781,6 +1917,24 @@ class NixlConnectorWorker:

1917 dst_engine_id: str,

1918 request_id: str,

1919 ):

1920 + block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)

1921 + if block_size_ratio > 1:

1922 + local_block_ids = self.get_mapped_blocks(

1923 + np.asarray(local_block_ids), block_size_ratio

1924 + )

1925 + if len(local_block_ids) > len(remote_block_ids):

1926 + # NOTE:

1927 + # get_mapped_blocks will always expand block_ids for n times.

1928 + # ex:

1929 + # prefill block_ids with block_size as 4:

1930 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

1931 + # Local decode block_ids with block_size as 16: [1, 2, 3]

1932 + # expland ecode block_ids with get_mapped_blocks from [1, 2, 3] to

1933 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

1934 + # Then we clip local to align with prefill

1935 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to

1936 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

1937 + local_block_ids = local_block_ids[: len(remote_block_ids)]

1938 # NOTE(rob): having the staging blocks be on the READER side is

1939 # not going to work well (since we will have to call rearrange tensors).

1940 # after we detect the txn is complete (which means we cannot make the

@@ -1823,7 +1977,10 @@ class NixlConnectorWorker:

1977 remote_block_ids = remote_block_ids[-num_local_blocks:]

1978

1979 # Get side handles.

1826 - local_xfer_side_handle = self.src_xfer_side_handle

1980 + remote_block_size = self.kv_topo.remote_block_size[dst_engine_id]

1981 + local_xfer_side_handle = self.src_xfer_side_handles.get(

1982 + remote_block_size, self.src_xfer_side_handle

1983 + )

1984 remote_xfer_side_handle = self.dst_xfer_side_handles[dst_engine_id]

1985

1986 # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from

@@ -1833,13 +1990,17 @@ class NixlConnectorWorker:

1990 # Get descs ids.

1991 local_block_descs_ids: np.ndarray

1992 remote_block_descs_ids: np.ndarray

1993 +

1994 if not self.block_window_per_layer:

1995 # Default case: assume global attention

1996 remote_block_descs_ids = self._get_block_descs_ids(

1839 - dst_engine_id, remote_block_ids

1997 + dst_engine_id,

1998 + remote_block_ids,

1999 )

2000 local_block_descs_ids = self._get_block_descs_ids(

1842 - self.engine_id, local_block_ids

2001 + self.engine_id,

2002 + local_block_ids,

2003 + block_size_ratio=block_size_ratio,

2004 )

2005 else:

2006 # TODO(mgoin): remove this once we have hybrid memory allocator

@@ -1860,10 +2021,15 @@ class NixlConnectorWorker:

2021

2022 # Get descs ids for the layer.

2023 layer_local_desc_ids = self._get_block_descs_ids(

1863 - self.engine_id, layer_local_block_ids, layer_idx

2024 + dst_engine_id,

2025 + layer_local_block_ids,

2026 + layer_idx,

2027 )

2028 layer_remote_desc_ids = self._get_block_descs_ids(

1866 - dst_engine_id, layer_remote_block_ids, layer_idx

2029 + self.engine_id,

2030 + layer_remote_block_ids,

2031 + layer_idx,

2032 + block_size_ratio=block_size_ratio,

2033 )

2034

2035 local_descs_list.append(layer_local_desc_ids)

@@ -1905,8 +2071,31 @@ class NixlConnectorWorker:

2071 self.nixl_wrapper.release_xfer_handle(handle)

2072 self._failed_recv_reqs.add(request_id)

2073

2074 + def get_mapped_blocks(self, block_ids, block_size_ratio):

2075 + """

2076 + Calculates the new set of block IDs by mapping every element

2077 + in the (potentially sparse) input array.

2078 + Example: block_ids=[0, 2], block_size_ratio=2

2079 + get_mapped_blocks 0 1 [2 3] 4 5

2080 + # remote is |h0-b0|h1-b0||h0-b1|h1-b1||h0-b1|h1-b1||

2081 + # local is |h0-b0......||h1-b0......||h2-b0........

2082 + local_block_ids 0 [1] 2

2083 + """

2084 + if block_ids.size == 0:

2085 + return np.array([], dtype=np.int64)

2086 +

2087 + start_ids = block_ids * block_size_ratio

2088 + offsets = np.arange(block_size_ratio)

2089 + mapped_2d = start_ids[:, None] + offsets[None, :]

2090 +

2091 + return mapped_2d.flatten().astype(np.int64)

2092 +

2093 def _get_block_descs_ids(

1909 - self, engine_id: str, block_ids: list[int], layer_idx: int | None = None

2094 + self,

2095 + engine_id: str,

2096 + block_ids: list[int],

2097 + layer_idx: int | None = None,

2098 + block_size_ratio: float | None = None,

2099 ) -> np.ndarray:

2100 """

2101 Get the descs ids for a set of block ids.

@@ -1929,6 +2118,8 @@ class NixlConnectorWorker:

2118 region_ids = np.arange(layer_idx, layer_idx + 1)

2119

2120 num_blocks = self.dst_num_blocks[engine_id]

2121 + if block_size_ratio is not None:

2122 + num_blocks = int(num_blocks * block_size_ratio)

2123

2124 # Compute the desc ids for each block.

2125 region_ids = region_ids[:, None]