[Bugfix][DSv4.1] Skip SWA bounded replay for KV loads that carry the window (#59197)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/distributed/aux_output_connector/test_store.py

@@ -96,6 +96,10 @@ class _SchedulerRequest:
         default_factory=lambda: SimpleNamespace(routed_experts_prompt_start=0)
     )
 
+    @property
+    def num_prompt_tokens(self) -> int:
+        return self.num_tokens - self.num_output_tokens
+
     def is_finished(self) -> bool:
         return self.finished
 
@@ -1439,6 +1443,18 @@ def test_scheduler_starts_worker_output_at_requested_prompt_token():
 
     assert metadata.requests == {request.request_id: 2}
 
+    # A P/D prefiller cuts its prompt short of the last token, which can pass
+    # the requested start; the rows then start at the end of the cut prompt.
+    connector = _make_connector()
+    request = _scheduler_request(
+        "cut", [b"a" * 32], num_tokens=3, num_output_tokens=0, prompt_start=4
+    )
+    metadata = connector.build_connector_meta(
+        _step_output([request.request_id], [0], [3]),
+        {request.request_id: request},
+    )
+    assert metadata.requests == {request.request_id: 3}
+
 
 def test_scheduler_connector_preserves_request_finish_order():
     connector = _make_connector()

tests/v1/core/test_prefix_caching.py

@@ -594,8 +594,12 @@ def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens):
         failed_recving_kv_req_ids=set(),
         finished_recving_kv_req_ids={request.request_id},
         prefix_replay_tokens=0,
+        prefix_replay_group_ids=(),
     )
     scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler)
+    scheduler._load_restores_replay_window = MethodType(
+        Scheduler._load_restores_replay_window, scheduler
+    )
     Scheduler._update_waiting_for_remote_kv(scheduler, request)
     assert request.num_tokens - request.num_computed_tokens == 1
     assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None

tests/v1/core/test_prefix_replay.py

