diff --git a/opt_prime/examples/pd/http_matrix/STATUS.md b/opt_prime/examples/pd/http_matrix/STATUS.md index 3464a402..268b0f05 100644 --- a/opt_prime/examples/pd/http_matrix/STATUS.md +++ b/opt_prime/examples/pd/http_matrix/STATUS.md @@ -5,3 +5,11 @@ State: review; full goal remains active. GPU evidence remains as recorded. ATOM ## Production EOS correction (#229) The common serve_batched factory now unions model generation EOS and tokenizer turn-end EOS. KONI config128001/tokenizer128009 previously continued into role delimiters; real GPU HTTP negative control ends length, fixed/stream/reuse end stop, ignore_eos remains length.22related CPU tests pass after3failed factory regressions. GPU and ATOM hardware both pass5/5 EOS/override/stream/reuse checks; NPU uses existing B1/context9216/k1 artifact. See serving-eos229 record/evidence229. + +## Fresh hybrid PD (#231) + +Gemma full-prefix KV imports now run through ATOM PD scheduling: 33 CPU tests and 18 real HTTP normal/reuse responses pass, with 20 remote / 0 local prefills including cancellations. Producer data is recorded RNGD KV replay. Live prefill, overlap, hybrid cache restoration and hardware batch >1 remain unverified. SIGPIPE log issue still depends on #214/#215. See hybrid-pd231 experiment and evidence231. + +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. diff --git a/opt_prime/examples/pd/http_matrix/evidence231/events.jsonl b/opt_prime/examples/pd/http_matrix/evidence231/events.jsonl new file mode 100644 index 00000000..fc3bf9e3 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/evidence231/events.jsonl @@ -0,0 +1,3 @@ +{"id": "2e762726f70745af91068ca493f9cdb6", "ts": "2026-10-23T21:56:32+09:00", "task": "Fresh hybrid PD scheduler (aicomp #231)", "result": "33 CPU tests pass; Gemma ATOM HTTP replay18 responses exact,20 remote/0 local prefills; cancellation and disconnect reuse pass. Known SIGPIPE logs remain #214.", "next": "Publish scoped PR; integrate live producer and measure overlap after long-context validation.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {"model": "Gemma3-1B", "batch": "1", "worker": "recorded_rngd_kv_replay"}, "artifacts": ["opt_prime/examples/pd/http_matrix/evidence231/gemma-http-final.json"], "experiments": ["opt_prime/examples/pd/http_matrix/experiments/2026-10-23-hybrid-pd231.md"], "labels": ["experiment"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "RNGD+ATOM host", "branch": "main", "commit": "351d39d-dirty"} +{"id": "c5fa00d2033c4c358b49d3d43194e104", "ts": "2026-10-23T22:01:44+09:00", "task": "Fresh hybrid PD scheduler (aicomp #231)", "result": "Reproduced wrong-head/KV-length import poisoning (2failed); request-scoped validation fix passes37tests, real device failures still poison. PR232 follow-up.", "next": "Update PR232 and finish live long-context queue before live heterogeneous PD.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {"validation": "cpu_recovery"}, "artifacts": [], "experiments": ["opt_prime/examples/pd/http_matrix/experiments/2026-10-23-hybrid-pd231.md"], "labels": ["bug"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "validation host", "branch": "main", "commit": "351d39d-dirty"} +{"id": "c0b7d30efdee4c99b1cae96743f5f6a0", "ts": "2026-10-23T22:12:20+09:00", "task": "EOS-aware four-model validation (aicomp #227)", "result": "Gemma6.4K GPU3/4,ATOM4/4; same equations GPU BF16/FP32 choose3 while RNGD2; growing layer KV error tracked#233. Long PD HTTP replay10responses/12remote verified on1e54786. Qwen4B PID131672 compiling; idleATOM released.", "next": "Finish Qwen4B semantic GPU/ATOM checks; isolate Gemma per-layer execution drift#233 and integrate live producers.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {"model": "Gemma3-1B", "width": "8192"}, "artifacts": [], "experiments": ["opt_prime/examples/pd/http_matrix/experiments/2026-10-23-hybrid-pd231.md"], "labels": ["analysis"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "validation host", "branch": "main", "commit": "351d39d-dirty"} diff --git a/opt_prime/examples/pd/http_matrix/evidence231/gemma-http-final.json b/opt_prime/examples/pd/http_matrix/evidence231/gemma-http-final.json new file mode 100644 index 00000000..bfc46cba --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/evidence231/gemma-http-final.json @@ -0,0 +1,493 @@ +{ + "worker": "recorded_rngd_kv_replay_not_live_prefill", + "rows": [ + { + "case": "capital-0", + "request": "normal-0-0", + "response": { + "id": "normal-0-0", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 23, + "completion_tokens": 2, + "total_tokens": 25 + } + } + }, + { + "case": "capital-100", + "request": "normal-0-1", + "response": { + "id": "normal-0-1", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1623, + "completion_tokens": 2, + "total_tokens": 1625 + } + } + }, + { + "case": "arithmetic-0", + "request": "normal-0-2", + "response": { + "id": "normal-0-2", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "42\n", + "token_ids": [ + 236812, + 236778, + 107, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 24, + "completion_tokens": 4, + "total_tokens": 28 + } + } + }, + { + "case": "arithmetic-100", + "request": "normal-0-3", + "response": { + "id": "normal-0-3", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "40\n", + "token_ids": [ + 236812, + 236771, + 107, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1624, + "completion_tokens": 4, + "total_tokens": 1628 + } + } + }, + { + "case": "korean-0", + "request": "normal-0-4", + "response": { + "id": "normal-0-4", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "서울특별시", + "token_ids": [ + 222340, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 24, + "completion_tokens": 2, + "total_tokens": 26 + } + } + }, + { + "case": "korean-100", + "request": "normal-0-5", + "response": { + "id": "normal-0-5", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "서울", + "token_ids": [ + 98665, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1624, + "completion_tokens": 2, + "total_tokens": 1626 + } + } + }, + { + "case": "retrieval-0", + "request": "normal-0-6", + "response": { + "id": "normal-0-6", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 28487, + 1618, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 28, + "completion_tokens": 3, + "total_tokens": 31 + } + } + }, + { + "case": "retrieval-100", + "request": "normal-0-7", + "response": { + "id": "normal-0-7", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 28487, + 1618, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1628, + "completion_tokens": 3, + "total_tokens": 1631 + } + } + }, + { + "case": "capital-0", + "request": "normal-1-0", + "response": { + "id": "normal-1-0", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 23, + "completion_tokens": 2, + "total_tokens": 25 + } + } + }, + { + "case": "capital-100", + "request": "normal-1-1", + "response": { + "id": "normal-1-1", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1623, + "completion_tokens": 2, + "total_tokens": 1625 + } + } + }, + { + "case": "arithmetic-0", + "request": "normal-1-2", + "response": { + "id": "normal-1-2", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "42\n", + "token_ids": [ + 236812, + 236778, + 107, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 24, + "completion_tokens": 4, + "total_tokens": 28 + } + } + }, + { + "case": "arithmetic-100", + "request": "normal-1-3", + "response": { + "id": "normal-1-3", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "40\n", + "token_ids": [ + 236812, + 236771, + 107, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1624, + "completion_tokens": 4, + "total_tokens": 1628 + } + } + }, + { + "case": "korean-0", + "request": "normal-1-4", + "response": { + "id": "normal-1-4", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "서울특별시", + "token_ids": [ + 222340, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 24, + "completion_tokens": 2, + "total_tokens": 26 + } + } + }, + { + "case": "korean-100", + "request": "normal-1-5", + "response": { + "id": "normal-1-5", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "서울", + "token_ids": [ + 98665, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1624, + "completion_tokens": 2, + "total_tokens": 1626 + } + } + }, + { + "case": "retrieval-0", + "request": "normal-1-6", + "response": { + "id": "normal-1-6", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 28487, + 1618, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 28, + "completion_tokens": 3, + "total_tokens": 31 + } + } + }, + { + "case": "retrieval-100", + "request": "normal-1-7", + "response": { + "id": "normal-1-7", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 28487, + 1618, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1628, + "completion_tokens": 3, + "total_tokens": 1631 + } + } + }, + { + "case": "capital-0", + "request": "after-cancel", + "response": { + "id": "after-cancel", + "object": "text_completion", + "created": 1792759826, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 23, + "completion_tokens": 2, + "total_tokens": 25 + } + } + }, + { + "case": "capital-0", + "request": "after-disconnect", + "response": { + "id": "after-disconnect", + "object": "text_completion", + "created": 1792759827, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 23, + "completion_tokens": 2, + "total_tokens": 25 + } + } + } + ], + "stats_pd": { + "remote": 20, + "local": 0 + }, + "free": [ + 0 + ], + "cancellation_reuse": true +} diff --git a/opt_prime/examples/pd/http_matrix/evidence231/gemma-http-long.json b/opt_prime/examples/pd/http_matrix/evidence231/gemma-http-long.json new file mode 100644 index 00000000..fd156171 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/evidence231/gemma-http-long.json @@ -0,0 +1,277 @@ +{ + "worker": "recorded_rngd_kv_replay_not_live_prefill", + "rows": [ + { + "case": "capital-400", + "request": "normal-0-0", + "response": { + "id": "normal-0-0", + "object": "text_completion", + "created": 1792760865, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6423, + "completion_tokens": 2, + "total_tokens": 6425 + } + } + }, + { + "case": "arithmetic-400", + "request": "normal-0-1", + "response": { + "id": "normal-0-1", + "object": "text_completion", + "created": 1792760865, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "25", + "token_ids": [ + 236778, + 236810, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6424, + "completion_tokens": 3, + "total_tokens": 6427 + } + } + }, + { + "case": "korean-400", + "request": "normal-0-2", + "response": { + "id": "normal-0-2", + "object": "text_completion", + "created": 1792760865, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "서울", + "token_ids": [ + 98665, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6424, + "completion_tokens": 2, + "total_tokens": 6426 + } + } + }, + { + "case": "retrieval-400", + "request": "normal-0-3", + "response": { + "id": "normal-0-3", + "object": "text_completion", + "created": 1792760866, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 28487, + 1618, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6428, + "completion_tokens": 3, + "total_tokens": 6431 + } + } + }, + { + "case": "capital-400", + "request": "normal-1-0", + "response": { + "id": "normal-1-0", + "object": "text_completion", + "created": 1792760866, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6423, + "completion_tokens": 2, + "total_tokens": 6425 + } + } + }, + { + "case": "arithmetic-400", + "request": "normal-1-1", + "response": { + "id": "normal-1-1", + "object": "text_completion", + "created": 1792760866, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "25", + "token_ids": [ + 236778, + 236810, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6424, + "completion_tokens": 3, + "total_tokens": 6427 + } + } + }, + { + "case": "korean-400", + "request": "normal-1-2", + "response": { + "id": "normal-1-2", + "object": "text_completion", + "created": 1792760866, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "서울", + "token_ids": [ + 98665, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6424, + "completion_tokens": 2, + "total_tokens": 6426 + } + } + }, + { + "case": "retrieval-400", + "request": "normal-1-3", + "response": { + "id": "normal-1-3", + "object": "text_completion", + "created": 1792760866, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "MAPLE", + "token_ids": [ + 28487, + 1618, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6428, + "completion_tokens": 3, + "total_tokens": 6431 + } + } + }, + { + "case": "capital-400", + "request": "after-cancel", + "response": { + "id": "after-cancel", + "object": "text_completion", + "created": 1792760867, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6423, + "completion_tokens": 2, + "total_tokens": 6425 + } + } + }, + { + "case": "capital-400", + "request": "after-disconnect", + "response": { + "id": "after-disconnect", + "object": "text_completion", + "created": 1792760867, + "model": "unsloth/gemma-3-1b-it", + "choices": [ + { + "index": 0, + "text": "Paris", + "token_ids": [ + 50429, + 106 + ], + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 6423, + "completion_tokens": 2, + "total_tokens": 6425 + } + } + } + ], + "stats_pd": { + "remote": 12, + "local": 0 + }, + "free": [ + 0 + ], + "cancellation_reuse": true +} diff --git a/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-hybrid-pd231.md b/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-hybrid-pd231.md new file mode 100644 index 00000000..e21da98a --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-hybrid-pd231.md @@ -0,0 +1,47 @@ +# Fresh hybrid KV imports through ATOM PD scheduling + +- **Date**: 2026-10-23 +- **Task**: aicomp #231; parent #201/#227 +- **Host / device**: RNGD+ATOM host, exclusively allocated one ATOM, Gemma3-1B B1/context9216/k3; existing corrected global-causal artifact from #225. +- **Commit**: base f84823a plus the #231 working tree; RNGD export producer 5e19f0e. + +## Question + +Can the production ATOM PD scheduler import a fresh full hybrid KV prefix, decode correctly and recycle its lane after cancellation? + +## Setup + +Global attention slots use the assigned cache block. Sliding slots use the request lane and retain the latest window. Physical allocation size determines staging dtype because source float32 metadata can describe two-byte DLFloat16 storage. Reset the request lane's decode mask and reject partial hybrid prefills. Prefix/host-cache restoration remains unsupported. + +Hardware validation uses recorded RNGD KV exports for four questions at short and medium lengths (23–28 and 1623–1628 tokens), replayed through PrefillWorker into the actual production batcher, owner and HTTP server. This is not live RNGD prefill or a transport benchmark. Checkpoint revision 5b11413a10db4e486ef16a20101fd028f8f2499c. The existing artifact and isolated SDK bootstrap are those recorded in #225/#227; no shared SDK installation was modified. + +## Command + +```bash +PYTHONPATH=opt_prime python -m pytest -q opt_prime/tests/test_atom_hybrid.py opt_prime/tests/test_atom_batching.py opt_prime/tests/test_pd_recovery.py opt_prime/tests/test_serve_batched.py +python opt_prime/examples/pd/validate_hybrid_pd.py --batching opt_prime/opt_prime/accel/platforms/atom/batching.py --cache-dir "$CORRECTED_GEMMA_ARTIFACT_ROOT" --exports semantic227-gemma-rngd128-0.pt semantic227-gemma-rngd2048-0.pt semantic227-gemma-rngd128-1.pt semantic227-gemma-rngd2048-1.pt semantic227-gemma-rngd128-2.pt semantic227-gemma-rngd2048-2.pt semantic227-gemma-rngd128-3.pt semantic227-gemma-rngd2048-3.pt --baseline semantic227-gemma-atom.jsonl --output hybrid231-gemma-http-final.json +``` + +Run the second command through the single-device bootstrap when the container exposes a remapped ATOM allocation. Exports are trusted outputs of #227; they are not downloaded executable artifacts. + +## Results + +- Initial regressions: 2 failed / 4 passed. Partial-result recovery subsequently failed until validation was moved into per-request error handling. +- Final related CPU suite: 33 passed, one existing unfused attention warning. Includes two-lane reuse, physical staging geometry, late cancelled future and malformed partial prefix recovery. +- Final hardware rerun: 18 normal/reuse responses match native tokens and stop at EOS. Counters: 20 remote, 0 local prefills including cancellation requests. Explicit cancellation and socket disconnection both drain requests and allow immediate reuse; free lane list ends [0]. +- Vendor native SIGPIPE backtraces still appear on disconnected streams, although the process survives and reuse checks pass. This is tracked separately by #214 / PR #215, whose fix is not included in this stack. Logs are not clean. +- Ruff on added validator and changed tests, and git diff --check: pass. + +Raw local evidence: analysis/serving-matrix-201/hybrid231-{red,partial-red,tests-full,gemma-http-final}.log and hybrid231-gemma-http-final.json. Selected public result: ../evidence231/gemma-http-final.json. + +## Conclusion + +Fresh full-prefix hybrid PD is now exercised through the production scheduler with recorded RNGD exports. This does not establish live producer integration, chunked overlap, hardware batch >1, cache restoration or serving speedup. The four-model goal remains active. + +## Geometry-error recovery follow-up + +Two new regressions (wrong head count and mismatched K/V length) failed because import validation escaped the per-request handler and poisoned the scheduler. Distinguish InvalidKVImport, raised before device writes, from real device/runtime failures. Reject only the malformed request and retain lane reuse. RuntimeError and ValueError from device import still poison the scheduler. Final selected suite:37passed, one existing warning; scoped Ruff and diff checks pass. No additional hardware run was performed for this exception-only follow-up; the hardware evidence above belongs to c76e03a. Local logs: hybrid231-geometry-red.log (2failed/1passed), hybrid231-geometry-final.log (37passed). + +## Long-input hardware follow-up (1e54786) + +Latest exception-recovery code was copied into the isolated runtime and validated with the same HTTP harness, replacing the export list by semantic227-gemma-rngd8192-{0,1,2,3}.pt and output by hybrid231-gemma-http-long.json. Four inputs6423–6428tokens, two normal passes, cancellation/disconnect and immediate reuse:10normal/reuse responses match native ATOM through EOS;12remote/0local prefills including cancellation requests; free lane[0]. Native and imported arithmetic both yield25 (wrong), so this is behavior parity, not answer quality. Vendor SIGPIPE traces remain the separately tracked #214/#215 issue. Producer remains recorded RNGD KV replay; no live transfer/overlap claim. diff --git a/opt_prime/examples/pd/validate_hybrid_pd.py b/opt_prime/examples/pd/validate_hybrid_pd.py new file mode 100644 index 00000000..a5b3c1db --- /dev/null +++ b/opt_prime/examples/pd/validate_hybrid_pd.py @@ -0,0 +1,154 @@ +"""Real ATOM scheduler/HTTP hybrid PD validation using recorded RNGD KV, not live prefill.""" + +import argparse +import importlib.util +import json +import sys +import threading +import time +import urllib.request +from pathlib import Path + +import torch + +p = argparse.ArgumentParser(description=__doc__) +p.add_argument("--batching", required=True) +p.add_argument("--cache-dir", required=True) +p.add_argument("--exports", nargs="+", required=True) +p.add_argument("--baseline", required=True) +p.add_argument("--output", required=True) +a = p.parse_args() +name = "opt_prime.accel.platforms.atom.batching" +spec = importlib.util.spec_from_file_location(name, a.batching) +module = importlib.util.module_from_spec(spec) +sys.modules[name] = module +spec.loader.exec_module(module) +from opt_prime.accel.pd.pd_batching import PrefillOut, PrefillWorker +from opt_prime.accel.platforms.atom.engine import AtomRblnEngine +from opt_prime.accel.serve.batched import BatchedOwner, BatchedService +from opt_prime.accel.serve.llmd import serve_stage +from transformers import AutoTokenizer + +exports = [torch.load(f, weights_only=True) for f in a.exports] +lookup = {tuple(r["ids"]): r for r in exports} +baseline = { + r["case"]: r["tokens"] + for r in map(json.loads, Path(a.baseline).read_text().splitlines()) +} + + +class Replay: + def prefill(self, request, send_from=0): + assert send_from == 0 + r = lookup[request.input_ids] + return PrefillOut(int(r["logits"].argmax()), 0, r["keys"], r["values"]) + + +nodes = list(Path("/dev").glob("rbln[0-9]*")) +assert len(nodes) == 1 +engine = AtomRblnEngine( + devices=[int(nodes[0].name[4:])], + max_seq_len=9216, + batch_size=1, + kv_blocks=3, + cache_dir=a.cache_dir, +) +worker = PrefillWorker(Replay()) +owner = thread = server = http = None +rows = [] +try: + model = "unsloth/gemma-3-1b-it" + engine.start(model) + tok = AutoTokenizer.from_pretrained(model, local_files_only=True) + b = engine.batcher(prefix_cache=False, prefill_worker=worker, local_prefill=False) + service = BatchedService( + b, + model=model, + vocab_size=engine.model.config.vocab_size, + max_output=64, + render=lambda ids: tok.decode(ids, skip_special_tokens=True), + tokenizer=tok, + eos_token_ids=tuple(exports[0]["eos_ids"]), + ) + owner = BatchedOwner(service) + thread = threading.Thread(target=owner.run, daemon=True) + thread.start() + server, http = serve_stage(owner, "decode", 0) + + def request(path, body, rid): + req = urllib.request.Request( + f"http://127.0.0.1:{server.server_port}{path}", + data=json.dumps(body).encode(), + headers={"x-request-id": rid}, + ) + return urllib.request.urlopen(req, timeout=60) + + def complete(r, rid): + with request( + "/v1/completions", + {"model": model, "prompt": r["ids"], "max_tokens": 64}, + rid, + ) as response: + out = json.load(response) + assert out["choices"][0]["finish_reason"] == "stop" + assert out["choices"][0]["token_ids"] == baseline[r["metadata"]["case"]], out + rows.append({"case": r["metadata"]["case"], "request": rid, "response": out}) + + # Alternate long and short prefixes while blocks rotate under the same local lane. + for repeat in range(2): + for i, r in enumerate(exports): + complete(r, f"normal-{repeat}-{i}") + for mode in ["cancel", "disconnect"]: + r = exports[-1] + with request( + "/v1/completions", + { + "model": model, + "prompt": r["ids"], + "max_tokens": 64, + "stream": True, + "ignore_eos": True, + }, + mode, + ) as response: + for raw in response: + if raw.startswith(b"data: ") and b"[DONE]" not in raw: + event = json.loads(raw[6:]) + if event.get("choices") and event["choices"][0].get("text"): + break + else: + raise AssertionError("no streamed text") + if mode == "cancel": + with request("/v1/cancel", {}, mode) as cancelled: + assert cancelled.status == 202 + end = time.monotonic() + 10 + while mode in b.by_id and time.monotonic() < end: + time.sleep(0.01) + assert mode not in b.by_id + complete(exports[0], "after-" + mode) + assert b.stats_pd["local"] == 0 and sorted(b.free) == [0] + result = { + "worker": "recorded_rngd_kv_replay_not_live_prefill", + "rows": rows, + "stats_pd": b.stats_pd, + "free": b.free, + "cancellation_reuse": True, + } + Path(a.output).write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n") + print( + json.dumps( + {"requests": len(rows), "stats_pd": b.stats_pd, "cancellation_reuse": True} + ), + flush=True, + ) +finally: + if owner: + owner.submit("stop") + if thread: + thread.join(15) + if server: + server.shutdown() + server.server_close() + http.join(5) + worker.close() + engine.close() diff --git a/opt_prime/opt_prime/accel/platforms/atom/batching.py b/opt_prime/opt_prime/accel/platforms/atom/batching.py index 2be40067..42f04cb7 100644 --- a/opt_prime/opt_prime/accel/platforms/atom/batching.py +++ b/opt_prime/opt_prime/accel/platforms/atom/batching.py @@ -26,6 +26,9 @@ - a host-DRAM tier: device entries evicted for space are copied to host memory (per-layer `get_kv_cache_for_layer` through a 2 MiB-aligned staging buffer that mirrors the layer's block layout) and written back on a later hit; - KV events (BlockStored / BlockRemoved, device and CPU tiers) and prefix-hit counters, as for the GPU pool. +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. 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. """ @@ -39,6 +42,10 @@ from opt_prime.accel.pd.sampling import select_token +class InvalidKVImport(ValueError): + """Invalid remote KV rejected before any device write; the lane remains reusable.""" + + _DLFLOAT = None @@ -74,6 +81,7 @@ def __init__(self, model, lanes, *, prefix_cache=True, host_bytes=0, medium_devi from opt_prime.accel.pd.paged import block_hashes self._hashes = block_hashes rc = model.rbln_config + 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: @@ -87,7 +95,7 @@ def __init__(self, model, lanes, *, prefix_cache=True, host_bytes=0, medium_devi self.entries = OrderedDict() # block -> prompt unit hashes (unreferenced, least recent first) self.index = {} # unit hash -> a block holding that prefix self.copy_rt, self.io_rt = model.model[0], model.model[1] - self.metas = [m if isinstance(m, dict) else {k: getattr(m, k) for k in ("name", "shape", "dtype") + self.metas = [m if isinstance(m, dict) else {k: getattr(m, k) for k in ("name", "shape", "dtype", "layer_type") if hasattr(m, k)} for m in rc.kvcache_metas] self.host_bytes, self.host_used = host_bytes, 0 self.host = OrderedDict() # last unit hash -> (hashes, tokens, [per-layer arrays]) @@ -106,7 +114,13 @@ def _staging(self): self._stage = [] for m in self.metas: n = self.io_rt.get_kv_cache_size_for_layer(m["name"]) - dt = np.uint16 if m.get("dtype", "bfloat16") == "bfloat16" else np.float32 + count = int(np.prod(m["shape"])) + if n == count * 2: + dt = np.uint16 # physical DLFloat16, even when source metadata says float32 + elif n == count * 4: + dt = np.float32 + else: + raise ValueError("unsupported physical KV element size") self._stage.append(aligned_tensor(n // np.dtype(dt).itemsize, dtype=dt, alignment=2 << 20, tensor_type="np").reshape(m["shape"])) return self._stage @@ -229,8 +243,27 @@ def import_span(self, lane, start, keys, values): """PD: write KV computed elsewhere ([layers, kv_heads, n, head_dim], any float dtype) into the lane's block at positions start..start+n; the SDK's KV tensors are (key, value) per layer in bfloat16""" import numpy as np + 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 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() + if len({buf.dtype for buf in stage}) != 1: + raise ValueError("mixed physical KV formats are unsupported") + destinations = [] + for m, buf in zip(self.metas, stage): + if (keys.shape[1], keys.shape[3]) != (buf.shape[1], buf.shape[3]): + raise InvalidKVImport("KV export heads/dimensions do not match compiled cache") + local = self.hybrid_cache and m.get("layer_type") == "sliding_attention" + if self.hybrid_cache and m.get("layer_type") not in ("sliding_attention", "full_attention"): + raise ValueError("unknown hybrid KV layer type") + dest, offset, kept = (lane, 0, min(n, buf.shape[2])) if local else (block, start, n) + if dest is None or not 0 <= dest < buf.shape[0] or offset < 0 or offset + kept > buf.shape[2]: + raise InvalidKVImport("KV import exceeds assigned cache capacity") + destinations.append((dest, offset, kept)) # all layers converted at once ([layers, 2, kv_heads, n, head_dim], K then V), then one copy per KV tensor kv = torch.stack([keys, values], dim=1).detach().to("cpu") if stage[0].dtype == np.uint16: # the device holds DLFloat16 @@ -238,8 +271,9 @@ def import_span(self, lane, start, keys, values): else: data = kv.float().numpy() for i, (m, buf) in enumerate(zip(self.metas, stage)): - buf[block, :, start:start + n] = data[i // 2, i % 2] - self.io_rt.set_kv_cache_for_layer(m["name"], block, start, n, buf) + dest, offset, kept = destinations[i] + buf[dest, :, offset:offset + kept] = data[i // 2, i % 2, :, n - kept:] + self.io_rt.set_kv_cache_for_layer(m["name"], dest, offset, kept, buf) def drain_events(self): out, self.events = self.events, [] @@ -263,8 +297,8 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib= local_prefill=False keeps every prefill remote (decode never stalls on a prompt; TPOT first).""" rc = model.rbln_config self.hybrid_cache = getattr(rc, "cache_impl", None) == "hybrid" - if self.hybrid_cache and (prefix_cache or host_kv_gib or prefill_worker is not None): - raise ValueError("hybrid sliding cache restoration is not implemented for prefix/host cache or PD") + if self.hybrid_cache and (prefix_cache or host_kv_gib): + raise ValueError("hybrid sliding cache restoration is not implemented for prefix/host cache") self.lanes = rc.batch_size decoders = getattr(model, "decoders", None) or {} if self.lanes not in decoders: @@ -446,6 +480,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") except BaseException as error: # this request fails; the lanes go on with self._lock: self._finish(r, error) @@ -453,8 +489,16 @@ def _pd_step(self): r.pending = None ids = r.request.input_ids t0 = time.perf_counter() - if out.keys.shape[2]: - self.blocks.import_span(r.slot, out.start, out.keys, out.values) + try: + if out.keys.shape[2]: + self.blocks.import_span(r.slot, out.start, out.keys, out.values) + except InvalidKVImport as error: + with self._lock: + self._finish(r, error) + continue + if self.hybrid_cache: + self.decoder.dec_attn_mask[r.slot].zero_() + self.decoder.dec_attn_mask[r.slot, :len(ids)] = 1 # remote cost per prompt: the prefill itself, its transfer and the import here self._remote_s = _ema(self._remote_s, (out.prefill_s or 0.0) + (out.transfer_s or 0.0) + time.perf_counter() - t0) diff --git a/opt_prime/tests/test_atom_hybrid.py b/opt_prime/tests/test_atom_hybrid.py index 0d3b32dc..723db291 100644 --- a/opt_prime/tests/test_atom_hybrid.py +++ b/opt_prime/tests/test_atom_hybrid.py @@ -4,7 +4,7 @@ import torch from opt_prime.accel.pd import PDRequest from opt_prime.accel.platforms.atom.batching import AtomBatcher -from test_atom_batching import FakeBlocks, _greedy, _llama, _prompts +from test_atom_batching import FakeBlocks, _greedy, _llama, _prompts, fake_rebel # noqa: F401 class HybridFake(FakeBlocks): @@ -12,6 +12,8 @@ def __init__(self): super().__init__(_llama(), batch=2, blocks=4) self.rbln_config.cache_impl = "hybrid" self.rbln_config.use_position_ids = True + for meta in self.rbln_config.kvcache_metas: + meta["layer_type"] = "full_attention" self.dec_attn_mask = torch.ones(2, 128) self.decoders = {2: self} self.local_calls = [] @@ -50,7 +52,7 @@ def test_hybrid_refill_reinitializes_only_request_lane(): @pytest.mark.parametrize( "options", - [{"prefix_cache": True}, {"host_kv_gib": 1}, {"prefill_worker": object()}], + [{"prefix_cache": True}, {"host_kv_gib": 1}], ) def test_hybrid_unrestored_cache_modes_fail_before_sdk(options): model = HybridFake() @@ -77,3 +79,204 @@ def test_hybrid_staggered_cancel_and_refill_preserves_other_lane(): assert second.result() == _greedy(model.m, q, 8) assert third.result() == _greedy(model.m, r, 4) assert sorted(b.free) == [0, 1] + + +def test_hybrid_remote_prefill_resets_masks_and_reuses_lanes(fake_rebel): # noqa: F811 + import copy + + from opt_prime.accel.pd.pd_batching import PrefillWorker, TorchPrefillEngine + + model = HybridFake() + fake_rebel["model"] = model + worker = PrefillWorker( + TorchPrefillEngine( + copy.deepcopy(model.m), max_context=128, chunk=16, page_size=16 + ) + ) + try: + batcher = AtomBatcher(model, prefill_worker=worker, local_prefill=False) + prompts = _prompts(5, 20, 35) + futures = [ + batcher.submit(PDRequest(str(i), p, 5)) for i, p in enumerate(prompts) + ] + batcher.run_until_idle() + assert [f.result() for f in futures] == [ + _greedy(model.m, p, 5) for p in prompts + ] + assert batcher.stats_pd["remote"] == 5 and batcher.stats_pd["local"] == 0 + assert not model.local_calls + assert sorted(batcher.free) == [0, 1] + finally: + worker.close() + + +def test_hybrid_import_uses_lane_tail_and_global_block_with_physical_dtype(fake_rebel): # noqa: F811 + from types import SimpleNamespace + + import numpy as np + from opt_prime.accel.platforms.atom.batching import AtomKVBlocks, bf16_to_dlfloat + + metas = [ + {"name": f"kv{i}", "shape": shape, "dtype": "float32", "layer_type": kind} + for i, (shape, kind) in enumerate( + [ + ([2, 1, 4, 2], "sliding_attention"), + ([2, 1, 4, 2], "sliding_attention"), + ([4, 1, 8, 2], "full_attention"), + ([4, 1, 8, 2], "full_attention"), + ] + ) + ] + writes = [] + rt = SimpleNamespace( + get_kv_cache_size_for_layer=lambda name: ( + int(np.prod(metas[int(name[2:])]["shape"])) * 2 + ), + set_kv_cache_for_layer=lambda name, block, start, n, buf: writes.append( + (name, block, start, n, buf.copy()) + ), + ) + rc = SimpleNamespace( + cache_impl="hybrid", + prefill_chunk_size=2, + kvcache_num_blocks=4, + kvcache_metas=metas, + ) + model = SimpleNamespace(rbln_config=rc, model=[rt, rt]) + blocks = AtomKVBlocks(model, 2, prefix_cache=False) + blocks.admit(1, tuple(range(6))) + assert blocks.lane_block[1] == 0 + keys = torch.arange(24).reshape(2, 1, 6, 2).float() + values = keys + 30 + blocks.import_span(1, 0, keys, values) + assert len(writes) == 4 + for i, (name, block, start, n, buf) in enumerate(writes): + assert buf.dtype == np.uint16 + assert (block, start, n) == ((1, 0, 4) if i < 2 else (0, 0, 6)) + source = (keys if i % 2 == 0 else values)[i // 2, :, -n:] + expected = bf16_to_dlfloat( + source.to(torch.bfloat16) + .contiguous() + .view(torch.int16) + .numpy() + .view(np.uint16) + ) + np.testing.assert_array_equal(buf[block, :, :n], expected) + with pytest.raises(ValueError, match="full"): + blocks.import_span(1, 1, keys, values) + assert len(writes) == 4 + + +@pytest.mark.parametrize("invalid", ["partial", "heads", "values"]) +def test_invalid_hybrid_result_fails_request_without_poisoning_lane( + fake_rebel, invalid # noqa: F811 +): + import copy + + from opt_prime.accel.pd.pd_batching import PrefillWorker, TorchPrefillEngine + + model = HybridFake() + fake_rebel["model"] = model + engine = TorchPrefillEngine( + copy.deepcopy(model.m), 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 calls == 1: + if invalid == "partial": + out.start = 1 + out.keys, out.values = out.keys[:, :, 1:], out.values[:, :, 1:] + elif invalid == "heads": + out.keys = out.keys.repeat(1, 2, 1, 1) + out.values = out.values.repeat(1, 2, 1, 1) + else: + out.values = out.values[:, :, :-1] + return out + + engine.prefill = prefill + worker = PrefillWorker(engine) + try: + batcher = AtomBatcher(model, prefill_worker=worker, local_prefill=False) + prompt = _prompts(1, 20, 25)[0] + bad = batcher.submit(PDRequest("partial", prompt, 3)) + batcher.run_until_idle() + with pytest.raises(ValueError): + bad.result() + good = batcher.submit(PDRequest("next", prompt, 3)) + batcher.run_until_idle() + assert good.result() == _greedy(model.m, prompt, 3) + assert not batcher.poisoned + finally: + worker.close() + + +def test_hybrid_cancel_pending_then_reuse_ignores_late_result(fake_rebel): # noqa: F811 + import copy + from concurrent.futures import Future + + from opt_prime.accel.pd.pd_batching import TorchPrefillEngine + from opt_prime.accel.pd.runtime import RequestCancelled + + class Worker: + def __init__(self): + self.pending = [] + + def submit(self, request, send_from=0, slot=None): + future = Future() + self.pending.append((request, future)) + return future + + model = HybridFake() + fake_rebel["model"] = model + worker = Worker() + engine = TorchPrefillEngine( + copy.deepcopy(model.m), max_context=128, chunk=16, page_size=16 + ) + batcher = AtomBatcher(model, prefill_worker=worker, local_prefill=False) + p, q = _prompts(2, 20, 35) + gone = batcher.submit(PDRequest("gone", p, 5)) + batcher.step() + assert batcher.cancel("gone") + batcher.step() + with pytest.raises(RequestCancelled): + gone.result() + good = batcher.submit(PDRequest("reuse", q, 5)) + batcher.step() + for request, future in worker.pending: + if not future.done(): + future.set_result(engine.prefill(request)) + batcher.run_until_idle() + assert good.result() == _greedy(model.m, q, 5) + assert not batcher.by_id and sorted(batcher.free) == [0, 1] + + +@pytest.mark.parametrize("error_type", [RuntimeError, ValueError]) +def test_hybrid_device_import_failure_still_poisons(fake_rebel, error_type): # noqa: F811 + import copy + + from opt_prime.accel.pd.pd_batching import PrefillWorker, TorchPrefillEngine + + model = HybridFake() + fake_rebel["model"] = model + engine = TorchPrefillEngine( + copy.deepcopy(model.m), max_context=128, chunk=16, page_size=16 + ) + worker = PrefillWorker(engine) + try: + batcher = AtomBatcher(model, prefill_worker=worker, local_prefill=False) + + def broken_import(*args): + raise error_type("device write failed") + + batcher.blocks.import_span = broken_import + batcher.submit(PDRequest("device-failure", _prompts(1, 20, 25)[0], 3)) + with pytest.raises(error_type, match="device write failed"): + batcher.run_until_idle() + assert batcher.poisoned + finally: + worker.close()