diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index b64fa4ff9..b78154453 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -137,9 +137,9 @@ class DecodeReqToTokenPool: # Indices of reqs that already have a req_pool_idx and will reuse # their existing slot (e.g. chunked prefill continuing across chunks). reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None] - assert ( - len(reusing) <= 1 - ), "only one chunked request may reuse req_pool_idx in a batch" + assert len(reusing) <= 1, ( + "only one chunked request may reuse req_pool_idx in a batch" + ) assert all( reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in reusing ), "reusing request must be chunked or have committed KV" @@ -166,7 +166,6 @@ class DecodeReqToTokenPool: class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): - def __init__( self, size: int, @@ -793,9 +792,9 @@ class DecodePreallocQueue: """Pre-allocate the memory for req_to_token and token_kv_pool""" req_pool_indices = self.req_to_token_pool.alloc([req]) - assert ( - req_pool_indices is not None - ), "req_pool_indices is full! There is a bug in memory estimation." + assert req_pool_indices is not None, ( + "req_pool_indices is full! There is a bug in memory estimation." + ) # Alloc all tokens for the prebuilt req (except for the reserved input token for decoding) fill_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0) @@ -814,9 +813,9 @@ class DecodePreallocQueue: extend_num_tokens=fill_len, ) - assert ( - kv_loc is not None - ), "KV cache is full! There is a bug in memory estimation." + assert kv_loc is not None, ( + "KV cache is full! There is a bug in memory estimation." + ) self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc) @@ -1008,7 +1007,6 @@ class DecodeTransferQueue: class SchedulerDisaggregationDecodeMixin: - @torch.no_grad() def event_loop_normal_disagg_decode(self: Scheduler): """A normal scheduler loop for decode worker in disaggregation mode.""" @@ -1023,6 +1021,18 @@ class SchedulerDisaggregationDecodeMixin: # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch + disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) + + def pop_and_process(): + # Process the results of the last batch + tmp_batch, tmp_result = self.result_queue.popleft() + self.process_batch_result(tmp_batch, tmp_result) + + # If we need grammar sync (spec + grammar), process the last batch + # results first so that the grammar state is up-to-date before + # generating the bitmask for the current batch. + if disable_overlap_for_batch: + pop_and_process() # Launch the current batch if batch: @@ -1040,6 +1050,11 @@ class SchedulerDisaggregationDecodeMixin: self.result_queue = deque() self.last_batch: Optional[ScheduleBatch] = None + def pop_and_process(): + # Process the results of the last batch + tmp_batch, tmp_result = self.result_queue.popleft() + self.process_batch_result(tmp_batch, tmp_result) + while True: # Receive requests recv_reqs = self.recv_requests() @@ -1050,6 +1065,13 @@ class SchedulerDisaggregationDecodeMixin: # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch + disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) + + # If we need grammar sync (spec + grammar), process the last batch + # results first so that the grammar state is up-to-date before + # generating the bitmask for the current batch. + if self.last_batch and disable_overlap_for_batch: + pop_and_process() # Launch the current batch if batch: @@ -1060,8 +1082,8 @@ class SchedulerDisaggregationDecodeMixin: # Process the last batch if self.last_batch: - tmp_batch, tmp_result = self.result_queue.popleft() - self.process_batch_result(tmp_batch, tmp_result) + if not disable_overlap_for_batch: + pop_and_process() elif batch is None: self.self_check_during_idle() diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index ccfca17e6..766c105ad 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -75,7 +75,10 @@ class AttentionBackend(ABC): Here, we need to redo the computation of all metadata of the attention backend that depends on tree mask and position buffers. """ - raise NotImplementedError() + # if self.topk <= 1: + # return + # raise NotImplementedError() + pass @debug_kernel_api def forward( diff --git a/test/registered/unit/disaggregation/test_overlap_disagg_decode_event_loop.py b/test/registered/unit/disaggregation/test_overlap_disagg_decode_event_loop.py new file mode 100644 index 000000000..3dc07a446 --- /dev/null +++ b/test/registered/unit/disaggregation/test_overlap_disagg_decode_event_loop.py @@ -0,0 +1,85 @@ +"""Unit tests for overlap disaggregation decode event loop ordering.""" + +import unittest +from types import SimpleNamespace +from typing import Any, cast + +from sglang.srt.disaggregation.decode import SchedulerDisaggregationDecodeMixin +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=6, suite="stage-a-test-cpu") + + +class FakeBatch: + def __init__(self, label: str): + self.label = label + + def copy(self): + return self + + +class TestOverlapDisaggDecodeEventLoop(CustomTestCase): + def test_processes_previous_result_before_disabled_overlap_launch(self): + class StopLoop(Exception): + pass + + batch1 = FakeBatch("batch-1") + batch2 = FakeBatch("batch-2") + batches = iter([batch1, batch2]) + events = [] + + def get_next_batch(): + return next(batches) + + def run_batch(batch): + result = f"result-{batch.label}" + events.append(f"run:{batch.label}") + return result + + def process_batch_result(batch, result): + events.append(f"process:{batch.label}") + + def launch_batch_sample_if_needed(result): + events.append(f"sample:{result}") + if result == "result-batch-2": + raise StopLoop() + + scheduler = cast( + Any, + SimpleNamespace( + recv_requests=lambda: [], + process_input_requests=lambda reqs: None, + process_decode_queue=lambda: None, + get_next_disagg_decode_batch_to_run=get_next_batch, + run_batch=run_batch, + process_batch_result=process_batch_result, + is_disable_overlap_for_batch=lambda batch: batch is batch2, + launch_batch_sample_if_needed=launch_batch_sample_if_needed, + self_check_during_idle=lambda: None, + ), + ) + + with self.assertRaises(StopLoop): + SchedulerDisaggregationDecodeMixin.event_loop_overlap_disagg_decode( + scheduler + ) + + self.assertEqual( + events, + [ + "run:batch-1", + "sample:result-batch-1", + "process:batch-1", + "run:batch-2", + "sample:result-batch-2", + ], + ) + self.assertEqual(len(scheduler.result_queue), 1) + queued_batch, queued_result = scheduler.result_queue[0] + self.assertIs(queued_batch, batch2) + self.assertEqual(queued_result, "result-batch-2") + + +if __name__ == "__main__": + unittest.main()