diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 341d33bcf..e84aa2c13 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -141,6 +141,18 @@ class BaseKVReceiver(ABC): """ ... + def clear(self): + """ + Clear any internal states. + """ + pass + + def abort(self): + """ + Abort the current transfer. + """ + pass + class BaseKVBootstrapServer(ABC): @abstractmethod diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 6625b1750..48f570e0f 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -693,6 +693,50 @@ class DecodeTransferQueue: def extend(self, decode_reqs: List[DecodeRequest]) -> None: self.queue.extend(decode_reqs) + def _commit_transfer_to_req(self, decode_req: DecodeRequest) -> None: + idx = decode_req.metadata_buffer_index + ( + output_id, + cached_tokens, + output_token_logprobs_val, + output_token_logprobs_idx, + output_top_logprobs_val, + output_top_logprobs_idx, + output_topk_p, + output_topk_index, + output_hidden_states, + ) = self.metadata_buffers.get_buf(idx) + + decode_req.req.output_ids.append(output_id[0].item()) + decode_req.req.cached_tokens = cached_tokens[0].item() + if not self.spec_algorithm.is_none(): + decode_req.req.output_topk_p = output_topk_p + decode_req.req.output_topk_index = output_topk_index + decode_req.req.hidden_states_tensor = output_hidden_states + + if decode_req.req.return_logprob: + decode_req.req.output_token_logprobs_val.append( + output_token_logprobs_val[0].item() + ) + decode_req.req.output_token_logprobs_idx.append( + output_token_logprobs_idx[0].item() + ) + decode_req.req.output_top_logprobs_val.append( + output_top_logprobs_val[: decode_req.req.top_logprobs_num].tolist() + ) + decode_req.req.output_top_logprobs_idx.append( + output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist() + ) + + decode_req.kv_receiver.clear() + decode_req.kv_receiver = None + trace_slice_end( + RequestStage.DECODE_TRANSFERRED, + decode_req.req.rid, + auto_next_anon=True, + ) + decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter() + def pop_transferred(self) -> List[Req]: if not self.queue: return [] @@ -725,58 +769,9 @@ class DecodeTransferQueue: self.scheduler.metrics_collector.increment_transfer_failed_reqs() continue elif poll == KVPoll.Success: - - idx = decode_req.metadata_buffer_index - ( - output_id, - cached_tokens, - output_token_logprobs_val, - output_token_logprobs_idx, - output_top_logprobs_val, - output_top_logprobs_idx, - output_topk_p, - output_topk_index, - output_hidden_states, - ) = self.metadata_buffers.get_buf(idx) - - decode_req.req.output_ids.append(output_id[0].item()) - decode_req.req.cached_tokens = cached_tokens[0].item() - if not self.spec_algorithm.is_none(): - decode_req.req.output_topk_p = output_topk_p - decode_req.req.output_topk_index = output_topk_index - decode_req.req.hidden_states_tensor = output_hidden_states - - if decode_req.req.return_logprob: - decode_req.req.output_token_logprobs_val.append( - output_token_logprobs_val[0].item() - ) - decode_req.req.output_token_logprobs_idx.append( - output_token_logprobs_idx[0].item() - ) - decode_req.req.output_top_logprobs_val.append( - output_top_logprobs_val[ - : decode_req.req.top_logprobs_num - ].tolist() - ) - decode_req.req.output_top_logprobs_idx.append( - output_top_logprobs_idx[ - : decode_req.req.top_logprobs_num - ].tolist() - ) - - if hasattr(decode_req.kv_receiver, "clear"): - decode_req.kv_receiver.clear() - decode_req.kv_receiver = None - + self._commit_transfer_to_req(decode_req) indices_to_remove.add(i) - decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter() transferred_reqs.append(decode_req.req) - trace_slice_end( - RequestStage.DECODE_TRANSFERRED, - decode_req.req.rid, - auto_next_anon=True, - ) - elif poll in [ KVPoll.Bootstrapping, KVPoll.WaitingForInput, diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 044234e85..eb304f3b8 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -310,6 +310,37 @@ class SchedulerDisaggregationPrefillMixin: Mixin for Scheduler to handle disaggregation prefill """ + def get_next_disagg_prefill_batch_to_run( + self: Scheduler, + ) -> Optional[ScheduleBatch]: + if self.last_batch and self.last_batch.forward_mode.is_extend(): + if self.chunked_req: + # Move the chunked request out of the batch so that we can merge + # only finished requests to running_batch. + self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req) + self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) + if self.enable_overlap: + # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved + self.chunked_req.tmp_end_idx = min( + len(self.chunked_req.fill_ids), + len(self.chunked_req.origin_input_ids), + ) + else: + self.send_kv_chunk(self.chunked_req) + # chunked request keeps its rid but will get a new req_pool_idx + self.req_to_token_pool.free(self.chunked_req.req_pool_idx) + self.running_batch.batch_is_full = False + + batch = self.get_new_batch_prefill() + if self.require_mlp_sync: + batch = self.prepare_mlp_sync_batch(batch) + + if batch: + attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()} + trace_event_batch("schedule", batch.reqs, attrs=attrs) + + return batch + @torch.no_grad() def event_loop_normal_disagg_prefill(self: Scheduler) -> None: """A normal scheduler loop for prefill worker in disaggregation mode.""" @@ -320,26 +351,17 @@ class SchedulerDisaggregationPrefillMixin: self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) - self.process_prefill_chunk() - batch = self.get_new_batch_prefill() - if batch: - attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()} - trace_event_batch("schedule", batch.reqs, attrs=attrs) - - if self.require_mlp_sync: - batch = self.prepare_mlp_sync_batch(batch) + batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch if batch: result = self.run_batch(batch) self.process_batch_result_disagg_prefill(batch, result) - - if len(self.disagg_prefill_inflight_queue) > 0: - self.process_disagg_prefill_inflight_queue() - - if batch is None and len(self.disagg_prefill_inflight_queue) == 0: + else: self.self_check_during_idle() + self.process_disagg_prefill_inflight_queue() + self.last_batch = batch # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it # Otherwise, it hangs under high concurrency @@ -355,14 +377,7 @@ class SchedulerDisaggregationPrefillMixin: self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) - self.process_prefill_chunk() - batch = self.get_new_batch_prefill() - if batch: - attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()} - trace_event_batch("schedule", batch.reqs, attrs=attrs) - - if self.require_mlp_sync: - batch = self.prepare_mlp_sync_batch(batch) + batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch batch_result = None @@ -373,15 +388,13 @@ class SchedulerDisaggregationPrefillMixin: if self.last_batch: tmp_batch, tmp_result = self.result_queue.popleft() self.process_batch_result_disagg_prefill(tmp_batch, tmp_result) + elif batch is None: + self.self_check_during_idle() - if len(self.disagg_prefill_inflight_queue) > 0: - self.process_disagg_prefill_inflight_queue() + self.process_disagg_prefill_inflight_queue() self.launch_batch_sample_if_needed(batch_result) - if batch is None and len(self.disagg_prefill_inflight_queue) == 0: - self.self_check_during_idle() - self.last_batch = batch # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it # Otherwise, it hangs under high concurrency @@ -599,25 +612,6 @@ class SchedulerDisaggregationPrefillMixin: return transferred_rids - def process_prefill_chunk(self: Scheduler) -> None: - if self.last_batch and self.last_batch.forward_mode.is_extend(): - if self.chunked_req: - # Move the chunked request out of the batch so that we can merge - # only finished requests to running_batch. - self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req) - self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) - if self.enable_overlap: - # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved - self.chunked_req.tmp_end_idx = min( - len(self.chunked_req.fill_ids), - len(self.chunked_req.origin_input_ids), - ) - else: - self.send_kv_chunk(self.chunked_req) - # chunked request keeps its rid but will get a new req_pool_idx - self.req_to_token_pool.free(self.chunked_req.req_pool_idx) - self.running_batch.batch_is_full = False - def send_kv_chunk( self: Scheduler, req: Req, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 39a323dd5..aaa6a4ce4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2443,15 +2443,13 @@ class Scheduler( for decode_req in self.disagg_decode_prealloc_queue.queue: if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}") - if hasattr(decode_req.kv_receiver, "abort"): - decode_req.kv_receiver.abort() + decode_req.kv_receiver.abort() # Abort requests waiting for kvcache to release tree cache for decode_req in self.disagg_decode_transfer_queue.queue: if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}") - if hasattr(decode_req.kv_receiver, "abort"): - decode_req.kv_receiver.abort() + decode_req.kv_receiver.abort() # Delete requests in the running batch if self.cur_batch is self.running_batch or self.cur_batch is None: diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index c97090bcb..e6f59e5b0 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -217,7 +217,10 @@ class SchedulerRuntimeCheckerMixin: self.tree_cache.sanity_check() def self_check_during_idle(self: Scheduler): - if self.disaggregation_mode == DisaggregationMode.DECODE: + if self.disaggregation_mode == DisaggregationMode.PREFILL: + if len(self.disagg_prefill_inflight_queue) > 0: + return + elif self.disaggregation_mode == DisaggregationMode.DECODE: queue_size = ( len(self.waiting_queue) + len(self.disagg_decode_transfer_queue.queue)