Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
<string>--cache-link-mbps</string><string>10000</string>
<string>--cache-default-rtt-ms</string><string>0.55</string>
<string>--remote-prefill-min-tokens</string><string>0</string>
<string>--prefill-worker-timeout-s</string><string>900</string>
<string>--prefill-worker-timeout-s</string><string>3600</string>
<string>--prefill-policy</string><string>remote-required</string>
<string>--network-http-host</string><string>127.0.0.1</string>
<string>--network-http-port</string><string>8090</string>
Expand Down
12 changes: 7 additions & 5 deletions docs/ops/distributed-prefill-kv-network.md
Original file line number Diff line number Diff line change
Expand Up @@ -342,11 +342,13 @@ Generator completion status and must not penalize an honest statement that an
open problem has no accepted proof.
The REPL ignores external `SIGTERM`; its shell supervisor restarts signal-based
exits. Only `/quit`, `/exit`, or EOF is treated as approval to stop.
Generator output is always streamed in full to Terminal. To keep the 16GB
allens Critic Prefill interactive, Critic receives a labeled extractive evidence
window (default 64 Generator tokens: beginning + conclusion) with the omitted
token count and EOS status. It must not interpret evidence-window omission as
Generator truncation. Long Prefill operations emit a heartbeat every 30 seconds.
Generator output is always streamed in full to Terminal and passed verbatim to
the Gemma Critic. Sampling, truncation, summarization, independent chunk scores,
and semantic fallback are forbidden. A global Critic score is valid only when
`review_scope=full`, `critic_context_tokens=generator_full_tokens`, and
`critic_omitted_tokens=0`. Long Prefill operations emit a heartbeat every 30
seconds; on the 16GB allens worker, full-context Critic Prefill may take 15–25
minutes.
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.

