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
5 changes: 5 additions & 0 deletions docs/ops/distributed-prefill-kv-network.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
8 changes: 8 additions & 0 deletions inference_engine/bench/prefill_fleet_report.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions inference_engine/network/dashboard.py
Original file line number Diff line number Diff line change
Expand Up @@ -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=>`<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">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>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>${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>`}
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'}<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);
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+'%';
Expand Down
55 changes: 43 additions & 12 deletions scripts/agent_gan_inference_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Expand All @@ -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",
}


Expand All @@ -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()
Expand Down Expand Up @@ -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},
Expand All @@ -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."
),
}]

Expand All @@ -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,
},
},
)
Expand All @@ -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"]
Expand All @@ -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"]
Expand All @@ -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:
Expand All @@ -208,17 +237,19 @@ 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})
critic_history.append({
"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({
Expand Down
21 changes: 17 additions & 4 deletions scripts/agent_gan_repl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand All @@ -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()

Expand Down Expand Up @@ -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}."
),
},
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -176,15 +181,22 @@ 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}."
),
},
{
"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']}"
),
},
]
Expand All @@ -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(
Expand Down
9 changes: 8 additions & 1 deletion tests/inference_engine/bench/test_prefill_fleet_report.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
64 changes: 64 additions & 0 deletions tests/inference_engine/bridge/test_agent_gan_demo.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from scripts.agent_gan_inference_demo import (
_agent_cache_gate,
_infer,
_output_metadata,
)

Expand All @@ -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
8 changes: 8 additions & 0 deletions tests/inference_engine/bridge/test_agent_gan_repl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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"]
Loading