Batch exact Frontier EP lane predictions

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,

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,