Expand Down
2 changes: 1 addition & 1 deletion inference_engine/network/dashboard.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def dashboard_html() -> str:
$('createRegistration').onclick=async()=>{let r=await fetch('/v1/network/nodes/register',{method:'POST',headers:writeHeaders(),body:JSON.stringify({alias:$('alias').value,address:$('address').value,region:$('region').value,role:'hybrid'})});let j=await r.json();$('pairing').textContent=r.ok?`Pairing token: ${j.pairing_token}\nExpires: ${new Date(j.expires_at*1000).toLocaleTimeString()}`:`Error: ${j.detail||r.status}`;$('pairing').classList.remove('hidden');load()};
$('createGroup').onclick=async()=>{await fetch('/v1/network/groups',{method:'POST',headers:writeHeaders(),body:JSON.stringify({name:$('groupName').value,node_ids:$('groupNodes').value.split(',').map(x=>x.trim()).filter(Boolean)})});load()};
function nodePosition(i,total){let a=(i/Math.max(total,1))*Math.PI*2;return {x:50+38*Math.cos(a),y:53+35*Math.sin(a)}}
function phaseCards(stages){$('benchmarkPhases').innerHTML=stages.map(x=>`<div class="card"><b>${esc(x.agent?`${x.agent} R${x.round}`:x.name)}</b><p>${esc(x.hit_source)} · ${x.ok?'PASS':'FAIL'}</p><div class="muted">Stop ${esc(x.stop_reason||'n/a')} · ${x.complete===false?'INCOMPLETE':'complete'}${x.critic_omitted_tokens?`<br>Critic evidence ${x.critic_evidence_tokens}/${x.generator_full_tokens} tokens · omitted ${x.critic_omitted_tokens}`:''}<br>TTFT ${num(x.ttft_s)}s<br>Prefill/restore ${num(x.prefill_or_restore_tok_s)} tok/s<br>Decode ${num(x.decode_tok_s)} tok/s<br>Generation ${num(x.generation_latency_ms_per_token)} ms/token<br>E2E ${num(x.e2e_tok_s)} tok/s</div></div>`).join('')||'<div class="card muted">No stages yet.</div>'}
function phaseCards(stages){$('benchmarkPhases').innerHTML=stages.map(x=>`<div class="card"><b>${esc(x.agent?`${x.agent} R${x.round}`:x.name)}</b><p>${esc(x.hit_source)} · ${x.ok?'PASS':'FAIL'}</p><div class="muted">Stop ${esc(x.stop_reason||'n/a')} · ${x.complete===false?'INCOMPLETE':'complete'}${x.review_scope?`<br>Critic review ${esc(x.review_scope)} · context ${x.critic_context_tokens}/${x.generator_full_tokens} tokens · omitted ${x.critic_omitted_tokens||0}`:''}<br>TTFT ${num(x.ttft_s)}s<br>Prefill/restore ${num(x.prefill_or_restore_tok_s)} tok/s<br>Decode ${num(x.decode_tok_s)} tok/s<br>Generation ${num(x.generation_latency_ms_per_token)} ms/token<br>E2E ${num(x.e2e_tok_s)} tok/s</div></div>`).join('')||'<div class="card muted">No stages yet.</div>'}
async function showBenchmark(id){let r=await fetch('/v1/network/benchmarks/'+encodeURIComponent(id)).then(x=>x.json());phaseCards(r.stages||[]);$('benchmarkDetail').innerHTML=`<b>${esc(r.id)} · ${esc(r.status)}</b><p class="muted">${esc(r.kind)} · ${new Date(r.started_at*1000).toLocaleString()}</p><table><thead><tr><th>Phase</th><th>Source</th><th>Stop</th><th>TTFT</th><th>Prefill/restore</th><th>Decode</th><th>Latency/token</th><th>E2E</th></tr></thead><tbody>${(r.stages||[]).map(x=>`<tr><td>${esc(x.agent?`${x.agent} R${x.round}`:x.name)}</td><td>${esc(x.hit_source)}</td><td>${esc(x.stop_reason||'n/a')}</td><td>${num(x.ttft_s)}s</td><td>${num(x.prefill_or_restore_tok_s)}</td><td>${num(x.decode_tok_s)}</td><td>${num(x.generation_latency_ms_per_token)}ms</td><td>${num(x.e2e_tok_s)}</td></tr>`).join('')}</tbody></table>`}
async function load(){let [s,n,g,live,runs]=await Promise.all([fetch('/v1/network/summary').then(r=>r.json()),fetch('/v1/network/nodes').then(r=>r.json()),fetch('/v1/network/groups').then(r=>r.json()),fetch('/v1/network/benchmarks/live').then(r=>r.json()),fetch('/v1/network/benchmarks?limit=20').then(r=>r.json())]);
$('online').textContent=s.online_nodes;$('groupCount').textContent=s.groups;$('tokens').textContent=fmt(s.completed_tokens);$('hitRate').textContent=(s.kv_hit_rate*100).toFixed(0)+'%';$('cache').textContent=gb(s.cache_bytes_used+s.cache_bytes_free)+' GB';let p=s.prefill||{};$('remoteJobs').textContent=fmt(p.remote_jobs);$('remoteHits').textContent=fmt(p.remote_hits);$('reusedTokens').textContent=fmt(p.tokens_reused);$('evictions').textContent=fmt(s.cache_evictions);$('publishFailures').textContent=fmt(p.publish_failures);
Expand Down
53 changes: 18 additions & 35 deletions scripts/agent_gan_inference_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,33 +34,13 @@ def _output_metadata(text: str) -> dict:
}


def build_critic_evidence(tokenizer, text: str, max_tokens: int) -> tuple[str, dict]:
if max_tokens <= 0:
raise ValueError("critic evidence token budget must be > 0")
def build_critic_context(tokenizer, text: str) -> tuple[str, dict]:
full_ids = tokenizer.encode(text, add_special_tokens=False)
if len(full_ids) <= max_tokens:
return text, {
"generator_full_tokens": len(full_ids),
"critic_evidence_tokens": len(full_ids),
"critic_omitted_tokens": 0,
}
head_count = max_tokens // 2
tail_count = max_tokens - head_count
head = tokenizer.decode(full_ids[:head_count], skip_special_tokens=True)
tail = tokenizer.decode(full_ids[-tail_count:], skip_special_tokens=True)
omitted = len(full_ids) - max_tokens
evidence = (
"[BEGIN GENERATOR EVIDENCE]\n"
f"{head}\n"
f"[... {omitted} generator tokens omitted from Critic context ...]\n"
f"{tail}\n"
"[END GENERATOR EVIDENCE]"
)
evidence_tokens = len(tokenizer.encode(evidence, add_special_tokens=False))
return evidence, {
return text, {
"generator_full_tokens": len(full_ids),
"critic_evidence_tokens": evidence_tokens,
"critic_omitted_tokens": omitted,
"critic_context_tokens": len(full_ids),
"critic_omitted_tokens": 0,
"review_scope": "full",
}


Expand Down Expand Up @@ -154,12 +134,11 @@ def main() -> int:
default=0,
help="Optional client response cap; 0 means generate until model EOS.",
)
parser.add_argument("--critic-evidence-tokens", type=int, default=128)
parser.add_argument("--report", default="/tmp/kakeya-agent-gan-demo.json")
parser.add_argument("--skip-ensure", action="store_true")
args = parser.parse_args()
if min(args.rounds, args.output_tokens, args.critic_evidence_tokens) <= 0:
raise SystemExit("rounds, output-tokens and critic evidence must be > 0")
if min(args.rounds, args.output_tokens) <= 0:
raise SystemExit("rounds and output-tokens must be > 0")

