fix(mamba): skip prefill donation when an in-flight verify aliases the keep slot

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -369,18 +369,25 @@ def process_batch_result_prefill(
                         and not batch.spec_algorithm.is_none()
                     ):
                         req.kv.kv_committed_len += 1
+                    keep_may_be_written_in_flight = (
+                        self._mamba_lazy_keep_may_be_written_in_flight(req, batch)
+                    )
                     if req.finished():
                         if sampling_mask_finish_reason is None:
                             self._maybe_collect_routed_experts(req)
                             self._maybe_collect_indexer_topk(req)
                         release_kv_cache(
                             req,
                             self.tree_cache,
-                            is_insert=sampling_mask_finish_reason is None,
+                            is_insert=(
+                                sampling_mask_finish_reason is None
+                                and not keep_may_be_written_in_flight
+                            ),
                         )
                         req.time_stats.set_completion_time()
                     elif not batch.decoding_reqs or req not in batch.decoding_reqs:
-                        maybe_cache_unfinished_req(req, self.tree_cache)
+                        if not keep_may_be_written_in_flight:
+                            maybe_cache_unfinished_req(req, self.tree_cache)
                         if get_memory().enable_hisparse:
                             self.hisparse_coordinator.admit_request_into_staging(req)
 
@@ -1419,6 +1426,47 @@ def _mamba_prefix_cache_update(
                 batch.req_to_token_pool.get_mamba_ping_pong_other_idx(track_idx)
             )
 
+    def _mamba_lazy_keep_may_be_written_in_flight(
+        self, req: Req, batch: ScheduleBatch
+    ) -> bool:
+        """Whether a spec verify overlapping this prefill may scatter into the
+        keep slot of the request.
+
+        Prefill-side twin of the decode guard keep_may_be_written_in_flight
+        recompute in _mamba_lazy_spec_update. Under lazy extra-buffer + spec,
+        when the pending ping-pong position is empty while a track-interval
+        crossing is reachable, the verify plan (mamba_lazy_spec_prepare) fell
+        back -- or will fall back -- to the keep slot as its scatter
+        destination. Publishing that physical slot to the radix tree now would
+        record it under the prefill tracked depth C while the verify commits
+        the crossing state T > C into the same slot, mislabeling the newer
+        state as C. The caller must skip the donation/insert while this holds;
+        the first verify result repairs the request-side label.
+        """
+        if not get_exec().mamba.enable_mamba_extra_buffer_lazy:
+            return False
+        # getattr: partial test fixtures model the batch without a spec bag.
+        spec_algorithm = getattr(batch, "spec_algorithm", None)
+        if spec_algorithm is None or spec_algorithm.is_none():
+            return False
+        buf = getattr(req.kv, "mamba_ping_pong_track_buffer", None)
+        next_idx = req.kv.mamba_next_track_idx
+        if buf is None or next_idx is None:
+            return False
+        other_idx = 1 - next_idx
+        if other_idx >= buf.numel():
+            # Single-slot buffer (non-overlap scheduler): donation substitutes
+            # the replacement slot before any verify plan can capture the keep
+            # slot, so the published slot is never the captured destination.
+            return False
+        if buf[other_idx].item() != -1:
+            return False
+        return mamba_lazy_spec_in_window(
+            req,
+            mamba_track_grid(self.tree_cache.page_size),
+            max_speculative_num_draft_tokens(),
+        )
+
     def _mamba_lazy_spec_update(
         self, req: Req, batch: ScheduleBatch, i: int, crossed: bool, track_seqlen: int
     ) -> None:

test/registered/unit/managers/test_batch_result_processor_mamba_lazy_donation.py

