[KV Connector] Fix NIXL mamba state pairing for multi-slot block tables (#50153)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 29, 2026 at 09:45 UTC
82553692b899059d2ec967ece7677a76bf2daab2
6 files changed
+163
-34
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+85
@@ -365,6 +365,28 @@ def test_apply_prefix_caching_mamba_hybrid(
365
[[6, 7, 8, 9], [99]],
366
id="fa_prefix_hit_and_ssm_trim",
367
),
368
+ # Multi-slot SSM ("all" mode): a local prefix hit leaves fewer local
369
+ # slots; the earlier remote slots are covered locally → remote tail.
370
+ pytest.param(
371
+ 10,
372
+ 10,
373
+ [list(range(10)), [5, 6]],
374
+ [list(range(10)), [1, 2, 3]],
375
+ [list(range(10)), [5, 6]],
376
+ [list(range(10)), [2, 3]],
377
+ id="ssm_multi_block_local_hit_tail",
378
+ ),
379
+ # Multi-slot SSM ("all" mode): the one trailing local position holds
380
+ # the token D recomputes itself → local head-clip.
381
+ pytest.param(
382
+ 10,
383
+ 10,
384
+ [list(range(10)), [4, 5, 6]],
385
+ [list(range(10)), [8, 9]],
386
+ [list(range(10)), [4, 5]],
387
+ [list(range(10)), [8, 9]],
388
+ id="ssm_multi_block_local_extra_head_clip",
389
+ ),
390
],
391
)
392
def test_apply_prefix_caching_ssm_prefix_cache_hit(
@@ -402,6 +424,28 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit(
424
)
425
426
427
+@pytest.mark.cpu_test
428
+def test_apply_prefix_caching_ssm_unpairable_slots_rejected():
429
+ """Local SSM slots can only exceed the remote ones by the position D
430
+ recomputes itself. A larger excess means the lists aren't
431
+ position-aligned: fail loudly rather than transfer into wrong slots."""
432
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
433
+ NixlConnectorWorker,
434
+ )
435
+ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
436
+
437
+ worker = object.__new__(NixlConnectorWorker)
438
+ worker._has_mamba = True
439
+ worker._physical_blocks_per_logical_kv_block = 10
440
+ worker._group_spec_types = (FullAttentionSpec, MambaSpec)
441
+ worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
442
+
443
+ with pytest.raises(AssertionError, match="unpairable SSM state slots"):
444
+ worker._apply_prefix_caching(
445
+ [list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10
446
+ )
447
+
448
+
449
@pytest.mark.cpu_test
450
@pytest.mark.parametrize(
451
"local_physical_per_logical,remote_physical_per_logical,"
@@ -1382,6 +1426,47 @@ def test_logical_to_kernel_block_ids_with_remote_ratio(
1426
)
1427
1428
1429
+@pytest.mark.cpu_test
1430
+def test_exchange_clipped_blocks_ssm_single_state():
1431
+ """In single-state cache modes, SSM lists are reduced to the running
1432
+ state slot: speculative scratch slots, null placeholders and the previous
1433
+ step's state carry nothing. Attention groups pass through untouched."""
1434
+ sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
1435
+ sched.blocks_per_sw = [0, 0]
1436
+ sched._ssm_spec_blocks = [None, 2]
1437
+ sched._ssm_state_slots_are_positional = False
1438
+
1439
+ # Align-mode list: null placeholders, state block, 2 speculative slots.
1440
+ clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 0, 7, 8, 9]))
1441
+ assert clipped == ([1, 2, 3], [7])
1442
+
1443
+ # Same, still holding the previous step's state block (freed a step later).
1444
+ assert sched.get_exchange_clipped_blocks(([1], [0, 6, 7, 8, 9]))[1] == [7]
1445
+
1446
+ # Default (mamba_block_size=max_model_len): state block, 2 scratch slots.
1447
+ assert sched.get_exchange_clipped_blocks(([1], [7, 8, 9]))[1] == [7]
1448
+
1449
+ # Scratch slots not allocated: the state slot still survives.
1450
+ assert sched.get_exchange_clipped_blocks(([1], [5]))[1] == [5]
1451
+
1452
+ # Non-mamba models pass through unchanged.
1453
+ fa_sched = make_nixl_scheduler(has_mamba=False)
1454
+ assert fa_sched.get_exchange_clipped_blocks(([1, 2],)) == ([1, 2],)
1455
+
1456
+
1457
+@pytest.mark.cpu_test
1458
+def test_exchange_clipped_blocks_ssm_positional_states():
1459
+ """In "all" mode every position holds a state, so only the speculative
1460
+ slots go; placeholders stay to keep the list position-indexed."""
1461
+ sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
1462
+ sched.blocks_per_sw = [0, 0]
1463
+ sched._ssm_spec_blocks = [None, 2]
1464
+ sched._ssm_state_slots_are_positional = True
1465
+
1466
+ clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 5, 6, 7, 8, 9]))
1467
+ assert clipped == ([1, 2, 3], [0, 5, 6, 7])
1468
+
1469
+
1470
# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
1471
1472
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+3
-3
@@ -106,9 +106,9 @@ class _BlocksMock:
106
107
108
def _stub_sw_clipping(scheduler) -> None:
109
- """Make ``get_sw_clipped_blocks`` a passthrough so tests don't need
110
- the full sliding-window machinery."""
111
- scheduler.get_sw_clipped_blocks = lambda block_ids: block_ids
109
+ """Make ``get_exchange_clipped_blocks`` a passthrough so tests don't
110
+ need the full sliding-window machinery."""
111
+ scheduler.get_exchange_clipped_blocks = lambda block_ids, clip_ssm=True: block_ids
112
113
114
# ----------------------------------------------------------------- #
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+52
-17
@@ -135,6 +135,20 @@ class NixlBaseConnectorScheduler:
135
for n_tokens, block_size in sw_sizes_tokens
136
]
137
138
+ # Trailing scratch slots that mamba managers co-allocate per request
139
+ # for speculative decoding; None for non-SSM groups.
140
+ self._ssm_spec_blocks = [
141
+ g.kv_cache_spec.num_speculative_blocks
142
+ if isinstance(g.kv_cache_spec, MambaSpec)
143
+ else None
144
+ for g in kv_cache_config.kv_cache_groups
145
+ ]
146
+ # Only "all" mode keeps a state per block position; the other modes
147
+ # keep a single running state in the last non-speculative slot.
148
+ self._ssm_state_slots_are_positional = (
149
+ vllm_config.cache_config.mamba_cache_mode == "all"
150
+ )
151
+
152
# Threshold to decide whether to compute kv cache locally
153
# or pull from a remote node: minimum number of remote
154
# tokens to amortize the xfer latencies
@@ -218,13 +232,25 @@ class NixlBaseConnectorScheduler:
232
# Clean up empty engines so we don't leak a key when remote dies.
233
del self._heartbeat_by_engine[engine_id]
234
221
- def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
222
- """
223
- Clip the number of blocks to the sliding window size for each kv cache group
224
- that employs SWA.
225
- This is necessary because the KV Cache manager initially allocates blocks for
226
- the entire sequence length, and successively cleans up blocks that are outside
227
- the window prior to the `request_finished_all_groups` hook.
235
+ def get_exchange_clipped_blocks(
236
+ self, block_ids: BlockIds, clip_ssm: bool = True
237
+ ) -> BlockIds:
238
+ """Clip a request's block lists down to the transferable blocks.
239
+
240
+ Sliding-window groups keep only the in-window tail: the KV cache
241
+ manager allocates blocks for the entire sequence length and cleans up
242
+ out-of-window blocks only prior to the `request_finished_all_groups`
243
+ hook.
244
+
245
+ SSM groups keep only their state-bearing slots: the trailing
246
+ speculative scratch slots always go, and in single-state cache modes
247
+ so does everything before the running state (null placeholders and
248
+ the previous step's superseded state). "all" mode keeps its remaining
249
+ slots, which the worker pairs position-wise.
250
+
251
+ Use this at every block-id exchange point. Pass ``clip_ssm=False``
252
+ for per-step partial lists (host-buffer save), where the SSM strip
253
+ does not apply.
254
"""
255
if len(block_ids) == 0 or not self._is_hma_required:
256
# No blocks to clip eg Full prefix cache hit or not a hybrid model.
@@ -235,15 +261,22 @@ class NixlBaseConnectorScheduler:
261
assert len(block_ids) == len(self.blocks_per_sw), (
262
"Number of KV cache groups must match"
263
)
238
- # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
239
- return tuple(
240
- [
241
- blocks[-self.blocks_per_sw[i] :]
242
- if self.blocks_per_sw[i] > 0
243
- else blocks
244
- for i, blocks in enumerate(block_ids)
245
- ]
246
- )
264
+ clipped = []
265
+ for i, blocks in enumerate(block_ids):
266
+ if n_sw := self.blocks_per_sw[i]:
267
+ blocks = blocks[-n_sw:]
268
+ elif (
269
+ clip_ssm
270
+ and blocks
271
+ and (n_spec_blocks := self._ssm_spec_blocks[i]) is not None
272
+ ):
273
+ if n_spec := min(n_spec_blocks, len(blocks) - 1):
274
+ blocks = blocks[:-n_spec]
275
+ if not self._ssm_state_slots_are_positional:
276
+ # Never empty: downstream reads that as a full prefix hit.
277
+ blocks = blocks[-1:]
278
+ clipped.append(blocks)
279
+ return tuple(clipped)
280
281
def set_xfer_handshake_metadata(
282
self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
@@ -380,7 +413,9 @@ class NixlBaseConnectorScheduler:
413
req = req_to_save
414
415
assert req.kv_transfer_params is not None
383
- clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
416
+ clipped_block_id_groups = self.get_exchange_clipped_blocks(
417
+ new_block_id_groups, clip_ssm=False
418
+ )
419
meta.add_new_req_to_save(
420
request_id=req_id,
421
local_block_ids=clipped_block_id_groups,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+19
-10
@@ -2439,16 +2439,25 @@ class NixlBaseConnectorWorker:
2439
for i, remote_group in enumerate(remote_block_ids):
2440
num_local_blocks = len(local_block_ids[i])
2441
num_remote_blocks = len(remote_group)
2442
- if (
2443
- _is_ssm_spec(self._group_spec_types[i])
2444
- and num_local_blocks < num_remote_blocks
2445
- ):
2446
- # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks
2447
- # prior to the last one are placeholders (null blocks). Mind that
2448
- # this doesn't really impact transfer, as we only still care about
2449
- # the last "block", the full in-place state.
2450
- assert num_local_blocks == 1, "SSM can only have one local block"
2451
- remote_block_ids[i] = remote_group[-num_local_blocks:]
2442
+ if _is_ssm_spec(self._group_spec_types[i]):
2443
+ if num_local_blocks == num_remote_blocks:
2444
+ continue
2445
+ # Only state-bearing slots reach here, single-state modes
2446
+ # just one (see get_exchange_clipped_blocks), so differing
2447
+ # counts mean position-indexed "all"-mode lists. A longer
2448
+ # remote list carries earlier positions the local side
2449
+ # already has (prefix hit) -> read its tail; a longer local
2450
+ # list holds the position D recomputes itself, which gets
2451
+ # no remote state.
2452
+ assert num_local_blocks - num_remote_blocks <= 1, (
2453
+ f"Group {i}: unpairable SSM state slots, "
2454
+ f"local={num_local_blocks} remote={num_remote_blocks}"
2455
+ )
2456
+ num_blocks = min(num_local_blocks, num_remote_blocks)
2457
+ if num_local_blocks < num_remote_blocks:
2458
+ remote_block_ids[i] = remote_group[-num_blocks:]
2459
+ else:
2460
+ local_block_ids[i] = local_block_ids[i][:num_blocks]
2461
elif (
2462
self._physical_blocks_per_logical_kv_block
2463
== remote_physical_per_logical
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+2
-2
@@ -155,7 +155,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
155
if num_external_tokens > 0
156
else ()
157
)
158
- local_block_ids = self.get_sw_clipped_blocks(
158
+ local_block_ids = self.get_exchange_clipped_blocks(
159
unhashed_local_block_ids
160
)
161
@@ -262,7 +262,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
262
# trimming down after allocating for the whole sequence length. Empty
263
# blocks are always at the start of the list.
264
# Here we "unpad" blocks to send the actual remote blocks to be read.
265
- block_ids = self.get_sw_clipped_blocks(block_ids)
265
+ block_ids = self.get_exchange_clipped_blocks(block_ids)
266
267
remote_num_tokens = request.num_computed_tokens
268
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
+2
-2
@@ -170,7 +170,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
170
request.request_id,
171
)
172
local_block_ids: BlockIds = blocks.get_unhashed_block_ids_all_groups()
173
- local_block_ids = self.get_sw_clipped_blocks(local_block_ids)
173
+ local_block_ids = self.get_exchange_clipped_blocks(local_block_ids)
174
175
# ``remote_*`` fields are P's coordinates (from D's perspective).
176
# ``decode_*`` fields are D's own info that P needs for the
@@ -271,7 +271,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
271
time.perf_counter() + self._kv_lease_duration
272
)
273
274
- block_ids = self.get_sw_clipped_blocks(block_ids)
274
+ block_ids = self.get_exchange_clipped_blocks(block_ids)
275
remote_num_tokens = request.num_computed_tokens
276
277
# Store finished blocks for worker-level matching with D