from kakeya import Client
from transformers import AutoTokenizer
Expand Down Expand Up @@ -199,9 +178,8 @@ def main() -> int:
"incomplete merely because it refuses to fabricate a solution to "
"an open problem. Claim truncation only when completion_status is "
"not EOS or the text is syntactically cut off."
" You receive an explicitly bounded evidence window; omitted "
"middle tokens are a transport constraint, not evidence that the "
"Generator itself failed to complete."
" Review the complete Generator response as one semantic argument. "
"Do not sample, summarize, or infer claims from partial text."
),
}]

Expand Down Expand Up @@ -293,16 +271,21 @@ def execute_agent(client, name, round_index, history, extra_metrics=None):
client, "generator", round_index, generator_history,
)
generator_history.append({"role": "assistant", "content": proposal})
evidence, evidence_metrics = build_critic_evidence(
critic_context, context_metrics = build_critic_context(
tokenizer,
proposal,
args.critic_evidence_tokens,
)
if (
critic_context != proposal
or context_metrics["critic_omitted_tokens"] != 0
or context_metrics["review_scope"] != "full"
):
raise RuntimeError("Critic full-context invariant violated")
critic_history.append({
"role": "user",
"content": (
f"Architecture task:\n{task}\n\n"
f"Generator evidence window:\n{evidence}"
f"Complete Generator response:\n{critic_context}"
f"\n\ncompletion_status={generator_stage['stop_reason']}; "
f"complete={generator_stage['complete']}"
),
Expand All @@ -312,7 +295,7 @@ def execute_agent(client, name, round_index, history, extra_metrics=None):
"critic",
round_index,
critic_history,
extra_metrics=evidence_metrics,
extra_metrics=context_metrics,
)
critic_history.append({
"role": "assistant",
Expand Down
28 changes: 16 additions & 12 deletions scripts/agent_gan_repl.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from scripts.agent_gan_inference_demo import (
_agent_cache_gate,
_infer,
build_critic_evidence,
build_critic_context,
)
from scripts.benchmark_prefill_architecture import (
_ensure_services,
Expand Down Expand Up @@ -61,7 +61,7 @@ def build_generator_messages(prompt: str) -> list[dict[str, str]]:

def build_critic_messages(
prompt: str,
evidence: str,
generator_response: str,
*,
stop_reason: str,
complete: bool,
Expand All @@ -72,14 +72,14 @@ def build_critic_messages(
"content": (
"Score the answer 0-10, identify false claims, and give "
"specific corrections. Do not penalize an honest statement "
"that an open problem is unsolved. Evidence is intentionally "
"bounded; omitted tokens do not imply truncation."
"that an open problem is unsolved. Review the complete response "
"as one semantic argument; do not sample or summarize it."
),
},
{
"role": "user",
"content": (
f"Task:\n{prompt}\n\nEvidence:\n{evidence}\n\n"
f"Task:\n{prompt}\n\nComplete response:\n{generator_response}\n\n"
f"Completion: {stop_reason}; complete={complete}"
),
},
Expand Down Expand Up @@ -172,11 +172,10 @@ def main() -> int:
default=0,
help="Optional client response cap; 0 means generate until model EOS.",
)
parser.add_argument("--critic-evidence-tokens", type=int, default=64)
parser.add_argument("--skip-ensure", action="store_true")
args = parser.parse_args()
if min(args.output_tokens, args.critic_evidence_tokens) <= 0:
raise SystemExit("output-tokens and critic-evidence-tokens must be > 0")
if args.output_tokens <= 0:
raise SystemExit("output-tokens must be > 0")

from kakeya import Client
from transformers import AutoTokenizer
Expand Down Expand Up @@ -284,14 +283,19 @@ def get_stats():
body={"stages": [generator_stage]},
)

evidence, evidence_metrics = build_critic_evidence(
critic_context, context_metrics = build_critic_context(
tokenizer,
generator_text,
args.critic_evidence_tokens,
)
if (
critic_context != generator_text
or context_metrics["critic_omitted_tokens"] != 0
or context_metrics["review_scope"] != "full"
):
raise RuntimeError("Critic full-context invariant violated")
critic_messages = build_critic_messages(
prompt,
evidence,
critic_context,
stop_reason=generator_actual["stop_reason"],
complete=generator_actual["complete"],
)
Expand Down Expand Up @@ -330,7 +334,7 @@ def get_stats():
critic_warm,
critic_actual,
critic_text,
extra_metrics=evidence_metrics,
extra_metrics=context_metrics,
)
if not critic_stage["ok"] and not telemetry_state["degraded"]:
raise RuntimeError("Critic KV gate failed")
Expand Down
22 changes: 7 additions & 15 deletions tests/inference_engine/bridge/test_agent_gan_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
_agent_cache_gate,
_infer,
_output_metadata,
build_critic_evidence,
build_critic_context,
)


Expand Down Expand Up @@ -98,18 +98,10 @@ def decode(self, token_ids, **_kwargs):
return "".join(chr(token) for token in token_ids)


def test_critic_evidence_is_bounded_and_explicit_about_omission():
evidence, metrics = build_critic_evidence(CharTokenizer(), "abcdefghij", 4)
assert "ab" in evidence and "ij" in evidence
assert "6 generator tokens omitted" in evidence
def test_critic_context_preserves_complete_generator_response():
context, metrics = build_critic_context(CharTokenizer(), "abcdefghij")
assert context == "abcdefghij"
assert metrics["generator_full_tokens"] == 10
assert metrics["critic_omitted_tokens"] == 6
full, full_metrics = build_critic_evidence(CharTokenizer(), "abc", 4)
assert full == "abc"
assert full_metrics["critic_omitted_tokens"] == 0
try:
build_critic_evidence(CharTokenizer(), "abc", 0)
except ValueError:
pass
else:
raise AssertionError("expected evidence budget validation")
assert metrics["critic_context_tokens"] == 10
assert metrics["critic_omitted_tokens"] == 0
assert metrics["review_scope"] == "full"
20 changes: 14 additions & 6 deletions tests/inference_engine/bridge/test_agent_gan_repl.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def test_prefill_heartbeat_reports_elapsed_progress(capsys):
assert "Critic Prefill still running" in output


def test_stage_includes_evidence_window_metrics():
def test_stage_includes_full_context_metrics():
warm = {
"prefix_tokens": 10,
"e2e_s": 1,
Expand All @@ -124,9 +124,16 @@ def test_stage_includes_evidence_window_metrics():
warm,
actual,
"ok",
extra_metrics={"critic_omitted_tokens": 100},
extra_metrics={
"generator_full_tokens": 100,
"critic_context_tokens": 100,
"critic_omitted_tokens": 0,
"review_scope": "full",
},
)
assert stage["critic_omitted_tokens"] == 100
assert stage["critic_context_tokens"] == 100
assert stage["critic_omitted_tokens"] == 0
assert stage["review_scope"] == "full"


def test_telemetry_timeout_warns_without_stopping_inference(
Expand All @@ -148,13 +155,13 @@ def test_interactive_prompts_are_deterministic_for_kv_reuse():
generator_b = build_generator_messages("prove RH")
critic_a = build_critic_messages(
"prove RH",
"bounded evidence",
"complete generator response",
stop_reason="eos",
complete=True,
)
critic_b = build_critic_messages(
"prove RH",
"bounded evidence",
"complete generator response",
stop_reason="eos",
complete=True,
)
Expand All @@ -163,4 +170,5 @@ def test_interactive_prompts_are_deterministic_for_kv_reuse():
combined = repr(generator_a + critic_a)
assert "Internal run" not in combined
assert "open problem" in combined
assert "omitted tokens do not imply truncation" in combined
assert "Review the complete response" in combined
assert "do not sample or summarize" in combined
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def test_two_mac_deployment_uses_allens_as_prefill_only():
in plist
)
assert (
"<string>--prefill-worker-timeout-s</string><string>900</string>"
"<string>--prefill-worker-timeout-s</string><string>3600</string>"
in plist
)
assert (
Expand Down
Loading