sglang/test/registered/unit/disaggregation/test_decode_queue_compaction.py

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