diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 5b7e014f..d2fc598e 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -302,6 +302,24 @@ curl -fsS https://kakeya.ai/v1/network/benchmarks/ The `Benchmarks` dashboard tab shows live progress, phase comparison, history, and complete redacted stage details. +### Generator/Critic Agent GAN demo + +Run two logical agents through real multi-round model inference: + +```bash +bash scripts/run_agent_gan_demo.sh \ + --rounds 2 \ + --output-tokens 64 \ + --report /tmp/kakeya-agent-gan-demo.json +``` + +For every Generator and Critic turn the task first asks allens to Prefill the +complete agent context, then performs the actual inference against the promoted +Primary hot snapshot. Terminal output contains the real proposal/critique; +persisted reports contain only output length/hash and metrics. The report +separates inference-only KV hit rate from whole-workload hit rate including +warmup, plus per-agent and aggregate token throughput/latency. + ## Rollback The cache is an optimization; inference correctness does not depend on it. diff --git a/inference_engine/bench/prefill_fleet_report.py b/inference_engine/bench/prefill_fleet_report.py index f4b69898..6ffdd354 100644 --- a/inference_engine/bench/prefill_fleet_report.py +++ b/inference_engine/bench/prefill_fleet_report.py @@ -4,7 +4,13 @@ import statistics from typing import Any, Sequence -PHASES = ("remote_compute", "primary_hot_hit", "allens_cold_restore") +PHASES = ( + "remote_compute", + "primary_hot_hit", + "allens_cold_restore", + "agent_generator", + "agent_critic", +) HIT_SOURCES = ("remote_worker", "primary_hot", "allens_offload", "unknown") _PRIVATE_KEYS = { "prompt", @@ -60,6 +66,20 @@ def summarize_stages(stages: Sequence[dict[str, Any]]) -> dict[str, Any]: for stage in normalized: sources[stage["hit_source"]] += 1 decode = [stage["decode_tok_s"] for stage in normalized] + prefix_tokens = sum(stage["prefix_tokens"] for stage in normalized) + hit_tokens = sum( + int(stage.get("delta", {}).get("tokens_reused", 0)) + for stage in normalized + ) + warmup_prefix_tokens = sum( + int(stage.get("warmup_prefix_tokens", 0)) for stage in normalized + ) + warmup_hit_tokens = sum( + int(stage.get("warmup_tokens_reused", 0)) for stage in normalized + ) + decode_tokens = sum(stage["output_tokens"] for stage in normalized) + decode_seconds = sum(stage["decode_s"] for stage in normalized) + e2e_seconds = sum(stage["e2e_s"] for stage in normalized) return { "stages_total": len(normalized), "stages_failed": sum(not stage.get("ok", False) for stage in normalized), @@ -72,6 +92,20 @@ def summarize_stages(stages: Sequence[dict[str, Any]]) -> dict[str, Any]: "generation_latency_ms_p50": _median( stage["generation_latency_ms_per_token"] for stage in normalized ), + "inference_kv_token_hit_rate": ( + hit_tokens / prefix_tokens if prefix_tokens else 0.0 + ), + "workload_kv_token_hit_rate": ( + (hit_tokens + warmup_hit_tokens) + / (prefix_tokens + warmup_prefix_tokens) + if prefix_tokens + warmup_prefix_tokens else 0.0 + ), + "aggregate_decode_tok_s": ( + decode_tokens / decode_seconds if decode_seconds else 0.0 + ), + "aggregate_e2e_tok_s": ( + decode_tokens / e2e_seconds if e2e_seconds else 0.0 + ), "bytes_received": sum( int(stage.get("delta", {}).get("bytes_received", 0)) for stage in normalized diff --git a/inference_engine/network/dashboard.py b/inference_engine/network/dashboard.py index bfcb79f6..25243441 100644 --- a/inference_engine/network/dashboard.py +++ b/inference_engine/network/dashboard.py @@ -38,7 +38,7 @@ def dashboard_html() -> str: - +

Exact IPs, raw prompt hashes and cache keys are administrator-only. Region is operator-selected and coarse.

""" diff --git a/scripts/agent_gan_inference_demo.py b/scripts/agent_gan_inference_demo.py new file mode 100644 index 00000000..1171618a --- /dev/null +++ b/scripts/agent_gan_inference_demo.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""Real Generator/Critic multi-agent inference over the two-Mac KV architecture.""" +from __future__ import annotations + +import argparse +import hashlib +import json +import time +import uuid +from pathlib import Path + +from scripts.benchmark_prefill_architecture import ( + _delta, + _ensure_services, + _json_request, +) + + +def _agent_cache_gate(warm_delta: dict, actual_delta: dict) -> bool: + return ( + warm_delta["remote_jobs"] >= 1 + and warm_delta["remote_hits"] >= 1 + and actual_delta["local_hits"] >= 1 + and actual_delta["remote_jobs"] == 0 + and actual_delta["tokens_computed"] == 0 + and actual_delta["fallbacks"] == 0 + ) + + +def _output_metadata(text: str) -> dict: + return { + "output_chars": len(text), + "output_hash": hashlib.sha256(text.encode()).hexdigest(), + } + + +def _infer(client, eos_ids, token_ids, output_tokens: int, get_stats): + before = get_stats() + started = time.perf_counter() + with client.create_session(eos_token_ids=eos_ids, client_label="agent-gan") as s: + append_started = time.perf_counter() + s.append(token_ids) + append_done = time.perf_counter() + first_at = None + generated = [] + for token in s.generate(max_tokens=output_tokens): + generated.append(int(token)) + if first_at is None: + first_at = time.perf_counter() + done = time.perf_counter() + after = get_stats() + first_at = first_at or done + return generated, { + "prefix_tokens": len(token_ids), + "output_tokens": len(generated), + "append_s": append_done - append_started, + "ttft_s": first_at - started, + "decode_s": done - append_done, + "e2e_s": done - started, + "delta": _delta(before, after), + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--worker-ssh", default="allens") + parser.add_argument("--address", default="127.0.0.1:51051") + parser.add_argument("--dashboard", default="http://127.0.0.1:8090") + parser.add_argument("--api-key-file", default="~/.kakeya/network_api_key") + parser.add_argument("--tokenizer-id", required=True) + parser.add_argument("--rounds", type=int, default=2) + parser.add_argument("--output-tokens", type=int, default=64) + parser.add_argument("--report", default="/tmp/kakeya-agent-gan-demo.json") + parser.add_argument("--skip-ensure", action="store_true") + args = parser.parse_args() + if args.rounds <= 0 or args.output_tokens <= 0: + raise SystemExit("rounds and output-tokens must be > 0") + + from kakeya import Client + from transformers import AutoTokenizer + from scripts.chat_grpc import _resolve_eos_token_ids + + if not args.skip_ensure: + _ensure_services(args.worker_ssh) + api_key = Path(args.api_key_file).expanduser().read_text().strip() + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id) + eos_ids = _resolve_eos_token_ids(tokenizer) + run_nonce = uuid.uuid4().hex + task = ( + "Evaluate and improve the current two-Mac architecture where Primary " + "is decode-only, allens performs prefill, Primary keeps hot KV, and " + "allens provides cold KV offload. Produce concrete correctness, " + "throughput, memory, and failure-mode recommendations. " + f"Evaluation run {run_nonce}." + ) + generator_history = [ + { + "role": "system", + "content": ( + "You are the Generator agent. Propose a technically precise " + "architecture improvement. Respond with actionable reasoning." + ), + }, + {"role": "user", "content": task}, + ] + critic_history = [{ + "role": "system", + "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." + ), + }] + + run = _json_request( + f"{args.dashboard}/v1/network/benchmarks", + api_key=api_key, + method="POST", + body={ + "kind": "agent_gan_inference_demo", + "config": { + "model_id": "gemma-4-26B-A4B-it-mlx-4bit", + "topology": "primary-decode-allens-prefill", + "agents": ["generator", "critic"], + "rounds": args.rounds, + "output_tokens": args.output_tokens, + }, + }, + ) + run_id = run["id"] + all_stages = [] + + def get_stats(): + return _json_request(f"{args.dashboard}/v1/network/prefill") + + def execute_agent(client, name, round_index, history): + token_ids = tokenizer.apply_chat_template( + history, + add_generation_prompt=True, + tokenize=True, + return_dict=False, + enable_thinking=False, + ) + warm_tokens, warm = _infer(client, eos_ids, token_ids, 1, get_stats) + del warm_tokens + generated, actual = _infer( + client, + eos_ids, + token_ids, + args.output_tokens, + get_stats, + ) + text = tokenizer.decode(generated, skip_special_tokens=True) + delta = actual["delta"] + ok = _agent_cache_gate(warm["delta"], delta) + stage = { + **actual, + "name": f"agent_{name}", + "agent": name, + "round": round_index, + "hit_source": "primary_hot" if delta["local_hits"] else "unknown", + "ok": ok, + "warmup_prefix_tokens": warm["prefix_tokens"], + "warmup_tokens_reused": ( + warm["delta"]["tokens_reused"] + if warm["delta"]["remote_jobs"] == 0 else 0 + ), + "warmup_wall_s": warm["e2e_s"], + "warmup_remote_jobs": warm["delta"]["remote_jobs"], + **_output_metadata(text), + } + if not ok: + raise RuntimeError(f"{name} round {round_index} cache gate failed") + _json_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={"stages": [stage]}, + ) + all_stages.append(stage) + print(f"\n[{name.upper()} round {round_index}]\n{text}\n", flush=True) + return text + + try: + with Client(args.address) as client: + critic_feedback = "" + for round_index in range(1, args.rounds + 1): + if critic_feedback: + generator_history.append({ + "role": "user", + "content": ( + "Revise the architecture using this critic feedback:\n" + + critic_feedback + ), + }) + proposal = execute_agent( + client, "generator", round_index, generator_history, + ) + generator_history.append({"role": "assistant", "content": proposal}) + critic_history.append({ + "role": "user", + "content": ( + f"Architecture task:\n{task}\n\nGenerator proposal:\n{proposal}" + ), + }) + critic_feedback = execute_agent( + client, "critic", round_index, critic_history, + ) + critic_history.append({ + "role": "assistant", + "content": critic_feedback, + }) + completed = _json_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={ + "status": "completed", + "finished_at": time.time(), + }, + ) + except Exception: + _json_request( + f"{args.dashboard}/v1/network/benchmarks/{run_id}", + api_key=api_key, + method="PATCH", + body={"status": "failed", "finished_at": time.time()}, + ) + raise + Path(args.report).write_text(json.dumps(completed, indent=2)) + print(json.dumps({ + "ok": True, + "run_id": run_id, + "report": args.report, + "summary": completed["summary"], + }, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_agent_gan_demo.sh b/scripts/run_agent_gan_demo.sh new file mode 100644 index 00000000..fb8e670e --- /dev/null +++ b/scripts/run_agent_gan_demo.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)" +PYTHON="${KAKEYA_BENCH_PYTHON:-$HOME/.venv-distwan/bin/python}" +MODEL="${KAKEYA_BENCH_MODEL:-$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit}" + +exec env PYTHONPATH="$REPO_ROOT:$REPO_ROOT/sdks/python" \ + "$PYTHON" "$REPO_ROOT/scripts/agent_gan_inference_demo.py" \ + --tokenizer-id "$MODEL" \ + "$@" diff --git a/tests/inference_engine/bench/test_prefill_fleet_report.py b/tests/inference_engine/bench/test_prefill_fleet_report.py index dff0a5af..b7d35d10 100644 --- a/tests/inference_engine/bench/test_prefill_fleet_report.py +++ b/tests/inference_engine/bench/test_prefill_fleet_report.py @@ -18,7 +18,9 @@ def _stage(name="remote_compute", source="remote_worker"): "ttft_s": 5.2, "decode_s": 2.0, "e2e_s": 7.0, - "delta": {"bytes_received": 1000}, + "delta": {"bytes_received": 1000, "tokens_reused": 100}, + "warmup_prefix_tokens": 100, + "warmup_tokens_reused": 0, } @@ -41,12 +43,19 @@ def test_summary_aggregates_sources_and_medians(): assert summary["hit_source_counts"]["remote_worker"] == 1 assert summary["hit_source_counts"]["primary_hot"] == 1 assert summary["bytes_received"] == 3000 + assert summary["inference_kv_token_hit_rate"] == 1.0 + assert summary["workload_kv_token_hit_rate"] == 0.5 + assert summary["aggregate_decode_tok_s"] == 5.0 + assert summary["aggregate_e2e_tok_s"] == 10 / 7 assert summarize_stages([])["decode_tok_s_p50"] == 0 def test_schema_rejects_unknown_and_private_fields(): with pytest.raises(ValueError, match="unknown benchmark phase"): normalize_stage(_stage("bad")) + assert normalize_stage(_stage("agent_generator", "primary_hot"))["name"] == ( + "agent_generator" + ) with pytest.raises(ValueError, match="unknown hit_source"): normalize_stage(_stage(source="peer:1")) with pytest.raises(ValueError, match="non-negative"): diff --git a/tests/inference_engine/bridge/test_agent_gan_demo.py b/tests/inference_engine/bridge/test_agent_gan_demo.py new file mode 100644 index 00000000..d17337ad --- /dev/null +++ b/tests/inference_engine/bridge/test_agent_gan_demo.py @@ -0,0 +1,25 @@ +from scripts.agent_gan_inference_demo import ( + _agent_cache_gate, + _output_metadata, +) + + +def test_agent_gate_requires_remote_warmup_and_primary_hot_inference(): + warm = {"remote_jobs": 1, "remote_hits": 1} + actual = { + "local_hits": 1, + "remote_jobs": 0, + "tokens_computed": 0, + "fallbacks": 0, + } + assert _agent_cache_gate(warm, actual) + assert not _agent_cache_gate({**warm, "remote_jobs": 0}, actual) + assert not _agent_cache_gate(warm, {**actual, "local_hits": 0}) + assert not _agent_cache_gate(warm, {**actual, "fallbacks": 1}) + + +def test_agent_output_report_is_redacted(): + result = _output_metadata("private model output") + assert result["output_chars"] == 20 + assert len(result["output_hash"]) == 64 + assert "output" not in result