From df0222383519a9dcd5ec00243d017ecc400f34d7 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 22:56:26 +0900 Subject: [PATCH] feat: enable cache-off single-block ATOM PD --- opt_prime/examples/pd/http_matrix/STATUS.md | 4 + .../pd/http_matrix/evidence236/events.jsonl | 1 + .../evidence236/koni-live-http.json | 415 ++++++++++++++++++ .../experiments/2026-10-23-live-http234.md | 44 ++ .../2026-10-23-single-block-pd236.md | 30 ++ .../accel/platforms/atom/batching.py | 21 +- opt_prime/tests/test_atom_batching.py | 85 +++- 7 files changed, 592 insertions(+), 8 deletions(-) create mode 100644 opt_prime/examples/pd/http_matrix/evidence236/events.jsonl create mode 100644 opt_prime/examples/pd/http_matrix/evidence236/koni-live-http.json create mode 100644 opt_prime/examples/pd/http_matrix/experiments/2026-10-23-live-http234.md create mode 100644 opt_prime/examples/pd/http_matrix/experiments/2026-10-23-single-block-pd236.md diff --git a/opt_prime/examples/pd/http_matrix/STATUS.md b/opt_prime/examples/pd/http_matrix/STATUS.md index 268b0f05..9f065c3a 100644 --- a/opt_prime/examples/pd/http_matrix/STATUS.md +++ b/opt_prime/examples/pd/http_matrix/STATUS.md @@ -13,3 +13,7 @@ Gemma full-prefix KV imports now run through ATOM PD scheduling: 33 CPU tests an Geometry recovery follow-up: malformed remote heads/KV lengths fail only their request; real device import failures still poison the scheduler.37related tests pass. Hardware evidence remains c76e03a. Long hardware follow-up on1e54786:6423–6428token inputs,10normal/reuse HTTP responses match nativeATOM;12remote/0local prefills including cancellation. Recorded RNGD KV replay, arithmetic wrong on both native/imported paths; no quality or live-transfer claim. + +## Single-block static PD (#236) + +Static B1 with one KV block now supports cache-off full-prefix PD. KONI live RNGD->ATOM HTTP14/14 output checks and16remote/0local prefills pass, including cancel/disconnect reuse.43CPU tests cover queued reuse, late cancellation and partial-prefix error recovery; cache/multi-lane guards remain. No new artifact, hardware batch or cache-restoration claim. See single-block-pd236 record and evidence236. diff --git a/opt_prime/examples/pd/http_matrix/evidence236/events.jsonl b/opt_prime/examples/pd/http_matrix/evidence236/events.jsonl new file mode 100644 index 00000000..00e3b45a --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/evidence236/events.jsonl @@ -0,0 +1 @@ +{"id": "3352f09cfbc74d869aefd49b13ae2cc7", "ts": "2026-10-23T22:53:49+09:00", "task": "Single-block ATOM PD (aicomp #236)", "result": "Static B1 cache-off KONI live PD enabled without extra KV block:14/14 HTTP outputs exact,16remote/0local; cancel/disconnect reuse pass.43CPU tests incl partial/late-cancel recovery; unsupported cache/multi-lane still rejected.", "next": "Publish scoped scheduler PR; continue medium live matrix and wider performance/cache coverage.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {"model": "KONI-8B", "batch": "1", "kv_blocks": "1"}, "artifacts": ["opt_prime/examples/pd/http_matrix/evidence236/koni-live-http.json"], "experiments": ["opt_prime/examples/pd/http_matrix/experiments/2026-10-23-single-block-pd236.md"], "labels": ["experiment"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "RNGD+ATOM host", "branch": "main", "commit": "351d39d-dirty"} diff --git a/opt_prime/examples/pd/http_matrix/evidence236/koni-live-http.json b/opt_prime/examples/pd/http_matrix/evidence236/koni-live-http.json new file mode 100644 index 00000000..b3a13143 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/evidence236/koni-live-http.json @@ -0,0 +1,415 @@ +{ + "worker": "live_rngd_subprocess_prefill", + "rows": [ + { + "case": "capital-0", + "request": "normal-0-0", + "response": { + "id": "normal-0-0", + "object": "text_completion", + "created": 1792763172, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 60704, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 138, + "completion_tokens": 2, + "total_tokens": 140 + } + }, + "latency_s": 0.4229255300015211 + }, + { + "case": "arithmetic-0", + "request": "normal-0-1", + "response": { + "id": "normal-0-1", + "object": "text_completion", + "created": 1792763173, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "17 + 25 = 42.", + "token_ids": [ + 1114, + 489, + 220, + 914, + 284, + 220, + 2983, + 13, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 137, + "completion_tokens": 9, + "total_tokens": 146 + } + }, + "latency_s": 0.7874407280469313 + }, + { + "case": "korean-0", + "request": "normal-0-2", + "response": { + "id": "normal-0-2", + "object": "text_completion", + "created": 1792763174, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "서울", + "token_ids": [ + 115978, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 139, + "completion_tokens": 2, + "total_tokens": 141 + } + }, + "latency_s": 0.3736360300099477 + }, + { + "case": "retrieval-0", + "request": "normal-0-3", + "response": { + "id": "normal-0-3", + "object": "text_completion", + "created": 1792763174, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 18082, + 877, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 143, + "completion_tokens": 3, + "total_tokens": 146 + } + }, + "latency_s": 0.41646945395041257 + }, + { + "case": "capital-0", + "request": "normal-1-0", + "response": { + "id": "normal-1-0", + "object": "text_completion", + "created": 1792763174, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 60704, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 138, + "completion_tokens": 2, + "total_tokens": 140 + } + }, + "latency_s": 0.3646502789342776 + }, + { + "case": "arithmetic-0", + "request": "normal-1-1", + "response": { + "id": "normal-1-1", + "object": "text_completion", + "created": 1792763175, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "17 + 25 = 42.", + "token_ids": [ + 1114, + 489, + 220, + 914, + 284, + 220, + 2983, + 13, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 137, + "completion_tokens": 9, + "total_tokens": 146 + } + }, + "latency_s": 0.8413291039178148 + }, + { + "case": "korean-0", + "request": "normal-1-2", + "response": { + "id": "normal-1-2", + "object": "text_completion", + "created": 1792763176, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "서울", + "token_ids": [ + 115978, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 139, + "completion_tokens": 2, + "total_tokens": 141 + } + }, + "latency_s": 0.3697471059858799 + }, + { + "case": "retrieval-0", + "request": "normal-1-3", + "response": { + "id": "normal-1-3", + "object": "text_completion", + "created": 1792763176, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 18082, + 877, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 143, + "completion_tokens": 3, + "total_tokens": 146 + } + }, + "latency_s": 0.44573811802547425 + }, + { + "case": "capital-0", + "request": "concurrent-0", + "response": { + "id": "concurrent-0", + "object": "text_completion", + "created": 1792763176, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 60704, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 138, + "completion_tokens": 2, + "total_tokens": 140 + } + }, + "latency_s": 0.36879166797734797 + }, + { + "case": "arithmetic-0", + "request": "concurrent-1", + "response": { + "id": "concurrent-1", + "object": "text_completion", + "created": 1792763177, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "17 + 25 = 42.", + "token_ids": [ + 1114, + 489, + 220, + 914, + 284, + 220, + 2983, + 13, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 137, + "completion_tokens": 9, + "total_tokens": 146 + } + }, + "latency_s": 1.2019858930725604 + }, + { + "case": "korean-0", + "request": "concurrent-2", + "response": { + "id": "concurrent-2", + "object": "text_completion", + "created": 1792763178, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "서울", + "token_ids": [ + 115978, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 139, + "completion_tokens": 2, + "total_tokens": 141 + } + }, + "latency_s": 1.5741613419959322 + }, + { + "case": "retrieval-0", + "request": "concurrent-3", + "response": { + "id": "concurrent-3", + "object": "text_completion", + "created": 1792763178, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 18082, + 877, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 143, + "completion_tokens": 3, + "total_tokens": 146 + } + }, + "latency_s": 2.007744933012873 + }, + { + "case": "capital-0", + "request": "after-cancel", + "response": { + "id": "after-cancel", + "object": "text_completion", + "created": 1792763179, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 60704, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 138, + "completion_tokens": 2, + "total_tokens": 140 + } + }, + "latency_s": 0.3682253049919382 + }, + { + "case": "capital-0", + "request": "after-disconnect", + "response": { + "id": "after-disconnect", + "object": "text_completion", + "created": 1792763180, + "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 60704, + 128009 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 138, + "completion_tokens": 2, + "total_tokens": 140 + } + }, + "latency_s": 0.3731464589945972 + } + ], + "stats_pd": { + "remote": 16, + "local": 0 + }, + "free": [ + 0 + ], + "cancellation_reuse": true +} diff --git a/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-live-http234.md b/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-live-http234.md new file mode 100644 index 00000000..6431c29c --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-live-http234.md @@ -0,0 +1,44 @@ +# Four-model live RNGD to ATOM HTTP PD + +- **Date**: 2026-10-23 +- **Task**: aicomp #234, single-block support #236 +- **Host / device**: owned paired RNGD+ATOM allocation, one of each accelerator; hardware batch1, CPU threads4. +- **Commit**: adapter branch3855050 plus current HTTP/KONI changes; AtomBatcher1e54786 for Gemma/Qwen, #236 working tree based69bca73 for KONI; HTTP SIGPIPE fix9e442cd loaded process-locally. +- **SDKs**: producer furiosa-torch2026.4.0 / torch2.10.0+cpu / transformers5.19.0; decoder optimum-rbln0.11.0.post1 / rebel-compiler0.11.0 / torch2.11.0+cpu / transformers5.8.1. + +## Question + +Can each validated producer handle real HTTP-triggered prefill through SubprocessPrefillWorker, shared-memory KV import and the production ATOM scheduler/owner/server, including queued clients and cancellation/reuse? + +## Setup + +No saved KV is loaded by check_live_http.py. It reads only input/EOS fixtures and native output baselines. Each request performs live RNGD prefill, returns KV through the existing shared-memory protocol, and decodes on ATOM. Four semantic questions, two sequential passes, one four-client concurrent pass, two streaming cancellation/disconnect requests and a successful reuse after each.14normal/reuse output checks and16remote prefills per model. Concurrency queues on a B1 lane; this is not hardware batch4. + +Gemma/Qwen artifacts use context9216/k3; KONI uses its existing static context9216/k1 artifact with the #236 scheduler fix. Prefix and host caches disabled. Corrected Gemma global-causal artifact from #225; pinned revisions are unchanged from #227. Token lists and EOS finish reason must equal native ATOM baselines. This criterion does not establish answer quality (known shared wrong answers remain in #227/#233). + +## Command + +```bash +bash run-live-http234.sh --batching atom_batching231.py --http-module llmd214.py --cache-dir atom-artifacts-global225 --fixtures semantic227-gemma-fixtures.json --cases capital-0 arithmetic-0 korean-0 retrieval-0 --worker-python "$WORKER_PYTHON" --engine-factory gemma_prefill:Prefill --width 128 --baseline semantic227-gemma-atom.jsonl --output live234-gemma-http.json +``` + +The isolated launcher runs check_live_http.py under the decoder interpreter/bootstrap. For Qwen use --model Qwen/Qwen3-0.6B orQwen/Qwen3-4B, --cache-dir atom-artifacts, --engine-factory prefill:Prefill and matching qwen06b/qwen4b fixtures/baselines/output. For KONI use --model KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024, --batching atom_batching236.py, --cache-dir atom-artifacts --kv-blocks 1, --engine-factory koni_prefill:Prefill --width 512, semantic227-koni-fixtures-eos.json and semantic227-koni-atom.jsonl. No installed SDK/core files changed. + +## Results + +| Model | Prompt tokens | Compiled width | Exact normal/reuse responses | Remote/local prefills | Cancel/disconnect reuse | +|---|---:|---:|---:|---:|---| +| Gemma3-1B | 23–28 | 128 | 14/14 | 16/0 | pass | +| Qwen3-0.6B | 26–31 | 128 | 14/14 | 16/0 | pass | +| Qwen3-4B | 26–31 | 128 | 14/14 | 16/0 | pass | +| KONI-8B | 137–143 | 512 | 14/14 | 16/0 | pass | + +Total56exact responses and64actual RNGD prefills, zero local fallbacks. Each final free-lane list is[0]. Client latency observations are included in JSON but this is not a matched native-vs-PD benchmark or throughput claim. The previous vendor SIGPIPE backtraces do not appear with the separately tested #214 module; the ordinary tokenizer cleanup warning remains on KONI. + +Supporting CPU tests: #234 adapter/geometry22tests plus5subtests; #236 scheduler43tests, one existing unfused attention warning. Before #236, single-block PD fails at construction (1failed/3passed); cache/multi-lane rejection tests remain green. KONI raw config omits head_dim: worker infers128 only when absent and preserves explicit Qwen/Gemma head dimensions. + +## Conclusion + +All four models now have short-input real RNGD->ATOM HTTP evidence, replacing replay-only coverage for this condition. Remaining: live medium/long prefixes, prefill-time cancellation/fault recovery, prefix/host-cache restoration, hardware batch coverage, true compute-transfer overlap and controlled serving performance comparison. A width2048 queue for the four *-100 fixtures has started in the owned paired allocation; results are pending. Overall goal remains active. + +Raw evidence: analysis/serving-matrix-201/live234-{gemma,qwen06b,qwen4b,koni}-http.json/.log and live234-http-summary.json. Pending queue: live234-mid-queue.log. diff --git a/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-single-block-pd236.md b/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-single-block-pd236.md new file mode 100644 index 00000000..2badb0a4 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-single-block-pd236.md @@ -0,0 +1,30 @@ +# Static single-block ATOM PD + +- **Date**: 2026-10-23 +- **Task**: aicomp #236 +- **Host / device**: owned RNGD+ATOM allocation, KONI-8B static B1/context9216/k1 artifact; SDK versions in live-http234 record. +- **Commit**: feat/236-atom-single-pd working tree based69bca73; final changes preserve the hardware-tested behavior, with only explanatory text changed after the run. + +## Question + +Does static B1 cache-off remote PD require a separate idle-lane scratch block? + +## Setup + +For exactly one block and one non-hybrid lane, with prefix/host caches disabled, allow block0 to be assigned and released per request. No active/idle lane pair exists, so there is no concurrent scratch access. All other scratch/cache requirements remain. Require a full prefix from position0; invalid partial results fail per request. + +## Command + +```bash +PYTHONPATH=opt_prime python -m pytest -q opt_prime/tests/test_atom_batching.py opt_prime/tests/test_atom_hybrid.py opt_prime/tests/test_pd_recovery.py opt_prime/tests/test_serve_batched.py +``` + +Hardware command and pinned checkpoint are in live-http234; the KONI row uses this candidate and its existing k1 artifact. + +## Results + +Initial regression1failed/3passed at the existing PD capacity guard. Final43tests pass (one existing unfused attention warning), including queued request parity/reuse, cancellation while prefill pending and ignoring its late result, malformed partial result recovery, and rejection of cache/multi-lane combinations. Real KONI HTTP:14normal/reuse responses token-identical through EOS,16remote/0local prefills, cancellation/disconnect followed by successful reuse; free lane[0]. No new artifact or extra KV block was required. + +## Conclusion + +Static B1 cache-off PD is supported with one block. This is not prefix-cache support, hybrid single-block support or hardware batch>1. The broader four-model goal remains active. Evidence: single236-red.log, single236-final.log and live234-koni-http.json in the local serving-matrix-201 analysis directory. diff --git a/opt_prime/opt_prime/accel/platforms/atom/batching.py b/opt_prime/opt_prime/accel/platforms/atom/batching.py index 42f04cb7..1fa67d0b 100644 --- a/opt_prime/opt_prime/accel/platforms/atom/batching.py +++ b/opt_prime/opt_prime/accel/platforms/atom/batching.py @@ -29,6 +29,8 @@ Hybrid fresh PD imports global layers into their assigned block and the latest sliding window into the request lane. Source dtype metadata may differ from physical storage, so staging validates allocation size. Hybrid prefix and host-cache restoration remain unsupported; remote hybrid requests must return their full prompt KV. +Static B1 cache-off PD may use its sole block: it never decodes an idle lane alongside an active lane. +This path requires full prompt KV and releases the block after each request; it retains no prefix entries. Idle lanes then decode on a reserved scratch block. Measured on ATOM (Llama-3.2-1B, 600-token prompt): full prefill 85.5 ms; device copy of a 512-token prefix 0.8 ms + 17.5 ms suffix prefill; host reload 8.9 ms; tokens identical. """ @@ -84,12 +86,14 @@ def __init__(self, model, lanes, *, prefix_cache=True, host_bytes=0, medium_devi self.hybrid_cache = getattr(rc, "cache_impl", None) == "hybrid" self.unit, self.lanes = rc.prefill_chunk_size, lanes self.blocks = rc.kvcache_num_blocks - if self.blocks < lanes + 1: + self.single_block = self.blocks == lanes == 1 and not (prefix_cache or host_bytes or self.hybrid_cache) + if self.blocks < lanes + 1 and not self.single_block: raise ValueError("own block tables need kv_blocks >= lanes + 1 (a scratch block for idle lanes)") if prefix_cache and self.blocks < lanes + 2: # else a prefix hit may find only its own source to evict raise ValueError("prefix caching needs kv_blocks >= lanes + 2 (a scratch block and an evictable entry)") self.scratch = self.blocks - 1 - self.free = list(range(self.blocks - 1)) + # B1 cache-off PD never decodes an idle lane beside an active one. + self.free = list(range(self.blocks if self.single_block else self.blocks - 1)) self.lane_block = [None] * lanes self.prefix_cache = prefix_cache self.entries = OrderedDict() # block -> prompt unit hashes (unreferenced, least recent first) @@ -246,8 +250,8 @@ def import_span(self, lane, start, keys, values): if keys.ndim != 4 or values.ndim != 4: raise InvalidKVImport("KV export must have layer/head/token/dimension axes") block, n = self.lane_block[lane], keys.shape[2] - if self.hybrid_cache and start: - raise InvalidKVImport("hybrid PD requires a full prefix from position zero") + if (self.hybrid_cache or self.single_block) and start: + raise InvalidKVImport("hybrid or single-block PD requires a full prefix from position zero") if keys.shape != values.shape or len(self.metas) != keys.shape[0] * 2: raise InvalidKVImport("KV export geometry does not match compiled cache") stage = self._staging() @@ -325,7 +329,10 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib= self._remote_s = self._local_s = None # moving averages (seconds per prompt) self.local_prefill = local_prefill self.blocks = None - if (getattr(rc, "kvcache_num_blocks", self.lanes) or self.lanes) > self.lanes: + blocks = getattr(rc, "kvcache_num_blocks", self.lanes) or self.lanes + single_pd = blocks == self.lanes == 1 and prefill_worker is not None and not ( + prefix_cache or host_kv_gib or self.hybrid_cache) + if blocks > self.lanes or single_pd: self.blocks = AtomKVBlocks(model, self.lanes, prefix_cache=prefix_cache, host_bytes=int(host_kv_gib * 2 ** 30)) self.block_size = self.blocks.unit @@ -480,8 +487,8 @@ def _pd_step(self): self.prefilling.remove(r) try: out = r.pending.result() - if self.hybrid_cache and (out.start != 0 or out.keys.shape[2] != len(r.request.input_ids)): - raise ValueError("hybrid PD requires the full prompt KV") + if (self.hybrid_cache or self.blocks.single_block) and (out.start != 0 or out.keys.shape[2] != len(r.request.input_ids)): + raise ValueError("hybrid or single-block PD requires the full prompt KV") except BaseException as error: # this request fails; the lanes go on with self._lock: self._finish(r, error) diff --git a/opt_prime/tests/test_atom_batching.py b/opt_prime/tests/test_atom_batching.py index 8359c822..2434a808 100644 --- a/opt_prime/tests/test_atom_batching.py +++ b/opt_prime/tests/test_atom_batching.py @@ -180,7 +180,7 @@ def _decode(self, input_ids, cache_position, block_tables): logits = torch.zeros(B, 1, self.config.vocab_size) for lane in range(B): block, pos = int(block_tables[lane, 0]), int(cache_position[lane, 0]) - if block == self.rbln_config.kvcache_num_blocks - 1: + if self.rbln_config.kvcache_num_blocks > 1 and block == self.rbln_config.kvcache_num_blocks - 1: continue # scratch block: idle lane with torch.inference_mode(): out = self.m(input_ids[lane:lane + 1], past_key_values=self._cache(block, pos), use_cache=True, @@ -351,3 +351,86 @@ def test_pd_remote_only_with_host_tier_events_and_metrics(fake_rebel): flat = [e for batch in events.batches for e in batch[1]] assert any(e["type"] == "BlockStored" and e["medium"] == "CPU" for e in flat) worker.close() + + +@pytest.mark.parametrize("partial", [False, True]) +def test_single_block_pd_reuses_full_prefix_without_cache(fake_rebel, partial): + import copy + from opt_prime.accel.pd.pd_batching import PrefillWorker, TorchPrefillEngine + model = _llama() + fake = _blocks_model(fake_rebel, model, batch=1, blocks=1) + engine = TorchPrefillEngine(copy.deepcopy(model), max_context=128, chunk=16, page_size=16) + original = engine.prefill + calls = 0 + def prefill(request, send_from=0): + nonlocal calls + out = original(request, send_from) + calls += 1 + if partial and calls == 1: + out.start = 1 + out.keys, out.values = out.keys[:, :, 1:], out.values[:, :, 1:] + return out + engine.prefill = prefill + worker = PrefillWorker(engine) + try: + b = AtomBatcher(fake, prefix_cache=False, prefill_worker=worker, local_prefill=False) + prompts = _prompts(4, 20, 70, seed=236) + fs = [b.submit(PDRequest(str(i), p, 5)) for i,p in enumerate(prompts)] + b.run_until_idle() + for i,(p,f) in enumerate(zip(prompts,fs)): + if partial and i == 0: + with pytest.raises(ValueError, match="full prompt"): + f.result() + else: + assert f.result() == _greedy(model,p,5) + assert b.stats_pd == {"remote":4,"local":0} + assert b.blocks.free == [0] and b.blocks.lane_block == [None] + assert not fake.calls and not b.by_id and not b.poisoned + finally: + worker.close() + + +@pytest.mark.parametrize("batch,blocks,prefix,host", [(1,1,True,0),(1,1,False,1),(2,2,False,0)]) +def test_single_block_pd_keeps_cache_and_multilane_guards(fake_rebel,batch,blocks,prefix,host): + model = _llama() + fake = _blocks_model(fake_rebel,model,batch=batch,blocks=blocks) + with pytest.raises(ValueError): + AtomBatcher(fake,prefix_cache=prefix,host_kv_gib=host,prefill_worker=object(),local_prefill=False) + + +def test_single_block_pd_cancel_late_result_and_reuse(fake_rebel): + import copy + from concurrent.futures import Future + from opt_prime.accel.pd.pd_batching import TorchPrefillEngine + + class Worker: + def __init__(self): + self.pending = [] + + def submit(self, request, send_from=0, slot=None): + assert send_from == slot == 0 + f = Future() + self.pending.append((request, f)) + return f + + model = _llama() + fake = _blocks_model(fake_rebel, model, batch=1, blocks=1) + worker = Worker() + engine = TorchPrefillEngine(copy.deepcopy(model), max_context=128, chunk=16, page_size=16) + b = AtomBatcher(fake, prefix_cache=False, prefill_worker=worker, local_prefill=False) + p,q = _prompts(2,20,35) + gone = b.submit(PDRequest('gone',p,5)) + b.step() + assert b.cancel('gone') + b.step() + with pytest.raises(RequestCancelled): + gone.result() + good = b.submit(PDRequest('next',q,5)) + b.step() + for request,f in worker.pending: + if not f.done(): + f.set_result(engine.prefill(request)) + b.run_until_idle() + assert good.result() == _greedy(model,q,5) + assert b.blocks.free == [0] and b.blocks.lane_block == [None] + assert not b.by_id and not b.poisoned