perf(prefill): segment long jobs for autoresearch

autoresearch/prefill/candidate.py

@@ -0,0 +1,7 @@
+"""The only Prefill strategy file the autoresearch agent may edit."""
+
+PREFILL_COMPUTE_CHUNK_TOKENS = 256
+SNAPSHOT_MODE = "final_only"
+MAX_SEGMENT_SECONDS = 300.0
+REQUIRE_FULL_CONTEXT = True
+ALLOW_FALLBACK = False

autoresearch/prefill/prepare.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""Fixed evaluation harness for Karpathy-style Prefill autoresearch."""
+from __future__ import annotations
+
+import argparse
+import csv
+import importlib.util
+import json
+import time
+from pathlib import Path
+
+
+def _load_candidate(path: Path):
+    spec = importlib.util.spec_from_file_location("prefill_candidate", path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError("cannot load candidate")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def evaluate(report: dict, candidate) -> dict:
+    stages = report.get("stages", [])
+    critic = next(
+        (stage for stage in stages if stage.get("name") == "agent_critic"),
+        None,
+    )
+    if critic is None:
+        raise ValueError("report has no Critic stage")
+    prefix_tokens = int(critic.get("prefix_tokens", 0))
+    warmup_s = float(critic.get("warmup_wall_s", 0))
+    measured_tps = prefix_tokens / warmup_s if warmup_s > 0 else 0.0
+    estimated_max_segment_s = (
+        candidate.PREFILL_COMPUTE_CHUNK_TOKENS / measured_tps
+        if measured_tps > 0 else float("inf")
+    )
+    delta = critic.get("delta", {})
+    constraints = {
+        "stage_ok": bool(critic.get("ok")),
+        "complete": bool(critic.get("complete")),
+        "full_context": (
+            critic.get("review_scope") == "full"
+            and int(critic.get("critic_omitted_tokens", -1)) == 0
+            and int(critic.get("critic_context_tokens", -1))
+            == int(critic.get("generator_full_tokens", -2))
+        ),
+        "recursive_protocol": (
+            critic.get("critic_protocol")
+            == "recursive_proof_decomposition_v2"
+        ),
+        "no_fallback": int(delta.get("fallbacks", 0)) == 0,
+        "no_job_failure": int(delta.get("remote_job_failures", 0)) == 0,
+        "segment_under_budget": (
+            estimated_max_segment_s <= candidate.MAX_SEGMENT_SECONDS
+        ),
+        "final_only_snapshot": candidate.SNAPSHOT_MODE == "final_only",
+    }
+    return {
+        "accepted": all(constraints.values()),
+        "metric_cold_critic_prefill_s": warmup_s,
+        "measured_prefill_tps": measured_tps,
+        "estimated_max_segment_s": estimated_max_segment_s,
+        "compute_chunk_tokens": candidate.PREFILL_COMPUTE_CHUNK_TOKENS,
+        "constraints": constraints,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--report", type=Path, required=True)
+    parser.add_argument(
+        "--candidate",
+        type=Path,
+        default=Path(__file__).with_name("candidate.py"),
+    )
+    parser.add_argument(
+        "--results",
+        type=Path,
+        default=Path(__file__).with_name("results.tsv"),
+    )
+    args = parser.parse_args()
+    candidate = _load_candidate(args.candidate)
+    result = evaluate(json.loads(args.report.read_text()), candidate)
+    write_header = not args.results.exists()
+    with args.results.open("a", newline="") as handle:
+        writer = csv.DictWriter(
+            handle,
+            fieldnames=(
+                "timestamp",
+                "accepted",
+                "metric_cold_critic_prefill_s",
+                "measured_prefill_tps",
+                "estimated_max_segment_s",
+                "compute_chunk_tokens",
+            ),
+            delimiter="\t",
+        )
+        if write_header:
+            writer.writeheader()
+        writer.writerow({"timestamp": time.time(), **{
+            key: result[key] for key in writer.fieldnames if key != "timestamp"
+        }})
+    print(json.dumps(result, indent=2, sort_keys=True))
+    return 0 if result["accepted"] else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

autoresearch/prefill/program.md

@@ -0,0 +1,41 @@
+# Prefill AutoResearch Program
+
+You are optimizing the two-Mac full-context Prefill system.
+
+## Ownership
+
+- You may edit only `candidate.py`.
+- Never edit `prepare.py`, benchmark reports, tests, or production metrics.
+- The human owns this file.
+
+## Objective
+
+Minimize `metric_cold_critic_prefill_s`. Lower is better.
+
+## Hard constraints
+
+- Every compute segment must remain at or below 300 seconds.
+- Critic must receive the complete Generator response.
+- `critic_omitted_tokens` must equal zero.
+- Protocol must be `recursive_proof_decomposition_v2`.
+- Snapshot mode must remain `final_only`.
+- No fallback, local Primary Prefill, failed remote job, sampling, summary, or
+  semantic simplification is allowed.
+- Primary remains decode-only and allens remains Prefill-only.
+
+## Experiment loop
+
+1. Read `candidate.py` and `results.tsv`.
+2. State one concrete performance hypothesis.
+3. Modify only `candidate.py`.
+4. Deploy the candidate to allens.
+5. Clear Primary and allens caches.
+6. Run the fixed full-context acceptance workload.
+7. Run `prepare.py` against the resulting report.
+8. Keep the candidate only if every hard constraint passes and cold Critic
+   Prefill time improves. Otherwise restore the previous candidate.
+9. Append the result and repeat.
+
+Do not optimize output wording, scores, prizes, or other proof-irrelevant
+content. Optimize only measured Prefill execution while preserving the complete
+semantic contract.

deploy/install_prefill_worker_launchd.sh

@@ -15,6 +15,7 @@ CACHE_GB="${KAKEYA_WORKER_CACHE_GB:-4}"
 CACHE_MIN_GB="${KAKEYA_WORKER_CACHE_MIN_GB:-1}"
 MEMORY_RESERVE_GB="${KAKEYA_WORKER_MEMORY_RESERVE_GB:-0.5}"
 SNAPSHOT_BYTES_PER_TOKEN="${KAKEYA_SNAPSHOT_BYTES_PER_TOKEN:-400000}"
+COMPUTE_CHUNK_TOKENS="${KAKEYA_PREFILL_COMPUTE_CHUNK_TOKENS:-256}"
 ADAPTIVE_CACHE="${KAKEYA_WORKER_ADAPTIVE_CACHE:-0}"
 PSK_FILE="${KAKEYA_FLEET_PSK_FILE:-}"
 CACHE_MODEL_ID="${KAKEYA_CACHE_MODEL_ID:-$KAKEYA_WORKER_MODEL}"
@@ -78,6 +79,7 @@ cat > "$PLIST" <<EOF
     <string>--cache-min-gb</string><string>$CACHE_MIN_GB</string>
     <string>--memory-reserve-gb</string><string>$MEMORY_RESERVE_GB</string>
     <string>--estimated-snapshot-bytes-per-token</string><string>$SNAPSHOT_BYTES_PER_TOKEN</string>
+    <string>--prefill-compute-chunk-tokens</string><string>$COMPUTE_CHUNK_TOKENS</string>
     $adaptive_xml
     <string>--sink</string><string>$SINK</string>
     <string>--window</string><string>$WINDOW</string>

deploy/launchd/ai.kakeya.prefill-worker-peer.plist

@@ -27,6 +27,7 @@
     <string>--adaptive-cache</string>
     <string>--memory-reserve-gb</string><string>0.5</string>
     <string>--estimated-snapshot-bytes-per-token</string><string>400000</string>
+    <string>--prefill-compute-chunk-tokens</string><string>256</string>
     <string>--prefill-tps</string><string>1</string>
     <string>--max-concurrent-jobs</string><string>1</string>
     <string>--network</string><string>thunderbolt</string>

docs/ops/distributed-prefill-kv-network.md

@@ -365,6 +365,17 @@ The MLX worker exports and compresses only the final chained-prefix snapshot.
 The final hash commits every preceding token block, so exporting a growing full
 snapshot at every 64-token boundary is redundant and creates quadratic
 serialization/compression work for long Critic contexts.
+Model compute is segmented independently from the 64-token content-addressed
+hash boundary. The default segment is 256 tokens, keeping each allens step below
+the five-minute research budget at the measured Prefill rate. Worker job status
+updates `tokens_computed` after every segment; Terminal heartbeats display
+tokens, percentage, and ETA.
+
+Karpathy-style optimization lives in `autoresearch/prefill/`. Humans edit
+`program.md`, the research agent edits only `candidate.py`, and immutable
+`prepare.py` evaluates full-context correctness plus cold Critic Prefill time.
+Candidates are retained only when every semantic/topology constraint passes and
+the metric improves; `results.tsv` records experiments.
 Interactive prompt templates are deterministic and contain no per-run nonce, so
 repeating the same task can reuse allens cold-tier and Primary hot-tier KV.
 

inference_engine/backends/mlx/prefill_worker.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 import threading
-from typing import Sequence
+from typing import Callable, Sequence
 
 from inference_engine.backends.mlx.prefill_snapshot import (
     export_mlx_prefill_snapshot,
@@ -18,11 +18,31 @@
 class MLXPrefillComputeEngine:
     """Serially runs prefill with a loaded MLX verifier and exports one snapshot."""
 
-    def __init__(self, verifier, compatibility: CacheCompatibility) -> None:
+    def __init__(
+        self,
+        verifier,
+        compatibility: CacheCompatibility,
+        *,
+        compute_chunk_tokens: int = 256,
+    ) -> None:
+        if compute_chunk_tokens <= 0:
+            raise ValueError("compute_chunk_tokens must be > 0")
         self.verifier = verifier
         self.compatibility = compatibility
+        self.compute_chunk_tokens = int(compute_chunk_tokens)
+        self._progress_callback: Callable[[int], None] | None = None
         self._lock = threading.Lock()
 
+    def set_progress_callback(
+        self,
+        callback: Callable[[int], None] | None,
+    ) -> None:
+        self._progress_callback = callback
+
+    def _report_progress(self, token_count: int) -> None:
+        if self._progress_callback is not None:
+            self._progress_callback(int(token_count))
+
     def compute_prefill(
         self,
         token_ids: Sequence[int],
@@ -43,18 +63,24 @@ def compute_prefill(
         with self._lock:
             if cancelled.is_set():
                 raise InterruptedError("prefill job cancelled")
-            first_end = min(size, len(tokens))
+            first_end = min(self.compute_chunk_tokens, len(tokens))
             self.verifier.prefill(tokens[:first_end])
-            for start in range(first_end, len(tokens), size):
+            self._report_progress(first_end)
+            for start in range(
+                first_end,
+                len(tokens),
+                self.compute_chunk_tokens,
+            ):
                 if cancelled.is_set():
                     raise InterruptedError("prefill job cancelled")
-                block = tokens[start:start + size]
+                block = tokens[start:start + self.compute_chunk_tokens]
                 logits = self.verifier.forward_block(block)
                 self.verifier.commit_or_truncate(
                     forwarded=len(block),
                     accepted=len(block),
                 )
                 self.verifier.next_token_logits = logits[-1].clone()
+                self._report_progress(min(start + len(block), len(tokens)))
             # Export exactly once. Intermediate full snapshots make encoding
             # and compression quadratic in prompt length and are not required
             # for correctness because the final chained hash commits every

inference_engine/distributed/prefill_cache_runtime.py

@@ -72,6 +72,8 @@ class PrefillReuseStats:
     bytes_received: int = 0
     remote_jobs: int = 0
     remote_job_failures: int = 0
+    remote_job_tokens_total: int = 0
+    remote_job_tokens_computed: int = 0
     fallbacks: int = 0
     last_fallback_reason: str = ""
     publish_attempts: int = 0
@@ -356,6 +358,8 @@ def _compute_remote(
                 auth=self.auth,
             )
             self.stats.remote_jobs += 1
+            self.stats.remote_job_tokens_total = len(tokens)
+            self.stats.remote_job_tokens_computed = 0
             deadline = time.monotonic() + self.worker_timeout_s
             while time.monotonic() < deadline:
                 status_request = distributed_pb2.GetPrefillJobStatusRequest(
@@ -368,6 +372,10 @@ def _compute_remote(
                     timeout_s=self.lookup_timeout_s,
                     auth=self.auth,
                 )
+                self.stats.remote_job_tokens_computed = min(
+                    len(tokens),
+                    int(status.tokens_computed),
+                )
                 if status.status == int(PrefillJobState.COMPLETED):
                     return _Hit(
                         source=status.cache_address or target.address,

inference_engine/distributed/prefill_worker.py

@@ -258,6 +258,7 @@ def _run(self, job_id: str) -> None:
             job.state = PrefillJobState.RUNNING
         started = time.perf_counter()
         timer = None
+        engine = None
         if job.deadline_at:
             remaining = job.deadline_at - time.time()
             if remaining <= 0:
@@ -274,7 +275,13 @@ def _run(self, job_id: str) -> None:
                 job.job_id,
                 len(job.token_ids) * self.estimated_snapshot_bytes_per_token,
             )
-            blocks = tuple(self._engine_for_current_thread().compute_prefill(
+            engine = self._engine_for_current_thread()
+            set_progress = getattr(engine, "set_progress_callback", None)
+            if callable(set_progress):
+                set_progress(
+                    lambda count: self._update_job_progress(job.job_id, count),
+                )
+            blocks = tuple(engine.compute_prefill(
                 job.token_ids,
                 job.block_hashes,
                 compression=job.compression,
@@ -313,13 +320,25 @@ def _run(self, job_id: str) -> None:
                 job.state = PrefillJobState.FAILED
                 job.failure_reason = f"{type(exc).__name__}: {exc}"
         finally:
+            if engine is not None:
+                set_progress = getattr(engine, "set_progress_callback", None)
+                if callable(set_progress):
+                    set_progress(None)
             self.cache_store.release_reservation(job.job_id)
             if timer is not None:
                 timer.cancel()
             with self._lock:
                 job.compute_ms = (time.perf_counter() - started) * 1000.0
                 job.finished_at = time.time()
 
+    def _update_job_progress(self, job_id: str, token_count: int) -> None:
+        with self._lock:
+            job = self._jobs[job_id]
+            job.tokens_computed = min(
+                len(job.token_ids),
+                max(job.tokens_computed, int(token_count)),
+            )
+
     def _engine_for_current_thread(self) -> PrefillComputeEngine:
         if self.engine is not None:
             return self.engine