diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 0124f6516e..c3058eaddb 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/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 ───────────────────── diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index adf0063701..69887de4a0 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/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 # ----------------------------------------------------------------- # diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 7a3629c722..e53fb5673e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/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, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 12b9f0d937..13f45dfcc4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 51a54d06f0..9beceeabad 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index 413ebdcdb0..a02491e487 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/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