Fix log_decode_stats_every_iteration when having TP in attention (#16871)

This commit is contained in:
fzyzcjy
2026-01-10 20:10:23 +08:00
committed by GitHub
parent 641830c1c2
commit a0899bdbd8
2 changed files with 9 additions and 11 deletions

View File

@@ -440,11 +440,12 @@ class SchedulerMetricsMixin:
def log_decode_stats_every_iteration(
self: Scheduler, batch: ScheduleBatch, num_accepted_tokens: int
):
self.metrics_collector.increment_realtime_tokens(
# TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator
decode_tokens=batch.batch_size() + num_accepted_tokens,
dp_cooperation_info=batch.dp_cooperation_info,
)
if self.enable_metrics:
self.metrics_collector.increment_realtime_tokens(
# TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator
decode_tokens=batch.batch_size() + num_accepted_tokens,
dp_cooperation_info=batch.dp_cooperation_info,
)
def log_batch_result_stats(
self: Scheduler,

View File

@@ -470,12 +470,9 @@ class SchedulerOutputProcessorMixin:
self.token_to_kv_pool_allocator.free_group_end()
self.forward_ct_decode = (self.forward_ct_decode + 1) % (1 << 30)
if (
self.current_scheduler_metrics_enabled
and self.forward_ct_decode % self.server_args.decode_log_interval == 0
):
self.log_decode_stats(can_run_cuda_graph, running_batch=batch)
if self.enable_metrics:
if self.current_scheduler_metrics_enabled:
if self.forward_ct_decode % self.server_args.decode_log_interval == 0:
self.log_decode_stats(can_run_cuda_graph, running_batch=batch)
self.log_decode_stats_every_iteration(
batch, num_accepted_tokens=result.num_accepted_tokens
)