@@ -247,6 +247,33 @@ def test_remote_kv_hit_is_taken_in_whole_blocks():
     assert swa_blocks[new_req.replay_start // BLOCK_SIZE] is not swa_manager._null_block
 
 
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.parametrize("matched", [WINDOW, 4 * BLOCK_SIZE - 1])
+def test_kv_load_carrying_window_does_not_replay(is_async, matched):
+    """A load of the request's own blocks (P/D) restores every transfer group,
+    window included, so the hit is taken as is: no replay, no cut to whole
+    blocks, and a hit no longer than the window is kept."""
+    scheduler = _replay_scheduler(
+        use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=is_async)
+    )
+    request = create_requests(
+        num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
+    )[0]
+    scheduler.connector.get_loaded_kv_cache_group_ids = lambda request: (FULL, SWA)
+    scheduler.add_request(request)
+    out = scheduler.schedule()
+    if is_async:
+        assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+        scheduler.update_from_output(
+            out, create_model_runner_output([], finished_recving={request.request_id})
+        )
+        out = scheduler.schedule()
+    new_req = _new_req_data(out, request)
+    assert new_req.num_computed_tokens == matched
+    assert new_req.replay_start == 0
+    assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS - matched
+
+
 @pytest.mark.parametrize("via_connector", [False, True])
 def test_hit_no_longer_than_window_is_ignored(via_connector):
     """Such a hit would be recomputed in full anyway. It is not adopted (an

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -1559,17 +1559,22 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks():
 
 @pytest.mark.cpu_test
 @pytest.mark.parametrize(
-    "has_mamba,is_hma_required,expected_count",
+    "has_mamba,is_hma_required,bounded_replay,expected_count",
     [
-        (True, True, 9),
-        (False, False, 10),
-        (False, True, 10),
+        (True, True, False, 9),
+        (False, False, False, 10),
+        (False, True, False, 10),
+        (False, True, True, 9),
     ],
-    ids=["mamba", "fa_only", "swa_only"],
+    ids=["mamba", "fa_only", "swa_only", "swa_bounded_replay"],
 )
-def test_mamba_n1_d_side(has_mamba, is_hma_required, expected_count):
-    """D-side: Mamba gets N-1 matched tokens, non-Mamba gets N."""
-    sched = make_nixl_scheduler(has_mamba=has_mamba, is_hma_required=is_hma_required)
+def test_mamba_n1_d_side(has_mamba, is_hma_required, bounded_replay, expected_count):
+    """D-side: Mamba and SWA bounded replay get N-1 matched tokens, others N."""
+    sched = make_nixl_scheduler(
+        has_mamba=has_mamba,
+        is_hma_required=is_hma_required,
+        bounded_replay=bounded_replay,
+    )
     req = create_request(num_tokens=10, do_remote_prefill=True)
 
     count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
@@ -1638,6 +1643,20 @@ def test_mamba_n1_p_side_truncation():
     fa_sched.on_new_request(fa_req)
     assert len(fa_req.prompt_token_ids) == fa_original
 
+    # SWA bounded replay: the window is not replayed after the load, so the
+    # prefiller stops short of the last token too.
+    swa_sched = make_nixl_scheduler(is_hma_required=True, bounded_replay=True)
+    swa_req = create_request(num_tokens=10, do_remote_decode=True)
+    swa_sched.on_new_request(swa_req)
+    assert len(swa_req.prompt_token_ids) == 9
+
+    # A request that skips reading the prefix cache (prompt logprobs) is not
+    # cut: the decoder loads nothing and recomputes its whole prompt.
+    logprobs_req = create_request(num_tokens=10, do_remote_decode=True)
+    logprobs_req.sampling_params.skip_reading_prefix_cache = True
+    swa_sched.on_new_request(logprobs_req)
+    assert len(logprobs_req.prompt_token_ids) == 10
+
 
 @pytest.mark.cpu_test
 @pytest.mark.parametrize(

tests/v1/kv_connector/unit/utils.py

@@ -532,6 +532,7 @@ def make_kv_cache_config(
 def make_nixl_scheduler(
     has_mamba: bool = False,
     is_hma_required: bool = False,
+    bounded_replay: bool = False,
     heartbeat: bool = False,
     kv_lease_duration: int = 30,
 ):
@@ -548,6 +549,7 @@ def make_nixl_scheduler(
 
     sched = object.__new__(NixlConnectorScheduler)
     sched._has_mamba = has_mamba
+    sched._bounded_replay = bounded_replay
     sched._is_hma_required = is_hma_required
     sched.kv_cache_config = make_kv_cache_config(
         block_size=16,
@@ -615,6 +617,7 @@ def make_nixl_push_scheduler(
     sched.side_channel_port = 5600
     sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled
     sched._has_mamba = has_mamba
+    sched._bounded_replay = False
     sched.kv_cache_config = make_kv_cache_config(
         block_size=16,
         mamba_enabled=has_mamba,

vllm/distributed/aux_output_connector/connector.py

@@ -67,8 +67,12 @@ class AuxOutputSchedulerConnector:
                 block_hashes_by_request[request_id] = packed
             self._sent_hash_counts[request_id] = len(request.block_hashes)
             assert request.sampling_params is not None
+            # A P/D prefiller may cut the prompt short of the requested start.
             scheduled_requests[request_id] = max(
-                request.sampling_params.routed_experts_prompt_start,
+                min(
+                    request.sampling_params.routed_experts_prompt_start,
+                    request.num_prompt_tokens,
+                ),
                 0 if request.num_output_tokens == 0 else request.num_tokens - 1,
             )
         # A settled token can complete a hash block after the next async schedule

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -897,6 +897,10 @@ class MooncakeConnectorScheduler:
         # GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec
         # path is currently tested with GDN; Mamba2 is not validated yet.
         self._has_mamba = kv_cache_config.has_mamba_layers
+        self._bounded_replay = any(
+            g.kv_cache_spec.prefix_replay_tokens > 0
+            for g in kv_cache_config.kv_cache_groups
+        )
 
         # Requests that need to start recv/send.
         # New requests are added by update_state_after_alloc in
@@ -941,8 +945,10 @@ class MooncakeConnectorScheduler:
 
     def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
         """D-side only. Returns N-1 for Mamba models since the decoder
-        always recomputes the last token and must start from h(N-1)."""
-        if self._has_mamba and num_prompt_tokens > 1:
+        always recomputes the last token and must start from h(N-1), and for
+        DSV41 SWA bounded replay, whose window is not replayed after a P/D load
+        and must be laid out for that token."""
+        if (self._has_mamba or self._bounded_replay) and num_prompt_tokens > 1:
             return num_prompt_tokens - 1
         return num_prompt_tokens
 
