[Refactor] scheduler: Hold wrapped requests in request lists and abort through each list

python/sglang/srt/disaggregation/decode_hicache_mixin.py

@@ -169,7 +169,7 @@ def _hicache_pending_restore_tokens(self) -> int:
             return 0
         return sum(
             dr.prefix_match.restore_token_count
-            for dr in self.transfer_queue.queue
+            for dr in self.transfer_queue.queue.items()
             if dr.prefix_match is not None and dr.hicache_restore_lock is None
         )
 

python/sglang/srt/disaggregation/utils.py

@@ -182,7 +182,7 @@ def busy() -> bool:
             return bool(
                 any(
                     not decode_req.host_staged
-                    for decode_req in scheduler.disagg_decode_transfer_queue.queue
+                    for decode_req in scheduler.disagg_decode_transfer_queue.queue.items()
                 )
                 or scheduler.disagg_decode_prealloc_queue.has_published_destinations
                 or (

python/sglang/srt/managers/hisparse_coordinator.py

@@ -2,7 +2,15 @@
 
 import logging
 import weakref
-from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    List,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import torch
 
@@ -39,7 +47,11 @@ def load_blocks_to_device_buffer_mha(*args, **kwargs):
 from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import Req
-from sglang.srt.managers.scheduler_components.request_list import ListKind
+from sglang.srt.managers.scheduler_components.request_list import (
+    ListKind,
+    RequestList,
+    _RequestList,
+)
 from sglang.srt.mem_cache.allocator.hisparse import (
     DeepSeekV4HiSparseTokenToKVPoolAllocator,
     HiSparseTokenToKVPoolAllocator,
@@ -666,7 +678,9 @@ def __init__(
 
         self.write_staging_stream = device_module.Stream()
         self.decode_backup_stream = device_module.Stream()
-        self.ack_staging_queue: List[HiSparseAct] = []
+        self.ack_staging_queue = _RequestList(
+            ListKind.HISPARSE_STAGING, req_of=lambda x: x.req
+        )
         self.decode_producer_stream = None
         self._backup_done_event = device_module.Event()
         self._has_pending_backup = False
@@ -874,7 +888,7 @@ def admit_request_into_staging(self, req: Req) -> None:
             if device_indices.is_cuda:
                 device_indices.record_stream(self.write_staging_stream)
 
-        self.ack_staging_queue.append(HiSparseAct(start_event, finish_event, req))
+        self.ack_staging_queue.join([HiSparseAct(start_event, finish_event, req)])
 
     def admit_request_direct(self, req: Req) -> None:
         """Direct-to-host path: KV data already resides in host pool via RDMA.
@@ -1058,19 +1072,17 @@ def _grow_device_buffers(
         reserved_positions = (seq_lens - 1).clamp(max=self.device_buffer_size)
         return self.req_to_device_buffer[req_pool_indices, reserved_positions]
 
-    def request_lists(self) -> Dict[ListKind, List[Req]]:
-        return {ListKind.HISPARSE_STAGING: [act.req for act in self.ack_staging_queue]}
-
-    def has_ongoing_staging(self) -> bool:
-        return len(self.ack_staging_queue) > 0
+    def request_lists(self) -> Dict[ListKind, RequestList]:
+        return {ListKind.HISPARSE_STAGING: self.ack_staging_queue}
 
     def collect_ready_reqs(self) -> List[Req]:
         ready_reqs: List[Req] = []
         if len(self.ack_staging_queue) == 0:
             return ready_reqs
 
+        staged = list(self.ack_staging_queue.items())
         finish_count = 0
-        for _, finish_event, _ in self.ack_staging_queue:
+        for _, finish_event, _ in staged:
             if not finish_event.query():
                 break
             finish_count += 1
@@ -1083,13 +1095,12 @@ def collect_ready_reqs(self) -> List[Req]:
                 group=self.tp_group,
             )
         finish_count = int(queue_size.item())
-        while finish_count > 0:
-            _, _, req = self.ack_staging_queue.pop(0)
+        for _, _, req in staged[:finish_count]:
+            self.ack_staging_queue.leave([req])
             # prepare device buffer and update req
             self.alloc_device_buffer(req)
             self._skip_first_backup[req.kv.req_pool_idx] = True
             req.hisparse_staging = False
-            finish_count -= 1
             ready_reqs.append(req)
         return ready_reqs
 
@@ -1371,9 +1382,7 @@ def abort_staging_request(self, req: Req) -> None:
         but has not yet completed (i.e. req.hisparse_staging is True).
         """
         # Remove from staging queue
-        self.ack_staging_queue = [
-            act for act in self.ack_staging_queue if act.req is not req
-        ]
+        self.ack_staging_queue.leave([req])
         # Wait for any in-flight staging DMA to complete before freeing
         self.write_staging_stream.synchronize()
 

python/sglang/srt/managers/scheduler.py

@@ -2582,6 +2582,7 @@ def maybe_init_uocr_runtime(self) -> None:
             decode_graph_io=lambda: resolve_uocr_decode_graph_io(
                 runner.decode_cuda_graph_runner
             ),
+            abort_windows=self._abort_uocr_doc_windows,
         )
         self.resolve_forward_inputs = self.uocr_runtime.resolve_forward_inputs
 
@@ -5698,6 +5699,7 @@ def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool:
                 ListKind.PD_HELD,
                 # Counted by uocr blocks_flush below.
                 ListKind.UOCR_PARKED,
+                ListKind.UOCR_RELEASED,
                 # Finished, and held only on the output rank; all ranks must agree.
                 ListKind.DEFERRED,
             }
@@ -5989,12 +5991,10 @@ def record_weight_version_change(self, new_version: Optional[str]) -> None:
             *self.collect_inflight_reqs(),
             *self.waiting_queue,
             *([self.chunked_req] if self.chunked_req is not None else []),
-            *self.deferred_responses.held_reqs(),
+            *self.deferred_responses.held,
         }
         if self.hisparse_coordinator is not None:
-            live_reqs.update(
-                act.req for act in self.hisparse_coordinator.ack_staging_queue
-            )
+            live_reqs.update(self.hisparse_coordinator.ack_staging_queue)
         num_recorded = record_weight_version_events(live_reqs, old_version=old_version)
         logger.info(
             f"Weight version changed. {old_version=} {new_version=} {num_recorded=}"
@@ -6019,7 +6019,6 @@ def _collect_inflight_batches(self) -> list[ScheduleBatch | None]:
         return inflight_batches
 
     def abort_request(self, recv_req: AbortReq):
-        self.deferred_responses.abort(recv_req)
         if (chunked_req := self.chunked_req) is not None:
             if recv_req.matches(chunked_req.rid):
                 self._pending_chunked_abort_req = chunked_req
@@ -6035,15 +6034,6 @@ def abort_request(self, recv_req: AbortReq):
             if recv_req.matches(req.rid):
                 to_del.append(i)
 
-        if self.uocr_runtime is not None:
-            # Parked document windows are off the waiting queue; abort them alike.
-            self._abort_uocr_doc_windows(
-                self.uocr_runtime.drop_parked(
-                    "" if recv_req.abort_all else recv_req.rid
-                ),
-                finished_reason=recv_req.finished_reason,
-            )
-
         # Sort in reverse order to avoid index issues when deleting
         for i in reversed(to_del):
             # Abort method 1: directly pop from the queue
@@ -6082,29 +6072,6 @@ def abort_request(self, recv_req: AbortReq):
             if isinstance(request_list, RequestList):
                 request_list.abort(recv_req)
 
-        # Delete requests not in the waiting queue when PD disaggregation is enabled
-        if self.disaggregation_mode == DisaggregationMode.DECODE:
-            # Abort requests that have not yet finished preallocation
-            for decode_req in self.disagg_decode_prealloc_queue.queue:
-                if recv_req.matches(decode_req.req.rid):
-                    logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}")
-                    decode_req.kv_receiver.abort()
-                    if get_parallel().pp_size > 1:
-                        prepare_abort(decode_req.req, "Aborted by AbortReq.")
-
-            # Abort requests waiting for kvcache to release tree cache
-            for decode_req in self.disagg_decode_transfer_queue.queue:
-                if recv_req.matches(decode_req.req.rid):
-                    logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}")
-                    if decode_req.host_staged:
-                        # Keep the host destination alive until prefill stops writing.
-                        prepare_abort(decode_req.req, "Aborted by AbortReq.")
-                        continue
-                    # The receiver arms drain-ack accounting before sending the
-                    # ABORT (see CommonKVReceiver._send_abort_notification), so
-                    # an ack racing back is never dropped.
-                    decode_req.kv_receiver.abort()
-
         # Delete requests in the running batch
         for req in self.collect_inflight_reqs():
             if not req.finished() and recv_req.matches(req.rid):

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

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Callable, Dict, Optional
 
 from sglang.srt.disaggregation.mode import DisaggregationMode
 from sglang.srt.managers.req_types import FINISH_ABORT
