← back to Exo
Capture energy in prefill and ageneration separately (#2124)
051a64e3b4956da788cbedf4caf21319df06fc2d · 2026-05-28 14:42:36 -0700 · ciaranbor
## Motivation
Energy was reported as a single aggregate. Split into prefill vs.
generation so each phase can be analysed independently.
## Changes
- `PowerSampler`: `mark_prefill_done()` + `trapezoidal_energy_range()`
helper; `result()` now emits per-phase splits.
- `PowerUsage` / `NodePowerStats`: optional `prefill_*` / `generation_*`
fields (back-compat: `None` if unmarked).
- API marks the boundary on the first non-`PrefillProgressChunk`.
- `bench/exo_bench.py` surfaces the split in the log line and persists
`power_usage` to JSON.
- METHODOLOGY: one sentence + one bullet.
## Why It Works
First non-prefill chunk *is* the boundary. Anchoring a sample there and
interpolating power at the boundary makes phase energies sum exactly to
the unsplit total.
## Test Plan
### Manual Testing
`eco`-reserved nodes:
- M3 Ultra, Qwen3-VL-4B, pp=8192/tg=1024: server 1940 J vs client 1931 J
(+0.5 %)
- M4 Pro, Qwen3.6-27B, pp=16384/tg=2048: server 20,292 J vs client
20,221 J (+0.35 %)
### Automated Testing
5 new tests in `test_power_sampler.py` (range integrator,
splits-sum-to-total, `None`-when-unmarked, idempotency). 14/14 pass.
Files touched
M bench/METHODOLOGY.mdM bench/exo_bench.pyM src/exo/api/main.pyM src/exo/api/types/api.pyM src/exo/utils/power_sampler.pyM src/exo/utils/tests/test_power_sampler.py
Diff
commit 051a64e3b4956da788cbedf4caf21319df06fc2d
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Thu May 28 14:42:36 2026 -0700
Capture energy in prefill and ageneration separately (#2124)
## Motivation
Energy was reported as a single aggregate. Split into prefill vs.
generation so each phase can be analysed independently.
## Changes
- `PowerSampler`: `mark_prefill_done()` + `trapezoidal_energy_range()`
helper; `result()` now emits per-phase splits.
- `PowerUsage` / `NodePowerStats`: optional `prefill_*` / `generation_*`
fields (back-compat: `None` if unmarked).
- API marks the boundary on the first non-`PrefillProgressChunk`.
- `bench/exo_bench.py` surfaces the split in the log line and persists
`power_usage` to JSON.
- METHODOLOGY: one sentence + one bullet.
## Why It Works
First non-prefill chunk *is* the boundary. Anchoring a sample there and
interpolating power at the boundary makes phase energies sum exactly to
the unsplit total.
## Test Plan
### Manual Testing
`eco`-reserved nodes:
- M3 Ultra, Qwen3-VL-4B, pp=8192/tg=1024: server 1940 J vs client 1931 J
(+0.5 %)
- M4 Pro, Qwen3.6-27B, pp=16384/tg=2048: server 20,292 J vs client
20,221 J (+0.35 %)
### Automated Testing
5 new tests in `test_power_sampler.py` (range integrator,
splits-sum-to-total, `None`-when-unmarked, idempotency). 14/14 pass.
---
bench/METHODOLOGY.md | 3 +-
bench/exo_bench.py | 26 ++++++
src/exo/api/main.py | 2 +
src/exo/api/types/api.py | 16 ++++
src/exo/utils/power_sampler.py | 122 +++++++++++++++++++++++++++
src/exo/utils/tests/test_power_sampler.py | 136 ++++++++++++++++++++++++++++++
6 files changed, 304 insertions(+), 1 deletion(-)
diff --git a/bench/METHODOLOGY.md b/bench/METHODOLOGY.md
index d42624c0..be589110 100644
--- a/bench/METHODOLOGY.md
+++ b/bench/METHODOLOGY.md
@@ -125,7 +125,7 @@ A background thread polls each node at 1 Hz, collecting:
- System power draw (W)
- CPU cluster usage (performance and efficiency cores)
-**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`.
+**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`. The server additionally returns a `power_usage` block in each non-stream `/bench/chat/completions` response that splits energy into prefill and generation phases, with the boundary anchored to the first non-`PrefillProgressChunk` from the runner.
---
@@ -136,6 +136,7 @@ Results are written as JSON with three top-level keys:
- **`runs`**: Array of per-request result objects, each containing:
- `elapsed_s`, `output_text_preview` (first 200 chars)
- `stats`: `{ prompt_tps, generation_tps, prompt_tokens, generation_tokens, peak_memory_usage }`
+ - `power_usage`: server-side total + prefill/generation split, per-node breakdown (non-stream requests only)
- Placement metadata: `model_id`, `placement_sharding`, `placement_instance_meta`, `placement_nodes`
- Run metadata: `pp_tokens`, `tg`, `repeat_index`, `concurrency`, `concurrent_index`
- `download_duration_s` (if model was freshly downloaded)
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 3322402b..68c52fa8 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -295,6 +295,7 @@ def run_one_completion(
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
+ power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
content = message.get("content") or ""
@@ -330,6 +331,7 @@ def run_one_completion(
elapsed = time.perf_counter() - t0
preview = "".join(text_parts)[:200]
+ power_usage = None
if not stats:
ttft = (first_token_time - t0) if first_token_time else elapsed
@@ -348,6 +350,7 @@ def run_one_completion(
"elapsed_s": elapsed,
"output_text_preview": preview,
"stats": stats,
+ "power_usage": power_usage,
}, pp_tokens
@@ -764,6 +767,7 @@ def main() -> int:
out = c.post_bench_chat_completions(_payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
+ power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = (
choices[0].get("message", {}) if choices else {}
@@ -773,6 +777,7 @@ def main() -> int:
"elapsed_s": elapsed,
"output_text_preview": text[:200],
"stats": stats,
+ "power_usage": power_usage,
}, _actual_pp
inf_t0 = time.monotonic()
@@ -868,6 +873,27 @@ def main() -> int:
inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
avg_watts = joules / inf_seconds if inf_seconds > 0 else 0
summary += f" energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
+
+ # mean() not sum() across concurrent runs: each
+ # request's PowerSampler observes the same shared
+ # cluster state, so they all report the same figure.
+ prefill_energies = [
+ (x.get("power_usage") or {}).get("prefill_energy_joules")
+ for x in runs
+ ]
+ gen_energies = [
+ (x.get("power_usage") or {}).get("generation_energy_joules")
+ for x in runs
+ ]
+ prefill_vals = [e for e in prefill_energies if e is not None]
+ gen_vals = [e for e in gen_energies if e is not None]
+ if prefill_vals and gen_vals:
+ avg_pref = mean(prefill_vals)
+ avg_gen = mean(gen_vals)
+ summary += (
+ f" prefill_energy={avg_pref:.1f}J "
+ f"gen_energy={avg_gen:.1f}J"
+ )
logger.info(f"{summary}\n")
time.sleep(2)
finally:
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index d15882b1..fcd54c93 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -806,6 +806,8 @@ class API:
if isinstance(chunk, PrefillProgressChunk):
continue
+ sampler.mark_prefill_done()
+
if chunk.finish_reason == "error":
raise HTTPException(
status_code=500,
diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py
index 45782267..02e397ad 100644
--- a/src/exo/api/types/api.py
+++ b/src/exo/api/types/api.py
@@ -186,6 +186,12 @@ class NodePowerStats(BaseModel, frozen=True):
node_id: NodeId
samples: int
avg_sys_power: float
+ # Per-phase breakdown. Populated only when the caller marks a phase
+ # boundary (e.g. prefill -> generation); None otherwise.
+ prefill_avg_sys_power: float | None = None
+ generation_avg_sys_power: float | None = None
+ prefill_energy_joules: float | None = None
+ generation_energy_joules: float | None = None
class PowerUsage(BaseModel, frozen=True):
@@ -193,6 +199,16 @@ class PowerUsage(BaseModel, frozen=True):
nodes: list[NodePowerStats]
total_avg_sys_power_watts: float
total_energy_joules: float
+ # Split between the prefill (prompt-processing) phase and the
+ # generation/decode phase. Populated only when the caller marks a phase
+ # boundary; None otherwise. The two phase energies should sum to
+ # approximately `total_energy_joules` (modulo interpolation rounding).
+ prefill_seconds: float | None = None
+ generation_seconds: float | None = None
+ prefill_energy_joules: float | None = None
+ generation_energy_joules: float | None = None
+ prefill_avg_sys_power_watts: float | None = None
+ generation_avg_sys_power_watts: float | None = None
class BenchChatCompletionResponse(ChatCompletionResponse):
diff --git a/src/exo/utils/power_sampler.py b/src/exo/utils/power_sampler.py
index c6e61b41..223d3ed3 100644
--- a/src/exo/utils/power_sampler.py
+++ b/src/exo/utils/power_sampler.py
@@ -24,6 +24,7 @@ class PowerSampler:
] = defaultdict(list)
self._start_time: float | None = None
self._stopped = False
+ self._prefill_done_at: float | None = None
def _take_sample(self, t_rel: float | None = None) -> None:
assert self._start_time is not None
@@ -38,14 +39,35 @@ class PowerSampler:
await anyio.sleep(self._interval)
self._take_sample()
+ def mark_prefill_done(self) -> None:
+ """Anchor the prefill→generation boundary on a fresh sample.
+ Idempotent. Safe to call before `run()`; boundary then lands at t=0.
+ """
+ if self._prefill_done_at is not None:
+ return
+ if self._start_time is None:
+ self._prefill_done_at = 0.0
+ return
+ t_rel = time.perf_counter() - self._start_time
+ self._take_sample(t_rel=t_rel)
+ self._prefill_done_at = t_rel
+
def result(self) -> PowerUsage:
self._stopped = True
assert self._start_time is not None, "result() called before run()"
elapsed = time.perf_counter() - self._start_time
self._take_sample(t_rel=elapsed)
+ # Clamp the split point to [0, elapsed] in case timing is weird (e.g.
+ # mark called after result, or sampler ran for < the prefill window).
+ split = self._prefill_done_at
+ if split is not None:
+ split = max(0.0, min(elapsed, split))
+
node_stats: list[NodePowerStats] = []
total_energy_j = 0.0
+ total_prefill_energy_j = 0.0
+ total_generation_energy_j = 0.0
for node_id, ts_profiles in self._samples.items():
n = len(ts_profiles)
if n == 0:
@@ -53,20 +75,68 @@ class PowerSampler:
node_energy_j = trapezoidal_energy(ts_profiles, elapsed)
avg_power_w = node_energy_j / elapsed if elapsed > 0 else 0.0
total_energy_j += node_energy_j
+
+ prefill_e: float | None = None
+ generation_e: float | None = None
+ prefill_avg: float | None = None
+ generation_avg: float | None = None
+ if split is not None:
+ prefill_e = trapezoidal_energy_range(ts_profiles, 0.0, split)
+ generation_e = trapezoidal_energy_range(ts_profiles, split, elapsed)
+ total_prefill_energy_j += prefill_e
+ total_generation_energy_j += generation_e
+ prefill_dt = split
+ generation_dt = elapsed - split
+ prefill_avg = prefill_e / prefill_dt if prefill_dt > 0 else 0.0
+ generation_avg = (
+ generation_e / generation_dt if generation_dt > 0 else 0.0
+ )
+
node_stats.append(
NodePowerStats(
node_id=node_id,
samples=n,
avg_sys_power=avg_power_w,
+ prefill_avg_sys_power=prefill_avg,
+ generation_avg_sys_power=generation_avg,
+ prefill_energy_joules=prefill_e,
+ generation_energy_joules=generation_e,
)
)
total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
+
+ prefill_seconds: float | None = None
+ generation_seconds: float | None = None
+ prefill_energy_joules: float | None = None
+ generation_energy_joules: float | None = None
+ prefill_avg_w: float | None = None
+ generation_avg_w: float | None = None
+ if split is not None:
+ prefill_seconds = split
+ generation_seconds = elapsed - split
+ prefill_energy_joules = total_prefill_energy_j
+ generation_energy_joules = total_generation_energy_j
+ prefill_avg_w = (
+ total_prefill_energy_j / prefill_seconds if prefill_seconds > 0 else 0.0
+ )
+ generation_avg_w = (
+ total_generation_energy_j / generation_seconds
+ if generation_seconds > 0
+ else 0.0
+ )
+
return PowerUsage(
elapsed_seconds=elapsed,
nodes=node_stats,
total_avg_sys_power_watts=total_avg_sys_w,
total_energy_joules=total_energy_j,
+ prefill_seconds=prefill_seconds,
+ generation_seconds=generation_seconds,
+ prefill_energy_joules=prefill_energy_joules,
+ generation_energy_joules=generation_energy_joules,
+ prefill_avg_sys_power_watts=prefill_avg_w,
+ generation_avg_sys_power_watts=generation_avg_w,
)
@@ -89,3 +159,55 @@ def trapezoidal_energy(
continue
energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
return energy_j
+
+
+def trapezoidal_energy_range(
+ ts_profiles: list[tuple[float, SystemPerformanceProfile]],
+ t_start: float,
+ t_end: float,
+) -> float:
+ """Integrate sys_power(t) over [t_start, t_end] using the trapezoidal rule.
+
+ Linearly interpolates power at the endpoints when they fall between
+ existing samples, so callers can integrate over arbitrary sub-windows
+ (e.g. the prefill segment) without losing accuracy. Returns 0 for an
+ empty or zero-length window. Falls back to constant-power assumption
+ when only one sample exists.
+ """
+ if t_end <= t_start:
+ return 0.0
+ if len(ts_profiles) == 0:
+ return 0.0
+ if len(ts_profiles) == 1:
+ return ts_profiles[0][1].sys_power * (t_end - t_start)
+
+ def power_at(t: float) -> float:
+ if t <= ts_profiles[0][0]:
+ return ts_profiles[0][1].sys_power
+ if t >= ts_profiles[-1][0]:
+ return ts_profiles[-1][1].sys_power
+ for i in range(1, len(ts_profiles)):
+ t_cur, p_cur = ts_profiles[i]
+ if t_cur >= t:
+ t_prev, p_prev = ts_profiles[i - 1]
+ span = t_cur - t_prev
+ if span <= 0:
+ return p_cur.sys_power
+ frac = (t - t_prev) / span
+ return p_prev.sys_power + frac * (p_cur.sys_power - p_prev.sys_power)
+ return ts_profiles[-1][1].sys_power
+
+ p_start = power_at(t_start)
+ p_end = power_at(t_end)
+ in_range: list[tuple[float, float]] = [
+ (t, profile.sys_power) for t, profile in ts_profiles if t_start < t < t_end
+ ]
+ seq: list[tuple[float, float]] = [(t_start, p_start)] + in_range + [(t_end, p_end)]
+
+ energy_j = 0.0
+ for i in range(1, len(seq)):
+ dt = seq[i][0] - seq[i - 1][0]
+ if dt <= 0:
+ continue
+ energy_j += (seq[i - 1][1] + seq[i][1]) / 2.0 * dt
+ return energy_j
diff --git a/src/exo/utils/tests/test_power_sampler.py b/src/exo/utils/tests/test_power_sampler.py
index 7880936c..cbd8e702 100644
--- a/src/exo/utils/tests/test_power_sampler.py
+++ b/src/exo/utils/tests/test_power_sampler.py
@@ -141,6 +141,142 @@ def test_trapezoidal_unit_single_sample() -> None:
assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
+def test_trapezoidal_range_interpolation() -> None:
+ """Sub-window integration should linearly interpolate at the boundaries."""
+ from exo.utils.power_sampler import trapezoidal_energy_range
+
+ # Two samples: t=0 W=10, t=10 W=20 -> power(t) = 10 + t
+ samples = [
+ (0.0, _make_profile(10.0)),
+ (10.0, _make_profile(20.0)),
+ ]
+ # Integral from t=4 to t=6: power goes 14 -> 16, mean 15, dt=2 -> 30 J
+ assert abs(trapezoidal_energy_range(samples, 4.0, 6.0) - 30.0) < 1e-9
+ # Integral over the full window matches the full trapezoidal integral.
+ full = trapezoidal_energy_range(samples, 0.0, 10.0)
+ assert abs(full - 150.0) < 1e-9
+
+
+def test_trapezoidal_range_zero_window() -> None:
+ """Zero-length or reversed windows integrate to zero."""
+ from exo.utils.power_sampler import trapezoidal_energy_range
+
+ samples = [(0.0, _make_profile(10.0)), (5.0, _make_profile(20.0))]
+ assert trapezoidal_energy_range(samples, 3.0, 3.0) == 0.0
+ assert trapezoidal_energy_range(samples, 5.0, 3.0) == 0.0
+
+
+def test_trapezoidal_range_splits_sum_to_full() -> None:
+ """Energy split at an arbitrary boundary should sum back to the full integral."""
+ from exo.utils.power_sampler import (
+ trapezoidal_energy,
+ trapezoidal_energy_range,
+ )
+
+ samples = [
+ (0.0, _make_profile(10.0)),
+ (1.0, _make_profile(20.0)),
+ (3.0, _make_profile(15.0)),
+ (5.0, _make_profile(25.0)),
+ ]
+ full = trapezoidal_energy(samples, elapsed=5.0)
+ # Split at t=2.5 (between samples) — interpolation should be exact.
+ left = trapezoidal_energy_range(samples, 0.0, 2.5)
+ right = trapezoidal_energy_range(samples, 2.5, 5.0)
+ assert abs((left + right) - full) < 1e-9
+
+
+async def test_prefill_generation_split() -> None:
+ """When mark_prefill_done() is called, the result should split energy."""
+ state: dict[NodeId, SystemPerformanceProfile] = {
+ NODE_A: _make_profile(10.0),
+ }
+ sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
+
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ # "Prefill" phase: power = 10 W
+ await anyio.sleep(0.1)
+ # Mark the boundary BEFORE changing state — this matches what
+ # _collect_text_generation_with_stats does in production: the mark
+ # fires on the first non-prefill chunk, so the boundary sample is
+ # the genuine end-of-prefill reading rather than the new phase's.
+ sampler.mark_prefill_done()
+ state[NODE_A] = _make_profile(30.0)
+ # "Generation" phase: power = 30 W
+ await anyio.sleep(0.1)
+ tg.cancel_scope.cancel()
+
+ result = sampler.result()
+ assert result.prefill_seconds is not None
+ assert result.generation_seconds is not None
+ assert result.prefill_energy_joules is not None
+ assert result.generation_energy_joules is not None
+ assert result.prefill_avg_sys_power_watts is not None
+ assert result.generation_avg_sys_power_watts is not None
+
+ # Phase durations should sum to the elapsed seconds.
+ assert (
+ abs(
+ (result.prefill_seconds + result.generation_seconds)
+ - result.elapsed_seconds
+ )
+ < 1e-6
+ )
+ # Phase energies should sum to (approximately) the total.
+ assert (
+ abs(
+ (result.prefill_energy_joules + result.generation_energy_joules)
+ - result.total_energy_joules
+ )
+ < 1e-6
+ )
+ # With the boundary sample anchored at the genuine end-of-prefill (10 W),
+ # prefill avg should converge tightly on 10 W and generation on 30 W.
+ # 15 W cleanly separates the two and would catch any cross-contamination.
+ assert result.prefill_avg_sys_power_watts < 15.0
+ assert result.generation_avg_sys_power_watts > 15.0
+ assert result.nodes[0].prefill_avg_sys_power is not None
+ assert result.nodes[0].generation_avg_sys_power is not None
+
+
+async def test_no_split_when_unmarked() -> None:
+ """If mark_prefill_done() is never called, phase fields stay None."""
+ state: dict[NodeId, SystemPerformanceProfile] = {
+ NODE_A: _make_profile(10.0),
+ }
+ sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ await anyio.sleep(0.05)
+ tg.cancel_scope.cancel()
+
+ result = sampler.result()
+ assert result.prefill_seconds is None
+ assert result.generation_seconds is None
+ assert result.prefill_energy_joules is None
+ assert result.generation_energy_joules is None
+ assert result.nodes[0].prefill_energy_joules is None
+ assert result.nodes[0].generation_energy_joules is None
+
+
+async def test_mark_prefill_done_is_idempotent() -> None:
+ """Only the first call to mark_prefill_done() should take effect."""
+ state: dict[NodeId, SystemPerformanceProfile] = {
+ NODE_A: _make_profile(10.0),
+ }
+ sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ await anyio.sleep(0.05)
+ sampler.mark_prefill_done()
+ first_prefill_at = sampler._prefill_done_at # pyright: ignore[reportPrivateUsage]
+ await anyio.sleep(0.05)
+ sampler.mark_prefill_done()
+ assert sampler._prefill_done_at == first_prefill_at # pyright: ignore[reportPrivateUsage]
+ tg.cancel_scope.cancel()
+
+
async def test_result_stops_sampling() -> None:
"""Calling result() should stop the sampler's run loop."""
state: dict[NodeId, SystemPerformanceProfile] = {
← a8602ea6 fix(bug): no longer repeated _trigger_notify_user_to_downloa
·
back to Exo
·
fix: make app builds work again (#2127) f9f8cbb3 →