@@ -958,6 +964,9 @@ class MooncakeConnectorScheduler:
             params is not None
             and not params.get("_p_side_truncated")
             and request.num_prompt_tokens > 1
+            # A request that skips reading the prefix cache (prompt logprobs)
+            # loads nothing on the decoder, which recomputes its whole prompt.
+            and not request.get_skip_reading_prefix_cache()
         ):
             if request.prompt_token_ids is not None:
                 request.prompt_token_ids.pop()
@@ -973,7 +982,11 @@ class MooncakeConnectorScheduler:
 
     def on_new_request(self, request: "Request") -> None:
         params = request.kv_transfer_params
-        if params is not None and params.get("do_remote_decode") and self._has_mamba:
+        if (
+            params is not None
+            and params.get("do_remote_decode")
+            and (self._has_mamba or self._bounded_replay)
+        ):
             self._truncate_mamba_request_for_prefill(request)
 
     def get_num_new_matched_tokens(

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

@@ -101,6 +101,10 @@ class NixlBaseConnectorScheduler:
             )
         )
         self._has_mamba = kv_cache_config.has_mamba_layers
+        self._bounded_replay = any(
+            g.kv_cache_spec.prefix_replay_tokens > 0
+            for g in kv_cache_config.kv_cache_groups
+        )
 
         logger.info("Initializing NIXL Scheduler %s", engine_id)
         if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
@@ -376,14 +380,16 @@ class NixlBaseConnectorScheduler:
         """Trailing prompt tokens the prefiller must not compute; the decoder
         recomputes them locally.
 
-        Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module
+        Mamba needs h(N-1) so the decoder can derive h(N) itself. DSV41's sliding
+        window under bounded replay is not replayed after a P/D load, so it must
+        be laid out for the token the decoder recomputes. Multi-module
         MTP needs to keep its whole lookahead window off of the prefiller, which
         would otherwise embed the unverified drafts in the MTP layer's KV cache. The
         decoder would never rebuild them, because the update is sized by the rejection
         count, which is zero for the first decode.
         """
         return max(
-            1 if self._has_mamba else 0,
+            1 if self._has_mamba or self._bounded_replay else 0,
             self.vllm_config.num_prefill_lookahead_tokens - 1,
         )
 
@@ -412,6 +418,9 @@ class NixlBaseConnectorScheduler:
             # Guard against repeated truncation after preemption/reschedule.
             and not params.get("_p_side_truncated")
             and request.num_prompt_tokens > backoff
+            # A request that skips reading the prefix cache (prompt logprobs)
+            # loads nothing on the decoder, which recomputes its whole prompt.
+            and not request.get_skip_reading_prefix_cache()
         ):
             if request.prompt_token_ids is not None:
                 del request.prompt_token_ids[-backoff:]

vllm/v1/core/sched/scheduler.py

@@ -290,10 +290,14 @@ class Scheduler(SchedulerInterface):
         # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt
         # after a prefix hit by recomputing its trailing tokens. One window
         # for all such groups, so the rewind matches every group's allocation.