-from sglang.srt.managers.scheduler_components.request_list import ListKind
+from sglang.srt.managers.scheduler_components.request_list import (
+    ListKind,
+    RequestList,
+    _RequestList,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.managers.deferred_output import ResponseCompletion
@@ -32,38 +36,38 @@ def __init__(
         self.disaggregation_mode = disaggregation_mode
         self.output_streamer = output_streamer
         self.send_abort = send_abort
-        self.held: Dict[str, Tuple[Req, ResponseCompletion]] = {}
+        # (req, completion) per held response.
+        self.held = _RequestList(
+            ListKind.DEFERRED, abort=self._abort, req_of=lambda x: x[0]
+        )
 
-    def held_reqs(self) -> List[Req]:
-        return [req for req, _ in self.held.values()]
-
-    def request_lists(self) -> Dict[ListKind, List[Req]]:
-        return {ListKind.DEFERRED: self.held_reqs()}
+    def request_lists(self) -> Dict[ListKind, RequestList]:
+        return {ListKind.DEFERRED: self.held}
 
     def defer(self, req: Req, completion: ResponseCompletion) -> None:
         if self.disaggregation_mode != DisaggregationMode.NULL:
             raise ValueError("deferred responses do not support PD disaggregation")
         if not req.finished() or req.finished_output:
             raise ValueError("only finished, unstreamed responses can be deferred")
-        if req.rid in self.held:
+        if any(held.rid == req.rid for held in self.held):
             raise ValueError(f"response {req.rid!r} is already deferred")
-        self.held[req.rid] = (req, completion)
+        self.held.join([(req, completion)])
         req.defer_output = True
         self.output_streamer.defer_outputs = True
 
     def stream_released(self) -> None:
         if not self.held:
             return
         ready = []
-        for rid, (req, completion) in list(self.held.items()):
+        for req, completion in self.held.items():
             if not completion.done():
                 continue
-            del self.held[rid]
+            self.held.leave([req])
             req.defer_output = False
             try:
                 completion.result()
             except Exception as exc:
-                logger.exception("Deferred response failed for %s", rid)
+                logger.exception("Deferred response failed for %s", req.rid)
                 self._abort_response(
                     req, FINISH_ABORT(str(exc), status_code=500).to_json()
                 )
@@ -84,12 +88,12 @@ def _abort_response(
         req.finished_output = True
         self.send_abort(req, finished_reason)
 
-    def abort(self, recv_req: AbortReq) -> None:
-        for rid, (req, completion) in list(self.held.items()):
-            if recv_req.matches(rid):
-                del self.held[rid]
+    def _abort(self, recv_req: AbortReq) -> None:
+        for req, completion in self.held.items():
+            if recv_req.matches(req.rid):
+                self.held.leave([req])
                 try:
                     completion.cancel()
                 except Exception:
-                    logger.exception("Failed to cancel deferred work for %s", rid)
+                    logger.exception("Failed to cancel deferred work for %s", req.rid)
                 self._abort_response(req, recv_req.finished_reason)

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

@@ -199,7 +199,7 @@ def get_loads(self) -> LoadSnapshot:
             )
             ready_reqs = [
                 decode_req.req
-                for decode_req in self.get_disagg_decode_prealloc_queue().queue
+                for decode_req in self.get_disagg_decode_prealloc_queue().queue.items()
                 if decode_req.waiting_for_input
             ]
             decode_prealloc_ready = len(ready_reqs)

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

@@ -1449,7 +1449,7 @@ def _update_decode_queue_stats(self, priority_enabled: bool) -> None:
             self.scheduler.disagg_decode_transfer_queue.num_pending_deferred_releases()
         )
         host_reqs = (
-            [req for req in transfer_queue if req.host_staged]
+            [d.req for d in transfer_queue.items() if d.host_staged]
             if get_disagg().disaggregation_decode_host_receive_threshold < 1
             else []
         )