Fix single-request generation being killed early on hybrid-SWA models under DFLASH/DSPARK speculative decoding

python/sglang/srt/speculative/spec_utils.py

@@ -1037,6 +1037,17 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
             max_speculative_num_draft_tokens(),
         )
     if batch.spec_algorithm.is_dflash_family():
+        # Mirror eagle_prepare_for_decode: evict out-of-window SWA KV and
+        # advance decode_batch_idx (the eviction gate requires >= 1). Without
+        # this, hybrid-SWA models retain SWA KV for every generated token and
+        # a single request exhausts the SWA pool after
+        # swa_full_tokens_ratio * max_total_num_tokens generated tokens, at
+        # which point the scheduler retracts it. Pool sizing already assumes
+        # this eviction runs (pool_configurator: trailing_tokens =
+        # window + eviction_interval * draft_tokens + page_size).
+        batch.maybe_evict_swa()
+        for r in batch.reqs:
+            r.decode_batch_idx += 1
         batch.spec_info.prepare_for_decode(batch)
     else:
         from sglang.srt.speculative.eagle_utils import eagle_prepare_for_decode

scripts/lint/check_decode_bookkeeping_ownership.py

@@ -50,6 +50,9 @@
     ("mem_cache/allocation.py", "alloc_for_extend", "evict"): 1,
     ("mem_cache/allocation.py", "alloc_for_extend", "kv_allocated_len"): 1,
     ("mem_cache/allocation.py", "alloc_for_decode", "evict"): 1,
+    # spec: the dflash-family branch mirrors eagle_prepare_for_decode.
+    ("speculative/spec_utils.py", "spec_prepare_for_decode", "decode_batch_idx"): 1,
+    ("speculative/spec_utils.py", "spec_prepare_for_decode", "evict"): 1,
     ("mem_cache/allocation.py", "alloc_for_decode", "kv_allocated_len"): 1,
     # spec v2: no pre-claim; resolve commits the full accepted run uniformly.
     # kv_allocated_len for spec v2 draft decode (eagle + dflash) is settled

test/registered/unit/speculative/test_spec_prepare_swa_eviction.py

@@ -0,0 +1,82 @@
+"""spec_prepare_for_decode must run SWA eviction on the dflash-family path.
+
+The dflash-family branch (DFLASH, DSPARK) historically skipped both
+``maybe_evict_swa`` and the ``decode_batch_idx`` tick that
+``eagle_prepare_for_decode`` performs. On hybrid-SWA models that retained
+SWA KV for every generated token, so a single request exhausted the SWA pool
+after ``swa_full_tokens_ratio * max_total_num_tokens`` generated tokens and
+was retracted by the scheduler. Pool sizing assumes the eviction runs
+(``pool_configurator``: ``trailing_tokens = sliding_window +
+eviction_interval * draft_tokens + page_size``).
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+def _make_batch(spec_is_dflash_family: bool):
+    batch = MagicMock()
+    batch.spec_algorithm.is_dflash_family.return_value = spec_is_dflash_family
+    req = SimpleNamespace(decode_batch_idx=0)
+    batch.reqs = [req]
+    return batch, req
+
+
+class TestSpecPrepareSwaEviction(CustomTestCase):
+    def _run(self, batch):
+        from sglang.srt.speculative import spec_utils
+
+        with patch.object(
+            spec_utils, "mamba_extra_buffer_lazy_enabled", return_value=False
+        ):
+            spec_utils.spec_prepare_for_decode(batch)
+
+    def test_dflash_family_evicts_and_ticks(self):
+        batch, req = _make_batch(spec_is_dflash_family=True)
+        self._run(batch)
+        batch.maybe_evict_swa.assert_called_once_with()
+        self.assertEqual(req.decode_batch_idx, 1)
+        batch.spec_info.prepare_for_decode.assert_called_once_with(batch)
+
+    def test_tick_advances_every_iteration(self):
+        """decode_batch_idx is a clock, not a flag: it must keep advancing so
+        the SWA leaf-lock release gate (decode_batch_idx >= sliding_window_size)
+        can fire."""
+        batch, req = _make_batch(spec_is_dflash_family=True)
+        self._run(batch)
+        self._run(batch)
+        self.assertEqual(req.decode_batch_idx, 2)
+        self.assertEqual(batch.maybe_evict_swa.call_count, 2)
+
+    def test_dflash_family_evicts_before_tick(self):
+        """The overlap-scheduler gate (decode_batch_idx >= 1) must see the
+        pre-tick value, exactly as in eagle_prepare_for_decode."""
+        batch, req = _make_batch(spec_is_dflash_family=True)
+        seen_idx = []
+        batch.maybe_evict_swa.side_effect = lambda: seen_idx.append(
+            req.decode_batch_idx
+        )
+        self._run(batch)
+        self.assertEqual(seen_idx, [0])
+        self.assertEqual(req.decode_batch_idx, 1)
+
+    def test_eagle_path_unchanged(self):
+        batch, req = _make_batch(spec_is_dflash_family=False)
+        with patch(
+            "sglang.srt.speculative.eagle_utils.eagle_prepare_for_decode"
+        ) as eagle_prep:
+            self._run(batch)
+        eagle_prep.assert_called_once_with(batch)
+        # The dflash-branch tick must not run on the eagle path.
+        self.assertEqual(req.decode_batch_idx, 0)
+        batch.spec_info.prepare_for_decode.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()