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)