Stabilize disaggregated decode burst handoff
Burst decode can sit in disaggregation prealloc/transfer queues while health probes still need an immediate scheduler-alive response, and EAGLE prebuilt metadata must outlive the synthetic prebuilt step until the first real decode consume. The transfer path also needs to send final aux metadata even when there are no new KV pages in the final chunk. This keeps decode metadata slot ownership narrow instead of cloning EAGLE tensors, adds fail-safe cleanup for prebuilt exceptions/finished prebuilt requests, fixes final empty chunk metadata transfer, and records the investigation ledger for future debugging. Constraint: bs=1 historically worked, so decode compute/sampling behavior must not be changed without direct evidence. Rejected: Clone EAGLE metadata tensors on transfer commit | avoids lifetime issues but adds hot-path CPU/GPU memory traffic. Rejected: Treat prefill AbortReq logs as root cause | current evidence shows they can be downstream of decode/router aborts. Confidence: medium Scope-risk: moderate Directive: Do not move EAGLE metadata slot release earlier than first real decode result processing without proving the H2D/spec_info consume point changed. Tested: g0034 docker py_compile for touched scheduler/disagg/test files; PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_req_to_token_pool.py test/registered/unit/managers/test_scheduler_health_check.py -> 24 passed Not-tested: Fresh end-to-end burst traffic after restart; decode node_rank=1 persistent log capture.
This commit is contained in:
@@ -1482,9 +1482,10 @@ class SchedulerDisaggregationDecodeMixin:
|
||||
new_batch.prepare_for_prebuilt()
|
||||
try:
|
||||
new_batch.process_prebuilt(self.server_args, self.future_map)
|
||||
finally:
|
||||
except BaseException:
|
||||
for req in can_run_list:
|
||||
self._free_decode_metadata_index_if_held(req)
|
||||
raise
|
||||
|
||||
return new_batch
|
||||
|
||||
|
||||
@@ -946,11 +946,18 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
page_size,
|
||||
)
|
||||
|
||||
if len(page_indices) == 0:
|
||||
if len(page_indices) == 0 and not last_chunk:
|
||||
logger.info(
|
||||
f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
|
||||
f"Skip sending non-final kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
|
||||
)
|
||||
return
|
||||
if len(page_indices) == 0:
|
||||
logger.warning(
|
||||
"[CP_SHARED_KV_TRANSFER][final_empty_chunk] "
|
||||
"sending final aux metadata without new KV pages: "
|
||||
f"{req.rid=} {req.bootstrap_room=} {start_idx=} {end_idx=} "
|
||||
f"origin_len={len(req.origin_input_ids)} fill_len={len(req.fill_ids)}"
|
||||
)
|
||||
prefill_queue = getattr(self, "disagg_prefill_bootstrap_queue", None)
|
||||
has_draft_pool = (
|
||||
getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None
|
||||
|
||||
@@ -1567,8 +1567,15 @@ class Scheduler(
|
||||
if is_health_check_generate_req(recv_req) and not self.is_fully_idle(
|
||||
for_health_check=True
|
||||
):
|
||||
self.return_health_check_ipcs.append(
|
||||
getattr(recv_req, "http_worker_ipc", None)
|
||||
# Send the liveness signal immediately. Deferring it until
|
||||
# process_batch_result can false-fail disaggregated decode when
|
||||
# requests are waiting in prealloc/transfer queues and no decode
|
||||
# batch is currently producing output.
|
||||
self.send_to_tokenizer.send_output(
|
||||
HealthCheckOutput(
|
||||
http_worker_ipc=getattr(recv_req, "http_worker_ipc", None)
|
||||
),
|
||||
recv_req,
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -2902,11 +2909,11 @@ class Scheduler(
|
||||
return ClearHiCacheReqOutput(success=if_success)
|
||||
|
||||
def is_fully_idle(self, for_health_check=False) -> bool:
|
||||
# Health check piggybacks on running requests in process_output.
|
||||
# Only running_batch + waiting_queue guarantee active GPU processing;
|
||||
# disagg queues (bootstrap/prealloc/transfer) may have items without
|
||||
# any request actually running on GPU — e.g. stuck handshake, full
|
||||
# KV cache, or stalled transfer — so they can't carry health info.
|
||||
# Health checks need a scheduler-liveness answer. Disaggregated queues
|
||||
# may wait for prefill/transfer without a local decode batch result, so
|
||||
# they must still count as busy; otherwise /health injects a generation
|
||||
# request that can sit behind the same queue and trigger a false
|
||||
# detokenizer-hang shutdown.
|
||||
# Batch running status
|
||||
idle = (
|
||||
self.running_batch.is_empty()
|
||||
@@ -2921,17 +2928,18 @@ class Scheduler(
|
||||
# Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode)
|
||||
idle &= len(self.waiting_queue) == 0
|
||||
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||
idle &= len(self.disagg_prefill_inflight_queue) == 0
|
||||
idle &= len(self.disagg_prefill_bootstrap_queue.queue) == 0
|
||||
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
idle &= len(self.disagg_decode_prealloc_queue.queue) == 0
|
||||
idle &= len(self.disagg_decode_transfer_queue.queue) == 0
|
||||
|
||||
if not for_health_check:
|
||||
# Grammar queue and prefill inflight queue may not produce batch
|
||||
# results instantly, but they still indicate the server is not idle.
|
||||
idle &= len(self.grammar_manager.grammar_queue) == 0
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||
idle &= len(self.disagg_prefill_inflight_queue) == 0
|
||||
idle &= len(self.disagg_prefill_bootstrap_queue.queue) == 0
|
||||
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
idle &= len(self.disagg_decode_prealloc_queue.queue) == 0
|
||||
idle &= len(self.disagg_decode_transfer_queue.queue) == 0
|
||||
|
||||
# HiCache: in-flight async ops (GPU↔Host↔L3) must drain before
|
||||
# destructive operations like attach/detach/flush_cache.
|
||||
@@ -3189,6 +3197,60 @@ class Scheduler(
|
||||
return RpcReqOutput(success, "" if not exec else str(exec))
|
||||
|
||||
def abort_request(self, recv_req: AbortReq):
|
||||
if self.disaggregation_mode != DisaggregationMode.NULL:
|
||||
running_reqs = (
|
||||
len(self.running_batch.reqs)
|
||||
if self.running_batch is not None and self.running_batch.reqs is not None
|
||||
else 0
|
||||
)
|
||||
prefill_bootstrap_reqs = (
|
||||
len(self.disagg_prefill_bootstrap_queue.queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL
|
||||
and getattr(self, "disagg_prefill_bootstrap_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
prefill_inflight_reqs = (
|
||||
len(self.disagg_prefill_inflight_queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL
|
||||
and getattr(self, "disagg_prefill_inflight_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
decode_prealloc_reqs = (
|
||||
len(self.disagg_decode_prealloc_queue.queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE
|
||||
and getattr(self, "disagg_decode_prealloc_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
decode_transfer_reqs = (
|
||||
len(self.disagg_decode_transfer_queue.queue)
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE
|
||||
and getattr(self, "disagg_decode_transfer_queue", None) is not None
|
||||
else 0
|
||||
)
|
||||
if (
|
||||
recv_req.abort_all
|
||||
or prefill_bootstrap_reqs
|
||||
or prefill_inflight_reqs
|
||||
or decode_prealloc_reqs
|
||||
or decode_transfer_reqs
|
||||
):
|
||||
logger.warning(
|
||||
"[DISAGG_ABORT_TRACE] scheduler abort_request rid=%s abort_all=%s "
|
||||
"mode=%s waiting=%s running=%s prefill_bootstrap=%s "
|
||||
"prefill_inflight=%s decode_prealloc=%s decode_transfer=%s "
|
||||
"abort_message=%s",
|
||||
recv_req.rid,
|
||||
recv_req.abort_all,
|
||||
self.disaggregation_mode,
|
||||
len(self.waiting_queue),
|
||||
running_reqs,
|
||||
prefill_bootstrap_reqs,
|
||||
prefill_inflight_reqs,
|
||||
decode_prealloc_reqs,
|
||||
decode_transfer_reqs,
|
||||
recv_req.abort_message,
|
||||
)
|
||||
|
||||
# todo hisparse, release resources for abort requests in hisparse coordinator
|
||||
# Delete requests in the waiting queue
|
||||
to_del = []
|
||||
|
||||
@@ -64,9 +64,11 @@ class SchedulerOutputProcessorMixin:
|
||||
allocator.free(idx)
|
||||
req.metadata_buffer_index = -1
|
||||
# The EAGLE handoff tensors are views into the reusable metadata slot.
|
||||
# process_prebuilt consumes them into the batch before this helper is
|
||||
# called; keeping the views afterwards makes the request observe a
|
||||
# future transfer that reuses the same slot.
|
||||
# They are safe to release only after either the request finishes in the
|
||||
# synthetic prebuilt step or the first real decode forward consumes the
|
||||
# spec_info built from those views. Releasing immediately after
|
||||
# process_prebuilt lets a burst transfer overwrite pinned CPU metadata
|
||||
# before pending H2D copies have been consumed.
|
||||
req.output_topk_p = None
|
||||
req.output_topk_index = None
|
||||
req.hidden_states_tensor = None
|
||||
@@ -184,6 +186,7 @@ class SchedulerOutputProcessorMixin:
|
||||
if req.finished():
|
||||
req.time_stats.set_quick_finish_time()
|
||||
release_kv_cache(req, self.tree_cache)
|
||||
self._free_decode_metadata_index_if_held(req)
|
||||
|
||||
# Note: Logprobs should be handled on the prefill engine.
|
||||
self.stream_output(batch.reqs, batch.return_logprob)
|
||||
@@ -459,6 +462,13 @@ class SchedulerOutputProcessorMixin:
|
||||
if result.copy_done is not None:
|
||||
result.copy_done.synchronize()
|
||||
|
||||
# Requests entering decode through disaggregated prebuilt EAGLE hold
|
||||
# reusable metadata slots until the first real decode forward consumes
|
||||
# the initial draft state. Free here, not in get_new_prebuilt_batch(),
|
||||
# so burst transfers cannot overwrite the source views too early.
|
||||
for req in batch.reqs:
|
||||
self._free_decode_metadata_index_if_held(req)
|
||||
|
||||
logits_output, next_token_ids, can_run_cuda_graph = (
|
||||
result.logits_output,
|
||||
result.next_token_ids,
|
||||
|
||||
@@ -2187,7 +2187,14 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
||||
continue
|
||||
|
||||
is_stream = getattr(state.obj, "stream", False)
|
||||
already_sent_data = state.response_sent_to_client_ts > 0
|
||||
time_stats = getattr(state, "time_stats", None)
|
||||
already_sent_data = bool(
|
||||
getattr(state, "response_sent_to_client_ts", 0) > 0
|
||||
or (
|
||||
time_stats is not None
|
||||
and getattr(time_stats, "response_sent_to_client_time", 0) > 0
|
||||
)
|
||||
)
|
||||
|
||||
if is_stream and already_sent_data:
|
||||
# Streaming request that already sent partial data to client.
|
||||
|
||||
Reference in New Issue
Block a user