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
[Bugfix][Nixl] Fix kernel physical<>logical block_size issue (#28677)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Nov 14, 2025 at 15:40 UTC
96b23b8e3b5cd5d05345489a304e65f7ab53ef8e
3 files changed
+73
-17
tests/v1/worker/test_gpu_model_runner.py
+4
-2
@@ -985,8 +985,10 @@ def test_hybrid_block_table_initialization():
985
req_index = 0
986
block_table.append_row(kvcache_manager_blocks, req_index)
987
# Get expected kernel blocks from the implementation for verification.
988
- expected_kernel_blocks = block_table._map_to_kernel_blocks(
989
- np.array(kvcache_manager_blocks)
988
+ expected_kernel_blocks = block_table.map_to_kernel_blocks(
989
+ np.array(kvcache_manager_blocks),
990
+ block_table.blocks_per_kv_block,
991
+ block_table._kernel_block_arange,
992
)
993
# Verify block table state
994
assert block_table.num_blocks_per_row[req_index] == len(expected_kernel_blocks)
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+57
-10
@@ -49,6 +49,7 @@ from vllm.platforms import current_platform
49
from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
50
from vllm.v1.attention.backends.utils import get_kv_cache_layout
51
from vllm.v1.core.sched.output import SchedulerOutput
52
+from vllm.v1.worker.block_table import BlockTable
53
54
if TYPE_CHECKING:
55
from vllm.attention.backends.abstract import AttentionMetadata
@@ -112,6 +113,8 @@ class NixlAgentMetadata(KVConnectorHandshakeMetadata):
113
@dataclass
114
class ReqMeta:
115
local_block_ids: list[int]
116
+ # To be used when logical block size does not match the kernel block size
117
+ local_physical_block_ids: list[int]
118
remote_block_ids: list[int]
119
remote_host: str
120
remote_port: int
@@ -139,6 +142,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
142
assert load_remote_cache ^ save_to_host
143
_req = ReqMeta(
144
local_block_ids=local_block_ids,
145
+ local_physical_block_ids=local_block_ids,
146
remote_block_ids=kv_transfer_params["remote_block_ids"],
147
remote_engine_id=kv_transfer_params["remote_engine_id"],
148
remote_host=kv_transfer_params["remote_host"],
@@ -935,6 +939,7 @@ class NixlConnectorWorker:
939
attn_backend=backend,
940
)
941
self._use_pallas = self.kv_topo._use_pallas
942
+ self._physical_blocks_per_logical_kv_block = 1
943
944
def _nixl_handshake(
945
self,
@@ -1133,6 +1138,22 @@ class NixlConnectorWorker:
1138
if base_addr in seen_base_addresses:
1139
continue
1140
1141
+ # TODO (NickLucche): Get kernel_block_size in a cleaner way
1142
+ # NHD default "view" for non-MLA cache
1143
+ kernel_block_size = cache.shape[-2] if self.use_mla else cache.shape[-3]
1144
+
1145
+ if self.block_size != kernel_block_size:
1146
+ logger.info_once(
1147
+ "User-specified logical block size (%s) does not match"
1148
+ " physical kernel block size (%s). Using the latter. ",
1149
+ self.block_size,
1150
+ kernel_block_size,
1151
+ )
1152
+ self._physical_blocks_per_logical_kv_block = (
1153
+ self.block_size // kernel_block_size
1154
+ )
1155
+ self.block_size = kernel_block_size
1156
+
1157
seen_base_addresses.append(base_addr)
1158
curr_tensor_size_bytes = cache.numel() * cache.element_size()
1159
@@ -1479,7 +1500,7 @@ class NixlConnectorWorker:
1500
assert self.use_host_buffer
1501
assert self.copy_blocks is not None
1502
1482
- local_block_ids = meta.local_block_ids
1503
+ local_block_ids = meta.local_physical_block_ids
1504
self.copy_blocks(
1505
self.host_xfer_buffers,
1506
self.device_kv_caches,
@@ -1492,7 +1513,7 @@ class NixlConnectorWorker:
1513
"synced recved kv of request[%s] to device kv buffer,"
1514
"local_block_ids: %s. ",
1515
req_id,
1495
- ",".join(map(str, meta.local_block_ids)),
1516
+ ",".join(map(str, local_block_ids)),
1517
)
1518
1519
def save_kv_to_host(self, metadata: NixlConnectorMetadata):
@@ -1501,19 +1522,22 @@ class NixlConnectorWorker:
1522
assert self.copy_blocks is not None
1523
1524
for req_id, meta in metadata.reqs_to_save.items():
1525
+ meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
1526
+ meta.local_block_ids
1527
+ )
1528
if logger.isEnabledFor(logging.DEBUG):
1529
logger.debug(
1530
"save_load_kv for request[%s] to host xfer buffer."
1531
"local_block_ids: %s. ",
1532
req_id,
1509
- ",".join(map(str, meta.local_block_ids)),
1533
+ ",".join(map(str, meta.local_physical_block_ids)),
1534
)
1535
# blocking
1536
self.copy_blocks(
1537
self.device_kv_caches,
1538
self.host_xfer_buffers,
1515
- meta.local_block_ids,
1516
- meta.local_block_ids,
1539
+ meta.local_physical_block_ids,
1540
+ meta.local_physical_block_ids,
1541
"d2h",
1542
)
1543
@@ -1582,7 +1606,7 @@ class NixlConnectorWorker:
1606
if self.use_host_buffer:
1607
self.sync_recved_kv_to_device(req_id, meta)
1608
if self.enable_permute_local_kv:
1585
- block_ids_to_permute += meta.local_block_ids
1609
+ block_ids_to_permute += meta.local_physical_block_ids
1610
if len(block_ids_to_permute) > 0:
1611
self.permute_device_kv(block_ids_to_permute)
1612
@@ -1669,7 +1693,7 @@ class NixlConnectorWorker:
1693
req_id,
1694
xfer_state,
1695
)
1672
- # mark all blocks for this request as invalid
1696
+ # mark all (logical)blocks for this request as invalid
1697
if meta := self._recving_metadata.pop(req_id, None):
1698
self._invalid_block_ids.update(meta.local_block_ids)
1699
self._recving_metadata.pop(req_id, None)
@@ -1686,13 +1710,19 @@ class NixlConnectorWorker:
1710
We check for these trnxs to complete in each step().
1711
"""
1712
for req_id, meta in metadata.reqs_to_recv.items():
1713
+ meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
1714
+ meta.local_block_ids
1715
+ )
1716
+ meta.remote_block_ids = self._logical_to_kernel_block_ids(
1717
+ meta.remote_block_ids
1718
+ )
1719
remote_engine_id = meta.remote_engine_id
1720
logger.debug(
1721
"start_load_kv for request %s from remote engine %s. "
1722
"Num local_block_ids: %s. Num remote_block_ids: %s. ",
1723
req_id,
1724
remote_engine_id,
1695
- len(meta.local_block_ids),
1725
+ len(meta.local_physical_block_ids),
1726
len(meta.remote_block_ids),
1727
)
1728
# always store metadata for failure recovery
@@ -1740,7 +1770,7 @@ class NixlConnectorWorker:
1770
self._read_blocks(
1771
request_id=req_id,
1772
dst_engine_id=meta.remote_engine_id,
1743
- local_block_ids=meta.local_block_ids,
1773
+ local_block_ids=meta.local_physical_block_ids,
1774
remote_block_ids=meta.remote_block_ids,
1775
)
1776
@@ -1867,7 +1897,7 @@ class NixlConnectorWorker:
1897
"Marking blocks as invalid.",
1898
request_id,
1899
)
1870
- # mark all blocks for this request as invalid
1900
+ # mark all (logical) blocks for this request as invalid
1901
if meta := self._recving_metadata.get(request_id):
1902
self._invalid_block_ids.update(meta.local_block_ids)
1903
self.xfer_stats.record_failed_transfer()
@@ -1906,6 +1936,23 @@ class NixlConnectorWorker:
1936
descs_ids = region_ids * num_blocks + block_ids
1937
return descs_ids.flatten()
1938
1939
+ def _logical_to_kernel_block_ids(self, block_ids: list[int]) -> list[int]:
1940
+ """
1941
+ Convert logical block ids to kernel physical block ids.
1942
+ This is required when the logical block size (the one set by the user)
1943
+ does not match the one required by the attn backend.
1944
+ """
1945
+ if self._physical_blocks_per_logical_kv_block == 1:
1946
+ # Noop when physical and logical block sizes are the same
1947
+ return block_ids
1948
+ block_ids_np = np.array(block_ids)
1949
+ block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape(
1950
+ 1, -1
1951
+ )
1952
+ return BlockTable.map_to_kernel_blocks(
1953
+ block_ids_np, self._physical_blocks_per_logical_kv_block, block_arange
1954
+ ).tolist()
1955
+
1956
def get_backend_aware_kv_block_len(self, layer_idx: int):
1957
"""
1958
Get the block length for one K/V element (K and V have the same size).
vllm/v1/worker/block_table.py
+12
-5
@@ -98,7 +98,9 @@ class BlockTable:
98
return
99
100
if self.use_hybrid_blocks:
101
- block_ids = self._map_to_kernel_blocks(np.array(block_ids))
101
+ block_ids = self.map_to_kernel_blocks(
102
+ np.array(block_ids), self.blocks_per_kv_block, self._kernel_block_arange
103
+ )
104
105
num_blocks = len(block_ids)
106
start = self.num_blocks_per_row[row_idx]
@@ -188,7 +190,12 @@ class BlockTable:
190
self.block_table.gpu.fill_(0)
191
self.block_table.cpu.fill_(0)
192
191
- def _map_to_kernel_blocks(self, kv_manager_block_ids: np.ndarray) -> np.ndarray:
193
+ @staticmethod
194
+ def map_to_kernel_blocks(
195
+ kv_manager_block_ids: np.ndarray,
196
+ blocks_per_kv_block: int,
197
+ kernel_block_arange: np.ndarray,
198
+ ) -> np.ndarray:
199
"""Convert kv_manager_block_id IDs to kernel block IDs.
200
201
Example:
@@ -203,12 +210,12 @@ class BlockTable:
210
# kv_manager_block_id 1 → kernel block id [2, 3]
211
# kv_manager_block_id 2 → kernel block id [4, 5]
212
"""
206
- if not self.use_hybrid_blocks:
213
+ if blocks_per_kv_block == 1:
214
return kv_manager_block_ids
215
216
kernel_block_ids = (
210
- kv_manager_block_ids.reshape(-1, 1) * self.blocks_per_kv_block
211
- + self._kernel_block_arange
217
+ kv_manager_block_ids.reshape(-1, 1) * blocks_per_kv_block
218
+ + kernel_block_arange
219
)
220
221
return kernel_block_ids.reshape(-1)