diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist
index 0bc8eb92..17c530e8 100644
--- a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist
+++ b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist
@@ -38,7 +38,7 @@
--cache-link-mbps10000
--cache-default-rtt-ms0.55
--remote-prefill-min-tokens0
- --prefill-worker-timeout-s300
+ --prefill-worker-timeout-s900
--prefill-policyremote-required
--network-http-host127.0.0.1
--network-http-port8090
diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md
index 7b55fae9..5f0623bf 100644
--- a/docs/ops/distributed-prefill-kv-network.md
+++ b/docs/ops/distributed-prefill-kv-network.md
@@ -335,10 +335,11 @@ 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.
+model's EOS. `--max-response-tokens` defaults to `0` (no client-side cap);
+operators may explicitly set a positive cap, which 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.
The REPL ignores external `SIGTERM`; its shell supervisor restarts signal-based
exits. Only `/quit`, `/exit`, or EOF is treated as approval to stop.
diff --git a/scripts/agent_gan_inference_demo.py b/scripts/agent_gan_inference_demo.py
index 8d894a3d..91720868 100644
--- a/scripts/agent_gan_inference_demo.py
+++ b/scripts/agent_gan_inference_demo.py
@@ -51,11 +51,19 @@ def _infer(
append_done = time.perf_counter()
first_at = None
generated = []
- response_limit = int(max_response_tokens or output_tokens)
+ response_limit = (
+ int(output_tokens)
+ if max_response_tokens is None
+ else int(max_response_tokens) or None
+ )
stop_reason = "unknown"
- while len(generated) < response_limit:
+ while response_limit is None or len(generated) < response_limit:
before_count = len(generated)
- chunk = min(output_tokens, response_limit - len(generated))
+ chunk = (
+ output_tokens
+ if response_limit is None
+ else 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:
@@ -73,7 +81,11 @@ def _infer(
if len(generated) == before_count:
stop_reason = "no_progress"
break
- if len(generated) >= response_limit and stop_reason == "max_tokens":
+ if (
+ response_limit is not None
+ and len(generated) >= response_limit
+ and stop_reason == "max_tokens"
+ ):
stop_reason = "client_safety_limit"
done = time.perf_counter()
after = get_stats()
@@ -106,7 +118,12 @@ 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(
+ "--max-response-tokens",
+ type=int,
+ default=0,
+ help="Optional client response cap; 0 means generate until model EOS.",
+ )
parser.add_argument("--report", default="/tmp/kakeya-agent-gan-demo.json")
parser.add_argument("--skip-ensure", action="store_true")
args = parser.parse_args()
diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py
index e9e8adcf..1c603e38 100644
--- a/scripts/agent_gan_repl.py
+++ b/scripts/agent_gan_repl.py
@@ -76,7 +76,12 @@ 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(
+ "--max-response-tokens",
+ type=int,
+ default=0,
+ help="Optional client response cap; 0 means generate until model EOS.",
+ )
parser.add_argument("--skip-ensure", action="store_true")
args = parser.parse_args()
diff --git a/tests/inference_engine/bridge/test_agent_gan_demo.py b/tests/inference_engine/bridge/test_agent_gan_demo.py
index 82edc4db..d339b295 100644
--- a/tests/inference_engine/bridge/test_agent_gan_demo.py
+++ b/tests/inference_engine/bridge/test_agent_gan_demo.py
@@ -67,7 +67,7 @@ def test_infer_continues_chunks_until_eos():
2,
lambda: {},
on_token=lambda values: streamed.append(list(values)),
- max_response_tokens=10,
+ max_response_tokens=0,
)
assert tokens == [1, 2, 3]
assert streamed == [[1], [1, 2], [1, 2, 3]]
diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py
index 3044b677..f305f806 100644
--- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py
+++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py
@@ -48,6 +48,10 @@ def test_two_mac_deployment_uses_allens_as_prefill_only():
"--prefill-policyremote-required"
in plist
)
+ assert (
+ "--prefill-worker-timeout-s900"
+ in plist
+ )
assert (
"--cache-tenant-idprivate-fleet"
in plist