Support EPLB balancedness prometheus metric without GPU->CPU synchronize (#15401)

2025-12-18 22:24:23 +08:00
parent 602fe3b296
commit 88a405cc10
8 changed files with 111 additions and 33 deletions
@@ -406,6 +406,7 @@ class TpModelWorker(BaseTpWorker):
            batch_result = GenerationBatchResult(
                logits_output=logits_output,
                can_run_cuda_graph=can_run_cuda_graph,
+                expert_distribution_metrics=out.expert_distribution_metrics,
            )

            if is_verify:
@@ -460,6 +461,7 @@ class TpModelWorker(BaseTpWorker):
            return GenerationBatchResult(
                pp_hidden_states_proxy_tensors=pp_proxy_tensors,
                can_run_cuda_graph=can_run_cuda_graph,
+                expert_distribution_metrics=out.expert_distribution_metrics,
            )

    def forward_batch_split_prefill(self, batch: ScheduleBatch):
@@ -482,6 +484,7 @@ class TpModelWorker(BaseTpWorker):
        batch_result = GenerationBatchResult(
            logits_output=logits_output,
            can_run_cuda_graph=can_run_cuda_graph,
+            expert_distribution_metrics=out.expert_distribution_metrics,
        )
        batch_result.next_token_ids = next_token_ids
        return batch_result