diff --git 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 new file mode 100644 index 0000000..cc6d4ee --- /dev/null +++ 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, ++ )