[Observability] Add pending token count to prefill log and get_load (#22480)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
c554dc5c64
commit
37107bee6f
@@ -1857,6 +1857,7 @@ class GetLoadReqOutput(BaseReq):
|
|||||||
num_reqs: int
|
num_reqs: int
|
||||||
num_waiting_reqs: int
|
num_waiting_reqs: int
|
||||||
num_tokens: int
|
num_tokens: int
|
||||||
|
num_pending_tokens: int
|
||||||
ts_tic: float
|
ts_tic: float
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -2579,9 +2579,18 @@ class Scheduler(
|
|||||||
|
|
||||||
new_batch.prepare_for_extend()
|
new_batch.prepare_for_extend()
|
||||||
|
|
||||||
# Record prefill stats for logging after forward
|
# Record prefill stats for logging after forward.
|
||||||
new_batch.prefill_stats = PrefillStats.from_adder(
|
new_batch.prefill_stats = PrefillStats.from_adder(
|
||||||
adder, self.running_batch.reqs, self.enable_priority_scheduling
|
adder,
|
||||||
|
self.running_batch.reqs,
|
||||||
|
self.enable_priority_scheduling,
|
||||||
|
num_pending_tokens=self._get_num_pending_tokens(
|
||||||
|
chunk_deduct=(
|
||||||
|
self.chunked_req.extend_input_len
|
||||||
|
if self.chunked_req is not None
|
||||||
|
else 0
|
||||||
|
)
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Mixed-style chunked prefill
|
# Mixed-style chunked prefill
|
||||||
|
|||||||
@@ -55,6 +55,7 @@ class PrefillStats:
|
|||||||
new_token_ratio: float
|
new_token_ratio: float
|
||||||
num_running_reqs: QueueCount
|
num_running_reqs: QueueCount
|
||||||
num_new_seqs: int # len(can_run_list)
|
num_new_seqs: int # len(can_run_list)
|
||||||
|
num_pending_tokens: int = 0
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_adder(
|
def from_adder(
|
||||||
@@ -62,6 +63,7 @@ class PrefillStats:
|
|||||||
adder: PrefillAdder,
|
adder: PrefillAdder,
|
||||||
running_reqs: List[Req],
|
running_reqs: List[Req],
|
||||||
enable_priority_scheduling: bool = False,
|
enable_priority_scheduling: bool = False,
|
||||||
|
num_pending_tokens: int = 0,
|
||||||
):
|
):
|
||||||
return cls(
|
return cls(
|
||||||
log_input_tokens=adder.log_input_tokens,
|
log_input_tokens=adder.log_input_tokens,
|
||||||
@@ -71,6 +73,7 @@ class PrefillStats:
|
|||||||
running_reqs, enable_priority_scheduling
|
running_reqs, enable_priority_scheduling
|
||||||
),
|
),
|
||||||
num_new_seqs=len(adder.can_run_list),
|
num_new_seqs=len(adder.can_run_list),
|
||||||
|
num_pending_tokens=num_pending_tokens,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -393,6 +396,7 @@ class SchedulerMetricsMixin:
|
|||||||
f"{token_usage_msg}"
|
f"{token_usage_msg}"
|
||||||
f"#running-req: {prefill_stats.num_running_reqs.total}, "
|
f"#running-req: {prefill_stats.num_running_reqs.total}, "
|
||||||
f"#queue-req: {len(self.waiting_queue)}, "
|
f"#queue-req: {len(self.waiting_queue)}, "
|
||||||
|
f"#pending-token: {prefill_stats.num_pending_tokens}, "
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||||
@@ -862,6 +866,26 @@ class SchedulerMetricsMixin:
|
|||||||
self.stats.token_usage / 0.9,
|
self.stats.token_usage / 0.9,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _get_num_pending_tokens(self: Scheduler, chunk_deduct: int = 0) -> int:
|
||||||
|
"""Get the total number of tokens pending prefill.
|
||||||
|
|
||||||
|
This includes tokens from waiting queue requests plus remaining tokens
|
||||||
|
from the currently chunked request.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
chunk_deduct: extra tokens to subtract from the chunked request's
|
||||||
|
remaining count. At batch-scheduling time the current chunk
|
||||||
|
has been planned but ``prefix_indices`` does not yet include it,
|
||||||
|
so callers pass ``extend_input_len`` here. At query time
|
||||||
|
(``get_load``) ``prefix_indices`` is already up-to-date, so
|
||||||
|
the default 0 is correct.
|
||||||
|
"""
|
||||||
|
num_pending_tokens = sum(req.seqlen for req in self.waiting_queue)
|
||||||
|
if self.chunked_req is not None:
|
||||||
|
req = self.chunked_req
|
||||||
|
num_pending_tokens += req.seqlen - len(req.prefix_indices) - chunk_deduct
|
||||||
|
return num_pending_tokens
|
||||||
|
|
||||||
def get_load(self: Scheduler, _: GetLoadReqInput = None) -> GetLoadReqOutput:
|
def get_load(self: Scheduler, _: GetLoadReqInput = None) -> GetLoadReqOutput:
|
||||||
if self.is_hybrid_swa:
|
if self.is_hybrid_swa:
|
||||||
full_num_used, swa_num_used, *_ = self._get_swa_token_info()
|
full_num_used, swa_num_used, *_ = self._get_swa_token_info()
|
||||||
@@ -871,7 +895,9 @@ class SchedulerMetricsMixin:
|
|||||||
else:
|
else:
|
||||||
num_tokens = self._get_token_info()[0]
|
num_tokens = self._get_token_info()[0]
|
||||||
|
|
||||||
# Tokens in waiting queue, bootstrap queue, prealloc queue
|
num_pending_tokens = self._get_num_pending_tokens()
|
||||||
|
|
||||||
|
# Tokens and request count in waiting queue, bootstrap queue, prealloc queue
|
||||||
waiting_queues = [self.waiting_queue]
|
waiting_queues = [self.waiting_queue]
|
||||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||||
waiting_queues.append(self.disagg_prefill_bootstrap_queue.queue)
|
waiting_queues.append(self.disagg_prefill_bootstrap_queue.queue)
|
||||||
@@ -888,6 +914,7 @@ class SchedulerMetricsMixin:
|
|||||||
num_reqs=len(self.running_batch.reqs) + num_waiting_reqs,
|
num_reqs=len(self.running_batch.reqs) + num_waiting_reqs,
|
||||||
num_waiting_reqs=num_waiting_reqs,
|
num_waiting_reqs=num_waiting_reqs,
|
||||||
num_tokens=num_tokens,
|
num_tokens=num_tokens,
|
||||||
|
num_pending_tokens=num_pending_tokens,
|
||||||
ts_tic=time.perf_counter(),
|
ts_tic=time.perf_counter(),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user