perf(disaggregation): compact prefill bootstrap queue in one pass

python/sglang/srt/disaggregation/prefill.py

@@ -260,7 +260,7 @@ class PrefillBootstrapQueue:
 
         bootstrapped_reqs = []
         failed_reqs = []
-        indices_to_remove = set()
+        remaining_queue = []
 
         if len(self.queue) == 0:
             if return_failed_reqs is False:
@@ -278,9 +278,11 @@ class PrefillBootstrapQueue:
             if rids_to_check is not None:
                 # if req not in reqs_info_to_check, skip
                 if req.rid not in rids_to_check:
+                    remaining_queue.append(req)
                     continue
 
             if poll == KVPoll.Bootstrapping:
+                remaining_queue.append(req)
                 continue
             elif poll == KVPoll.Failed:
                 error_message = f"Prefill bootstrap failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}"
@@ -294,7 +296,6 @@ class PrefillBootstrapQueue:
                     req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
                 )
                 self.scheduler.stream_output([req], req.return_logprob)
-                indices_to_remove.add(i)
                 failed_reqs.append(req)
                 if self.scheduler.enable_metrics:
                     self.scheduler.metrics_collector.increment_bootstrap_failed_reqs()
@@ -307,6 +308,8 @@ class PrefillBootstrapQueue:
             req.time_stats.set_bootstrap_done_time()
             num_kv_indices = len(req.origin_input_ids)
             if self.req_to_metadata_buffer_idx_allocator.available_size() == 0:
+                remaining_queue.append(req)
+                remaining_queue.extend(self.queue[i + 1 :])
                 break
 
             req.metadata_buffer_index = (
@@ -318,12 +321,9 @@ class PrefillBootstrapQueue:
             req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index)
 
             bootstrapped_reqs.append(req)
-            indices_to_remove.add(i)
             req.time_stats.set_wait_queue_entry_time()
 
-        self.queue = [
-            entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
-        ]
+        self.queue = remaining_queue
 
         if return_failed_reqs is False:
             return bootstrapped_reqs
@@ -563,7 +563,6 @@ class SchedulerDisaggregationPrefillMixin:
         undone_reqs: List[Req] = []
         # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue
         for req, poll in zip(self.disagg_prefill_inflight_queue, polls):
-
             if rids_to_check is not None:
                 if req.rid not in rids_to_check:
                     undone_reqs.append(req)

test/registered/unit/disaggregation/test_prefill_bootstrap_queue.py

@@ -0,0 +1,129 @@
+"""Unit tests for prefill bootstrap queue compaction."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from sglang.srt.disaggregation.base import KVPoll
+from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue
+from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=6, suite="stage-a-test-cpu")
+
+
+class FakeSender:
+    def __init__(self, should_fail=False):
+        self.should_fail = should_fail
+        self.init_calls = []
+
+    def init(self, num_pages, metadata_buffer_index):
+        self.init_calls.append((num_pages, metadata_buffer_index))
+
+    def failure_exception(self):
+        if self.should_fail:
+            raise RuntimeError("boom")
+
+
+class TestPrefillBootstrapQueue(CustomTestCase):
+    def _make_req(self, rid, bootstrap_room, origin_input_ids, sender):
+        req = MagicMock()
+        req.rid = rid
+        req.bootstrap_room = bootstrap_room
+        req.origin_input_ids = origin_input_ids
+        req.disagg_kv_sender = sender
+        req.return_logprob = False
+        req.metadata_buffer_index = -1
+        req.time_stats = SimpleNamespace(
+            set_bootstrap_done_time=MagicMock(),
+            set_wait_queue_entry_time=MagicMock(),
+            trace_ctx=SimpleNamespace(abort=MagicMock()),
+        )
+        return req
+
+    def _build_queue(self, allocator_size=4):
+        queue = PrefillBootstrapQueue.__new__(PrefillBootstrapQueue)
+        queue.req_to_metadata_buffer_idx_allocator = ReqToMetadataIdxAllocator(
+            size=allocator_size
+        )
+        queue.token_to_kv_pool = SimpleNamespace(page_size=4)
+        queue.scheduler = SimpleNamespace(
+            stream_output=MagicMock(),
+            enable_metrics=False,
+            enable_hicache_storage=False,
+            attn_cp_cpu_group=MagicMock(),
+            attn_tp_cpu_group=MagicMock(),
+        )
+        queue.tp_rank = 0
+        queue.queue = []
+        return queue
+
+    def test_pop_bootstrapped_compacts_queue_in_one_pass(self):
+        queue = self._build_queue(allocator_size=2)
+
+        waiting = self._make_req("wait", 1, list(range(8)), FakeSender())
+        ready = self._make_req("ready", 2, list(range(10)), FakeSender())
+        failed = self._make_req("failed", 3, list(range(6)), FakeSender(True))
+        queue.queue = [waiting, ready, failed]
+
+        with patch(
+            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
+            return_value=[
+                KVPoll.Bootstrapping,
+                KVPoll.WaitingForInput,
+                KVPoll.Failed,
+            ],
+        ):
+            bootstrapped, failed_reqs = queue.pop_bootstrapped(return_failed_reqs=True)
+
+        self.assertEqual(bootstrapped, [ready])
+        self.assertEqual(failed_reqs, [failed])
+        self.assertEqual(queue.queue, [waiting])
+        self.assertEqual(ready.metadata_buffer_index, 0)
+        self.assertEqual(ready.disagg_kv_sender.init_calls, [(3, 0)])
+
+    def test_pop_bootstrapped_stops_on_metadata_exhaustion(self):
+        queue = self._build_queue(allocator_size=1)
+
+        first_ready = self._make_req("ready-1", 1, list(range(8)), FakeSender())
+        second_ready = self._make_req("ready-2", 2, list(range(12)), FakeSender())
+        trailing = self._make_req("trailing", 3, list(range(4)), FakeSender())
+        queue.queue = [first_ready, second_ready, trailing]
+
+        with patch(
+            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
+            return_value=[
+                KVPoll.WaitingForInput,
+                KVPoll.WaitingForInput,
+                KVPoll.Bootstrapping,
+            ],
+        ):
+            bootstrapped = queue.pop_bootstrapped()
+
+        self.assertEqual(bootstrapped, [first_ready])
+        self.assertEqual(queue.queue, [second_ready, trailing])
+        self.assertEqual(first_ready.disagg_kv_sender.init_calls, [(2, 0)])
+        self.assertEqual(second_ready.disagg_kv_sender.init_calls, [])
+
+    def test_pop_bootstrapped_preserves_unchecked_entries(self):
+        queue = self._build_queue(allocator_size=2)
+
+        skipped = self._make_req("skip", 1, list(range(8)), FakeSender())
+        checked = self._make_req("check", 2, list(range(10)), FakeSender())
+        queue.queue = [skipped, checked]
+
+        with patch(
+            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
+            return_value=[KVPoll.WaitingForInput, KVPoll.WaitingForInput],
+        ):
+            bootstrapped = queue.pop_bootstrapped(rids_to_check=["check"])
+
+        self.assertEqual(bootstrapped, [checked])
+        self.assertEqual(queue.queue, [skipped])
+        self.assertEqual(skipped.disagg_kv_sender.init_calls, [])
+        self.assertEqual(checked.disagg_kv_sender.init_calls, [(3, 0)])
+
+
+if __name__ == "__main__":
+    unittest.main()