[KV Connector] Fix NIXL mamba state pairing for multi-slot block tables (#50153)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -365,6 +365,28 @@ def test_apply_prefix_caching_mamba_hybrid(
             [[6, 7, 8, 9], [99]],
             id="fa_prefix_hit_and_ssm_trim",
         ),
+        # Multi-slot SSM ("all" mode): a local prefix hit leaves fewer local
+        # slots; the earlier remote slots are covered locally → remote tail.
+        pytest.param(
+            10,
+            10,
+            [list(range(10)), [5, 6]],
+            [list(range(10)), [1, 2, 3]],
+            [list(range(10)), [5, 6]],
+            [list(range(10)), [2, 3]],
+            id="ssm_multi_block_local_hit_tail",
+        ),
+        # Multi-slot SSM ("all" mode): the one trailing local position holds
+        # the token D recomputes itself → local head-clip.
+        pytest.param(
+            10,
+            10,
+            [list(range(10)), [4, 5, 6]],
+            [list(range(10)), [8, 9]],
+            [list(range(10)), [4, 5]],
+            [list(range(10)), [8, 9]],
+            id="ssm_multi_block_local_extra_head_clip",
+        ),
     ],
 )
 def test_apply_prefix_caching_ssm_prefix_cache_hit(
@@ -402,6 +424,28 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit(
     )
 
 
+@pytest.mark.cpu_test
+def test_apply_prefix_caching_ssm_unpairable_slots_rejected():
+    """Local SSM slots can only exceed the remote ones by the position D
+    recomputes itself. A larger excess means the lists aren't
+    position-aligned: fail loudly rather than transfer into wrong slots."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
+
+    worker = object.__new__(NixlConnectorWorker)
+    worker._has_mamba = True
+    worker._physical_blocks_per_logical_kv_block = 10
+    worker._group_spec_types = (FullAttentionSpec, MambaSpec)
+    worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
+
+    with pytest.raises(AssertionError, match="unpairable SSM state slots"):
+        worker._apply_prefix_caching(
+            [list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10
+        )
+
+
 @pytest.mark.cpu_test
 @pytest.mark.parametrize(
     "local_physical_per_logical,remote_physical_per_logical,"
@@ -1382,6 +1426,47 @@ def test_logical_to_kernel_block_ids_with_remote_ratio(
     )
 
 
+@pytest.mark.cpu_test
+def test_exchange_clipped_blocks_ssm_single_state():
+    """In single-state cache modes, SSM lists are reduced to the running
+    state slot: speculative scratch slots, null placeholders and the previous
+    step's state carry nothing. Attention groups pass through untouched."""
+    sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
+    sched.blocks_per_sw = [0, 0]
+    sched._ssm_spec_blocks = [None, 2]
+    sched._ssm_state_slots_are_positional = False
+
+    # Align-mode list: null placeholders, state block, 2 speculative slots.
+    clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 0, 7, 8, 9]))
+    assert clipped == ([1, 2, 3], [7])
+
+    # Same, still holding the previous step's state block (freed a step later).
+    assert sched.get_exchange_clipped_blocks(([1], [0, 6, 7, 8, 9]))[1] == [7]
+
+    # Default (mamba_block_size=max_model_len): state block, 2 scratch slots.
+    assert sched.get_exchange_clipped_blocks(([1], [7, 8, 9]))[1] == [7]
+
+    # Scratch slots not allocated: the state slot still survives.
+    assert sched.get_exchange_clipped_blocks(([1], [5]))[1] == [5]
+
+    # Non-mamba models pass through unchanged.
+    fa_sched = make_nixl_scheduler(has_mamba=False)
+    assert fa_sched.get_exchange_clipped_blocks(([1, 2],)) == ([1, 2],)
+
+
+@pytest.mark.cpu_test
+def test_exchange_clipped_blocks_ssm_positional_states():
+    """In "all" mode every position holds a state, so only the speculative
+    slots go; placeholders stay to keep the list position-indexed."""
+    sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
+    sched.blocks_per_sw = [0, 0]
+    sched._ssm_spec_blocks = [None, 2]
+    sched._ssm_state_slots_are_positional = True
+
+    clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 5, 6, 7, 8, 9]))
+    assert clipped == ([1, 2, 3], [0, 5, 6, 7])
+
+
 # ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
 
 

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -106,9 +106,9 @@ class _BlocksMock:
 
 
 def _stub_sw_clipping(scheduler) -> None:
