[Fix] scheduler: keep chunked prefill within allocatable KV under memory pressure

python/sglang/srt/managers/schedule_policy.py

@@ -1635,10 +1635,7 @@ def _select_prefill_admission(
                 return AddReqResult.OTHER
             max_new_tokens = 0
         elif chunk_tokens_limit is not None and chunk_fit_tokens > chunk_tokens_limit:
-            if (
-                has_chunked_req
-                and get_schedule().schedule_policy == "shortest-prefill-first"
-            ):
+            if has_chunked_req:
                 # Only one unfinished chunked request can be tracked.
                 return AddReqResult.OTHER
             if self.exact_chunk_fill:

python/sglang/srt/managers/scheduler.py

@@ -4184,13 +4184,9 @@ def _get_new_batch_prefill_raw(
             for req in ready_grammar_requests:
                 self._add_request_to_queue(req)
 
-        if self.enable_priority_preemption or self.is_hybrid_swa:
-            # Reset batch_is_full to try preemption with a prefill adder.
-            running_batch.batch_is_full = False
-        if running_batch.is_empty():
-            # An empty batch has no request left to leave, and only a leaving
-            # request clears the flag elsewhere, so admission would stay shut.
-            running_batch.batch_is_full = False
+        # A latched flag is otherwise only cleared when the running batch shrinks,
+        # stalling the queue until an unrelated finish. PrefillAdder enforces limits.
+        running_batch.batch_is_full = False
 
         if (
             running_batch.batch_is_full or len(self.waiting_queue) == 0

python/sglang/srt/mem_cache/prefill_budget.py

@@ -169,12 +169,15 @@ def available_chunk_tokens(self, chunk_limit: int) -> int | None:
         available_tokens = int(self.remaining_total)
         if is_cp_sharded_padded():
             available_tokens = get_cp_sharded_padded_available_tokens(available_tokens)
-        available = min(
-            chunk_limit,
-            available_tokens,
-        )
-        # Single-pool continuation must make progress to release its KV.
-        return available if available > 0 else chunk_limit
+        # Cap by the KV alloc_extend can really get (it over-asks one page).
+        free = self._available_and_evictable() - self.current_offset - self.page_size
+        free = free // self.page_size * self.page_size
+        available = min(chunk_limit, available_tokens, free)
+        if available > 0:
+            return available
+        # Decode headroom is an estimate: progress with the free KV, else defer.
+        available = min(chunk_limit, free)
+        return available if available > 0 else None
 
     def fit_chunk(
         self,

python/sglang/srt/models/deepseek_common/attention_backend_handler.py

@@ -182,6 +182,13 @@ def handle_attention_trtllm_mla(attn, forward_batch):
 
 
 def handle_attention_aiter(attn, forward_batch):
+    # Over-capacity extends run eagerly (graph replay rejects them), but the HIP
+    # eager fallback still sits in the PCG context: keep them chunked so the
+    # cached prefix is never materialized in one piece.
+    if forward_batch.forward_mode.is_extend_without_speculative() and not (
+        _support_mha_one_shot(attn, forward_batch, "aiter")
+    ):
+        return AttnForwardMethod.MHA_CHUNKED_KV
     # During PCG/BCG capture on ROCm, aiter fp8 MLA prefill has no capture
     # kernels; route through the MHA path (radix_attention swaps attn_mqa for
     # its attn_mha companion) so capture/replay use valid head/dim metadata.

test/registered/unit/managers/test_prefill_adder.py

@@ -266,6 +266,22 @@ def test_shortest_prefill_rejects_second_unfinished_chunk(self):
         self.assertIsNone(req.extend_end)
         self.mock_tree_cache.init_load_back.assert_not_called()
 
+    def test_default_policy_rejects_second_unfinished_chunk(self):
+        self.mock_tree_cache.supports_mamba.return_value = False
+        self.mock_tree_cache.supports_prefix_sharing.return_value = False
+        self.mock_token_allocator.available_size.return_value = 32768
+        adder = self.create_adder(
+            self.create_running_batch(), page_size=256, rem_chunk_tokens=512
+        )
+        req = self.create_shared_req("second-chunk")
+        req.full_untruncated_fill_ids = list(range(1024))
+        self.assertEqual(
+            adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None),
+            AddReqResult.OTHER,
+        )
+        self.assertEqual(adder.can_run_list, [])
+        self.assertIsNone(adder.new_chunked_req)
+
     def test_shortest_prefill_rechecks_chunk_limit_after_host_miss(self):
         adder = self.create_shortest_prefill_adder(chunk_tokens=512)
         req = self.create_shared_req("host-miss")

test/registered/unit/managers/test_scheduler_batch_is_full_latch.py

@@ -0,0 +1,64 @@
+"""An empty running batch never keeps admission closed (#41841)."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.scheduler import Scheduler
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class _AdmissionAttempted(Exception):
+    pass
+
+
+class TestEmptyBatchFullLatch(CustomTestCase):
+    def _scheduler(self) -> Scheduler:
+        s = Scheduler.__new__(Scheduler)
+        s.grammar_manager = SimpleNamespace(has_waiting_grammars=lambda: False)
+        s.enable_priority_preemption = False
+        s.waiting_queue = [SimpleNamespace(rid="req-0")]
+        s.chunked_req = None
+        s.min_free_slots_delayer = None
+        s.get_num_allocatable_reqs = Mock(return_value=8)
+        s.processed_tokens_counter = 0
+        s.disagg_prefill_bootstrap_queue = None
+        # Reaching the policy means the admission gates were passed.
+        s.policy = SimpleNamespace(calc_priority=Mock(side_effect=_AdmissionAttempted))
+        return s
+
+    def test_latched_empty_batch_still_admits(self):
+        # Whichever pass set the flag on an empty batch (no KV room, or no free
+        # request slot), no later step shrinks the batch to clear it.
+        s = self._scheduler()
+        running_batch = ScheduleBatch(reqs=[], batch_is_full=True)
+
+        with self.assertRaises(_AdmissionAttempted):
+            s._get_new_batch_prefill_raw(
+                prefill_delayer_single_pass=None, running_batch=running_batch
+            )
+        self.assertFalse(running_batch.batch_is_full)
+
+    def test_latched_running_batch_still_admits(self):
+        # Clearing the flag only when the batch shrinks would stall the queue
+        # until an unrelated request finishes.
+        s = self._scheduler()
+        running_batch = ScheduleBatch(
+            reqs=[SimpleNamespace(rid="running-0")], batch_is_full=True
+        )
+
+        with self.assertRaises(_AdmissionAttempted):
+            s._get_new_batch_prefill_raw(
+                prefill_delayer_single_pass=None, running_batch=running_batch
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_prefill_memory_budget.py

@@ -21,7 +21,7 @@
 from sglang.srt.mem_cache.common import evict_from_tree_cache
 from sglang.srt.mem_cache.kv.alloc.chain import ChainAllocator
 from sglang.srt.mem_cache.kv.alloc.interim import _FloatChainPolicy
-from sglang.srt.mem_cache.prefill_budget import ChainPrefillBudget
+from sglang.srt.mem_cache.prefill_budget import ChainPrefillBudget, PrefillBudget
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.sampling.sampling_params import SamplingParams
@@ -277,6 +277,21 @@ def test_cached_admission_reserves_fixed_length_decode(self):
                     self.assertEqual(result, AddReqResult.NO_TOKEN)
 
 
+class TestContinuingChunkCap(CustomTestCase):
+    def _budget(self, free_tokens):
+        allocator = SimpleNamespace(page_size=16, available_size=lambda: free_tokens)
+        tree_cache = SimpleNamespace(
+            supports_mamba=lambda: False, evictable_size=lambda: 0
+        )
+        return PrefillBudget(allocator, tree_cache)
+
+    def test_chunk_leaves_the_page_alloc_extend_over_asks(self):
+        self.assertEqual(self._budget(64).available_chunk_tokens(1024), 48)
+
+    def test_chunk_is_deferred_without_free_kv(self):
+        self.assertIsNone(self._budget(16).available_chunk_tokens(1024))
+
+
 class TestFixedPrefillMemoryBudget(unittest.TestCase):
     def _allocator(self, cls=SWATokenToKVPoolAllocator):
         allocator = object.__new__(cls)

test/registered/unit/models/test_aiter_attention_backend_handler.py

@@ -0,0 +1,26 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.srt.models.deepseek_common import attention_backend_handler as handler
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class TestAiterOverCapacityExtend(CustomTestCase):
+    def test_over_capacity_extend_stays_chunked_in_graph_context(self):
+        batch = SimpleNamespace(
+            forward_mode=SimpleNamespace(is_extend_without_speculative=lambda: True)
+        )
+        with (
+            patch.object(handler, "_support_mha_one_shot", return_value=False),
+            patch.object(handler, "is_in_breakable_cuda_graph", return_value=True),
+        ):
+            method = handler.handle_attention_aiter(None, batch)
+        self.assertEqual(method, handler.AttnForwardMethod.MHA_CHUNKED_KV)
+
+
+if __name__ == "__main__":
+    unittest.main()