minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fzyzcjy
2026-05-19 09:21:04 +08:00
committed by GitHub
parent fa37b68653
commit 1f3e5aa1e0
1 changed files with 15 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+15 -11
View File
@@ -2923,20 +2923,24 @@ 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
): ):
# Get the tensors indicating rank activeness return
tp_active_ranks = self.tp_group.active_ranks.detach().cpu().numpy() # Get the tensors indicating rank activeness
tp_active_ranks_cpu = self.tp_group.active_ranks_cpu.detach().numpy() tp_active_ranks = self.tp_group.active_ranks.detach().cpu().numpy()
tp_active_ranks &= tp_active_ranks_cpu tp_active_ranks_cpu = self.tp_group.active_ranks_cpu.detach().numpy()
dp_active_ranks = tp_active_ranks.reshape(self.ps.dp_size, -1).prod(axis=1) tp_active_ranks &= tp_active_ranks_cpu
self.ipc_channels.send_to_tokenizer.send_output( dp_active_ranks = tp_active_ranks.reshape(self.ps.dp_size, -1).prod(axis=1)
ActiveRanksOutput(status=dp_active_ranks.tolist()) self.ipc_channels.send_to_tokenizer.send_output(
) 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