#!/usr/bin/env python3
"""
Analyze REAP saliency JSON from llama-imatrix (ReapTraceCollector v2), verify §2.5 invariants,
generate per-layer expert keep-maps for K in {320, 256, 224} across:
- reap_l2_sum: count * REAP (expected gated output magnitude per token)
- reap_shrink: zero-prior damped REAP L2 mean S_e / (c_e + m) with LCB guard (protects high-norm
specialists with real support while scoring 0-hit/noise experts at 0)
- reap_hybrid: rank-percentile average 0.5 * pct(reap_l2_sum) + 0.5 * pct(reap_shrink)
(outlier-proof across skewed layers like L47)
and compute offline top-10 substitution error on a Q4RT trace file (§3.3).
"""
import argparse
import json
import math
from pathlib import Path
def verify_invariants(data: dict) -> dict:
layers = data["layers"]
assert len(layers) == 48, f"Expected 48 layers, got {len(layers)}"
n_picks = sum(layers["0"]["count"])
assert n_picks > 0 and n_picks % 10 == 0, f"Invalid L0 pick count: {n_picks}"
n_tokens = n_picks // 10
for l_idx in range(48):
ld = layers[str(l_idx)]
c_sum = sum(ld["count"])
p_sum = sum(ld["prob_sum"])
r_sum = sum(ld["reap_l2_sum"])
assert c_sum == n_picks, f"Layer {l_idx} count sum {c_sum} != {n_picks}"
assert abs(p_sum - n_tokens) / max(1, n_tokens) < 1e-2, f"Layer {l_idx} prob_sum {p_sum} != {n_tokens}"
assert r_sum > 0.0, f"Layer {l_idx} reap_l2_sum is zero!"
for e in range(len(ld["count"])):
if ld["count"][e] == 0:
assert ld["reap_l2_sum"][e] == 0.0, f"Layer {l_idx} expert {e} count=0 but reap>0"
return {
"total_tokens_evaluated": n_tokens,
"total_picks_per_layer": n_picks,
"nonfinite_drops": data.get("nonfinite_drops", 0),
"trace_tokens_written": data.get("trace_tokens_written", 0),
}
def _rank_percentiles(vals: list[float]) -> list[float]:
n = len(vals)
order = sorted(range(n), key=lambda i: (vals[i], i))
pct = [0.0] * n
for rank, idx in enumerate(order):
pct[idx] = rank / max(1, n - 1)
return pct
def compute_layer_scores(ld: dict, metric: str, shrink_m: float = 64.0) -> list[float]:
counts = ld["count"]
reap_sum = ld["reap_l2_sum"]
reap_sq = ld.get("reap_l2_sq_sum", [0.0] * len(counts))
n_exp = len(counts)
if metric == "reap_l2_sum":
return list(reap_sum)
if metric == "prob_sum":
return list(ld["prob_sum"])
# Zero-prior damped mean with LCB guard:
# Dead experts (c == 0) score 0.0. Rare experts require genuine sample support (c >> m)
# and low variance (mean - 1.0 * SE) so 1-hit/5-hit noise never beats active experts.
shrink_scores = []
for c, r, sq in zip(counts, reap_sum, reap_sq):
if c <= 0:
shrink_scores.append(0.0)
else:
mean = r / c
var = max(0.0, (sq / c) - (mean * mean))
se = math.sqrt(var / c)
lcb = max(0.0, mean - 1.0 * se)
damped = (c / (c + shrink_m)) * lcb
shrink_scores.append(damped)
if metric == "reap_shrink":
return shrink_scores
if metric == "reap_hybrid":
pct_sum = _rank_percentiles(reap_sum)
pct_shr = _rank_percentiles(shrink_scores)
return [0.5 * ps + 0.5 * pr for ps, pr in zip(pct_sum, pct_shr)]
raise ValueError(f"Unknown metric: {metric}")
def score_trace_substitution(trace_path: str, keep_sets: dict[int, set[int]], k_used: int = 10) -> dict:
import struct
raw = Path(trace_path).read_bytes()
magic, ver, n_layers, n_experts, top_m, n_tokens = struct.unpack_from("<4sHHHH I", raw, 0)
rec_bytes = top_m * 4
actual_tokens = (len(raw) - 16) // (n_layers * rec_bytes)
if n_tokens == 0 or actual_tokens < n_tokens:
n_tokens = actual_tokens
sample_tokens = min(n_tokens, 2048)
ids_fmt = f"<{top_m}H"
logits_fmt = f"<{top_m}e"
id_bytes = top_m * 2
total_slots = 0
replaced_slots = 0
prob_displaced = 0.0
for t in range(sample_tokens):
base_t = 16 + t * n_layers * rec_bytes
for l in range(n_layers):
off = base_t + l * rec_bytes
ids = struct.unpack_from(ids_fmt, raw, off)
logits = struct.unpack_from(logits_fmt, raw, off + id_bytes)
orig_top = ids[:k_used]
orig_logits = logits[:k_used]
max_z = orig_logits[0]
exps = [math.exp(z - max_z) for z in orig_logits]
z_sum = sum(exps)
probs = [x / z_sum for x in exps]
kset = keep_sets[l]
for eid, p in zip(orig_top, probs):
total_slots += 1
if eid not in kset:
replaced_slots += 1
prob_displaced += p
return {
"slot_replacement_rate": round(replaced_slots / max(1, total_slots), 4),
"mean_top10_weight_displaced": round(prob_displaced / max(1, sample_tokens * n_layers), 4),
}
def generate_keep_maps(
saliency_path: str,
out_dir: str,
ks: list[int] = [320, 256, 224],
metric: str = "reap_l2_sum",
trace_path: str | None = None,
) -> dict:
data = json.loads(Path(saliency_path).read_text())
inv = verify_invariants(data)
layers = data["layers"]
n_layers = len(layers)
Path(out_dir).mkdir(parents=True, exist_ok=True)
report = {
"invariants": inv,
"metric": metric,
"variants": {},
}
for k in ks:
keep_map = {}
keep_sets = {}
reap_retained_per_layer = []
prob_retained_per_layer = []
count_retained_per_layer = []
min_kept_count_per_layer = []
zero_hit_kept = 0
for l_idx in range(n_layers):
l_data = layers[str(l_idx)]
scores = compute_layer_scores(l_data, metric=metric)
reap_l2 = l_data["reap_l2_sum"]
probs = l_data["prob_sum"]
counts = l_data["count"]
n_exp = len(scores)
ranked = sorted(range(n_exp), key=lambda e: (scores[e], reap_l2[e], probs[e]), reverse=True)
kept = sorted(ranked[:k])
keep_map[str(l_idx)] = kept
keep_sets[l_idx] = set(kept)
tot_reap = sum(reap_l2)
kept_reap = sum(reap_l2[e] for e in kept)
reap_retained_per_layer.append(kept_reap / tot_reap if tot_reap > 0 else 1.0)
tot_prob = sum(probs)
kept_prob = sum(probs[e] for e in kept)
prob_retained_per_layer.append(kept_prob / tot_prob if tot_prob > 0 else 1.0)
tot_cnt = sum(counts)
kept_cnt = sum(counts[e] for e in kept)
count_retained_per_layer.append(kept_cnt / tot_cnt if tot_cnt > 0 else 1.0)
min_c = min(counts[e] for e in kept)
min_kept_count_per_layer.append(min_c)
zero_hit_kept += sum(1 for e in kept if counts[e] == 0)
assert zero_hit_kept == 0, f"Metric {metric} at K={k} kept {zero_hit_kept} dead (0-hit) experts!"
out_file = Path(out_dir) / f"keep_{metric}_k{k}.json"
out_file.write_text(json.dumps(keep_map, indent=2))
v_info = {
"keep_json": str(out_file),
"mean_reap_l2_retained": round(sum(reap_retained_per_layer) / n_layers, 4),
"min_reap_l2_retained": round(min(reap_retained_per_layer), 4),
"min_reap_layer": int(min(range(n_layers), key=lambda i: reap_retained_per_layer[i])),
"mean_prob_retained": round(sum(prob_retained_per_layer) / n_layers, 4),
"WholeModel_slot_replacement": round(1.0 - sum(count_retained_per_layer) / n_layers, 4),
"min_kept_expert_hits": int(min(min_kept_count_per_layer)),
"zero_hit_experts_kept": zero_hit_kept,
}
if trace_path and Path(trace_path).exists():
v_info["trace_substitution"] = score_trace_substitution(trace_path, keep_sets)
report["variants"][str(k)] = v_info
return report
def main():
p = argparse.ArgumentParser()
p.add_argument("--saliency", required=True)
p.add_argument("--out-dir", default="/content/keep_maps")
p.add_argument("--metric", default="reap_l2_sum", choices=["reap_l2_sum", "reap_shrink", "reap_hybrid", "prob_sum"])
p.add_argument("--trace", default="/content/routing_trace_v2.q4rt")
args = p.parse_args()
rep = generate_keep_maps(args.saliency, args.out_dir, metric=args.metric, trace_path=args.trace)
print(json.dumps(rep, indent=2))
if __name__ == "__main__":
main()