257 lines
8.0 KiB
Python
257 lines
8.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Offline predictor ablation for EXP-ATTN-STRUCTURED.
|
|
|
|
This is deliberately profile-only: it decides whether the structured model is
|
|
good enough to justify the expensive 7-cell trace replay.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.ensemble import RandomForestRegressor
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
REPO = ROOT.parents[1]
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--profile",
|
|
type=Path,
|
|
default=REPO
|
|
/ "runs/frontier-prefill-kvgrowth-fix-v0/profiles/"
|
|
"profile-v5-kvgrowth/attention.csv",
|
|
)
|
|
parser.add_argument(
|
|
"--frontier-checkout",
|
|
type=Path,
|
|
default=Path("/tmp/frontier-attn-structured-v0"),
|
|
)
|
|
parser.add_argument("--output-root", type=Path, default=ROOT / "results")
|
|
return parser.parse_args()
|
|
|
|
|
|
def normalize_bool(series: pd.Series) -> pd.Series:
|
|
return series.astype(str).str.strip().str.lower().isin(
|
|
{"1", "true", "t", "yes", "y"}
|
|
)
|
|
|
|
|
|
def load_profile(path: Path) -> pd.DataFrame:
|
|
df = pd.read_csv(path).drop_duplicates()
|
|
for column in ("is_prefill", "is_true_mixed_batch"):
|
|
df[column] = normalize_bool(df[column])
|
|
df = df[
|
|
(df["n_embd"] == 2048)
|
|
& (df["n_q_head"] == 32)
|
|
& (df["n_kv_head"] == 4)
|
|
& (df["block_size"] == 16)
|
|
& df["is_prefill"]
|
|
& ~df["is_true_mixed_batch"]
|
|
& (df["prefill_chunk_size"] > 0)
|
|
].copy()
|
|
df["prefill_chunk_size_squared"] = df["prefill_chunk_size"] ** 2
|
|
return df
|
|
|
|
|
|
def mape(actual: np.ndarray, predicted: np.ndarray) -> float:
|
|
return float(np.mean(np.abs((predicted - actual) / actual)) * 100)
|
|
|
|
|
|
def make_rf() -> RandomForestRegressor:
|
|
# Exact best parameters selected by the current profile-v5 GridSearchCV.
|
|
return RandomForestRegressor(
|
|
random_state=0,
|
|
n_estimators=250,
|
|
max_depth=8,
|
|
min_samples_split=2,
|
|
)
|
|
|
|
|
|
def features(df: pd.DataFrame) -> pd.DataFrame:
|
|
return df[["kv_cache_size", "prefill_chunk_size_squared"]]
|
|
|
|
|
|
def score_model(
|
|
name: str,
|
|
estimator: Any,
|
|
train: pd.DataFrame,
|
|
single: pd.DataFrame,
|
|
grid: pd.DataFrame,
|
|
) -> dict[str, Any]:
|
|
target = "time_stats.attn_prefill.median"
|
|
estimator.fit(features(train), train[target])
|
|
grid_prediction = estimator.predict(features(grid))
|
|
single_prediction = estimator.predict(features(single))
|
|
|
|
heldout_actual: list[float] = []
|
|
heldout_prediction: list[float] = []
|
|
for context in sorted(grid["kv_cache_size"].unique()):
|
|
test = grid[grid["kv_cache_size"] == context]
|
|
fold_train = train.drop(index=test.index, errors="ignore")
|
|
fold_model = (
|
|
make_rf()
|
|
if name.startswith("rf")
|
|
else estimator.__class__()
|
|
)
|
|
fold_model.fit(features(fold_train), fold_train[target])
|
|
heldout_actual.extend(test[target].astype(float))
|
|
heldout_prediction.extend(fold_model.predict(features(test)))
|
|
|
|
q = np.arange(1, 8193, dtype=float)
|
|
q_deltas: list[float] = []
|
|
prediction_min: list[float] = []
|
|
for context in (0, 8192, 16384, 24576, 32768, 40912):
|
|
X = pd.DataFrame(
|
|
{
|
|
"kv_cache_size": np.full_like(q, context),
|
|
"prefill_chunk_size_squared": q**2,
|
|
}
|
|
)
|
|
prediction = estimator.predict(X)
|
|
prediction_min.append(float(prediction.min()))
|
|
q_deltas.append(float(np.diff(prediction).min()))
|
|
|
|
kv = np.arange(0, 40913, 64, dtype=float)
|
|
kv_deltas: list[float] = []
|
|
for query in (64, 512, 2048, 4096, 8192):
|
|
X = pd.DataFrame(
|
|
{
|
|
"kv_cache_size": kv,
|
|
"prefill_chunk_size_squared": np.full_like(kv, query**2),
|
|
}
|
|
)
|
|
kv_deltas.append(float(np.diff(estimator.predict(X)).min()))
|
|
|
|
heldout_actual_array = np.asarray(heldout_actual)
|
|
heldout_prediction_array = np.asarray(heldout_prediction)
|
|
return {
|
|
"candidate": name,
|
|
"training_rows": len(train),
|
|
"grid_fit_mape_pct": mape(
|
|
grid[target].to_numpy(), np.asarray(grid_prediction)
|
|
),
|
|
"single_fit_mape_pct": mape(
|
|
single[target].to_numpy(), np.asarray(single_prediction)
|
|
),
|
|
"heldout_context_mape_pct": mape(
|
|
heldout_actual_array, heldout_prediction_array
|
|
),
|
|
"heldout_context_max_abs_error_pct": float(
|
|
np.max(
|
|
np.abs(
|
|
(heldout_prediction_array - heldout_actual_array)
|
|
/ heldout_actual_array
|
|
)
|
|
)
|
|
* 100
|
|
),
|
|
"prediction_min_ms": min(prediction_min),
|
|
"q_min_delta_ms": min(q_deltas),
|
|
"kv_min_delta_ms": min(kv_deltas),
|
|
"monotone_and_nonnegative": (
|
|
min(prediction_min) >= 0
|
|
and min(q_deltas) >= -1e-12
|
|
and min(kv_deltas) >= -1e-12
|
|
),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
sys.path.insert(0, str(args.frontier_checkout))
|
|
from frontier.execution_time_predictor.structured_attention_prefill import (
|
|
StructuredAttentionPrefillRegressor,
|
|
)
|
|
|
|
df = load_profile(args.profile)
|
|
records: list[dict[str, Any]] = []
|
|
data_audit: dict[str, Any] = {}
|
|
for tp in (1, 2, 4):
|
|
tp_df = df[df["num_tensor_parallel_workers"] == tp].copy()
|
|
single = tp_df[tp_df["batch_size"] == 1].copy()
|
|
grid = single[
|
|
single["prefill_chunk_size"].isin((2048, 4096, 8192))
|
|
& (single["kv_cache_size"] > 0)
|
|
].copy()
|
|
duplicate_groups = (
|
|
tp_df.groupby(
|
|
["kv_cache_size", "prefill_chunk_size_squared"]
|
|
)
|
|
.size()
|
|
.gt(1)
|
|
.sum()
|
|
)
|
|
data_audit[f"tp{tp}"] = {
|
|
"standard_rows": len(tp_df),
|
|
"single_request_rows": len(single),
|
|
"target_grid_rows": len(grid),
|
|
"duplicate_feature_groups": int(duplicate_groups),
|
|
}
|
|
|
|
candidates = (
|
|
("rf_all", make_rf(), tp_df),
|
|
("rf_single", make_rf(), single),
|
|
(
|
|
"structured_single",
|
|
StructuredAttentionPrefillRegressor(),
|
|
single,
|
|
),
|
|
)
|
|
for name, model, train in candidates:
|
|
result = score_model(name, model, train, single, grid)
|
|
result["tp"] = tp
|
|
records.append(result)
|
|
|
|
structured = [r for r in records if r["candidate"] == "structured_single"]
|
|
checks = {
|
|
"heldout_context_mape_le_5pct": all(
|
|
r["heldout_context_mape_pct"] <= 5 for r in structured
|
|
),
|
|
"monotone_and_nonnegative": all(
|
|
r["monotone_and_nonnegative"] for r in structured
|
|
),
|
|
}
|
|
checks["profile_gate"] = all(checks.values())
|
|
payload = {
|
|
"schema": "frontier-attn-structured-ablation-v1",
|
|
"profile": str(args.profile.resolve()),
|
|
"frontier_checkout": str(args.frontier_checkout.resolve()),
|
|
"data_audit": data_audit,
|
|
"results": records,
|
|
"checks": checks,
|
|
}
|
|
|
|
args.output_root.mkdir(parents=True, exist_ok=True)
|
|
(args.output_root / "predictor-ablation.json").write_text(
|
|
json.dumps(payload, indent=2)
|
|
)
|
|
with (args.output_root / "predictor-ablation.csv").open(
|
|
"w", newline=""
|
|
) as stream:
|
|
writer = csv.DictWriter(stream, fieldnames=list(records[0]))
|
|
writer.writeheader()
|
|
writer.writerows(records)
|
|
print(json.dumps(checks, indent=2))
|
|
for row in records:
|
|
print(
|
|
f"TP{row['tp']} {row['candidate']:18s} "
|
|
f"grid={row['grid_fit_mape_pct']:.2f}% "
|
|
f"heldout={row['heldout_context_mape_pct']:.2f}% "
|
|
f"max={row['heldout_context_max_abs_error_pct']:.2f}% "
|
|
f"monotone={row['monotone_and_nonnegative']}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|