minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PD] Improve optimistic prefill (#30951)

Browse Source
This commit is contained in:
cctry
2026-07-12 15:31:53 -07:00
committed by GitHub
parent 6cc9352dfe
commit c616d5a55e
6 changed files with 93 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+71 -67
View File
@@ -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)
python/sglang/srt/managers/schedule_batch.py
+5
View File
@@ -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
+2
View File
@@ -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
+1 -5
View File
@@ -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(
python/sglang/srt/server_args.py
+6 -6
View File
@@ -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
+8 -4
View File
@@ -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",