[Misc][PD] Nixl cleanup `get_backend_aware_kv_block_len` and `virtually_split_kv_in_blocks` (#49988)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Jul 27, 2026 at 18:51 UTC
2b465b2c42e6f7d37fbbc67956dc9741e832dc29
2 files changed
+44
-85
tests/v1/kv_connector/unit/test_nixl_connector.py
-12
@@ -1115,18 +1115,6 @@ class TestNixlHandshake:
1115
block_lens=[remote_block_len],
1116
)
1117
1118
- assert worker.get_backend_aware_kv_block_len(0, mamba_view=False) == (
1119
- local_block_len
1120
- )
1121
- assert (
1122
- worker.get_backend_aware_kv_block_len(0, first_split=True, mamba_view=True)
1123
- == worker._mamba_ssm_size[0]
1124
- )
1125
- assert (
1126
- worker.get_backend_aware_kv_block_len(0, first_split=False, mamba_view=True)
1127
- == worker._mamba_ssm_size[1]
1128
- )
1129
-
1118
assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
1119
[0x1000 + local_block_len, local_block_len, 0]
1120
]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+44
-73
@@ -96,7 +96,6 @@ class NixlBaseConnectorWorker:
96
physical_blocks_per_logical: int,
97
) -> np.ndarray:
98
"""Compute NIXL descriptor IDs for given block IDs."""
99
- num_fa_regions = self.num_regions
99
num_ssm_regions = 0
100
if self._has_mamba:
101
assert self._conv_decomp is not None
@@ -108,7 +107,7 @@ class NixlBaseConnectorWorker:
107
num_blocks = dst_num_blocks
108
if block_size_ratio is not None:
109
num_blocks = int(num_blocks * block_size_ratio)
111
- num_fa_descs = num_fa_regions * num_blocks
110
+ num_fa_descs = self.num_regions * num_blocks
111
112
# All-attention fast path: single vectorized broadcast.
113
if num_ssm_regions == 0:
@@ -119,7 +118,7 @@ class NixlBaseConnectorWorker:
118
# always differ (different areas). Therefore we can just flatten the
119
# block_ids and compute the descs ids for all groups at once.
120
block_arr = np.concatenate(block_ids)[None, :]
122
- region_ids = np.arange(num_fa_regions)[:, None]
121
+ region_ids = np.arange(self.num_regions)[:, None]
122
return (region_ids * num_blocks + block_arr).flatten()
123
124
# Compute desc ids per group using the right stride: FA descs have
@@ -130,7 +129,7 @@ class NixlBaseConnectorWorker:
129
for i, group in enumerate(block_ids):
130
group_arr = np.asarray(group)
131
if _is_attention_spec(self._group_spec_types[i]):
133
- fa_region_ids = np.arange(num_fa_regions)[:, None]
132
+ fa_region_ids = np.arange(self.num_regions)[:, None]
133
all_descs.append(
134
(fa_region_ids * num_blocks + group_arr[None, :]).flatten()
135
)
@@ -210,7 +209,7 @@ class NixlBaseConnectorWorker:
209
210
def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]:
211
"""Per-FA-descriptor replicate flag, in _build_fa_local emission order
213
- (region-major; K then optional V per region). Length ``num_fa_descs``.
212
+ (region-major; one desc per block, with K/V packed). Length ``num_fa_descs``.
213
"""
214
assert self.transfer_topo is not None
215
n_regions = len(self.block_len_per_layer)
@@ -1082,21 +1081,15 @@ class NixlBaseConnectorWorker:
1081
# With hybrid allocator, layers can share a kv cache tensor
1082
seen_base_addresses = []
1083
1085
- # Note(tms): I modified this from the original region setup code.
1086
- # K and V are now in different regions. Advantage is that we can
1087
- # elegantly support MLA and any cases where the K and V tensors
1088
- # are non-contiguous (it's not locally guaranteed that they will be)
1089
- # Disadvantage is that the encoded NixlAgentMetadata is now larger
1090
- # (roughly 8KB vs 5KB).
1091
- # Conversely for FlashInfer, K and V are registered in the same region
1092
- # to better exploit the memory layout (ie num_blocks is the first dim).
1084
+ # K and V are packed into the content dim, so each attention layer is a
1085
+ # single NIXL region whose block transfers as one unit. Mamba layers instead
1086
+ # register separate conv/ssm sub-regions (see `_build_mamba_local`).
1087
tensor_size_bytes = None
1088
1089
for layer_name, cache in xfer_buffers.items():
1096
- # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
1097
- # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
1098
- # However, physical page_size may differ when kernel requires a specific
1099
- # block size. This leads to SSM and FA layers having different num_blocks.
1090
+ # NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when
1091
+ # kernel requires a specific block size. This leads to SSM and FA layers
1092
+ # having different num_blocks.
1093
# `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
1094
layer_spec = self._layer_specs.get(layer_name)
1095
if layer_spec is None:
@@ -1107,7 +1100,7 @@ class NixlBaseConnectorWorker:
1100
)
1101
continue
1102
if isinstance(layer_spec, UniformTypeKVCacheSpecs):
1110
- # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
1103
+ # DSA Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
1104
layer_spec = layer_spec.kv_cache_specs[layer_name]
1105
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
1106
# the page_size assuming constant `self._logical_num_blocks`.
@@ -1270,7 +1263,32 @@ class NixlBaseConnectorWorker:
1263
block_size_ratio: int,
1264
) -> np.ndarray:
1265
"""Build desc regions (conv sub-projections + ssm) per layer for
1273
- local mamba blocks with DS conv layout, as an Nx3 uint64 array."""
1266
+ local mamba blocks with DS conv layout, as an Nx3 uint64 array.
1267
+
1268
+ A Mamba block interleaves conv and SSM state, which crucially differ in
1269
+ size, so the two are indexed as separate sub-regions. Attention blocks
1270
+ instead pack K and V into the content dim and transfer as a single unit.
1271
+ Reference diagram:
1272
+ KVCacheTensor (Shared)
1273
+ / \\
1274
+ / \\
1275
+ / \\
1276
+ Attention (FlashInfer) View Mamba View
1277
+ | |
1278
+ | |
1279
+ +-------------------+ +-------------------+
1280
+ | KVCacheTensor | | KVCacheTensor |
1281
+ | | | |
1282
+ |<----- page ------>| |<----- page ------->|
1283
+ | size | | size |
1284
+ | Key 0 | Val 0 | |Conv 0 | SSM 0 |
1285
+ | Key 1 | Val 1 | |Conv 1 | SSM 1 |
1286
+ | ... | ... | | ... | ... |
1287
+ | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 |
1288
+ | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 |
1289
+ +-------------------+ +--------------------+
1290
+ |1st_split-2nd_split| |1st_split-2nd_split |
1291
+ """
1292
assert block_size_ratio == 1, (
1293
"Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
1294
f"Got block_size_ratio={block_size_ratio}."
@@ -1362,15 +1380,11 @@ class NixlBaseConnectorWorker:
1380
block_arange = np.arange(num_blocks, dtype=np.uint64)
1381
parts: list[np.ndarray] = []
1382
for i, base_addr in enumerate(base_addresses):
1365
- kv_block_len = (
1366
- self.get_backend_aware_kv_block_len(
1367
- layer_idx=i, first_split=True, mamba_view=False
1368
- )
1369
- // block_size_ratio
1370
- )
1371
- page_stride = self.block_len_per_layer[i] // block_size_ratio
1372
- addrs = base_addr + block_arange * page_stride
1373
- parts.append(self._stack_descs(addrs, kv_block_len, device_id))
1383
+ # K/V are packed into the content dim, so the whole block transfers
1384
+ # as one unit: desc length equals the block stride.
1385
+ block_len = self.block_len_per_layer[i] // block_size_ratio
1386
+ addrs = base_addr + block_arange * block_len
1387
+ parts.append(self._stack_descs(addrs, block_len, device_id))
1388
return np.concatenate(parts)
1389
1390
def _build_fa_remote(
@@ -1397,9 +1411,7 @@ class NixlBaseConnectorWorker:
1411
for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1412
replicated = self._is_region_replicated(i)
1413
# Read our whole local region size from remote..
1400
- local_block_len = self.get_backend_aware_kv_block_len(
1401
- layer_idx=i, first_split=True, mamba_view=False
1402
- )
1414
+ local_block_len = self.block_len_per_layer[i]
1415
remote_kv_block_len = local_block_len // block_size_ratio
1416
if block_size_ratio > 1:
1417
# ..using remote kv_block_len as transfer unit
@@ -1521,8 +1533,7 @@ class NixlBaseConnectorWorker:
1533
)
1534
return self._remote_agents[engine_id][(0, remote_tp_rank)]
1535
1524
- # Compare physical regions, not self.num_regions (doubled by
1525
- # FlashInfer's virtual K/V split).
1536
+ # Number of physical regions registered locally (one per layer/tensor).
1537
num_local_regions = len(self.block_len_per_layer)
1538
if (
1539
self.pp_size > 1
@@ -2301,46 +2312,6 @@ class NixlBaseConnectorWorker:
2312
remote_block_ids[i] = remote_group[:num_blocks]
2313
return local_block_ids, remote_block_ids
2314
2304
- def get_backend_aware_kv_block_len(
2305
- self, layer_idx: int, first_split: bool = True, mamba_view: bool = False
2306
- ) -> int:
2307
- """
2308
- Get the block length for one K/V element (K and V have the same size).
2309
-
2310
- For FA and other backends, this is equal to the length of the whole
2311
- block, as K and V are in separate regions.
2312
- For FlashInfer, this is half the length of the whole block, as K and V
2313
- share the same region.
2314
- Similarly, for SSM-based models, state and conv are interleaved, but crucially
2315
- the their size differs.
2316
- Reference diagram:
2317
- KVCacheTensor (Shared)
2318
- / \\
2319
- / \\
2320
- / \\
2321
- Attention (FlashInfer) View Mamba View
2322
- | |
2323
- | |
2324
- +-------------------+ +-------------------+
2325
- | KVCacheTensor | | KVCacheTensor |
2326
- | | | |
2327
- |<----- page ------>| |<----- page ------->|
2328
- | size | | size |
2329
- | Key 0 | Val 0 | |Conv 0 | SSM 0 |
2330
- | Key 1 | Val 1 | |Conv 1 | SSM 1 |
2331
- | ... | ... | | ... | ... |
2332
- | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 |
2333
- | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 |
2334
- +-------------------+ +--------------------+
2335
- |1st_split-2nd_split| |1st_split-2nd_split |
2336
- """
2337
- assert self.transfer_topo is not None
2338
- if self.transfer_topo.virtually_split_kv_in_blocks and mamba_view:
2339
- block_len = self._mamba_ssm_size[not first_split]
2340
- else:
2341
- block_len = self.block_len_per_layer[layer_idx]
2342
- return block_len
2343
-
2315
def get_kv_connector_stats(self) -> KVConnectorStats | None:
2316
"""
2317
Get the KV transfer stats for the connector.