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