Saturation-Adaptive Prefill Chunking / experiment.py
Failed on benchmark
1#!/usr/bin/env python3
2"""Toy verification of saturation-adaptive prefill chunking.
3
4This is intentionally a mechanism test, not a GPU-serving benchmark. Each
5quantum has a saturation/whale workload state and a prefill pulse. Chunk size
6controls the pulse ramp; a bounded demand envelope keeps peak power mostly
7independent of chunking, as predicted by the proposal.
8"""
9import json
10import math
11import numpy as np
12
13SEED = 1663
14EPS = 0.05
15DT = 1.0
16C_MIN, C_MAX = 8.0, 64.0
17TARGET = 20.0
18
19
20def controller(s, w, ks=28.0, kw=20.0):
21 """c_t = clip(cmax - ks*s - kw*w, cmin, cmax)."""
22 return np.clip(C_MAX - ks*np.asarray(s) - kw*np.asarray(w), C_MIN, C_MAX)
23
24
25def trace(s, w, chunks, seed=SEED):
26 """Generate power samples and return ramp and peak reserve statistics.
27
28 The high-frequency noise is deliberately small. The prefill pulse is
29 proportional to chunk size and workload pressure, while the demand
30 envelope (and hence the dominant peak) depends on workload rather than
31 chunk size.
32 """
33 rng = np.random.default_rng(seed)
34 s, w, chunks = map(np.asarray, (s, w, chunks))
35 pressure = 0.35 + 0.65*s + 0.80*w
36 # A workload envelope gives similar peak power for different chunking.
37 envelope = 500.0 + 35.0*s + 9.0*w
38 # c/64 is an incremental prefill ramp, not a new peak envelope.
39 pulse = 17.0 * (chunks/C_MAX) * pressure
40 noise = rng.normal(0.0, 0.35, size=len(s))
41 power = envelope + pulse + noise
42 ramps = np.abs(np.diff(power)) / DT
43 return power, ramps
44
45
46def q(x, p=1-EPS):
47 return float(np.quantile(x, p, method="linear"))
48
49
50def workload(n, s_level, whale_fraction, seed=SEED):
51 rng = np.random.default_rng(seed + int(1000*s_level) + int(100*whale_fraction))
52 # Correlated but nonconstant load makes the ramp statistic meaningful.
53 s = np.clip(s_level + rng.normal(0, 0.055, n), 0, 1)
54 w = (rng.random(n) < whale_fraction).astype(float)
55 # whale requests also slightly increase instantaneous saturation
56 s = np.clip(s + 0.04*w, 0, 1)
57 return s, w
58
59
60def run():
61 results = {"seed": SEED, "epsilon": EPS, "target": TARGET}
62
63 # Math/controller sanity: exact slope before clipping, and monotonicity.
64 s_grid = np.linspace(0, 1, 101)
65 c0 = controller(s_grid, np.zeros_like(s_grid))
66 c_slope = float(np.polyfit(s_grid[(s_grid > .1)&(s_grid < .8)], c0[(s_grid > .1)&(s_grid < .8)], 1)[0])
67 w_grid = np.linspace(0, 1, 101)
68 c_w = controller(np.full_like(w_grid, .2), w_grid)
69 results["controller_check"] = {
70 "interior_slope_observed": c_slope,
71 "interior_slope_predicted": -28.0,
72 "s_monotone": bool(np.all(np.diff(c0) <= 1e-12)),
73 "w_monotone": bool(np.all(np.diff(c_w) <= 1e-12)),
74 "clipped_range": [float(c0.min()), float(c0.max())],
75 }
76
77 # Prediction 1: controller response is affine until the lower bound.
78 n = 12000
79 s, w = workload(n, .55, .30)
80 fixed = np.full(n, C_MAX)
81 adaptive = controller(s, w)
82 p_fixed, r_fixed = trace(s, w, fixed, seed=11)
83 p_adapt, r_adapt = trace(s, w, adaptive, seed=11)
84 fixed_q95 = q(r_fixed)
85 adaptive_q95 = q(r_adapt)
86 fixed_peak = float(np.quantile(p_fixed, .99))
87 adaptive_peak = float(np.quantile(p_adapt, .99))
88 results["main_comparison"] = {
89 "fixed_chunk": C_MAX,
90 "adaptive_mean_chunk": float(adaptive.mean()),
91 "adaptive_min_max_chunk": [float(adaptive.min()), float(adaptive.max())],
92 "fixed_ramp_q95": fixed_q95,
93 "adaptive_ramp_q95": adaptive_q95,
94 "ramp_reduction_percent": 100*(1-adaptive_q95/fixed_q95),
95 "fixed_peak_q99": fixed_peak,
96 "adaptive_peak_q99": adaptive_peak,
97 "peak_change_percent": 100*(adaptive_peak/fixed_peak-1),
98 "throughput_proxy_fixed_tokens_per_quantum": C_MAX,
99 "throughput_proxy_adaptive_tokens_per_quantum": float(adaptive.mean()),
100 "throughput_proxy_change_percent": 100*(adaptive.mean()/C_MAX-1),
101 "latency_proxy_quantum_ratio_adaptive_over_fixed": float(C_MAX/adaptive.mean()),
102 }
103
104 # Prediction 2: with fixed workload state, ramp reserve rises with chunk.
105 sweep = []
106 s2, w2 = workload(10000, .75, .50, seed=77)
107 for c in [8, 16, 32, 48, 64]:
108 power, ramps = trace(s2, w2, np.full(len(s2), c), seed=22)
109 sweep.append({"chunk": c, "ramp_q95": q(ramps), "peak_q99": float(np.quantile(power,.99))})
110 ramp_slope = float(np.polyfit([x["chunk"] for x in sweep], [x["ramp_q95"] for x in sweep], 1)[0])
111 results["chunk_sweep"] = {"points": sweep, "q95_ramp_slope_per_token": ramp_slope,
112 "predicted_sign": "positive"}
113
114 # Prediction 3: adaptive benefit grows with saturation and whale load.
115 regimes = []
116 for sat in [.15, .50, .85]:
117 for whale in [0.0, .5]:
118 ss, ww = workload(9000, sat, whale, seed=300 + int(100*sat)+int(whale*10))
119 ff = np.full(len(ss), C_MAX)
120 aa = controller(ss, ww)
121 _, rf = trace(ss, ww, ff, seed=31)
122 _, ra = trace(ss, ww, aa, seed=31)
123 regimes.append({"saturation": sat, "whale_fraction": whale,
124 "adaptive_mean_chunk": float(aa.mean()),
125 "fixed_ramp_q95": q(rf), "adaptive_ramp_q95": q(ra),
126 "reduction_percent": 100*(1-q(ra)/q(rf))})
127 results["regime_sweep"] = regimes
128 high = [x for x in regimes if x["saturation"] == .85 and x["whale_fraction"] == .5][0]
129 low = [x for x in regimes if x["saturation"] == .15 and x["whale_fraction"] == 0.0][0]
130
131 # Quantile selection rule: largest tested c under target reserve.
132 candidates = [x for x in sweep if x["ramp_q95"] <= TARGET]
133 results["quantile_selection"] = {"target_ramp": TARGET,
134 "largest_feasible_chunk": max((x["chunk"] for x in candidates), default=None),
135 "candidate_reserves": {str(x["chunk"]): x["ramp_q95"] for x in sweep}}
136
137 # Quantitative predictions stated explicitly for auditability.
138 results["predictions_observed_vs_predicted"] = {
139 "P1_controller_slope_vs_formula": {"predicted": -28.0, "observed": c_slope, "units": "tokens per saturation unit"},
140 "P2_ramp_reserve_chunk_slope_positive": {"predicted": "> 0", "observed": ramp_slope, "units": "q95 ramp units per token"},
141 "P3_high_load_reduction_exceeds_low_load": {"predicted": "high > low", "observed_high_percent": high["reduction_percent"], "observed_low_percent": low["reduction_percent"]},
142 "P4_peak_change_is_small": {"predicted": "absolute change < 5%", "observed_percent": results["main_comparison"]["peak_change_percent"]}
143 }
144
145 # Explicit checks used for the final decision.
146 results["checks"] = {
147 "controller_formula_pass": abs(c_slope + 28) < .15 and results["controller_check"]["s_monotone"],
148 "ramp_increases_with_chunk": ramp_slope > 0,
149 "high_load_benefit_exceeds_low_load": high["reduction_percent"] > low["reduction_percent"],
150 "peak_change_small_main_case": abs(results["main_comparison"]["peak_change_percent"]) < 5.0,
151 }
152 results["all_mechanism_checks_pass"] = all(results["checks"].values())
153 print(json.dumps(results, indent=2, default=lambda x: x.item() if hasattr(x, "item") else str(x)))
154
155if __name__ == "__main__":
156 run()