From 1f8900a4ac64e45754b03d0aa7c1dddab65785cf Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Thu, 23 Jul 2026 15:27:40 +0800 Subject: [PATCH] Experiment with structured attention prefill predictor --- .../shared_prediction_model_manager.py | 16 +++- .../sklearn_execution_time_predictor.py | 16 +++- .../structured_attention_prefill.py | 79 +++++++++++++++++++ .../unit/test_structured_attention_prefill.py | 59 ++++++++++++++ 4 files changed, 165 insertions(+), 5 deletions(-) create mode 100644 frontier/execution_time_predictor/structured_attention_prefill.py create mode 100644 tests/unit/test_structured_attention_prefill.py diff --git a/frontier/execution_time_predictor/shared_prediction_model_manager.py b/frontier/execution_time_predictor/shared_prediction_model_manager.py index 8a65a49..4a21165 100644 --- a/frontier/execution_time_predictor/shared_prediction_model_manager.py +++ b/frontier/execution_time_predictor/shared_prediction_model_manager.py @@ -19,6 +19,9 @@ from frontier.execution_time_predictor.attention_tp_policy import ( from frontier.execution_time_predictor.attention_dataset_contract import ( enforce_mixed_attention_input_contract, ) +from frontier.execution_time_predictor.structured_attention_prefill import ( + StructuredAttentionPrefillRegressor, +) from frontier.logger import init_logger from frontier.moe_gating_runtime import ( DEFAULT_MOE_GATING_RUNTIME_CONTEXT, @@ -1254,7 +1257,10 @@ class ExecutionTimePredictionModelManager: raise ValueError( "Missing required column 'prefill_chunk_size' in attention profiling data." ) - standard_prefill_df = prefill_df[prefill_df["prefill_chunk_size"] > 0].copy() + standard_prefill_df = prefill_df[ + (prefill_df["prefill_chunk_size"] > 0) + & (prefill_df["batch_size"] == 1) + ].copy() prefill_model_signature = f"attn_prefill_{attention_signature}" if prefill_model_signature not in trained_model_signatures: @@ -1742,7 +1748,13 @@ class ExecutionTimePredictionModelManager: # initialization to generate missing cache files. # ============================================================ - estimator, grid_search_params = self._create_estimator_and_params(execution_time_predictor_config) + if model_name == "attn_prefill": + estimator = StructuredAttentionPrefillRegressor() + grid_search_params = {} + else: + estimator, grid_search_params = self._create_estimator_and_params( + execution_time_predictor_config + ) cv = min(execution_time_predictor_config.k_fold_cv_splits, len(df)) if len(df) >= 2 else 2 diff --git a/frontier/execution_time_predictor/sklearn_execution_time_predictor.py b/frontier/execution_time_predictor/sklearn_execution_time_predictor.py index 27b62bf..b7f5350 100644 --- a/frontier/execution_time_predictor/sklearn_execution_time_predictor.py +++ b/frontier/execution_time_predictor/sklearn_execution_time_predictor.py @@ -45,6 +45,9 @@ from frontier.execution_time_predictor.attention_tp_policy import ( from frontier.execution_time_predictor.attention_dataset_contract import ( enforce_mixed_attention_input_contract, ) +from frontier.execution_time_predictor.structured_attention_prefill import ( + StructuredAttentionPrefillRegressor, +) from frontier.logger import init_logger from frontier.moe_gating_runtime import get_moe_gating_base_model_name from frontier.profiling.cpu_overhead.schema import ( @@ -2573,8 +2576,12 @@ class SklearnExecutionTimePredictor(BaseExecutionTimePredictor): if cached_model: return cached_model - model = self._get_estimator() - grid_search_params = self._get_grid_search_params() + if model_name == "attn_prefill": + model = StructuredAttentionPrefillRegressor() + grid_search_params = {} + else: + model = self._get_estimator() + grid_search_params = self._get_grid_search_params() if len(df) < self._config.k_fold_cv_splits: cv = 2 @@ -2869,7 +2876,10 @@ class SklearnExecutionTimePredictor(BaseExecutionTimePredictor): raise ValueError( "Missing required column 'prefill_chunk_size' in attention profiling data." ) - standard_prefill_df = prefill_df[prefill_df["prefill_chunk_size"] > 0].copy() + standard_prefill_df = prefill_df[ + (prefill_df["prefill_chunk_size"] > 0) + & (prefill_df["batch_size"] == 1) + ].copy() if len(standard_prefill_df) == 0: raise ValueError( "No standard prefill rows (prefill_chunk_size > 0) found in eager attention profiling data." diff --git a/frontier/execution_time_predictor/structured_attention_prefill.py b/frontier/execution_time_predictor/structured_attention_prefill.py new file mode 100644 index 0000000..1829047 --- /dev/null +++ b/frontier/execution_time_predictor/structured_attention_prefill.py @@ -0,0 +1,79 @@ +"""Structured latency model for single-request chunked prefill attention.""" + +from typing import Any + +import numpy as np +from sklearn.base import BaseEstimator, RegressorMixin +from sklearn.isotonic import IsotonicRegression +from sklearn.linear_model import LinearRegression + + +class StructuredAttentionPrefillRegressor(RegressorMixin, BaseEstimator): + """Model attention as a monotone base curve plus continuous KV growth. + + Input columns retain the existing Frontier contract: + ``[kv_cache_size, prefill_chunk_size_squared]``. + """ + + def fit(self, X: Any, y: Any) -> "StructuredAttentionPrefillRegressor": + values = self._as_feature_array(X) + target = np.asarray(y, dtype=float) + kv_cache_size = values[:, 0] + prefill_chunk_size = np.sqrt(np.maximum(values[:, 1], 0.0)) + + base_mask = np.isclose(kv_cache_size, 0.0) + growth_mask = kv_cache_size > 0.0 + if not np.any(base_mask) or not np.any(growth_mask): + raise ValueError( + "structured attn_prefill training requires both KV=0 base rows " + "and KV>0 growth rows" + ) + + base_q = prefill_chunk_size[base_mask] + base_y = target[base_mask] + unique_q = np.unique(base_q) + grouped_y = np.asarray( + [np.mean(base_y[np.isclose(base_q, q)]) for q in unique_q], + dtype=float, + ) + self._base_model = IsotonicRegression( + increasing=True, + out_of_bounds="clip", + ).fit(unique_q, grouped_y) + + growth_q = prefill_chunk_size[growth_mask] + growth_kv = kv_cache_size[growth_mask] + growth_base = self._base_model.predict(growth_q) + growth_features = np.column_stack( + (growth_kv, growth_q * growth_kv) + ) + self._growth_model = LinearRegression( + fit_intercept=False, + positive=True, + ).fit(growth_features, target[growth_mask] - growth_base) + + self.n_features_in_ = 2 + self._frontier_base_q_min = float(unique_q.min()) + self._frontier_base_q_max = float(unique_q.max()) + self._frontier_growth_kv_max = float(growth_kv.max()) + return self + + def predict(self, X: Any) -> np.ndarray: + values = self._as_feature_array(X) + kv_cache_size = values[:, 0] + prefill_chunk_size = np.sqrt(np.maximum(values[:, 1], 0.0)) + base = self._base_model.predict(prefill_chunk_size) + growth_features = np.column_stack( + (kv_cache_size, prefill_chunk_size * kv_cache_size) + ) + return np.maximum(base + self._growth_model.predict(growth_features), 0.0) + + @staticmethod + def _as_feature_array(X: Any) -> np.ndarray: + values = np.asarray(X, dtype=float) + if values.ndim != 2 or values.shape[1] != 2: + raise ValueError( + "structured attn_prefill expects exactly two features: " + "kv_cache_size and prefill_chunk_size_squared" + ) + return values diff --git a/tests/unit/test_structured_attention_prefill.py b/tests/unit/test_structured_attention_prefill.py new file mode 100644 index 0000000..12c4247 --- /dev/null +++ b/tests/unit/test_structured_attention_prefill.py @@ -0,0 +1,59 @@ +import pickle +import unittest + +import numpy as np + +from frontier.execution_time_predictor.structured_attention_prefill import ( + StructuredAttentionPrefillRegressor, +) + + +class StructuredAttentionPrefillRegressorTest(unittest.TestCase): + def setUp(self) -> None: + q = np.asarray([64, 128, 256, 512, 1024, 2048, 4096, 8192], dtype=float) + base = 0.05 + 1e-4 * q + 4e-8 * q**2 + context_q = np.asarray([2048, 4096, 8192] * 3, dtype=float) + context_kv = np.repeat([8192, 16384, 24576], 3).astype(float) + context_y = ( + np.interp(context_q, q, base) + + 1.5e-5 * context_kv + + 3e-8 * context_q * context_kv + ) + self.X = np.column_stack( + ( + np.concatenate((np.zeros_like(q), context_kv)), + np.concatenate((q**2, context_q**2)), + ) + ) + self.y = np.concatenate((base, context_y)) + + def test_recovers_structured_curve(self) -> None: + model = StructuredAttentionPrefillRegressor().fit(self.X, self.y) + np.testing.assert_allclose(model.predict(self.X), self.y, rtol=1e-6) + + def test_prediction_is_nonnegative_and_monotone(self) -> None: + model = StructuredAttentionPrefillRegressor().fit(self.X, self.y) + q = np.arange(1, 8193, dtype=float) + for kv in (0, 8192, 32768, 40912): + X = np.column_stack((np.full_like(q, kv), q**2)) + prediction = model.predict(X) + self.assertTrue(np.all(prediction >= 0)) + self.assertTrue(np.all(np.diff(prediction) >= -1e-12)) + + kv = np.arange(0, 40913, 64, dtype=float) + for q_value in (64, 2048, 8192): + X = np.column_stack((kv, np.full_like(kv, q_value**2))) + self.assertTrue(np.all(np.diff(model.predict(X)) >= -1e-12)) + + def test_pickle_round_trip(self) -> None: + model = StructuredAttentionPrefillRegressor().fit(self.X, self.y) + restored = pickle.loads(pickle.dumps(model)) + np.testing.assert_allclose(restored.predict(self.X), self.y, rtol=1e-6) + + def test_requires_base_and_growth_rows(self) -> None: + with self.assertRaisesRegex(ValueError, "KV=0 base rows"): + StructuredAttentionPrefillRegressor().fit(self.X[:8], self.y[:8]) + + +if __name__ == "__main__": + unittest.main() -- 2.43.0