[VLM] retire aborted disaggregated prefill results (#36988)
This commit is contained in:
Mick
@@ -729,6 +729,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
result.indexer_topk_output = None
|
||||
|
||||
logprob_pt = 0
|
||||
aborted_reqs: List[Req] = []
|
||||
assert batch.spec_info is result.next_draft_input
|
||||
draft_input = result.next_draft_input
|
||||
draft_hidden_states_cpu = None
|
||||
@@ -763,6 +764,13 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
if req.inflight_middle_chunks <= 0:
|
||||
req.time_stats.set_prefill_finished_time()
|
||||
|
||||
if is_aborted(req):
|
||||
if self._retire_aborted_prefill_result(req):
|
||||
req.time_stats.set_completion_time()
|
||||
aborted_reqs.append(req)
|
||||
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)
|
||||
@@ -770,6 +778,24 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
continue
|
||||
|
||||
req.output_ids.append(next_token_id)
|
||||
if req.grammar is not None:
|
||||
try:
|
||||
req.grammar.accept_token(next_token_id)
|
||||
except ValueError as e:
|
||||
error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
|
||||
prepare_abort(
|
||||
req,
|
||||
error_message,
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
req.grammar.finished = req.finished()
|
||||
if is_aborted(req):
|
||||
if self._retire_aborted_prefill_result(req):
|
||||
req.time_stats.set_completion_time()
|
||||
aborted_reqs.append(req)
|
||||
advance_logprob_pt(i, req)
|
||||
continue
|
||||
|
||||
maybe_cache_unfinished_req(req, self.tree_cache)
|
||||
self.disagg_prefill_inflight_queue.append(req)
|
||||
if self.spec_algorithm.is_eagle() and draft_input is not None:
|
||||
@@ -808,17 +834,6 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
self.send_kv_chunk(req, last_chunk=True)
|
||||
req.time_stats.set_prefill_transfer_queue_entry_time()
|
||||
|
||||
if req.grammar is not None:
|
||||
try:
|
||||
req.grammar.accept_token(next_token_id)
|
||||
except ValueError as e:
|
||||
error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
|
||||
prepare_abort(
|
||||
req,
|
||||
error_message,
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
req.grammar.finished = req.finished()
|
||||
else:
|
||||
# being chunked reqs' prefill is not finished
|
||||
req.inflight_middle_chunks -= 1
|
||||
@@ -831,16 +846,18 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
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)
|
||||
# Abort is terminal. Do not requeue an aborted optimistic
|
||||
# request merely because bootstrap is still pending.
|
||||
if is_aborted(req):
|
||||
if not still_chunking and self._retire_aborted_prefill_result(req):
|
||||
req.time_stats.set_completion_time()
|
||||
aborted_reqs.append(req)
|
||||
advance_logprob_pt(i, req)
|
||||
req.time_stats.set_last_chunked_prefill_finish_time()
|
||||
continue
|
||||
|
||||
# Optimistic bootstrap can fail while this overlapped chunk is
|
||||
# already running. Drop aborted chunks instead of sending KV.
|
||||
if is_aborted(req):
|
||||
self.clear_pending_chunk_send(req)
|
||||
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
|
||||
@@ -876,6 +893,12 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
auxiliary_output_starts,
|
||||
)
|
||||
|
||||
if aborted_reqs:
|
||||
self.output_streamer.stream_output(
|
||||
aborted_reqs,
|
||||
any(req.return_logprob for req in aborted_reqs),
|
||||
)
|
||||
|
||||
can_run_cuda_graph = result.can_run_cuda_graph
|
||||
self.metrics_reporter.report_prefill_stats(
|
||||
batch=batch,
|
||||
@@ -1030,6 +1053,35 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
"""
|
||||
self.disagg_prefill_pending_chunk_rids.discard(req.rid)
|
||||
|
||||
def _retire_aborted_prefill_result(self: Scheduler, req: Req) -> bool:
|
||||
"""Release an aborted request when its last prefill result is safe."""
|
||||
self.clear_pending_chunk_send(req)
|
||||
owns_resources = (
|
||||
req.kv.holds_kv or req.kv.holds_mamba or req.metadata_buffer_index >= 0
|
||||
)
|
||||
if not owns_resources:
|
||||
# A bootstrap failure or earlier abort already retired it.
|
||||
return False
|
||||
|
||||
sender = req.disagg_kv_sender
|
||||
if sender is not None:
|
||||
try:
|
||||
sender.abort()
|
||||
except Exception:
|
||||
# Transport notification is best effort; local ownership must
|
||||
# still be released or the next idle invariant check will fail.
|
||||
logger.exception("Failed to notify KV sender of abort for %s", req.rid)
|
||||
|
||||
if req.to_finish is not None and not req.finished():
|
||||
req.update_finish_state()
|
||||
maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator)
|
||||
req.pending_bootstrap = False
|
||||
if self.enable_hicache_storage:
|
||||
self.tree_cache.release_aborted_request(req.rid)
|
||||
if req.kv.holds_kv or req.kv.holds_mamba:
|
||||
release_kv_cache(req, self.tree_cache, is_insert=False)
|
||||
return True
|
||||
|
||||
def handle_bootstrap_failure(self: Scheduler, req: Req) -> None:
|
||||
self.clear_pending_chunk_send(req)
|
||||
error_message = (
|
||||
|
||||
Reference in New Issue
Block a user