[PD] Improve optimistic prefill (#30951)

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
+                    req.prefill_attempt_count
-                    < self.scheduler.server_args.optimistic_prefill_retries
+                    < 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,38 +636,15 @@ 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
+                # Test hook: exercise the release/requeue retry path.
-                if i in optimistic_polls:
+                if req.pending_bootstrap and should_force_retry(req):
-                    if not self.handle_pending_bootstrap(
+                    self.optimistic_release_and_requeue(req)
-                        req, optimistic_polls[i], defer_release=False
-                    ):
                     advance_logprob_pt(i, req)
                     continue
 
@@ -688,6 +674,7 @@ class SchedulerDisaggregationPrefillMixin:
                         logits_output,
                     )
                     logprob_pt += num_input_logprobs
+                if not req.pending_bootstrap:
                     self.send_kv_chunk(req, last_chunk=True)
                 req.time_stats.set_prefill_transfer_queue_entry_time()
 
@@ -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
+                # Still chunking iff its next chunk was launched: either it is
-                if req.pending_bootstrap:
+                # still self.chunked_req, or its final chunk (extend_range
-                    advance_logprob_pt(i, req)
+                # 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,6 +824,9 @@ class SchedulerDisaggregationPrefillMixin:
                 )
                 done_reqs.append(req)
                 if self.metrics_reporter.enable_metrics:
+                    if req.pending_bootstrap:
+                        self.metrics_collector.increment_bootstrap_failed_reqs()
+                    else:
                         self.metrics_collector.increment_transfer_failed_reqs()
             else:
                 logger.warning_once(
@@ -914,22 +919,15 @@ class SchedulerDisaggregationPrefillMixin:
         if self.enable_hicache_storage:
             self.tree_cache.release_aborted_request(req.rid)
 
-    def handle_pending_bootstrap(
+    def handle_pending_bootstrap(self: Scheduler, req: Req, poll: KVPoll) -> bool:
-        self: Scheduler, req: Req, poll: KVPoll, defer_release: bool
-    ) -> 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 should_force_retry(req):  # test hook
-            if force_retry:
-                if not defer_release:
-                    self.optimistic_release_and_requeue(req)
                 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(
+        return self.handle_pending_bootstrap(req, polls[0])
-            req, polls[0], defer_release=self.enable_overlap
-        )
 
     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:
+        if (req := self.chunked_req) is not None:
-            chunked_req_to_exclude.add(self.chunked_req)
+            chunked_req_to_exclude.add(req)
-            maybe_cache_unfinished_req(self.chunked_req, self.tree_cache, chunked=True)
+            maybe_cache_unfinished_req(req, self.tree_cache, chunked=True)
 
-            if not self.check_bootstrap(self.chunked_req):
+            if not self.check_bootstrap(req):
-                self.chunked_req = None  # stop the current chunked prefill
+                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(
+                req.tmp_end_idx = min(
-                    self.chunked_req.extend_range.end,
+                    req.extend_range.end,
-                    len(self.chunked_req.origin_input_ids),
+                    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 {req.rid} optimistic prefill yielded "
-                f"{req.time_stats.prefill_retry_count}/{max_retries}"
+                f"({req.prefill_attempt_count}/{max_attempts} attempts used)"
             )
             if self.metrics_reporter.enable_metrics:
                 self.metrics_collector.increment_prefill_retries(1)

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

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

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"transfer_total={self.transfer_total_mb:.2f} MB"
-                f"#retries={self.prefill_retry_count}"
             )
         elif self.disagg_mode == DisaggregationMode.DECODE:
             prealloc_duration = self.duration_between(

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:

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",