Publish elastic-EP active ranks from a dedicated step (#25721)

This commit is contained in:
fzyzcjy
2026-05-19 09:21:04 +08:00
committed by GitHub
parent fa37b68653
commit 1f3e5aa1e0
+7 -3
View File
@@ -2923,10 +2923,16 @@ class Scheduler(
pooled_hidden_states=pooler_output.pooled_hidden_states, pooled_hidden_states=pooler_output.pooled_hidden_states,
) )
if ( self._maybe_report_active_ranks()
return ret
def _maybe_report_active_ranks(self) -> None:
if not (
self.server_args.enable_dp_attention self.server_args.enable_dp_attention
and self.server_args.elastic_ep_backend is not None and self.server_args.elastic_ep_backend is not None
): ):
return
# Get the tensors indicating rank activeness # Get the tensors indicating rank activeness
tp_active_ranks = self.tp_group.active_ranks.detach().cpu().numpy() tp_active_ranks = self.tp_group.active_ranks.detach().cpu().numpy()
tp_active_ranks_cpu = self.tp_group.active_ranks_cpu.detach().numpy() tp_active_ranks_cpu = self.tp_group.active_ranks_cpu.detach().numpy()
@@ -2936,8 +2942,6 @@ class Scheduler(
ActiveRanksOutput(status=dp_active_ranks.tolist()) ActiveRanksOutput(status=dp_active_ranks.tolist())
) )
return ret
def launch_batch_sample_if_needed( def launch_batch_sample_if_needed(
self, batch_result: GenerationBatchResult self, batch_result: GenerationBatchResult
) -> Union[GenerationBatchResult]: ) -> Union[GenerationBatchResult]: