diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index e601f4df..79f6a9f3 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -334,6 +334,11 @@ Type a prompt at `prompt>`. The terminal prints allens Prefill progress, then streams `generator>` and `critic>` tokens as they arrive, followed by KV hit rate, decode tok/s, generation latency and E2E tok/s. `/quit` exits; services remain running. Reports remain redacted and appear in the Benchmarks tab. +Generation uses 64-token streaming chunks and automatically continues to the +model's EOS. `--max-response-tokens` defaults to 512 as an explicit safety cap; +reaching it marks the stage incomplete instead of presenting a truncated answer +as successful. The Critic receives the Generator completion status and must not +penalize an honest statement that an open problem has no accepted proof. ## Rollback diff --git a/inference_engine/bench/prefill_fleet_report.py b/inference_engine/bench/prefill_fleet_report.py index 6ffdd354..75a91f31 100644 --- a/inference_engine/bench/prefill_fleet_report.py +++ b/inference_engine/bench/prefill_fleet_report.py @@ -65,6 +65,10 @@ def summarize_stages(stages: Sequence[dict[str, Any]]) -> dict[str, Any]: sources = {source: 0 for source in HIT_SOURCES} for stage in normalized: sources[stage["hit_source"]] += 1 + stop_reasons: dict[str, int] = {} + for stage in normalized: + reason = str(stage.get("stop_reason", "unknown")) + stop_reasons[reason] = stop_reasons.get(reason, 0) + 1 decode = [stage["decode_tok_s"] for stage in normalized] prefix_tokens = sum(stage["prefix_tokens"] for stage in normalized) hit_tokens = sum( @@ -83,6 +87,10 @@ def summarize_stages(stages: Sequence[dict[str, Any]]) -> dict[str, Any]: return { "stages_total": len(normalized), "stages_failed": sum(not stage.get("ok", False) for stage in normalized), + "incomplete_stages": sum( + not stage.get("complete", True) for stage in normalized + ), + "stop_reason_counts": stop_reasons, "ttft_p50_s": _median(stage["ttft_s"] for stage in normalized), "prefill_tok_s_p50": _median( stage["prefill_or_restore_tok_s"] for stage in normalized diff --git a/inference_engine/network/dashboard.py b/inference_engine/network/dashboard.py index 25243441..469049ae 100644 --- a/inference_engine/network/dashboard.py +++ b/inference_engine/network/dashboard.py @@ -49,8 +49,8 @@ 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=>`
${esc(x.agent?`${x.agent} R${x.round}`:x.name)}

${esc(x.hit_source)} · ${x.ok?'PASS':'FAIL'}

TTFT ${num(x.ttft_s)}s
Prefill/restore ${num(x.prefill_or_restore_tok_s)} tok/s
Decode ${num(x.decode_tok_s)} tok/s
Generation ${num(x.generation_latency_ms_per_token)} ms/token
E2E ${num(x.e2e_tok_s)} tok/s
`).join('')||'
No stages yet.
'} -async function showBenchmark(id){let r=await fetch('/v1/network/benchmarks/'+encodeURIComponent(id)).then(x=>x.json());phaseCards(r.stages||[]);$('benchmarkDetail').innerHTML=`${esc(r.id)} · ${esc(r.status)}

${esc(r.kind)} · ${new Date(r.started_at*1000).toLocaleString()}

${(r.stages||[]).map(x=>``).join('')}
PhaseSourceTTFTPrefill/restoreDecodeLatency/tokenE2E
${esc(x.agent?`${x.agent} R${x.round}`:x.name)}${esc(x.hit_source)}${num(x.ttft_s)}s${num(x.prefill_or_restore_tok_s)}${num(x.decode_tok_s)}${num(x.generation_latency_ms_per_token)}ms${num(x.e2e_tok_s)}
`} +function phaseCards(stages){$('benchmarkPhases').innerHTML=stages.map(x=>`
${esc(x.agent?`${x.agent} R${x.round}`:x.name)}

${esc(x.hit_source)} · ${x.ok?'PASS':'FAIL'}

Stop ${esc(x.stop_reason||'n/a')} · ${x.complete===false?'INCOMPLETE':'complete'}
TTFT ${num(x.ttft_s)}s
Prefill/restore ${num(x.prefill_or_restore_tok_s)} tok/s
Decode ${num(x.decode_tok_s)} tok/s
Generation ${num(x.generation_latency_ms_per_token)} ms/token
E2E ${num(x.e2e_tok_s)} tok/s
`).join('')||'
No stages yet.
'} +async function showBenchmark(id){let r=await fetch('/v1/network/benchmarks/'+encodeURIComponent(id)).then(x=>x.json());phaseCards(r.stages||[]);$('benchmarkDetail').innerHTML=`${esc(r.id)} · ${esc(r.status)}

${esc(r.kind)} · ${new Date(r.started_at*1000).toLocaleString()}

${(r.stages||[]).map(x=>``).join('')}
PhaseSourceStopTTFTPrefill/restoreDecodeLatency/tokenE2E
${esc(x.agent?`${x.agent} R${x.round}`:x.name)}${esc(x.hit_source)}${esc(x.stop_reason||'n/a')}${num(x.ttft_s)}s${num(x.prefill_or_restore_tok_s)}${num(x.decode_tok_s)}${num(x.generation_latency_ms_per_token)}ms${num(x.e2e_tok_s)}
`} 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); let total=s.cache_bytes_used+s.cache_bytes_free,pct=total?s.cache_bytes_used/total*100:0;$('capacityLabel').textContent=`${gb(s.cache_bytes_used)} / ${gb(total)} GB`;$('capacityBar').style.width=pct+'%'; diff --git a/scripts/agent_gan_inference_demo.py b/scripts/agent_gan_inference_demo.py index ed0f247a..8d894a3d 100644 --- a/scripts/agent_gan_inference_demo.py +++ b/scripts/agent_gan_inference_demo.py @@ -41,6 +41,7 @@ def _infer( output_tokens: int, get_stats, on_token=None, + max_response_tokens=None, ): before = get_stats() started = time.perf_counter() @@ -50,12 +51,30 @@ def _infer( append_done = time.perf_counter() first_at = None generated = [] - for token in s.generate(max_tokens=output_tokens): - generated.append(int(token)) - if on_token is not None: - on_token(generated) - if first_at is None: - first_at = time.perf_counter() + response_limit = int(max_response_tokens or output_tokens) + stop_reason = "unknown" + while len(generated) < response_limit: + before_count = len(generated) + chunk = min(output_tokens, response_limit - len(generated)) + for token in s.generate(max_tokens=chunk): + generated.append(int(token)) + if on_token is not None: + on_token(generated) + if first_at is None: + first_at = time.perf_counter() + stop_reason = { + 1: "max_tokens", + 2: "eos", + 3: "cancelled", + 4: "truncated", + }.get(s.last_stop_reason, "unknown") + if stop_reason != "max_tokens": + break + if len(generated) == before_count: + stop_reason = "no_progress" + break + if len(generated) >= response_limit and stop_reason == "max_tokens": + stop_reason = "client_safety_limit" done = time.perf_counter() after = get_stats() first_at = first_at or done @@ -67,6 +86,8 @@ def _infer( "decode_s": done - append_done, "e2e_s": done - started, "delta": _delta(before, after), + "stop_reason": stop_reason, + "complete": stop_reason == "eos", } @@ -85,6 +106,7 @@ def main() -> int: "reliably; larger values require more worker memory or timeout.", ) parser.add_argument("--output-tokens", type=int, default=64) + parser.add_argument("--max-response-tokens", type=int, default=512) parser.add_argument("--report", default="/tmp/kakeya-agent-gan-demo.json") parser.add_argument("--skip-ensure", action="store_true") args = parser.parse_args() @@ -113,7 +135,9 @@ def main() -> int: "role": "system", "content": ( "You are the Generator agent. Propose a technically precise " - "architecture improvement. Respond with actionable reasoning." + "architecture improvement. Respond with actionable reasoning. " + "For open or unsolved problems, state the accepted boundary " + "honestly and never fabricate a proof." ), }, {"role": "user", "content": task}, @@ -123,7 +147,10 @@ def main() -> int: "content": ( "You are the Critic/Discriminator agent. Attack the proposal, " "identify false assumptions and bottlenecks, score it from 0 to " - "10, and demand specific corrections." + "10, and demand specific corrections. Do not call a response " + "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." ), }] @@ -139,6 +166,7 @@ def main() -> int: "agents": ["generator", "critic"], "rounds": args.rounds, "output_tokens": args.output_tokens, + "max_response_tokens": args.max_response_tokens, }, }, ) @@ -164,6 +192,7 @@ def execute_agent(client, name, round_index, history): token_ids, args.output_tokens, get_stats, + max_response_tokens=args.max_response_tokens, ) text = tokenizer.decode(generated, skip_special_tokens=True) delta = actual["delta"] @@ -174,7 +203,7 @@ def execute_agent(client, name, round_index, history): "agent": name, "round": round_index, "hit_source": "primary_hot" if delta["local_hits"] else "unknown", - "ok": ok, + "ok": ok and actual["complete"], "warmup_prefix_tokens": warm["prefix_tokens"], "warmup_tokens_reused": ( warm["delta"]["tokens_reused"] @@ -194,7 +223,7 @@ def execute_agent(client, name, round_index, history): ) all_stages.append(stage) print(f"\n[{name.upper()} round {round_index}]\n{text}\n", flush=True) - return text + return text, stage try: with Client(args.address) as client: @@ -208,7 +237,7 @@ def execute_agent(client, name, round_index, history): + critic_feedback ), }) - proposal = execute_agent( + proposal, generator_stage = execute_agent( client, "generator", round_index, generator_history, ) generator_history.append({"role": "assistant", "content": proposal}) @@ -216,9 +245,11 @@ def execute_agent(client, name, round_index, history): "role": "user", "content": ( f"Architecture task:\n{task}\n\nGenerator proposal:\n{proposal}" + f"\n\ncompletion_status={generator_stage['stop_reason']}; " + f"complete={generator_stage['complete']}" ), }) - critic_feedback = execute_agent( + critic_feedback, _critic_stage = execute_agent( client, "critic", round_index, critic_history, ) critic_history.append({ diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index b1307582..1eec84fd 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -42,7 +42,7 @@ def _stage(name: str, warm: dict, actual: dict, text: str) -> dict: "agent": name, "round": 1, "hit_source": "primary_hot" if delta["local_hits"] else "unknown", - "ok": _agent_cache_gate(warm["delta"], delta), + "ok": _agent_cache_gate(warm["delta"], delta) and actual["complete"], "warmup_prefix_tokens": warm["prefix_tokens"], "warmup_tokens_reused": ( warm["delta"]["tokens_reused"] @@ -63,6 +63,7 @@ def main() -> int: parser.add_argument("--api-key-file", default="~/.kakeya/network_api_key") parser.add_argument("--tokenizer-id", required=True) parser.add_argument("--output-tokens", type=int, default=64) + parser.add_argument("--max-response-tokens", type=int, default=512) parser.add_argument("--skip-ensure", action="store_true") args = parser.parse_args() @@ -121,7 +122,10 @@ def get_stats(): "role": "system", "content": ( "You are the Generator agent. Produce a concrete, " - "technically rigorous answer. Internal run " + "technically rigorous answer. For open or unsolved " + "problems, state the accepted boundary honestly, " + "provide rigorous context, and never fabricate a " + "proof. Internal run " f"{run_nonce}." ), }, @@ -149,6 +153,7 @@ def get_stats(): args.output_tokens, get_stats, on_token=generator_printer, + max_response_tokens=args.max_response_tokens, ) generator_printer.finish() generator_text = tokenizer.decode( @@ -176,7 +181,11 @@ def get_stats(): "content": ( "You are the Critic/Discriminator. Score the answer " "0-10, identify false assumptions, and propose " - "specific corrections. Internal run " + "specific corrections. Do not penalize a correct " + "statement that an open problem has no accepted " + "proof. Call an answer incomplete only when its " + "completion status is not EOS or its syntax is " + "visibly cut off. Internal run " f"{run_nonce}." ), }, @@ -184,7 +193,10 @@ def get_stats(): "role": "user", "content": ( f"Original task:\n{prompt}\n\n" - f"Generator answer:\n{generator_text}" + f"Generator answer:\n{generator_text}\n\n" + "Generator completion status: " + f"{generator_actual['stop_reason']}; " + f"complete={generator_actual['complete']}" ), }, ] @@ -210,6 +222,7 @@ def get_stats(): args.output_tokens, get_stats, on_token=critic_printer, + max_response_tokens=args.max_response_tokens, ) critic_printer.finish() critic_text = tokenizer.decode( diff --git a/tests/inference_engine/bench/test_prefill_fleet_report.py b/tests/inference_engine/bench/test_prefill_fleet_report.py index b7d35d10..92b3b314 100644 --- a/tests/inference_engine/bench/test_prefill_fleet_report.py +++ b/tests/inference_engine/bench/test_prefill_fleet_report.py @@ -36,10 +36,17 @@ def test_summary_aggregates_sources_and_medians(): summary = summarize_stages([ _stage(), _stage("primary_hot_hit", "primary_hot"), - {**_stage("allens_cold_restore", "allens_offload"), "ok": False}, + { + **_stage("allens_cold_restore", "allens_offload"), + "ok": False, + "complete": False, + "stop_reason": "client_safety_limit", + }, ]) assert summary["stages_total"] == 3 assert summary["stages_failed"] == 1 + assert summary["incomplete_stages"] == 1 + assert summary["stop_reason_counts"]["client_safety_limit"] == 1 assert summary["hit_source_counts"]["remote_worker"] == 1 assert summary["hit_source_counts"]["primary_hot"] == 1 assert summary["bytes_received"] == 3000 diff --git a/tests/inference_engine/bridge/test_agent_gan_demo.py b/tests/inference_engine/bridge/test_agent_gan_demo.py index d17337ad..82edc4db 100644 --- a/tests/inference_engine/bridge/test_agent_gan_demo.py +++ b/tests/inference_engine/bridge/test_agent_gan_demo.py @@ -1,5 +1,6 @@ from scripts.agent_gan_inference_demo import ( _agent_cache_gate, + _infer, _output_metadata, ) @@ -23,3 +24,66 @@ def test_agent_output_report_is_redacted(): assert result["output_chars"] == 20 assert len(result["output_hash"]) == 64 assert "output" not in result + + +class Session: + def __init__(self, chunks): + self.chunks = list(chunks) + self.last_stop_reason = None + self.calls = 0 + + def __enter__(self): + return self + + def __exit__(self, *_args): + pass + + def append(self, token_ids): + self.appended = list(token_ids) + + def generate(self, *, max_tokens): + tokens, reason = self.chunks[self.calls] + self.calls += 1 + assert len(tokens) <= max_tokens + yield from tokens + self.last_stop_reason = reason + + +class Client: + def __init__(self, session): + self.session = session + + def create_session(self, **_kwargs): + return self.session + + +def test_infer_continues_chunks_until_eos(): + session = Session([([1, 2], 1), ([3], 2)]) + streamed = [] + tokens, metrics = _infer( + Client(session), + [], + [9], + 2, + lambda: {}, + on_token=lambda values: streamed.append(list(values)), + max_response_tokens=10, + ) + assert tokens == [1, 2, 3] + assert streamed == [[1], [1, 2], [1, 2, 3]] + assert metrics["stop_reason"] == "eos" + assert metrics["complete"] is True + + +def test_infer_reports_explicit_client_safety_limit(): + tokens, metrics = _infer( + Client(Session([([1, 2], 1)])), + [], + [9], + 2, + lambda: {}, + max_response_tokens=2, + ) + assert tokens == [1, 2] + assert metrics["stop_reason"] == "client_safety_limit" + assert metrics["complete"] is False diff --git a/tests/inference_engine/bridge/test_agent_gan_repl.py b/tests/inference_engine/bridge/test_agent_gan_repl.py index 19d80fb5..ee34cad8 100644 --- a/tests/inference_engine/bridge/test_agent_gan_repl.py +++ b/tests/inference_engine/bridge/test_agent_gan_repl.py @@ -31,6 +31,8 @@ def test_repl_stage_is_redacted_and_passes_cache_gate(): "ttft_s": 0.2, "decode_s": 0.3, "e2e_s": 0.4, + "stop_reason": "eos", + "complete": True, "delta": { "local_hits": 1, "remote_jobs": 0, @@ -43,3 +45,9 @@ def test_repl_stage_is_redacted_and_passes_cache_gate(): assert stage["output_chars"] == 14 assert len(stage["output_hash"]) == 64 assert "output" not in stage + assert not _stage( + "generator", + warm, + {**actual, "complete": False, "stop_reason": "client_safety_limit"}, + "cut off", + )["ok"]