-    """Make ``get_sw_clipped_blocks`` a passthrough so tests don't need
-    the full sliding-window machinery."""
-    scheduler.get_sw_clipped_blocks = lambda block_ids: block_ids
+    """Make ``get_exchange_clipped_blocks`` a passthrough so tests don't
+    need the full sliding-window machinery."""
+    scheduler.get_exchange_clipped_blocks = lambda block_ids, clip_ssm=True: block_ids
 
 
 # ----------------------------------------------------------------- #

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

@@ -135,6 +135,20 @@ class NixlBaseConnectorScheduler:
             for n_tokens, block_size in sw_sizes_tokens
         ]
 
+        # Trailing scratch slots that mamba managers co-allocate per request
+        # for speculative decoding; None for non-SSM groups.
+        self._ssm_spec_blocks = [
+            g.kv_cache_spec.num_speculative_blocks
+            if isinstance(g.kv_cache_spec, MambaSpec)
+            else None
+            for g in kv_cache_config.kv_cache_groups
+        ]
+        # Only "all" mode keeps a state per block position; the other modes
+        # keep a single running state in the last non-speculative slot.
+        self._ssm_state_slots_are_positional = (
+            vllm_config.cache_config.mamba_cache_mode == "all"
+        )
+
         # Threshold to decide whether to compute kv cache locally
         # or pull from a remote node: minimum number of remote
         # tokens to amortize the xfer latencies
@@ -218,13 +232,25 @@ class NixlBaseConnectorScheduler:
                     # Clean up empty engines so we don't leak a key when remote dies.
                     del self._heartbeat_by_engine[engine_id]
 
