c9a39ccdd2
Decode queue compaction receives req_to_token rows after the prefill side has already populated cached prefix slots. Cache-hit requests therefore need the extend/suffix slice, not the leading prefix slice, when building the prebuilt transfer chunk. Constraint: Prefill/decode disaggregation shares req_to_token rows across cached prefix and new suffix positions. Rejected: Keep slicing from zero | cache-hit requests would copy prefix KV locs into the prebuilt suffix chunk. Confidence: medium Scope-risk: narrow Directive: Do not change prepare_for_prebuilt slicing without testing cache-hit req_to_token layouts. Tested: python -m py_compile on changed runtime files. Not-tested: Local pytest blocked before collection by missing orjson dependency. (cherry picked from commit 416112b617fabe71e8cff7484794af73f3e84440)
931 lines
34 KiB
Python
931 lines
34 KiB
Python
"""Unit tests for decode queue one-pass compaction."""
|
|
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
from sglang.srt.disaggregation.base import KVPoll
|
|
from sglang.srt.disaggregation.decode import (
|
|
DecodePreallocQueue,
|
|
DecodeRequest,
|
|
DecodeTransferQueue,
|
|
SchedulerDisaggregationDecodeMixin,
|
|
)
|
|
from sglang.srt.disaggregation.decode_schedule_batch_mixin import (
|
|
ScheduleBatchDisaggregationDecodeMixin,
|
|
)
|
|
from sglang.srt.disaggregation.utils import (
|
|
DisaggregationMode,
|
|
ReqToMetadataIdxAllocator,
|
|
)
|
|
from sglang.srt.managers.scheduler_output_processor_mixin import (
|
|
SchedulerOutputProcessorMixin,
|
|
)
|
|
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.init_next_round_calls = []
|
|
self.sampling_params = SimpleNamespace(max_new_tokens=0)
|
|
self.return_hidden_states = False
|
|
self.grammar = None
|
|
self.token_ids_logprob = None
|
|
self.top_logprobs_num = 0
|
|
self.multimodal_inputs = None
|
|
self.mamba_ping_pong_track_buffer = None
|
|
self.to_finish = None
|
|
|
|
class FakeTimeStats:
|
|
def __init__(self):
|
|
self.forward_entry_time = None
|
|
|
|
def set_bootstrap_done_time(self):
|
|
return None
|
|
|
|
def set_decode_transfer_queue_entry_time(self):
|
|
return None
|
|
|
|
def set_wait_queue_entry_time(self):
|
|
return None
|
|
|
|
def set_forward_entry_time(self, ts):
|
|
self.forward_entry_time = ts
|
|
|
|
def set_decode_prebuilt_finish_time(self):
|
|
return None
|
|
|
|
def set_quick_finish_time(self):
|
|
return None
|
|
|
|
def set_last_decode_finish_time(self):
|
|
return None
|
|
|
|
def set_completion_time(self):
|
|
return None
|
|
|
|
self.time_stats = FakeTimeStats()
|
|
|
|
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))
|
|
|
|
def init_next_round_input(self, tree_cache):
|
|
self.init_next_round_calls.append(tree_cache)
|
|
|
|
def check_finished(self, *args):
|
|
return None
|
|
|
|
def finished(self):
|
|
return self.finished_reason is not None
|
|
|
|
|
|
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 FakeBatch:
|
|
def __init__(self):
|
|
self.prepared = False
|
|
self.processed = []
|
|
|
|
def prepare_for_prebuilt(self):
|
|
self.prepared = True
|
|
|
|
def process_prebuilt(self, server_args, future_map):
|
|
self.processed.append((server_args, future_map))
|
|
|
|
|
|
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 FakeTokenToKVAllocator:
|
|
def __init__(self):
|
|
self.begin_calls = 0
|
|
self.end_calls = 0
|
|
|
|
def free_group_begin(self):
|
|
self.begin_calls += 1
|
|
|
|
def free_group_end(self):
|
|
self.end_calls += 1
|
|
|
|
|
|
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_pop_transferred_holds_eagle_metadata_slot_until_prebuilt_consumes(self):
|
|
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
|
|
allocator = FakeAllocator()
|
|
output_topk_p = torch.arange(16, dtype=torch.float32)
|
|
output_topk_index = torch.arange(16, dtype=torch.int64)
|
|
output_hidden_states = torch.arange(8, dtype=torch.float32)
|
|
queue.gloo_group = None
|
|
queue.req_to_metadata_buffer_idx_allocator = allocator
|
|
queue.tp_rank = 0
|
|
queue.metadata_buffers = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
get_buf=lambda idx: (
|
|
torch.tensor([7], dtype=torch.int32),
|
|
torch.tensor([320], dtype=torch.int32),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
output_topk_p,
|
|
output_topk_index,
|
|
output_hidden_states,
|
|
torch.tensor([3], dtype=torch.int64),
|
|
)
|
|
),
|
|
)
|
|
queue.tree_cache = cast(Any, object())
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake"),
|
|
stream_output=lambda reqs, return_logprob: None,
|
|
enable_metrics=False,
|
|
),
|
|
)
|
|
queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: False))
|
|
|
|
receiver = FakeReceiver()
|
|
decode_req = DecodeRequest(
|
|
req=cast(Any, FakeReq("eagle", 3)),
|
|
kv_receiver=cast(Any, receiver),
|
|
metadata_buffer_index=9,
|
|
)
|
|
queue.queue = [decode_req]
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.poll_and_all_reduce",
|
|
return_value=[KVPoll.Success],
|
|
):
|
|
transferred = queue.pop_transferred()
|
|
|
|
self.assertEqual(transferred, [decode_req.req])
|
|
self.assertEqual(queue.queue, [])
|
|
self.assertEqual(allocator.freed, [])
|
|
|
|
self.assertEqual(decode_req.req.output_ids, [7])
|
|
self.assertEqual(decode_req.req.cached_tokens, 320)
|
|
self.assertEqual(decode_req.req.metadata_buffer_index, 9)
|
|
self.assertEqual(
|
|
decode_req.req.output_topk_p.data_ptr(), output_topk_p.data_ptr()
|
|
)
|
|
self.assertEqual(
|
|
decode_req.req.output_topk_index.data_ptr(), output_topk_index.data_ptr()
|
|
)
|
|
self.assertEqual(
|
|
decode_req.req.hidden_states_tensor.data_ptr(),
|
|
output_hidden_states.data_ptr(),
|
|
)
|
|
|
|
def test_free_decode_metadata_index_if_held_releases_once(self):
|
|
allocator = FakeAllocator()
|
|
req = FakeReq("eagle", 3)
|
|
req.metadata_buffer_index = 9
|
|
req.output_topk_p = torch.ones((1,), dtype=torch.float32)
|
|
req.output_topk_index = torch.ones((1,), dtype=torch.int64)
|
|
req.hidden_states_tensor = torch.ones((4,), dtype=torch.float32)
|
|
|
|
scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin)
|
|
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
|
|
|
scheduler._free_decode_metadata_index_if_held(req)
|
|
scheduler._free_decode_metadata_index_if_held(req)
|
|
|
|
self.assertEqual(allocator.freed, [9])
|
|
self.assertEqual(req.metadata_buffer_index, -1)
|
|
self.assertIsNone(req.output_topk_p)
|
|
self.assertIsNone(req.output_topk_index)
|
|
self.assertIsNone(req.hidden_states_tensor)
|
|
|
|
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]]),
|
|
write=lambda *args, **kwargs: None,
|
|
),
|
|
)
|
|
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.draft_token_to_kv_pool = None
|
|
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, **kwargs: (
|
|
setattr(req, "req_pool_idx", 0),
|
|
torch.arange(len(req.origin_input_ids), dtype=torch.int64),
|
|
)[1]
|
|
|
|
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]]),
|
|
write=lambda *args, **kwargs: None,
|
|
),
|
|
)
|
|
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.draft_token_to_kv_pool = None
|
|
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, **kwargs: (
|
|
setattr(req, "req_pool_idx", 0),
|
|
torch.arange(len(req.origin_input_ids), dtype=torch.int64),
|
|
)[1]
|
|
|
|
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])
|
|
|
|
def test_get_new_prebuilt_batch_slices_waiting_queue_prefix(self):
|
|
allocator = FakeAllocator()
|
|
scheduler = cast(Any, SimpleNamespace())
|
|
scheduler.grammar_manager = SimpleNamespace(
|
|
has_waiting_grammars=lambda: False,
|
|
get_ready_grammar_requests=lambda: [],
|
|
)
|
|
scheduler._add_request_to_queue = lambda req: None
|
|
scheduler.waiting_queue = [FakeReq(f"req-{i}", i) for i in range(5)]
|
|
scheduler.running_batch = SimpleNamespace(batch_size=lambda: 1)
|
|
scheduler.req_to_token_pool = SimpleNamespace(size=8)
|
|
scheduler.max_running_requests = 4
|
|
scheduler.tree_cache = object()
|
|
scheduler.token_to_kv_pool_allocator = object()
|
|
scheduler.model_config = object()
|
|
scheduler.enable_overlap = False
|
|
scheduler.spec_algorithm = object()
|
|
scheduler.server_args = object()
|
|
scheduler.future_map = object()
|
|
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
|
scheduler._free_decode_metadata_index_if_held = (
|
|
SchedulerOutputProcessorMixin._free_decode_metadata_index_if_held.__get__(
|
|
scheduler
|
|
)
|
|
)
|
|
for i, req in enumerate(scheduler.waiting_queue[:3]):
|
|
req.metadata_buffer_index = 20 + i
|
|
|
|
captured = {}
|
|
|
|
def fake_init_new(reqs, *args, **kwargs):
|
|
captured["reqs"] = list(reqs)
|
|
batch = FakeBatch()
|
|
captured["batch"] = batch
|
|
return batch
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.ScheduleBatch.init_new", fake_init_new
|
|
):
|
|
batch = SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch(scheduler)
|
|
|
|
self.assertIs(batch, captured["batch"])
|
|
self.assertEqual(
|
|
[req.rid for req in captured["reqs"]], ["req-0", "req-1", "req-2"]
|
|
)
|
|
self.assertEqual(
|
|
[req.rid for req in scheduler.waiting_queue], ["req-3", "req-4"]
|
|
)
|
|
for req in captured["reqs"]:
|
|
self.assertEqual(req.init_next_round_calls, [scheduler.tree_cache])
|
|
self.assertIsNotNone(req.time_stats.forward_entry_time)
|
|
self.assertTrue(captured["batch"].prepared)
|
|
self.assertEqual(
|
|
captured["batch"].processed,
|
|
[(scheduler.server_args, scheduler.future_map)],
|
|
)
|
|
# EAGLE metadata slots are intentionally held past process_prebuilt().
|
|
# The initial draft state is consumed by the first real decode forward,
|
|
# so releasing here would let burst transfers overwrite the pinned CPU
|
|
# source views before the GPU copy/consume is complete.
|
|
self.assertEqual(allocator.freed, [])
|
|
for req in captured["reqs"]:
|
|
self.assertGreaterEqual(req.metadata_buffer_index, 20)
|
|
|
|
def test_prepare_for_prebuilt_uses_suffix_cache_locs_after_cache_hit(self):
|
|
"""Decode prebuilt must write new-token locs, not prefix locs.
|
|
|
|
Prefill transfers the full prompt KV to decode. During the first decode
|
|
prebuilt step, `prefix_indices` covers the cached prompt prefix and
|
|
`extend_input_len` covers only the prompt suffix that still needs a
|
|
local decode forward. The output cache loc tensor must therefore slice
|
|
req_to_token at [pre_len : pre_len + extend_input_len].
|
|
"""
|
|
|
|
req = FakeReq("cache-hit", 11)
|
|
req.req_pool_idx = 0
|
|
req.prefix_indices = torch.arange(100, 104, dtype=torch.int64)
|
|
req.extend_input_len = 3
|
|
req.fill_ids = [10, 11, 12, 13, 14, 15, 16]
|
|
req.origin_input_ids = list(req.fill_ids)
|
|
req.output_ids = []
|
|
req.retracted_stain = False
|
|
req.already_computed = len(req.prefix_indices)
|
|
req.top_logprobs_num = 0
|
|
req.token_ids_logprob = None
|
|
|
|
batch = cast(Any, SimpleNamespace())
|
|
batch.reqs = [req]
|
|
batch.req_to_token_pool = SimpleNamespace(
|
|
req_to_token=torch.tensor(
|
|
[[1000, 1001, 1002, 1003, 2000, 2001, 2002]],
|
|
dtype=torch.int64,
|
|
)
|
|
)
|
|
batch.device = "cpu"
|
|
batch.return_logprob = False
|
|
batch.model_config = SimpleNamespace(vocab_size=32000)
|
|
batch.multimodal_inputs = None
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode_schedule_batch_mixin."
|
|
"SamplingBatchInfo.from_schedule_batch",
|
|
return_value=SimpleNamespace(),
|
|
):
|
|
ScheduleBatchDisaggregationDecodeMixin.prepare_for_prebuilt(batch)
|
|
|
|
self.assertEqual(batch.prefix_lens, [4])
|
|
self.assertEqual(batch.extend_lens, [3])
|
|
self.assertEqual(batch.out_cache_loc.tolist(), [2000, 2001, 2002])
|
|
|
|
def test_get_new_prebuilt_batch_frees_metadata_on_prebuilt_error(self):
|
|
allocator = FakeAllocator()
|
|
scheduler = cast(Any, SimpleNamespace())
|
|
scheduler.grammar_manager = SimpleNamespace(
|
|
has_waiting_grammars=lambda: False,
|
|
get_ready_grammar_requests=lambda: [],
|
|
)
|
|
scheduler._add_request_to_queue = lambda req: None
|
|
req0 = FakeReq("req-0", 0)
|
|
scheduler.waiting_queue = [req0]
|
|
scheduler.running_batch = SimpleNamespace(batch_size=lambda: 0)
|
|
scheduler.req_to_token_pool = SimpleNamespace(size=8)
|
|
scheduler.max_running_requests = 8
|
|
scheduler.tree_cache = object()
|
|
scheduler.token_to_kv_pool_allocator = object()
|
|
scheduler.model_config = object()
|
|
scheduler.enable_overlap = False
|
|
scheduler.spec_algorithm = object()
|
|
scheduler.server_args = object()
|
|
scheduler.future_map = object()
|
|
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
|
scheduler._free_decode_metadata_index_if_held = (
|
|
SchedulerOutputProcessorMixin._free_decode_metadata_index_if_held.__get__(
|
|
scheduler
|
|
)
|
|
)
|
|
req0.metadata_buffer_index = 42
|
|
|
|
class FailingBatch(FakeBatch):
|
|
def process_prebuilt(self, server_args, future_map):
|
|
raise RuntimeError("prebuilt failed")
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.ScheduleBatch.init_new",
|
|
lambda reqs, *args, **kwargs: FailingBatch(),
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, "prebuilt failed"):
|
|
SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch(scheduler)
|
|
|
|
self.assertEqual(allocator.freed, [42])
|
|
self.assertEqual(req0.metadata_buffer_index, -1)
|
|
|
|
def test_process_batch_result_prebuilt_frees_finished_metadata(self):
|
|
allocator = FakeAllocator()
|
|
scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin)
|
|
scheduler.disaggregation_mode = DisaggregationMode.DECODE
|
|
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
|
scheduler.tree_cache = object()
|
|
scheduler.stream_output = lambda reqs, return_logprob: None
|
|
|
|
req = FakeReq("finished", 0)
|
|
req.metadata_buffer_index = 43
|
|
req.output_topk_p = torch.ones((1,), dtype=torch.float32)
|
|
req.output_topk_index = torch.ones((1,), dtype=torch.int64)
|
|
req.hidden_states_tensor = torch.ones((4,), dtype=torch.float32)
|
|
req.finished_reason = FINISH_ABORT("done")
|
|
batch = SimpleNamespace(reqs=[req], return_logprob=False)
|
|
|
|
with patch(
|
|
"sglang.srt.managers.scheduler_output_processor_mixin.release_kv_cache",
|
|
lambda *args, **kwargs: None,
|
|
):
|
|
scheduler.process_batch_result_prebuilt(batch)
|
|
|
|
self.assertEqual(allocator.freed, [43])
|
|
self.assertEqual(req.metadata_buffer_index, -1)
|
|
|
|
def test_process_batch_result_decode_releases_prebuilt_metadata_after_consume(self):
|
|
allocator = FakeAllocator()
|
|
token_allocator = FakeTokenToKVAllocator()
|
|
scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin)
|
|
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
|
scheduler.server_args = SimpleNamespace(
|
|
disaggregation_decode_enable_offload_kvcache=False
|
|
)
|
|
scheduler.enable_hisparse = False
|
|
scheduler.enable_overlap = False
|
|
scheduler.enable_metrics = False
|
|
scheduler.token_to_kv_pool_allocator = token_allocator
|
|
scheduler.tree_cache = object()
|
|
scheduler.forward_ct_decode = 0
|
|
scheduler.num_generated_tokens = 0
|
|
scheduler.stream_output = lambda reqs, return_logprob: None
|
|
scheduler.report_decode_stats = lambda *args, **kwargs: None
|
|
scheduler.update_spec_metrics = lambda *args, **kwargs: None
|
|
scheduler._maybe_log_eagle_accept_debug = lambda *args, **kwargs: None
|
|
|
|
req = FakeReq("decode", 0)
|
|
req.metadata_buffer_index = 44
|
|
req.output_topk_p = torch.ones((1,), dtype=torch.float32)
|
|
req.output_topk_index = torch.ones((1,), dtype=torch.int64)
|
|
req.hidden_states_tensor = torch.ones((4,), dtype=torch.float32)
|
|
batch = SimpleNamespace(
|
|
reqs=[req],
|
|
return_logprob=False,
|
|
spec_algorithm=SimpleNamespace(is_none=lambda: True),
|
|
is_spec_v2=False,
|
|
)
|
|
result = SimpleNamespace(
|
|
copy_done=None,
|
|
logits_output=SimpleNamespace(hidden_states=None, customized_info=None),
|
|
next_token_ids=torch.tensor([5], dtype=torch.int64),
|
|
can_run_cuda_graph=True,
|
|
num_accepted_tokens=1,
|
|
)
|
|
|
|
scheduler.process_batch_result_decode(batch, result)
|
|
|
|
self.assertEqual(allocator.freed, [44])
|
|
self.assertEqual(req.metadata_buffer_index, -1)
|
|
self.assertEqual(req.output_ids, [5])
|
|
self.assertEqual(token_allocator.begin_calls, 1)
|
|
self.assertEqual(token_allocator.end_calls, 1)
|
|
|
|
def test_get_new_prebuilt_batch_keeps_waiting_queue_when_no_capacity(self):
|
|
scheduler = cast(Any, SimpleNamespace())
|
|
scheduler.grammar_manager = SimpleNamespace(
|
|
has_waiting_grammars=lambda: False,
|
|
get_ready_grammar_requests=lambda: [],
|
|
)
|
|
scheduler._add_request_to_queue = lambda req: None
|
|
scheduler.waiting_queue = [FakeReq(f"req-{i}", i) for i in range(3)]
|
|
scheduler.running_batch = SimpleNamespace(batch_size=lambda: 4)
|
|
scheduler.req_to_token_pool = SimpleNamespace(size=8)
|
|
scheduler.max_running_requests = 4
|
|
scheduler.tree_cache = object()
|
|
scheduler.token_to_kv_pool_allocator = object()
|
|
scheduler.model_config = object()
|
|
scheduler.enable_overlap = False
|
|
scheduler.spec_algorithm = object()
|
|
scheduler.server_args = object()
|
|
scheduler.future_map = object()
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.ScheduleBatch.init_new",
|
|
side_effect=AssertionError(
|
|
"init_new should not be called without capacity"
|
|
),
|
|
):
|
|
batch = SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch(scheduler)
|
|
|
|
self.assertIsNone(batch)
|
|
self.assertEqual(
|
|
[req.rid for req in scheduler.waiting_queue], ["req-0", "req-1", "req-2"]
|
|
)
|
|
for req in scheduler.waiting_queue:
|
|
self.assertEqual(req.init_next_round_calls, [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|