From 1fabea26ef9e65c3afaed1e3bea4a8aef6fce37b Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 15 Jul 2026 22:36:54 +0800 Subject: [PATCH] Batch exact Frontier EP lane predictions --- ..._moe_ep_prefill_batched_lane_predict.patch | 82 +++++++++++++++++++ ...ier_moe_ep_prefill_use_batched_lanes.patch | 56 +++++++++++++ 2 files changed, 138 insertions(+) create mode 100644 runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_batched_lane_predict.patch create mode 100644 runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_use_batched_lanes.patch diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_batched_lane_predict.patch b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_batched_lane_predict.patch new file mode 100644 index 0000000..088aae6 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_batched_lane_predict.patch @@ -0,0 +1,82 @@ +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 +@@ -385,6 +385,78 @@ + allocation_ratios=allocation_ratios, + ) + ++ def _get_on_demand_predictions_batch( ++ self, ++ model_name: str, ++ feature_rows: List[Dict[str, float]], ++ ) -> List[float]: ++ """Predict independent feature rows in one exact forest invocation. ++ ++ ``RandomForestRegressor.predict`` applies the same trees to every row. ++ Batching the eight EP lanes removes repeated pandas/joblib call overhead; ++ it does not combine lanes or alter their features. Results populate the ++ same runtime cache used by the scalar prediction path. ++ """ ++ if not feature_rows: ++ return [] ++ ++ model_info = self._predictions.get(model_name) ++ if ( ++ not isinstance(model_info, dict) ++ or not model_info.get("_on_demand_prediction", False) ++ ): ++ raise ValueError( ++ f"Model {model_name} is not configured for on-demand prediction" ++ ) ++ model = model_info.get("_model") ++ feature_names = model_info.get("_feature_names", []) ++ if model is None: ++ raise ValueError(f"Model {model_name} has no trained model available") ++ ++ family_name = self._measurement_family_name(self._active_measurement_type) ++ runtime_cache = self._runtime_cache[family_name][model_name] ++ results: List[Optional[float]] = [None] * len(feature_rows) ++ misses: Dict[tuple, tuple] = {} ++ for index, features in enumerate(feature_rows): ++ missing = [name for name in feature_names if name not in features] ++ if missing: ++ raise ValueError( ++ f"On-demand prediction missing required features for {model_name}: " ++ f"{missing}. Provided keys: {sorted(list(features.keys()))}" ++ ) ++ cache_key = tuple(features[key] for key in sorted(features.keys())) ++ if cache_key in runtime_cache: ++ results[index] = float(runtime_cache[cache_key]) ++ continue ++ if cache_key not in misses: ++ misses[cache_key] = (features, []) ++ misses[cache_key][1].append(index) ++ ++ if misses: ++ miss_keys = list(misses) ++ frame = pd.DataFrame( ++ [ ++ [misses[key][0][name] for name in feature_names] ++ for key in miss_keys ++ ], ++ columns=feature_names, ++ ) ++ try: ++ predictions = model.predict(frame) ++ except Exception as error: ++ raise ValueError( ++ f"Batched on-demand prediction failed for {model_name}: {error}" ++ ) from error ++ for cache_key, raw_prediction in zip(miss_keys, predictions): ++ prediction = max(0.0, float(raw_prediction)) ++ runtime_cache[cache_key] = prediction ++ for index in misses[cache_key][1]: ++ results[index] = prediction ++ ++ if any(result is None for result in results): ++ raise AssertionError("batched on-demand prediction left an empty result") ++ return [float(result) for result in results] ++ + @lru_cache(maxsize=None) + def _get_critical_prefill_ep_lane_tokens( + self, diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_use_batched_lanes.patch b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_use_batched_lanes.patch new file mode 100644 index 0000000..5b23cff --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_prefill_use_batched_lanes.patch @@ -0,0 +1,56 @@ +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 +@@ -494,27 +494,35 @@ + 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) ++ lane_features = [ ++ self._build_moe_load_imbalance_features(allocation) ++ for allocation in lane_tokens ++ ] ++ lane_scores = [0.0] * lane_count ++ for operation in ("moe_shuffling", "moe_grouped_gemm"): ++ prediction = self._predictions.get(operation) ++ if isinstance(prediction, dict) and prediction.get( ++ "_on_demand_prediction", False ++ ): ++ operation_predictions = self._get_on_demand_predictions_batch( ++ operation, lane_features ++ ) ++ lane_scores = [ ++ score + operation_prediction ++ for score, operation_prediction in zip( ++ lane_scores, operation_predictions + ) +- return score ++ ] + +- return max( +- lane_tokens, +- key=lambda allocation: ( +- _lane_compute_score(allocation), +- sum(allocation.values()), +- tuple(allocation.values()), ++ critical_lane_index = max( ++ range(lane_count), ++ key=lambda lane_index: ( ++ lane_scores[lane_index], ++ sum(lane_tokens[lane_index].values()), ++ tuple(lane_tokens[lane_index].values()), + ), + ) ++ return lane_tokens[critical_lane_index] + + def predict_monolithic_decode_shared_domain_lane_moe_times_ms( + self,