@@ -0,0 +1,243 @@
+"""Prefill-time mamba keep-slot donation guard tests.
+
+Under extra_buffer_lazy + speculative decoding with the overlap scheduler, the
+first verify after a lazy final prefill is launched before the prefill result
+is processed. When the pending ping-pong slot allocation fails at verify
+prepare, the verify plan falls back to the keep slot as its scatter
+destination. If prefill result processing then donates/inserts that physical
+slot to the radix tree under the prefill tracked depth C, the verify commits
+the crossing state T > C into the same slot and the tree mislabels newer
+state as C.
+
+These tests pin the prefill-side alias predicate
+(_mamba_lazy_keep_may_be_written_in_flight) and the two decisions it gates:
+the finishing release insert flag and the unfinished donation call.
+"""
+
+import unittest
+from array import array
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.managers.scheduler_components.batch_result_processor import (
+    SchedulerBatchResultProcessor,
+)
+from sglang.srt.managers.utils import GenerationBatchResult
+from sglang.srt.runtime_context import get_context
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=8, suite="base-a-test-cpu")
+
+# page = chunk = interval = 8 keeps every checkpoint grid at 8; with two draft
+# tokens the in-window predicate fires at seq_len 12 but not at seq_len 8.
+GRID = 8
+DRAFT_TOKENS = 2
+IN_WINDOW_LEN = 12
+OUT_OF_WINDOW_LEN = 8
+KEEP_SLOT = 5
+PENDING_SLOT = 9
+
+
+def _make_req(
+    *,
+    max_new_tokens: int,
+    committed_len: int = IN_WINDOW_LEN,
+    ping_pong=(KEEP_SLOT, -1),
+) -> Req:
+    sampling_params = SamplingParams(max_new_tokens=max_new_tokens)
+    sampling_params.normalize(None)
+    req = Req(
+        rid="req",
+        origin_input_text="",
+        origin_input_ids=array("q", [1, 2]),
+        sampling_params=sampling_params,
+        vocab_size=128,
+    )
+    req.kv.kv_committed_len = committed_len
+    req.kv.req_pool_idx = 0
+    req.kv.mamba_last_track_seqlen = GRID
+    if ping_pong is not None:
+        req.kv.mamba_ping_pong_track_buffer = torch.tensor(
+            list(ping_pong), dtype=torch.int64
+        )
+        req.kv.mamba_next_track_idx = 0
+    return req
+
+
+def _make_prefill_batch(req: Req, *, spec_active: bool = True) -> ScheduleBatch:
+    batch = ScheduleBatch(reqs=[req])
+    batch.tree_cache = SimpleNamespace(page_size=GRID)
+    batch.device = "cpu"
+    batch.model_config = SimpleNamespace(is_encoder_decoder=False)
+    batch.enable_overlap = True
+    batch.spec_algorithm = SimpleNamespace(is_none=lambda: not spec_active)
+    batch.sampling_info = SimpleNamespace(
+        penalizer_orchestrator=SimpleNamespace(is_required=False)
+    )
+    batch.hisparse_coordinator = None
+    batch.decoding_reqs = []
+    batch.prefill_stats = None
+    return batch
+
+
+def _make_processor() -> SchedulerBatchResultProcessor:
+    metrics_reporter = MagicMock()
+    metrics_reporter.num_generated_tokens = 0
+    processor = SchedulerBatchResultProcessor(
+        is_generation=True,
+        disaggregation_mode=None,
+        enable_overlap=True,
+        enable_overlap_mlx=False,
+        model_config=SimpleNamespace(think_end_ids=None),
+        token_to_kv_pool_allocator=MagicMock(),
+        tree_cache=SimpleNamespace(page_size=GRID),
+        hisparse_coordinator=None,
+        beam_coordinator=MagicMock(),
+        req_to_token_pool=None,
+        decode_offload_manager=None,
+        metrics_collector=None,
+        metrics_reporter=metrics_reporter,
+        draft_worker=None,
+        model_worker=MagicMock(),
+        logprob_result_processor=None,
+        output_streamer=MagicMock(),
+        abort_request=lambda *args, **kwargs: None,
+    )
+    return processor
+
+
+def _make_prefill_result(next_token: int = 4) -> GenerationBatchResult:
+    return GenerationBatchResult(
+        logits_output=SimpleNamespace(
+            hidden_states=None,
+            customized_info=None,
+            sampling_mask_output=None,
+            next_token_sampling_mask_status=None,
+        ),
+        next_token_ids=torch.tensor([next_token], dtype=torch.int64),
+        speculative_num_draft_tokens=0,
+    )
+
+
+def _publish(*, strategy: str = "extra_buffer_lazy"):
+    return get_context().override_server_args(
+        mamba_radix_cache_strategy=strategy,
+        mamba_track_interval=GRID,
+        _mamba_cache_chunk_size=GRID,
+        speculative_num_draft_tokens=DRAFT_TOKENS,
+    )
+
+
+class TestKeepMayBeWrittenInFlightPredicate(unittest.TestCase):
+    def _predicate(self, req, batch, processor):
+        with _publish():
+            return processor._mamba_lazy_keep_may_be_written_in_flight(req, batch)
+
+    def test_alias_when_pending_empty_and_in_window(self):
+        req = _make_req(max_new_tokens=32)
+        batch = _make_prefill_batch(req)
+        self.assertTrue(self._predicate(req, batch, _make_processor()))
+
+    def test_no_alias_when_pending_slot_allocated(self):
+        req = _make_req(max_new_tokens=32, ping_pong=(KEEP_SLOT, PENDING_SLOT))
+        batch = _make_prefill_batch(req)
+        self.assertFalse(self._predicate(req, batch, _make_processor()))
+
+    def test_no_alias_outside_crossing_window(self):
+        req = _make_req(max_new_tokens=32, committed_len=OUT_OF_WINDOW_LEN)
+        batch = _make_prefill_batch(req)
+        self.assertFalse(self._predicate(req, batch, _make_processor()))
+
+    def test_no_alias_without_speculation(self):
+        req = _make_req(max_new_tokens=32)
+        batch = _make_prefill_batch(req, spec_active=False)
+        self.assertFalse(self._predicate(req, batch, _make_processor()))
+
+    def test_no_alias_without_lazy_strategy(self):
+        req = _make_req(max_new_tokens=32)
+        batch = _make_prefill_batch(req)
+        processor = _make_processor()
+        with _publish(strategy="extra_buffer"):
+            self.assertFalse(
+                processor._mamba_lazy_keep_may_be_written_in_flight(req, batch)
+            )
+
+    def test_no_alias_without_ping_pong_buffer(self):
+        req = _make_req(max_new_tokens=32, ping_pong=None)
+        batch = _make_prefill_batch(req)
+        self.assertFalse(self._predicate(req, batch, _make_processor()))
+
+    def test_no_alias_with_single_slot_buffer(self):
+        # Non-overlap scheduling keeps one ping-pong slot: donation substitutes
+        # the replacement before any verify plan can capture the keep slot.
+        req = _make_req(max_new_tokens=32, ping_pong=(KEEP_SLOT,))
+        batch = _make_prefill_batch(req)
+        self.assertFalse(self._predicate(req, batch, _make_processor()))
+
+
+@patch(
+    "sglang.srt.managers.scheduler_components.batch_result_processor.maybe_cache_unfinished_req"
+)
+@patch(
+    "sglang.srt.managers.scheduler_components.batch_result_processor.release_kv_cache"
+)
+class TestPrefillDonationDecision(unittest.TestCase):
+    def _run_prefill(self, req, batch):
+        processor = _make_processor()
+        with _publish():
+            processor.process_batch_result_prefill(batch, _make_prefill_result())
+
+    def test_finishing_suppresses_insert_under_alias(
+        self, release_kv_cache, maybe_cache_unfinished_req
+    ):
+        req = _make_req(max_new_tokens=1)
+        batch = _make_prefill_batch(req)
+        self._run_prefill(req, batch)
+        release_kv_cache.assert_called_once()
+        self.assertFalse(release_kv_cache.call_args.kwargs["is_insert"])
+        maybe_cache_unfinished_req.assert_not_called()
+
+    def test_finishing_inserts_when_pending_allocated(
+        self, release_kv_cache, maybe_cache_unfinished_req
+    ):
+        req = _make_req(max_new_tokens=1, ping_pong=(KEEP_SLOT, PENDING_SLOT))
+        batch = _make_prefill_batch(req)
+        self._run_prefill(req, batch)
+        release_kv_cache.assert_called_once()
+        self.assertTrue(release_kv_cache.call_args.kwargs["is_insert"])
+        maybe_cache_unfinished_req.assert_not_called()
+
+    def test_unfinished_skips_donation_under_alias(
+        self, release_kv_cache, maybe_cache_unfinished_req
+    ):
+        req = _make_req(max_new_tokens=32)
+        batch = _make_prefill_batch(req)
+        self._run_prefill(req, batch)
+        maybe_cache_unfinished_req.assert_not_called()
+        release_kv_cache.assert_not_called()
+
+    def test_unfinished_donates_when_pending_allocated(
+        self, release_kv_cache, maybe_cache_unfinished_req
+    ):
+        req = _make_req(max_new_tokens=32, ping_pong=(KEEP_SLOT, PENDING_SLOT))
+        batch = _make_prefill_batch(req)
+        self._run_prefill(req, batch)
+        maybe_cache_unfinished_req.assert_called_once()
+        release_kv_cache.assert_not_called()
+
+    def test_unfinished_donates_outside_crossing_window(
+        self, release_kv_cache, maybe_cache_unfinished_req
+    ):
+        req = _make_req(max_new_tokens=32, committed_len=OUT_OF_WINDOW_LEN)
+        batch = _make_prefill_batch(req)
+        self._run_prefill(req, batch)
+        maybe_cache_unfinished_req.assert_called_once()
+        release_kv_cache.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()