153 lines
5.4 KiB
Python
153 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare captured vLLM expert loads with Frontier's fixed routing prior."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import statistics
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--routing", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
return parser.parse_args()
|
|
|
|
|
|
def gini(values: np.ndarray) -> float:
|
|
values = np.asarray(values, dtype=np.float64)
|
|
if values.sum() == 0:
|
|
return 0.0
|
|
ordered = np.sort(values)
|
|
n = len(ordered)
|
|
indices = np.arange(1, n + 1, dtype=np.float64)
|
|
return float((2 * np.sum(indices * ordered) / np.sum(ordered) - (n + 1)) / n)
|
|
|
|
|
|
def stats(values: np.ndarray) -> dict[str, float]:
|
|
values = np.asarray(values, dtype=np.float64)
|
|
mean = float(np.mean(values))
|
|
return {
|
|
"load_cv": float(np.std(values) / mean),
|
|
"load_gini": gini(values),
|
|
"max_load_ratio": float(np.max(values) / mean),
|
|
"expert_utilization": float(np.count_nonzero(values) / len(values)),
|
|
}
|
|
|
|
|
|
def proportional_counts(total: int, ratios: np.ndarray) -> np.ndarray:
|
|
exact = total * ratios / ratios.sum()
|
|
counts = np.floor(exact).astype(np.int64)
|
|
remainder = total - int(counts.sum())
|
|
order = sorted(
|
|
range(len(ratios)),
|
|
key=lambda i: (-(exact[i] - counts[i]), -(ratios[i] / ratios.sum()), i),
|
|
)
|
|
for index in range(remainder):
|
|
counts[order[index % len(order)]] += 1
|
|
assert int(counts.sum()) == total
|
|
return counts
|
|
|
|
|
|
def correlation(left: np.ndarray, right: np.ndarray) -> float:
|
|
value = float(np.corrcoef(left, right)[0, 1])
|
|
return value if math.isfinite(value) else 0.0
|
|
|
|
|
|
def distribution(values: list[float]) -> dict[str, float]:
|
|
return {
|
|
"min": min(values),
|
|
"median": statistics.median(values),
|
|
"max": max(values),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
routing = json.loads(args.routing.read_text())
|
|
phases = routing["phases"]
|
|
layer_count = len(phases["prefill"]["per_layer"])
|
|
expert_count = len(phases["prefill"]["per_layer"][0]["counts"])
|
|
|
|
rows: list[dict[str, Any]] = []
|
|
for phase in ("prefill", "decode"):
|
|
for layer, actual in enumerate(phases[phase]["per_layer"]):
|
|
rng = np.random.RandomState(args.seed + layer)
|
|
ratios = rng.uniform(0.1, 1.0, expert_count)
|
|
synthetic = proportional_counts(int(actual["total_routed_tokens"]), ratios)
|
|
synthetic_stats = stats(synthetic)
|
|
for name in synthetic_stats:
|
|
if not math.isclose(
|
|
stats(np.asarray(actual["counts"]))[name],
|
|
float(actual[name]),
|
|
rel_tol=0,
|
|
abs_tol=1e-12,
|
|
):
|
|
raise ValueError(f"captured metric mismatch: {phase} layer {layer} {name}")
|
|
rows.append(
|
|
{
|
|
"phase": phase,
|
|
"layer": layer,
|
|
"total_routed_tokens": int(actual["total_routed_tokens"]),
|
|
"actual": {name: float(actual[name]) for name in synthetic_stats},
|
|
"frontier_simulation": synthetic_stats,
|
|
"actual_vs_frontier_pearson": correlation(
|
|
np.asarray(actual["counts"], dtype=np.float64), synthetic
|
|
),
|
|
}
|
|
)
|
|
|
|
phase_summary: dict[str, Any] = {}
|
|
for phase in ("prefill", "decode"):
|
|
selected = [row for row in rows if row["phase"] == phase]
|
|
phase_summary[phase] = {
|
|
"token_count": int(phases[phase]["token_count"]),
|
|
"actual": {
|
|
name: distribution([row["actual"][name] for row in selected])
|
|
for name in selected[0]["actual"]
|
|
},
|
|
"frontier_simulation": {
|
|
name: distribution([row["frontier_simulation"][name] for row in selected])
|
|
for name in selected[0]["frontier_simulation"]
|
|
},
|
|
"actual_vs_frontier_pearson": distribution(
|
|
[row["actual_vs_frontier_pearson"] for row in selected]
|
|
),
|
|
}
|
|
|
|
phase_correlations = []
|
|
for layer in range(layer_count):
|
|
prefill = np.asarray(phases["prefill"]["per_layer"][layer]["counts"])
|
|
decode = np.asarray(phases["decode"]["per_layer"][layer]["counts"])
|
|
phase_correlations.append(correlation(prefill, decode))
|
|
|
|
output = {
|
|
"schema": "frontier-routing-mismatch.v1",
|
|
"source": str(args.routing.resolve()),
|
|
"frontier_contract": {
|
|
"mode": "simulation",
|
|
"seed": args.seed,
|
|
"allocation": "per-layer fixed Uniform(0.1, 1.0), normalized once and reused for every batch and phase",
|
|
"layer_count": layer_count,
|
|
"expert_count": expert_count,
|
|
},
|
|
"phase_summary": phase_summary,
|
|
"actual_prefill_vs_decode_pearson": distribution(phase_correlations),
|
|
"frontier_prefill_vs_decode_pearson": 1.0,
|
|
"rows": rows,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
|
|
print(args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|