-        replay_windows = {
-            group.kv_cache_spec.prefix_replay_tokens
-            for group in kv_cache_config.kv_cache_groups
+        self.prefix_replay_group_ids = tuple(
+            group_id
+            for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
             if group.kv_cache_spec.prefix_replay_tokens > 0
+        )
+        replay_windows = {
+            kv_cache_config.kv_cache_groups[group_id].kv_cache_spec.prefix_replay_tokens
+            for group_id in self.prefix_replay_group_ids
         }
         assert len(replay_windows) <= 1, (
             f"Prefix replay windows should agree: {sorted(replay_windows)}"
@@ -973,14 +977,17 @@ class Scheduler(SchedulerInterface):
                             skip_request(request_queue)
                             continue
 
-                        if self.prefix_replay_tokens:
-                            # SWA bounded replay recomputes the hit's last
+                        if self.prefix_replay_tokens and not (
+                            ext_tokens and self._load_restores_replay_window(request)
+                        ):
+                            # DSV41 SWA bounded replay recomputes the hit's last
                             # window, from the block holding its first
                             # token; the sliding-window groups retire
                             # whole blocks below the window of the hit's
                             # next token. Hits end on a block boundary, or
                             # a hit ending one token short of one would
-                            # retire the block the replay starts in.
+                            # retire the block the replay starts in. A load
+                            # that carries the window replays nothing.
                             ext_tokens -= ext_tokens % self.block_size
                             load_kv_async = load_kv_async and ext_tokens > 0
 
@@ -1026,9 +1033,17 @@ class Scheduler(SchedulerInterface):
                         num_new_local_computed_tokens + num_external_computed_tokens
                     )
                     assert num_computed_tokens <= request.num_tokens
-                    if 0 < num_computed_tokens <= self.prefix_replay_tokens:
-                        # SWA bounded replay: a hit no longer than the replayed
-                        # window would be recomputed in full and save nothing.
+                    # DSV41 SWA bounded replay: a KV load that carries the
+                    # window is taken as is; any other hit no longer than the
+                    # replayed window would be recomputed in full and save
+                    # nothing.
+                    window_loaded = bool(
+                        num_external_computed_tokens
+                    ) and self._load_restores_replay_window(request)
+                    if (
+                        not window_loaded
+                        and 0 < num_computed_tokens <= self.prefix_replay_tokens
+                    ):
                         new_computed_blocks = (
                             self.kv_cache_manager.empty_kv_cache_blocks
                         )
@@ -1091,7 +1106,7 @@ class Scheduler(SchedulerInterface):
                     if did_prefix_cache_lookup:
                         # Fresh admission: local and/or sync external hit.
                         num_replay_tokens = self._mark_prefix_replay(
-                            request, num_computed_tokens
+                            request, num_computed_tokens, window_loaded
                         )
                         num_computed_tokens -= num_replay_tokens
                     # Number of tokens to be scheduled.
@@ -3083,16 +3098,31 @@ class Scheduler(SchedulerInterface):
             self._request_remaining_blocks(req) for req in self._inflight_prefills
         )
 
-    def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int:
+    def _load_restores_replay_window(self, request: Request) -> bool:
+        """Whether the KV connector's load for ``request`` restores every DSV41
+        SWA bounded-replay group, as a P/D transfer of the request's own blocks
+        does; a prefix-cache store holds no window."""
+        assert self.connector is not None
+        return bool(self.prefix_replay_group_ids) and set(
+            self.prefix_replay_group_ids
+        ).issubset(self.connector.get_loaded_kv_cache_group_ids(request))
+
+    def _mark_prefix_replay(
+        self, request: Request, num_hit_tokens: int, window_loaded: bool
+    ) -> int:
         """Record where a prefix hit's replay starts and return the number of
         hit tokens to recompute; the caller rewinds the computed count by it.
 
         The replayed tokens are the hit's last window (see
         ``Request.replay_start``): the worker rebuilds their sliding-window KV
-        and leaves their cached KV alone.
+        and leaves their cached KV alone. A hit whose KV load carries the
+        window replays nothing.
         """
         if not self.prefix_replay_tokens:
             return 0
+        if window_loaded:
+            request.replay_start = 0
+            return 0
         num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens)
         request.replay_start = num_hit_tokens - num_replay_tokens
         return num_replay_tokens
@@ -3106,7 +3136,8 @@ class Scheduler(SchedulerInterface):
         """
         assert self.connector is not None
 
-        if request.request_id in self.failed_recving_kv_req_ids:
+        load_failed = request.request_id in self.failed_recving_kv_req_ids
+        if load_failed:
             # Request had KV load failures; num_computed_tokens was already
             # updated in _update_requests_with_invalid_blocks
             if request.num_computed_tokens:
@@ -3132,11 +3163,15 @@ class Scheduler(SchedulerInterface):
             # This will cache the blocks iff caching is enabled.
             self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
 
-        # SWA bounded replay recomputes the tail of the hit, which covers the
-        # last token; otherwise a full prompt hit re-computes that token so the
-        # next one can be sampled.
+        # DSV41 SWA bounded replay recomputes the tail of a hit loaded without
+        # its window, or of a failed load, which may have lost it; the replay
+        # covers the last token. Otherwise a full prompt hit re-computes that
+        # token so the next one can be sampled.
         num_replay_tokens = self._mark_prefix_replay(
-            request, request.num_computed_tokens
+            request,
+            request.num_computed_tokens,
+            window_loaded=not load_failed
+            and self._load_restores_replay_window(request),
         )
         if num_replay_tokens > 0:
             request.num_computed_tokens -= num_replay_tokens