diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 1e0fdef8b..ad72e31f1 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -78,7 +78,8 @@ logger = logging.getLogger(__name__) def should_force_retry(req: Req) -> bool: """Test hook to force a request into optimistic prefill retry.""" retry_prob = envs.SGLANG_TEST_FORCE_OPTIMISTIC_PREFILL_RETRY_PROB.get() - if retry_prob <= 0 or req.time_stats.prefill_retry_count > 0 or req.is_retracted: + # Force only before/during the first attempt (count is 1 while it runs). + if retry_prob <= 0 or req.prefill_attempt_count > 1 or req.is_retracted: return False digest = hashlib.sha256(str(req.rid).encode()).digest() @@ -374,12 +375,13 @@ class PrefillBootstrapQueue: failed_reqs.append(req) elif poll == KVPoll.Bootstrapping: if ( - req.time_stats.prefill_retry_count - < self.scheduler.server_args.optimistic_prefill_retries + req.prefill_attempt_count + < self.scheduler.server_args.optimistic_prefill_attempts and not req.is_retracted # engine paused ): if not self.ensure_metadata_buffer(req): continue # no more metadata buffer + req.prefill_attempt_count += 1 bootstrapped_reqs.append(req) indices_to_remove.add(i) req.time_stats.set_wait_queue_entry_time() @@ -387,6 +389,7 @@ class PrefillBootstrapQueue: if should_force_retry(req): # skip checking for testing if not self.ensure_metadata_buffer(req): continue # no more metadata buffer + req.prefill_attempt_count += 1 elif not self.finalize_bootstrap(req): continue bootstrapped_reqs.append(req) @@ -467,6 +470,12 @@ class SchedulerDisaggregationPrefillMixin: req for req in self.waiting_queue if req not in failed ] + def has_bootstrapped_waiting_req(self: Scheduler) -> bool: + return any( + not req.pending_bootstrap and not is_aborted(req) + for req in self.waiting_queue + ) + @scheduler_nvtx_method("scheduler.get_next_batch_to_run") def get_next_disagg_prefill_batch_to_run( self: Scheduler, @@ -479,10 +488,10 @@ class SchedulerDisaggregationPrefillMixin: # Otherwise, it hangs under high concurrency running_batch.batch_is_full = False - self.process_prefill_chunk(last_batch=last_batch, running_batch=running_batch) - self.resolve_waiting_queue_bootstrap() + self.process_prefill_chunk(last_batch=last_batch, running_batch=running_batch) + prefill_plan = self.get_new_batch_prefill(running_batch) batch = prefill_plan.batch_to_run running_batch = prefill_plan.running_batch @@ -627,40 +636,17 @@ class SchedulerDisaggregationPrefillMixin: if extend_logprob_start_len < extend_input_len: logprob_pt += extend_input_len - extend_logprob_start_len - # Poll optimistic prefill requests in this batch. - # Note: In overlap scheduling, a chunked request that was still pending - # during process_prefill_chunk is not checked again here. - # If it becomes ready in the gap, we still retry the request to keep - # chunked-prefill state management simple. - optimistic_polls = {} - optimistic_reqs = [ - (i, req) - for i, req in enumerate(batch.reqs) - if req.pending_bootstrap and req.inflight_middle_chunks <= 0 - ] - if optimistic_reqs: - polls = poll_and_all_reduce_attn_cp_tp_group( - [req.disagg_kv_sender for _, req in optimistic_reqs], - self.attn_cp_cpu_group, - self.attn_tp_cpu_group, - ) - optimistic_polls = { - idx: poll for (idx, _), poll in zip(optimistic_reqs, polls) - } - for i, (req, next_token_id) in enumerate( zip(batch.reqs, next_token_ids, strict=True) ): if req.inflight_middle_chunks <= 0: req.time_stats.set_prefill_finished_time() - # For optimistic requests, check bootstrap before side effects - if i in optimistic_polls: - if not self.handle_pending_bootstrap( - req, optimistic_polls[i], defer_release=False - ): - advance_logprob_pt(i, req) - continue + # Test hook: exercise the release/requeue retry path. + if req.pending_bootstrap and should_force_retry(req): + self.optimistic_release_and_requeue(req) + advance_logprob_pt(i, req) + continue req.output_ids.append(next_token_id) maybe_cache_unfinished_req(req, self.tree_cache) @@ -688,7 +674,8 @@ class SchedulerDisaggregationPrefillMixin: logits_output, ) logprob_pt += num_input_logprobs - self.send_kv_chunk(req, last_chunk=True) + if not req.pending_bootstrap: + self.send_kv_chunk(req, last_chunk=True) req.time_stats.set_prefill_transfer_queue_entry_time() if req.grammar is not None: @@ -707,10 +694,17 @@ class SchedulerDisaggregationPrefillMixin: # being chunked reqs' prefill is not finished req.inflight_middle_chunks -= 1 - # Overlap deferred release for optimistic requests stopped in process_prefill_chunk - if req.pending_bootstrap: - advance_logprob_pt(i, req) + # Still chunking iff its next chunk was launched: either it is + # still self.chunked_req, or its final chunk (extend_range + # reaching the end of the input) is in flight. A yielded req + # is neither, so do its deferred release here. + still_chunking = self.chunked_req is req or ( + req.extend_range is not None + and req.extend_range.end >= len(req.origin_input_ids) + ) + if req.pending_bootstrap and not still_chunking: self.optimistic_release_and_requeue(req) + advance_logprob_pt(i, req) req.time_stats.set_last_chunked_prefill_finish_time() continue @@ -736,7 +730,9 @@ class SchedulerDisaggregationPrefillMixin: ) logprob_pt += num_input_logprobs - if self.enable_overlap: + # In non-overlap-mode, KV is sent in process_prefill_chunk + # Only send when req's sender is initialized + if self.enable_overlap and not req.pending_bootstrap: assert ( req.metadata_buffer_index >= 0 ), f"Req {req.rid} does not have metadata buffer allocated" @@ -792,7 +788,13 @@ class SchedulerDisaggregationPrefillMixin: undone_reqs.append(req) continue - if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: + if req.pending_bootstrap and poll != KVPoll.Failed: + # prefill finished before bootstrap + if poll == KVPoll.WaitingForInput: + assert self.disagg_prefill_bootstrap_queue.finalize_bootstrap(req) + self.send_kv_chunk(req, last_chunk=True) + undone_reqs.append(req) + elif poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: undone_reqs.append(req) elif poll == KVPoll.Success: # transfer done release_kv_cache(req, self.tree_cache) # unlock the tree @@ -822,7 +824,10 @@ class SchedulerDisaggregationPrefillMixin: ) done_reqs.append(req) if self.metrics_reporter.enable_metrics: - self.metrics_collector.increment_transfer_failed_reqs() + if req.pending_bootstrap: + self.metrics_collector.increment_bootstrap_failed_reqs() + else: + self.metrics_collector.increment_transfer_failed_reqs() else: logger.warning_once( f"Unexpected polling state {poll} for rid {req.rid} in inflight queue; " @@ -914,22 +919,15 @@ class SchedulerDisaggregationPrefillMixin: if self.enable_hicache_storage: self.tree_cache.release_aborted_request(req.rid) - def handle_pending_bootstrap( - self: Scheduler, req: Req, poll: KVPoll, defer_release: bool - ) -> bool: + def handle_pending_bootstrap(self: Scheduler, req: Req, poll: KVPoll) -> bool: """Return True when bootstrap is finalized and KV transfer can proceed.""" if poll == KVPoll.Failed: self.handle_bootstrap_failure(req) return False elif poll == KVPoll.Bootstrapping: - if not defer_release: - self.optimistic_release_and_requeue(req) return False elif poll == KVPoll.WaitingForInput: - force_retry = should_force_retry(req) # test hook - if force_retry: - if not defer_release: - self.optimistic_release_and_requeue(req) + if should_force_retry(req): # test hook return False # Metadata buffer was allocated in pop_bootstrapped before # the request entered the waiting queue, so finalize should not fail. @@ -950,9 +948,7 @@ class SchedulerDisaggregationPrefillMixin: self.attn_cp_cpu_group, self.attn_tp_cpu_group, ) - return self.handle_pending_bootstrap( - req, polls[0], defer_release=self.enable_overlap - ) + return self.handle_pending_bootstrap(req, polls[0]) def process_prefill_chunk( self: Scheduler, @@ -960,20 +956,28 @@ class SchedulerDisaggregationPrefillMixin: running_batch: ScheduleBatch, ) -> None: chunked_req_to_exclude = set() - if self.chunked_req: - chunked_req_to_exclude.add(self.chunked_req) - maybe_cache_unfinished_req(self.chunked_req, self.tree_cache, chunked=True) + if (req := self.chunked_req) is not None: + chunked_req_to_exclude.add(req) + maybe_cache_unfinished_req(req, self.tree_cache, chunked=True) - if not self.check_bootstrap(self.chunked_req): - self.chunked_req = None # stop the current chunked prefill + if not self.check_bootstrap(req): + if is_aborted(req): + # bootstrap failed + self.chunked_req = None + elif self.has_bootstrapped_waiting_req(): + # optimistic request yields to waiting requests + self.chunked_req = None + if not self.enable_overlap: + self.optimistic_release_and_requeue(req) + # else: still bootstrapping, keep computing without sending elif 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( - self.chunked_req.extend_range.end, - len(self.chunked_req.origin_input_ids), + req.tmp_end_idx = min( + req.extend_range.end, + len(req.origin_input_ids), ) else: - self.send_kv_chunk(self.chunked_req) + self.send_kv_chunk(req) if self.chunked_req is not None: running_batch.batch_is_full = False @@ -1142,7 +1146,7 @@ class SchedulerDisaggregationPrefillMixin: def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None: """Release KV cache and requeue an optimistic prefill request.""" - max_retries = self.server_args.optimistic_prefill_retries + max_attempts = self.server_args.optimistic_prefill_attempts maybe_cache_unfinished_req(req, self.tree_cache) release_kv_cache(req, self.tree_cache) req.reset_for_retract() @@ -1152,19 +1156,19 @@ class SchedulerDisaggregationPrefillMixin: req.hidden_states_tensor = None req.pending_bootstrap = True req.time_stats.reset_prefill_retry_time() - if req.time_stats.prefill_retry_count >= max_retries: + if req.prefill_attempt_count >= max_attempts: logger.info( - f"Req {req.rid} exhausted optimistic prefill retries " + f"Req {req.rid} exhausted optimistic prefill attempts " "falling back to bootstrap queue" ) # Reset it so the next real bootstrap done can be recorded. req.time_stats.bootstrap_done_time = 0.0 self.disagg_prefill_bootstrap_queue.queue.append(req) else: - req.time_stats.prefill_retry_count += 1 + req.prefill_attempt_count += 1 logger.info( - f"Req {req.rid} optimistic prefill retry " - f"{req.time_stats.prefill_retry_count}/{max_retries}" + f"Req {req.rid} optimistic prefill yielded " + f"({req.prefill_attempt_count}/{max_attempts} attempts used)" ) if self.metrics_reporter.enable_metrics: self.metrics_collector.increment_prefill_retries(1) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4c9be3c9c..632e27f9f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1019,6 +1019,8 @@ class Req(ReqDllmMixin): # Used in overlap sequence to signal that an optimistic request should # abort chunking. Set in create_sender, consumed in process_batch_result. self.pending_bootstrap = False + # Number of optimistic prefill forward passes started. preserved across retracts. + self.prefill_attempt_count = 0 # For Matryoshka embeddings self.dimensions = dimensions @@ -1497,6 +1499,8 @@ class Req(ReqDllmMixin): self.input_token_logprobs = None self.temp_input_top_logprobs_val = None self.temp_input_top_logprobs_idx = None + self.temp_input_token_ids_logprobs_val = None + self.temp_input_token_ids_logprobs_idx = None self.inflight_middle_chunks = 0 self.mamba_pool_idx = None self.mamba_ping_pong_track_buffer = None @@ -1605,6 +1609,7 @@ class Req(ReqDllmMixin): f"input_len={len(self.origin_input_ids)}, " f"cached_input_len={self.cached_tokens}, " f"output_len={len(self.output_ids)}, " + f"attempts={self.prefill_attempt_count}, " f"type={self.time_stats.disagg_mode_str()})" ) logger.info(f"{prefix}: {self.time_stats.convert_to_duration()}") diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 7d6d767ef..104b279f1 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -570,6 +570,8 @@ class SchedulerMetricsReporter: msg += ( f"#inflight-req: {len(self.scheduler.disagg_prefill_inflight_queue)}, " ) + num_optimistic = sum(1 for r in batch.reqs if r.pending_bootstrap) + msg += f"#optimistic-req: {num_optimistic}, " if ( self.scheduler.server_args.language_only diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 27734d988..8e03bcaea 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -623,9 +623,6 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): transfer_speed_gb_s: float = 0.0 transfer_total_mb: float = 0.0 - # Number of prefill retries for this request - prefill_retry_count: int = 0 - def __getstate__(self) -> object: # send to detokenizer/tokenizer if not self.enable_metrics: @@ -1089,8 +1086,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): f"forward_duration={self.format_duration(forward_duration)}, " f"entry_time={self.format_wallclock(self.prefill_bootstrap_queue_entry_time)}, " f"transfer_speed={self.transfer_speed_gb_s:.2f} GB/s, " - f"transfer_total={self.transfer_total_mb:.2f} MB, " - f"#retries={self.prefill_retry_count}" + f"transfer_total={self.transfer_total_mb:.2f} MB" ) elif self.disagg_mode == DisaggregationMode.DECODE: prealloc_duration = self.duration_between( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d27e50d0e..437073a64 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2441,9 +2441,9 @@ class ServerArgs: int, "The interval to poll requests in decode server. Can be set to >1 to reduce the overhead of this.", ] = 1 - optimistic_prefill_retries: A[ + optimistic_prefill_attempts: A[ int, - "Number of optimistic prefill retries that will skip the bootstrap wait. ", + "Number of optimistic prefill forward passes that skip the bootstrap wait.", ] = 0 # ------------------------------------------------------------------------- @@ -6382,21 +6382,21 @@ class ServerArgs: # Handle optimistic prefill validation if ( - self.optimistic_prefill_retries > 0 + self.optimistic_prefill_attempts > 0 and self.disaggregation_mode == "prefill" ): if self.pp_size > 1: logger.warning("Optimistic prefill does not support pp_size > 1") - self.optimistic_prefill_retries = 0 + self.optimistic_prefill_attempts = 0 elif self.enable_hierarchical_cache: logger.warning("Optimistic prefill does not support hierarchical cache") - self.optimistic_prefill_retries = 0 + self.optimistic_prefill_attempts = 0 elif resolved_view(self).uses_mamba_radix_cache: logger.warning( "Optimistic prefill does not support models that use " "mamba radix cache." ) - self.optimistic_prefill_retries = 0 + self.optimistic_prefill_attempts = 0 # Handle model inference tensor dump. if self.debug_tensor_dump_output_folder is not None: diff --git a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py index c99b5aa6f..58fa9b472 100644 --- a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py +++ b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py @@ -29,7 +29,7 @@ def rid_that_forces_retry(prefix: str) -> str: req = SimpleNamespace( rid=rid, is_retracted=False, - time_stats=SimpleNamespace(prefill_retry_count=0), + prefill_attempt_count=0, ) if should_force_retry(req): return rid @@ -72,7 +72,7 @@ class TestOptimisticPrefill( envs.SGLANG_TEST_FORCE_OPTIMISTIC_PREFILL_RETRY_PROB.set(FORCE_RETRY_PROB) cls.model = DEFAULT_MODEL_NAME_FOR_TEST cls.extra_prefill_args = [ - "--optimistic-prefill-retries", + "--optimistic-prefill-attempts", "3", "--chunked-prefill-size", "128", @@ -129,7 +129,11 @@ class TestOptimisticPrefill( self.assertGreater(j["meta_info"]["prompt_tokens"], 512) assert len(output_logprobs) == completion_tokens - assert len(input_logprobs) > 0 + # Input logprobs must be complete: retried or pending chunks must not + # drop their logprobs. + self.assertGreaterEqual( + len(input_logprobs), j["meta_info"]["prompt_tokens"] - 1 + ) class TestOptimisticPrefillFailure(PDDisaggregationServerBase): @@ -150,7 +154,7 @@ class TestOptimisticPrefillFailure(PDDisaggregationServerBase): cls.model = DEFAULT_MODEL_NAME_FOR_TEST cls.extra_prefill_args = [ - "--optimistic-prefill-retries", + "--optimistic-prefill-attempts", "3", "--chunked-prefill-size", "128",