-    def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
-        """
-        Clip the number of blocks to the sliding window size for each kv cache group
-        that employs SWA.
-        This is necessary because the KV Cache manager initially allocates blocks for
-        the entire sequence length, and successively cleans up blocks that are outside
-        the window prior to the `request_finished_all_groups` hook.
+    def get_exchange_clipped_blocks(
+        self, block_ids: BlockIds, clip_ssm: bool = True
+    ) -> BlockIds:
+        """Clip a request's block lists down to the transferable blocks.
+
+        Sliding-window groups keep only the in-window tail: the KV cache
+        manager allocates blocks for the entire sequence length and cleans up
+        out-of-window blocks only prior to the `request_finished_all_groups`
+        hook.
+
+        SSM groups keep only their state-bearing slots: the trailing
+        speculative scratch slots always go, and in single-state cache modes
+        so does everything before the running state (null placeholders and
+        the previous step's superseded state). "all" mode keeps its remaining
+        slots, which the worker pairs position-wise.
+
+        Use this at every block-id exchange point. Pass ``clip_ssm=False``
+        for per-step partial lists (host-buffer save), where the SSM strip
+        does not apply.
         """
         if len(block_ids) == 0 or not self._is_hma_required:
             # No blocks to clip eg Full prefix cache hit or not a hybrid model.
@@ -235,15 +261,22 @@ class NixlBaseConnectorScheduler:
         assert len(block_ids) == len(self.blocks_per_sw), (
             "Number of KV cache groups must match"
         )
-        # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
-        return tuple(
-            [
-                blocks[-self.blocks_per_sw[i] :]
-                if self.blocks_per_sw[i] > 0
-                else blocks
-                for i, blocks in enumerate(block_ids)
-            ]
-        )
+        clipped = []
+        for i, blocks in enumerate(block_ids):
+            if n_sw := self.blocks_per_sw[i]:
+                blocks = blocks[-n_sw:]
+            elif (
+                clip_ssm
+                and blocks
+                and (n_spec_blocks := self._ssm_spec_blocks[i]) is not None
+            ):
+                if n_spec := min(n_spec_blocks, len(blocks) - 1):
+                    blocks = blocks[:-n_spec]
+                if not self._ssm_state_slots_are_positional:
+                    # Never empty: downstream reads that as a full prefix hit.
+                    blocks = blocks[-1:]
+            clipped.append(blocks)
+        return tuple(clipped)
 
     def set_xfer_handshake_metadata(
         self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
@@ -380,7 +413,9 @@ class NixlBaseConnectorScheduler:
             req = req_to_save
 
             assert req.kv_transfer_params is not None
-            clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
+            clipped_block_id_groups = self.get_exchange_clipped_blocks(
+                new_block_id_groups, clip_ssm=False
+            )
             meta.add_new_req_to_save(
                 request_id=req_id,
                 local_block_ids=clipped_block_id_groups,

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -2439,16 +2439,25 @@ class NixlBaseConnectorWorker:
             for i, remote_group in enumerate(remote_block_ids):
                 num_local_blocks = len(local_block_ids[i])
                 num_remote_blocks = len(remote_group)
-                if (
-                    _is_ssm_spec(self._group_spec_types[i])
-                    and num_local_blocks < num_remote_blocks
-                ):
-                    # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks
-                    # prior to the last one are placeholders (null blocks). Mind that
-                    # this doesn't really impact transfer, as we only still care about
-                    # the last "block", the full in-place state.
-                    assert num_local_blocks == 1, "SSM can only have one local block"
-                    remote_block_ids[i] = remote_group[-num_local_blocks:]
+                if _is_ssm_spec(self._group_spec_types[i]):
+                    if num_local_blocks == num_remote_blocks:
+                        continue
+                    # Only state-bearing slots reach here, single-state modes
+                    # just one (see get_exchange_clipped_blocks), so differing
+                    # counts mean position-indexed "all"-mode lists. A longer
+                    # remote list carries earlier positions the local side
+                    # already has (prefix hit) -> read its tail; a longer local
+                    # list holds the position D recomputes itself, which gets
+                    # no remote state.
+                    assert num_local_blocks - num_remote_blocks <= 1, (
+                        f"Group {i}: unpairable SSM state slots, "
+                        f"local={num_local_blocks} remote={num_remote_blocks}"
+                    )
+                    num_blocks = min(num_local_blocks, num_remote_blocks)
+                    if num_local_blocks < num_remote_blocks:
+                        remote_block_ids[i] = remote_group[-num_blocks:]
+                    else:
+                        local_block_ids[i] = local_block_ids[i][:num_blocks]
                 elif (
                     self._physical_blocks_per_logical_kv_block
                     == remote_physical_per_logical

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -155,7 +155,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
                         if num_external_tokens > 0
                         else ()
                     )
-                    local_block_ids = self.get_sw_clipped_blocks(
+                    local_block_ids = self.get_exchange_clipped_blocks(
                         unhashed_local_block_ids
                     )
 
@@ -262,7 +262,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
             # trimming down after allocating for the whole sequence length. Empty
             # blocks are always at the start of the list.
             # Here we "unpad" blocks to send the actual remote blocks to be read.
-            block_ids = self.get_sw_clipped_blocks(block_ids)
+            block_ids = self.get_exchange_clipped_blocks(block_ids)
 
             remote_num_tokens = request.num_computed_tokens
 

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py

@@ -170,7 +170,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
             request.request_id,
         )
         local_block_ids: BlockIds = blocks.get_unhashed_block_ids_all_groups()
-        local_block_ids = self.get_sw_clipped_blocks(local_block_ids)
+        local_block_ids = self.get_exchange_clipped_blocks(local_block_ids)
 
         # ``remote_*`` fields are P's coordinates (from D's perspective).
         # ``decode_*`` fields are D's own info that P needs for the
@@ -271,7 +271,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
                 time.perf_counter() + self._kv_lease_duration
             )
 
-            block_ids = self.get_sw_clipped_blocks(block_ids)
+            block_ids = self.get_exchange_clipped_blocks(block_ids)
             remote_num_tokens = request.num_computed_tokens
 
             # Store finished blocks for worker-level matching with D