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
@@ -440,6 +440,7 @@ class SchedulerMetricsMixin:
def log_decode_stats_every_iteration( def log_decode_stats_every_iteration(
self: Scheduler, batch: ScheduleBatch, num_accepted_tokens: int self: Scheduler, batch: ScheduleBatch, num_accepted_tokens: int
): ):
if self.enable_metrics:
self.metrics_collector.increment_realtime_tokens( self.metrics_collector.increment_realtime_tokens(
# TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator # TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator
decode_tokens=batch.batch_size() + num_accepted_tokens, decode_tokens=batch.batch_size() + num_accepted_tokens,
@@ -470,12 +470,9 @@ class SchedulerOutputProcessorMixin:
self.token_to_kv_pool_allocator.free_group_end() self.token_to_kv_pool_allocator.free_group_end()
self.forward_ct_decode = (self.forward_ct_decode + 1) % (1 << 30) self.forward_ct_decode = (self.forward_ct_decode + 1) % (1 << 30)
if ( if self.current_scheduler_metrics_enabled:
self.current_scheduler_metrics_enabled if self.forward_ct_decode % self.server_args.decode_log_interval == 0:
and 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(can_run_cuda_graph, running_batch=batch)
if self.enable_metrics:
self.log_decode_stats_every_iteration( self.log_decode_stats_every_iteration(
batch, num_accepted_tokens=result.num_accepted_tokens batch, num_accepted_tokens=result.num_accepted_tokens
) )