Model critical EP lane for Frontier prefill

2026-07-15 20:15:09 +08:00
parent 9275619984
commit 76e0de33ee
1 changed files with 76 additions and 0 deletions
--- a/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_critical_lane.patch
+++ b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_critical_lane.patch
@@ -0,0 +1,76 @@
+diff --git a/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py b/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py
+--- a/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py
+++ b/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py
+@@ -386,0 +387,58 @@
+    def _get_critical_prefill_ep_lane_tokens(
+        self,
+        total_routed_tokens: int,
+        layer_id: int,
+    ) -> Dict[int, int]:
+        """Return the routed-token allocation for the slowest prefill EP lane.
+
+        A monolithic prefill starts with a global ``tokens * topk`` routing
+        population, but each EP rank executes only the assignments for its local
+        experts. The ranks synchronize after the MoE block, so the layer cost is
+        governed by the slowest lane rather than by either the global population
+        or the average lane.
+        """
+        total_experts = int(self._replica_config.total_expert_num)
+        lane_count = int(self._moe_ep_size)
+        if total_experts <= 0 or lane_count <= 1 or total_experts % lane_count != 0:
+            raise ValueError(
+                "Prefill EP routing requires total_expert_num to be positive and "
+                f"divisible by moe_ep_size > 1; got total_expert_num={total_experts}, "
+                f"moe_ep_size={lane_count}"
+            )
+
+        experts_per_lane = total_experts // lane_count
+        global_tokens = self._get_global_per_expert_tokens(
+            total_routed_tokens=total_routed_tokens,
+            layer_id=layer_id,
+        )
+        lane_tokens: List[Dict[int, int]] = [
+            {local_expert_id: 0 for local_expert_id in range(experts_per_lane)}
+            for _ in range(lane_count)
+        ]
+        for global_expert_id, token_count in global_tokens.items():
+            lane_id = int(global_expert_id) // experts_per_lane
+            local_expert_id = int(global_expert_id) % experts_per_lane
+            lane_tokens[lane_id][local_expert_id] = int(token_count)
+
+        def _lane_compute_score(allocation: Dict[int, int]) -> float:
+            features = self._build_moe_load_imbalance_features(allocation)
+            score = 0.0
+            for operation in ("moe_shuffling", "moe_grouped_gemm"):
+                prediction = self._predictions.get(operation)
+                if isinstance(prediction, dict) and prediction.get(
+                    "_on_demand_prediction", False
+                ):
+                    score += float(
+                        self._get_on_demand_prediction(operation, features)
+                    )
+            return score
+
+        return max(
+            lane_tokens,
+            key=lambda allocation: (
+                _lane_compute_score(allocation),
+                sum(allocation.values()),
+                tuple(allocation.values()),
+            ),
+        )
+
+@@ -1023,0 +1082,13 @@
+
+        from frontier.entities import EPBatchGroup
+
+        is_monolithic_prefill_ep = (
+            self._moe_ep_size > 1
+            and not isinstance(batch, EPBatchGroup)
+            and not bool(getattr(batch, "is_pure_decode_batch", False))
+        )
+        if is_monolithic_prefill_ep:
+            return self._get_critical_prefill_ep_lane_tokens(
+                total_routed_tokens=num_tokens * self._router_topk,
+                layer_id=layer_id,
+            )