chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

perf(disaggregation): compact decode queues in one pass

Browse Source
This commit is contained in:
wxiwnd
2026-04-08 22:37:40 +08:00
parent 9470d6d8dc
commit 6237271fc9
2 changed files with 537 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
python/sglang/srt/disaggregation/decode.py
View File

@@ -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

495
test/registered/unit/disaggregation/test_decode_queue_compaction.py Normal file
View File

@@ -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()