250 lines
11 KiB
Diff
250 lines
11 KiB
Diff
From 1f8900a4ac64e45754b03d0aa7c1dddab65785cf Mon Sep 17 00:00:00 2001
|
|
From: Gahow Wang <gahow.wang@gmail.com>
|
|
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
|
|
|