[Misc][Nixl] Unify `_logical_to_remote_kernel_block_ids` (#48717)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Jul 16, 2026 at 12:36 UTC
a317bc573921f9c23861f0eba452277752bbdf97
5 files changed
+41
-64
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+7
-5
@@ -86,7 +86,9 @@ def test_logical_to_kernel_block_ids_with_hma():
86
87
# Test conversion: FA + SW group
88
logical_block_ids = [[0, 1, 2], [3, 4]]
89
- kernel_block_ids = worker._logical_to_kernel_block_ids(logical_block_ids)
89
+ kernel_block_ids = worker._logical_to_kernel_block_ids(
90
+ logical_block_ids, worker._physical_blocks_per_logical_kv_block
91
+ )
92
93
expected_kernel_block_ids = [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]]
94
assert kernel_block_ids == expected_kernel_block_ids, (
@@ -179,7 +181,7 @@ def test_read_blocks_for_req_expands_remote_ids(
181
"""_read_blocks_for_req must expand remote logical block IDs to kernel
182
block IDs when kernel block size != logical block size.
183
182
- The hot path always calls _logical_to_remote_kernel_block_ids with
184
+ The hot path always calls _logical_to_kernel_block_ids with
185
remote_info.remote_physical_blocks_per_logical (model-agnostic).
186
"""
187
from unittest.mock import MagicMock
@@ -1140,7 +1142,7 @@ def test_derive_mamba_conv_split(
1142
),
1143
],
1144
)
1143
-def test_logical_to_remote_kernel_block_ids(
1145
+def test_logical_to_kernel_block_ids_with_remote_ratio(
1146
mamba_enabled,
1147
swa_enabled,
1148
local_physical_per_logical,
@@ -1148,7 +1150,7 @@ def test_logical_to_remote_kernel_block_ids(
1150
logical_block_ids,
1151
expected_kernel_block_ids,
1152
):
1151
- """Verify _logical_to_remote_kernel_block_ids uses the remote
1153
+ """Verify _logical_to_kernel_block_ids uses the remote
1154
physical_per_logical for FA expansion, not the local one.
1155
1156
This was the root cause of silent accuracy corruption in Qwen3.5
@@ -1169,7 +1171,7 @@ def test_logical_to_remote_kernel_block_ids(
1171
swa_enabled=swa_enabled,
1172
)
1173
1172
- result = worker._logical_to_remote_kernel_block_ids(
1174
+ result = worker._logical_to_kernel_block_ids(
1175
logical_block_ids,
1176
remote_physical_per_logical,
1177
)
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+3
-3
@@ -487,7 +487,7 @@ class TestPushWriterStartLoadKv:
487
# path here.
488
w._send_heartbeats = lambda metadata: None
489
# Stub logical-to-kernel mapping used by reqs_to_recv.
490
- w._logical_to_kernel_block_ids = lambda x: x
490
+ w._logical_to_kernel_block_ids = lambda x, ratio: x
491
492
meta = NixlConnectorMetadata()
493
meta.push_registrations = {
@@ -775,7 +775,7 @@ class TestPushWriterNegative:
775
"""Empty metadata must not wake the writer or enqueue anything."""
776
w = _StubWriterWorker.fresh()
777
w._send_heartbeats = lambda metadata: None
778
- w._logical_to_kernel_block_ids = lambda x: x
778
+ w._logical_to_kernel_block_ids = lambda x, ratio: x
779
780
meta = NixlConnectorMetadata()
781
w.start_load_kv(meta)
@@ -944,7 +944,7 @@ class TestPushWriterMlaReplication:
944
rank_offset_factor=0,
945
)
946
}
947
- w._logical_to_remote_kernel_block_ids = lambda block_ids, ratio: block_ids
947
+ w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
948
w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}}
949
w.src_xfer_handles_by_block_size = {16: 2000}
950
w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}}
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+24
-51
@@ -1852,7 +1852,7 @@ class NixlBaseConnectorWorker:
1852
1853
for req_id, meta in metadata.reqs_to_save.items():
1854
meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
1855
- meta.local_block_ids
1855
+ meta.local_block_ids, self._physical_blocks_per_logical_kv_block
1856
)
1857
if logger.isEnabledFor(logging.DEBUG):
1858
logger.debug(
@@ -2207,30 +2207,36 @@ class NixlBaseConnectorWorker:
2207
2208
return mapped_2d.flatten().astype(np.int64)
2209
2210
- def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
2210
+ def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds:
2211
"""
2212
- Convert logical block ids to kernel physical block ids.
2212
+ Convert block ids to kernel physical block ids.
2213
This is required when the logical block size (the one set by the user)
2214
does not match the one required by the attn backend.
2215
+ `ratio` is the number of physical blocks per logical block.
2216
+ We always receive logical blocks from the engine, so we expand them here eg:
2217
+ logical block ids: [(SW-clipped) [1], (FA) [2, 3]], ratio=2
2218
+ physical block ids: [(SW-clipped) [2, 3], (FA) [4, 5, 6, 7]]
2219
"""
2216
- if self._physical_blocks_per_logical_kv_block == 1:
2220
+ if ratio == 1:
2221
# Noop when physical and logical block sizes are the same
2222
return block_ids
2219
- block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape(
2220
- 1, -1
2221
- )
2222
- # Mamba blocks have no logical<>physical discrepancy
2223
+ block_arange = np.arange(0, ratio).reshape(1, -1)
2224
+ # Mamba blocks have no logical<>physical discrepancy (block-size=1)
2225
group_specs = self.kv_cache_config.kv_cache_groups
2224
- return [
2225
- BlockTable.map_to_kernel_blocks(
2226
- np.array(group),
2227
- self._physical_blocks_per_logical_kv_block,
2228
- block_arange,
2229
- ).tolist()
2230
- if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
2231
- else group
2232
- for i, group in enumerate(block_ids)
2233
- ]
2226
+ physical_block_ids = []
2227
+ for i, group in enumerate(block_ids):
2228
+ spec = group_specs[i].kv_cache_spec
2229
+ if isinstance(spec, MambaSpec):
2230
+ physical_block_ids.append(group)
2231
+ else:
2232
+ physical_block_ids.append(
2233
+ BlockTable.map_to_kernel_blocks(
2234
+ np.array(group),
2235
+ ratio,
2236
+ block_arange,
2237
+ ).tolist()
2238
+ )
2239
+ return physical_block_ids
2240
2241
def _apply_prefix_caching(
2242
self,
@@ -2308,39 +2314,6 @@ class NixlBaseConnectorWorker:
2314
remote_block_ids[i] = remote_group[:num_blocks]
2315
return local_block_ids, remote_block_ids
2316
2311
- def _logical_to_remote_kernel_block_ids(
2312
- self, block_ids: BlockIds, remote_physical_per_logical: int
2313
- ) -> BlockIds:
2314
- """Map logical block IDs to physical kernel block IDs on the remote.
2315
-
2316
- Args:
2317
- block_ids: per-group lists of logical block IDs.
2318
- remote_physical_per_logical: remote engine's physical blocks
2319
- per logical block.
2320
-
2321
- Returns:
2322
- Same structure with FA groups expanded (each logical block L
2323
- becomes kernel blocks [L*remote_physical_per_logical, ..
2324
- L*remote_physical_per_logical +
2325
- remote_physical_per_logical - 1]).
2326
- Mamba groups are passed through unchanged.
2327
- """
2328
- if remote_physical_per_logical == 1:
2329
- return block_ids
2330
- remote_arange = np.arange(remote_physical_per_logical).reshape(1, -1)
2331
- group_specs = self.kv_cache_config.kv_cache_groups
2332
- result = [
2333
- BlockTable.map_to_kernel_blocks(
2334
- np.array(group),
2335
- remote_physical_per_logical,
2336
- remote_arange,
2337
- ).tolist()
2338
- if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
2339
- else group
2340
- for i, group in enumerate(block_ids)
2341
- ]
2342
- return result
2343
-
2317
def get_backend_aware_kv_block_len(
2318
self, layer_idx: int, first_split: bool = True, mamba_view: bool = False
2319
) -> int:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+2
-2
@@ -48,7 +48,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
48
"""
49
for req_id, meta in metadata.reqs_to_recv.items():
50
meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
51
- meta.local_block_ids
51
+ meta.local_block_ids, self._physical_blocks_per_logical_kv_block
52
)
53
assert meta.remote is not None
54
# Remote block IDs are kept logical here; expanded in
@@ -134,7 +134,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
134
remote_info = self.transfer_topo.get_engine_info(engine_id)
135
tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
136
137
- meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
137
+ meta.remote.block_ids = self._logical_to_kernel_block_ids(
138
meta.remote.block_ids,
139
remote_info.remote_physical_blocks_per_logical,
140
)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+5
-3
@@ -152,7 +152,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
152
# D-side: track reqs waiting for P to push.
153
for req_id, meta in metadata.reqs_to_recv.items():
154
meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
155
- meta.local_block_ids
155
+ meta.local_block_ids, self._physical_blocks_per_logical_kv_block
156
)
157
assert meta.remote is not None
158
remote_engine_id = meta.remote.engine_id
@@ -414,7 +414,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
414
# expands each side using the appropriate ratio.
415
logical_local = self._as_grouped_block_ids(local_block_ids)
416
logical_remote = self._as_grouped_block_ids(remote_block_ids)
417
- physical_local = self._logical_to_kernel_block_ids(logical_local)
417
+ physical_local = self._logical_to_kernel_block_ids(
418
+ logical_local, self._physical_blocks_per_logical_kv_block
419
+ )
420
421
push_meta = ReqMeta(
422
local_block_ids=logical_local,
@@ -501,7 +503,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
503
# Expand D's logical IDs using the ratio learned during the
504
# NIXL handshake. ``meta`` is freshly built by
505
# ``_do_start_push_kv`` so mutating it here is safe.
504
- meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
506
+ meta.remote.block_ids = self._logical_to_kernel_block_ids(
507
meta.remote.block_ids,
508
remote_info.remote_physical_blocks_per_logical,
509
)