From 6237271fc9a411207f90420007ca5ad18e2a87c7 Mon Sep 17 00:00:00 2001 From: wxiwnd Date: Wed, 8 Apr 2026 22:37:40 +0800 Subject: [PATCH] perf(disaggregation): compact decode queues in one pass --- python/sglang/srt/disaggregation/decode.py | 75 +-- .../test_decode_queue_compaction.py | 495 ++++++++++++++++++ 2 files changed, 537 insertions(+), 33 deletions(-) create mode 100644 test/registered/unit/disaggregation/test_decode_queue_compaction.py diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 90ea6117c..226515a04 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -433,14 +433,17 @@ class DecodePreallocQueue: # allocate memory resumed_reqs = [] - indices_to_remove = set() + remaining_retracted_queue = [] allocatable_tokens = self._allocatable_tokens(count_retracted=False) for i, req in enumerate(self.retracted_queue): if rids_to_check is not None and req.rid not in rids_to_check: + remaining_retracted_queue.append(req) continue if self.req_to_token_pool.available_size() <= 0: + remaining_retracted_queue.append(req) + remaining_retracted_queue.extend(self.retracted_queue[i + 1 :]) break required_tokens_for_request = ( @@ -449,10 +452,11 @@ class DecodePreallocQueue: + self.num_reserved_decode_tokens ) if required_tokens_for_request > allocatable_tokens: + remaining_retracted_queue.append(req) + remaining_retracted_queue.extend(self.retracted_queue[i + 1 :]) break resumed_reqs.append(req) - indices_to_remove.add(i) req.is_retracted = False self._pre_alloc(req) allocatable_tokens -= required_tokens_for_request @@ -460,11 +464,7 @@ class DecodePreallocQueue: # load from cpu, release the cpu copy req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator) - self.retracted_queue = [ - entry - for i, entry in enumerate(self.retracted_queue) - if i not in indices_to_remove - ] + self.retracted_queue = remaining_retracted_queue return resumed_reqs @@ -605,7 +605,8 @@ class DecodePreallocQueue: failed_reqs = [] preallocated_reqs = [] - indices_to_remove = set() + remaining_queue = [] + prealloc_blocked = False # We need to make sure that the sum of inflight tokens and allocatable tokens is greater than maximum input+output length of each inflight request # Otherwise it is possible for one request running decode out of memory, while all other requests are in the transfer queue that cannot be retracted. @@ -616,33 +617,35 @@ class DecodePreallocQueue: allocatable_tokens = self._allocatable_tokens( retractable_tokens=retractable_tokens, count_retracted=True ) - # First, remove all failed requests from the queue for i, decode_req in enumerate(self.queue): if rids_to_check is not None and decode_req.req.rid not in rids_to_check: + remaining_queue.append(decode_req) continue + if isinstance(decode_req.req.finished_reason, FINISH_ABORT): self.scheduler.stream_output( [decode_req.req], decode_req.req.return_logprob ) failed_reqs.append(decode_req) - indices_to_remove.add(i) - - # Then, preallocate the remaining requests if possible - for i, decode_req in enumerate(self.queue): - if rids_to_check is not None and decode_req.req.rid not in rids_to_check: continue - if i in indices_to_remove: + if prealloc_blocked: + remaining_queue.append(decode_req) continue if not decode_req.waiting_for_input: + remaining_queue.append(decode_req) continue if self.req_to_token_pool.available_size() <= 0: - break + remaining_queue.append(decode_req) + prealloc_blocked = True + continue if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0: - break + remaining_queue.append(decode_req) + prealloc_blocked = True + continue # Memory estimation: don't add if the projected memory cannot be met # TODO: add new_token ratio @@ -663,9 +666,13 @@ class DecodePreallocQueue: ) > allocatable_tokens ): - break + remaining_queue.append(decode_req) + prealloc_blocked = True + continue if required_tokens_for_request > allocatable_tokens: - break + remaining_queue.append(decode_req) + prealloc_blocked = True + continue allocatable_tokens -= required_tokens_for_request self._pre_alloc(decode_req.req) @@ -727,12 +734,9 @@ class DecodePreallocQueue: page_indices, decode_req.metadata_buffer_index, state_indices ) preallocated_reqs.append(decode_req) - indices_to_remove.add(i) decode_req.req.time_stats.set_decode_transfer_queue_entry_time() - self.queue = [ - entry for i, entry in enumerate(self.queue) if i not in indices_to_remove - ] + self.queue = remaining_queue return preallocated_reqs, failed_reqs @@ -998,9 +1002,11 @@ class DecodeTransferQueue: ) transferred_reqs = [] - indices_to_remove = set() - for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): + completed_decode_reqs = [] + remaining_queue = [] + for decode_req, poll in zip(self.queue, polls): if rids_to_check is not None and decode_req.req.rid not in rids_to_check: + remaining_queue.append(decode_req) continue if poll == KVPoll.Failed: error_message = f"Decode transfer failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}" @@ -1019,14 +1025,14 @@ class DecodeTransferQueue: ) # release pre-allocated kv cache, but don't insert into the tree since it's failed release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) - indices_to_remove.add(i) + completed_decode_reqs.append(decode_req) if self.scheduler.enable_metrics: self.scheduler.metrics_collector.increment_transfer_failed_reqs() continue elif poll == KVPoll.Success: should_remove = self._commit_transfer_to_req(decode_req) if should_remove: - indices_to_remove.add(i) + completed_decode_reqs.append(decode_req) # Check if request was aborted due to corruption if isinstance(decode_req.req.finished_reason, FINISH_ABORT): self.scheduler.stream_output( @@ -1039,23 +1045,26 @@ class DecodeTransferQueue: self.scheduler.metrics_collector.increment_transfer_failed_reqs() else: transferred_reqs.append(decode_req.req) + else: + remaining_queue.append(decode_req) elif poll in [ KVPoll.Bootstrapping, KVPoll.WaitingForInput, KVPoll.Transferring, ]: - pass + remaining_queue.append(decode_req) else: raise ValueError(f"Unexpected poll case: {poll}") - for i in indices_to_remove: - idx = self.queue[i].metadata_buffer_index + if len(polls) < len(self.queue): + remaining_queue.extend(self.queue[len(polls) :]) + + for decode_req in completed_decode_reqs: + idx = decode_req.metadata_buffer_index assert idx != -1 self.req_to_metadata_buffer_idx_allocator.free(idx) - self.queue = [ - entry for i, entry in enumerate(self.queue) if i not in indices_to_remove - ] + self.queue = remaining_queue return transferred_reqs diff --git a/test/registered/unit/disaggregation/test_decode_queue_compaction.py b/test/registered/unit/disaggregation/test_decode_queue_compaction.py new file mode 100644 index 000000000..928c44824 --- /dev/null +++ b/test/registered/unit/disaggregation/test_decode_queue_compaction.py @@ -0,0 +1,495 @@ +"""Unit tests for decode queue one-pass compaction.""" + +import unittest +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import patch + +from sglang.srt.disaggregation.base import KVPoll +from sglang.srt.disaggregation.decode import ( + DecodePreallocQueue, + DecodeRequest, + DecodeTransferQueue, +) +from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator +from sglang.srt.managers.schedule_batch import FINISH_ABORT +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=8, suite="stage-a-test-cpu") + + +class FakeReq: + def __init__(self, rid, bootstrap_room): + self.rid = rid + self.bootstrap_room = bootstrap_room + self.bootstrap_host = "host" + self.return_logprob = False + self.latencies = [] + self.origin_input_ids = [] + self.output_ids = [] + self.is_retracted = True + self.load_calls = [] + self.prealloc_done = False + self.finished_reason = cast(Any, None) + self.cached_tokens = 0 + self.sampling_params = SimpleNamespace(max_new_tokens=0) + self.time_stats = SimpleNamespace( + set_bootstrap_done_time=lambda: None, + set_decode_transfer_queue_entry_time=lambda: None, + set_wait_queue_entry_time=lambda: None, + ) + + def add_latency(self, stage): + self.latencies.append(stage) + + def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): + self.load_calls.append((req_to_token_pool, token_to_kv_pool_allocator)) + + +class FakeReceiver: + def __init__(self, should_fail=False): + self.should_fail = should_fail + self.init_calls = [] + self.clear_calls = 0 + + def init(self, *args): + self.init_calls.append(args) + + def failure_exception(self): + if self.should_fail: + raise RuntimeError("boom") + + def clear(self): + self.clear_calls += 1 + return None + + +class FakeItem: + def __init__(self, value): + self.value = value + + def item(self): + return self.value + + +class FakeTensor: + def __init__(self, values): + self.values = values + + def __getitem__(self, item): + if isinstance(item, tuple): + row, col = item + return FakeTensor(self.values[row][col]) + return FakeTensor(self.values[item]) + + def cpu(self): + return self + + def numpy(self): + return self.values + + +class FakeAllocator(ReqToMetadataIdxAllocator): + def __init__(self): + super().__init__(size=0) + self.freed = [] + + def free(self, free_index): + self.freed.append(free_index) + + +class TestDecodeQueueCompaction(CustomTestCase): + def test_decode_transfer_queue_compacts_in_one_pass(self): + streamed = [] + released = [] + committed = [] + allocator = FakeAllocator() + + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.gloo_group = None + queue.req_to_metadata_buffer_idx_allocator = allocator + queue.tp_rank = 0 + queue.metadata_buffers = cast(Any, object()) + queue.tree_cache = cast(Any, object()) + queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True)) + queue.scheduler = cast( + Any, + SimpleNamespace( + stream_output=lambda reqs, return_logprob: streamed.extend( + req.rid for req in reqs + ), + enable_metrics=False, + token_to_kv_pool_allocator=SimpleNamespace( + get_kvcache=lambda: SimpleNamespace() + ), + ), + ) + + keep = DecodeRequest( + req=cast(Any, FakeReq("keep", 1)), + kv_receiver=cast(Any, FakeReceiver()), + metadata_buffer_index=10, + ) + success = DecodeRequest( + req=cast(Any, FakeReq("success", 2)), + kv_receiver=cast(Any, FakeReceiver()), + metadata_buffer_index=11, + ) + failed = DecodeRequest( + req=cast(Any, FakeReq("failed", 3)), + kv_receiver=cast(Any, FakeReceiver(should_fail=True)), + metadata_buffer_index=12, + ) + skipped = DecodeRequest( + req=cast(Any, FakeReq("skip", 4)), + kv_receiver=cast(Any, FakeReceiver()), + metadata_buffer_index=13, + ) + queue.queue = [keep, success, failed, skipped] + + with ( + patch( + "sglang.srt.disaggregation.decode.poll_and_all_reduce", + return_value=[ + KVPoll.Transferring, + KVPoll.Success, + KVPoll.Failed, + KVPoll.Success, + ], + ), + patch( + "sglang.srt.disaggregation.decode.release_kv_cache", + lambda req, tree_cache, is_insert=False: released.append( + (req.rid, is_insert) + ), + ), + ): + queue._commit_transfer_to_req = lambda decode_req: ( + committed.append(decode_req.req.rid) or True + ) + transferred = queue.pop_transferred( + rids_to_check=["keep", "success", "failed"] + ) + + self.assertEqual([req.rid for req in transferred], ["success"]) + self.assertEqual(committed, ["success"]) + self.assertEqual(streamed, ["failed"]) + self.assertEqual(released, [("failed", False)]) + self.assertEqual(allocator.freed, [11, 12]) + self.assertEqual(queue.queue, [keep, skipped]) + + def test_decode_transfer_queue_keeps_metadata_waiters(self): + allocator = FakeAllocator() + + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.gloo_group = None + queue.req_to_metadata_buffer_idx_allocator = allocator + queue.tp_rank = 0 + queue.metadata_buffers = cast(Any, object()) + queue.tree_cache = cast(Any, object()) + queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True)) + queue.scheduler = cast( + Any, + SimpleNamespace( + stream_output=lambda reqs, return_logprob: None, + enable_metrics=False, + token_to_kv_pool_allocator=SimpleNamespace( + get_kvcache=lambda: SimpleNamespace() + ), + ), + ) + + waiting = DecodeRequest( + req=cast(Any, FakeReq("waiting", 1)), + kv_receiver=cast(Any, FakeReceiver()), + metadata_buffer_index=10, + ) + queue.queue = [waiting] + + with patch( + "sglang.srt.disaggregation.decode.poll_and_all_reduce", + return_value=[KVPoll.Success], + ): + queue._commit_transfer_to_req = lambda decode_req: False + transferred = queue.pop_transferred() + + self.assertEqual(transferred, []) + self.assertEqual(allocator.freed, []) + self.assertEqual(queue.queue, [waiting]) + + def test_commit_transfer_to_req_waits_for_real_metadata(self): + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.metadata_buffers = cast( + Any, + SimpleNamespace( + get_buf=lambda idx: ( + [FakeItem(0)], + [FakeItem(0)], + None, + None, + None, + None, + None, + None, + None, + [FakeItem(0)], + ) + ), + ) + queue.scheduler = cast( + Any, + SimpleNamespace( + server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake") + ), + ) + queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True)) + + receiver = FakeReceiver() + decode_req = DecodeRequest( + req=cast(Any, FakeReq("waiting", 3)), + kv_receiver=cast(Any, receiver), + metadata_buffer_index=9, + ) + + should_remove = queue._commit_transfer_to_req(decode_req) + + self.assertIs(should_remove, False) + self.assertIs(decode_req.kv_receiver, receiver) + self.assertEqual(receiver.clear_calls, 0) + self.assertEqual(decode_req.req.output_ids, []) + + def test_commit_transfer_to_req_aborts_on_room_mismatch(self): + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.metadata_buffers = cast( + Any, + SimpleNamespace( + get_buf=lambda idx: ( + [FakeItem(0)], + [FakeItem(0)], + None, + None, + None, + None, + None, + None, + None, + [FakeItem(99)], + ) + ), + ) + queue.scheduler = cast( + Any, + SimpleNamespace( + server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake") + ), + ) + queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True)) + + receiver = FakeReceiver() + decode_req = DecodeRequest( + req=cast(Any, FakeReq("corrupt", 3)), + kv_receiver=cast(Any, receiver), + metadata_buffer_index=9, + ) + + aborted = [] + with patch( + "sglang.srt.disaggregation.decode.prepare_abort", + lambda req, message, status_code: aborted.append( + (req.rid, message, status_code) + ), + ): + should_remove = queue._commit_transfer_to_req(decode_req) + + self.assertIs(should_remove, True) + self.assertEqual(receiver.clear_calls, 1) + self.assertIsNone(decode_req.kv_receiver) + self.assertEqual(len(aborted), 1) + self.assertEqual(aborted[0][0], "corrupt") + + def test_resume_retracted_reqs_compacts_queue_in_one_pass(self): + prealloc_queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + prealloc_queue.req_to_token_pool = cast( + Any, SimpleNamespace(available_size=lambda: 1) + ) + prealloc_queue.token_to_kv_pool_allocator = cast(Any, object()) + prealloc_queue.num_reserved_decode_tokens = 1 + prealloc_queue._allocatable_tokens = ( + lambda retractable_tokens=None, count_retracted=False: 8 + ) + prealloc_queue._pre_alloc = lambda req: setattr(req, "prealloc_done", True) + + first = FakeReq("resume", 1) + first.origin_input_ids = [1, 2] + first.output_ids = [3] + skipped = FakeReq("skip", 2) + skipped.origin_input_ids = [4] + blocked = FakeReq("blocked", 3) + blocked.origin_input_ids = [5, 6, 7, 8, 9, 10, 11, 12] + + prealloc_queue.retracted_queue = cast(Any, [first, skipped, blocked]) + + resumed = prealloc_queue.resume_retracted_reqs( + rids_to_check=["resume", "blocked"] + ) + + self.assertEqual(resumed, [first]) + self.assertIs(first.is_retracted, False) + self.assertIs(first.prealloc_done, True) + self.assertEqual(len(first.load_calls), 1) + self.assertEqual(prealloc_queue.retracted_queue, [skipped, blocked]) + + def test_pop_preallocated_still_removes_failed_reqs_after_block(self): + streamed = [] + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue._resolve_pending_reqs = lambda: None + queue._update_handshake_waiters = lambda rids_to_check=None: None + queue.req_to_token_pool = cast( + Any, + SimpleNamespace( + available_size=lambda: 1, + req_to_token=FakeTensor([[1, 2, 3, 4, 5, 6, 7, 8]]), + ), + ) + queue.req_to_metadata_buffer_idx_allocator = cast( + Any, SimpleNamespace(available_size=lambda: 1, alloc=lambda: 7) + ) + queue.token_to_kv_pool_allocator = cast(Any, SimpleNamespace(page_size=1)) + queue.token_to_kv_pool = cast(Any, object()) + queue.num_reserved_decode_tokens = 1 + queue.scheduler = cast( + Any, + SimpleNamespace( + stream_output=lambda reqs, return_logprob: streamed.extend( + req.rid for req in reqs + ), + enable_metrics=False, + running_batch=SimpleNamespace(reqs=[]), + sliding_window_size=4, + ), + ) + queue._allocatable_tokens = ( + lambda retractable_tokens=None, count_retracted=True: 2 + ) + queue._pre_alloc = lambda req: setattr(req, "req_pool_idx", 0) + + blocked_req = FakeReq("blocked", 1) + blocked_req.origin_input_ids = [1, 2, 3] + blocked_req.sampling_params = SimpleNamespace(max_new_tokens=8) + blocked = DecodeRequest( + req=cast(Any, blocked_req), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + failed_req = FakeReq("failed", 2) + failed_req.finished_reason = FINISH_ABORT("boom") + failed = DecodeRequest( + req=cast(Any, failed_req), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + tail = DecodeRequest( + req=cast(Any, FakeReq("tail", 3)), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + queue.queue = cast(Any, [blocked, failed, tail]) + + preallocated, failed_reqs = queue.pop_preallocated( + rids_to_check=["blocked", "failed", "tail"] + ) + + self.assertEqual(preallocated, []) + self.assertEqual([req.req.rid for req in failed_reqs], ["failed"]) + self.assertEqual(streamed, ["failed"]) + self.assertEqual(queue.queue, [blocked, tail]) + + def test_pop_preallocated_compacts_queue_in_one_pass(self): + streamed = [] + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue._resolve_pending_reqs = lambda: None + queue._update_handshake_waiters = lambda rids_to_check=None: None + queue.req_to_token_pool = cast( + Any, + SimpleNamespace( + available_size=lambda: 1, + req_to_token=FakeTensor([[1, 2, 3, 4, 5, 6, 7, 8]]), + ), + ) + queue.req_to_metadata_buffer_idx_allocator = cast( + Any, SimpleNamespace(available_size=lambda: 1, alloc=lambda: 7) + ) + queue.token_to_kv_pool_allocator = cast(Any, SimpleNamespace(page_size=1)) + queue.token_to_kv_pool = cast(Any, object()) + queue.num_reserved_decode_tokens = 1 + queue.scheduler = cast( + Any, + SimpleNamespace( + stream_output=lambda reqs, return_logprob: streamed.extend( + req.rid for req in reqs + ), + enable_metrics=False, + running_batch=SimpleNamespace(reqs=[]), + sliding_window_size=4, + ), + ) + queue._allocatable_tokens = ( + lambda retractable_tokens=None, count_retracted=True: 6 + ) + queue._pre_alloc = lambda req: setattr(req, "req_pool_idx", 0) + + skipped = DecodeRequest( + req=cast(Any, FakeReq("skip", 1)), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + failed_req = FakeReq("failed", 2) + failed_req.finished_reason = FINISH_ABORT("boom") + failed = DecodeRequest( + req=cast(Any, failed_req), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + waiting = DecodeRequest( + req=cast(Any, FakeReq("wait", 3)), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=False, + ) + success_req = FakeReq("success", 4) + success_req.origin_input_ids = [1, 2] + success_req.sampling_params = SimpleNamespace(max_new_tokens=2) + success = DecodeRequest( + req=cast(Any, success_req), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + blocked_req = FakeReq("blocked", 5) + blocked_req.origin_input_ids = [1, 2, 3, 4, 5, 6] + blocked_req.sampling_params = SimpleNamespace(max_new_tokens=8) + blocked = DecodeRequest( + req=cast(Any, blocked_req), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + tail = DecodeRequest( + req=cast(Any, FakeReq("tail", 6)), + kv_receiver=cast(Any, FakeReceiver()), + waiting_for_input=True, + ) + + queue.queue = cast(Any, [skipped, failed, waiting, success, blocked, tail]) + + preallocated, failed_reqs = queue.pop_preallocated( + rids_to_check=["failed", "wait", "success", "blocked"] + ) + + self.assertEqual([req.req.rid for req in preallocated], ["success"]) + self.assertEqual([req.req.rid for req in failed_reqs], ["failed"]) + self.assertEqual(streamed, ["failed"]) + self.assertEqual(success.metadata_buffer_index, 7) + self.assertEqual(queue.queue, [skipped, waiting, blocked, tail]) + + +if __name__ == "__main__": + unittest.main()