diff --git a/opt_prime/docs/validation/batching/ATOM.md b/opt_prime/docs/validation/batching/ATOM.md index 3bda29e8..5cfebc15 100644 --- a/opt_prime/docs/validation/batching/ATOM.md +++ b/opt_prime/docs/validation/batching/ATOM.md @@ -100,3 +100,7 @@ shared (load average 75–135 from other users' jobs), so absolute values differ With one cache block the host tier quadruples the reused prefix (2048 → 11264 tokens). Under that run's host load the reloads did not turn into throughput: host staging is CPU-side work. PD mode (KV prefilled elsewhere, ATOM decodes only) is measured in `PD.md`. + +## Hybrid attention (#207) + +Gemma3 standalone external-cache serving uses global block ownership and lane-local sliding cache ownership. Fresh prefill resets the owning lane’s decode attention mask. Validation: [runtime integration](runtime207/experiments/2026-10-23-atom-hybrid-runtime.md). Prefix caching, host caching and remote prefill are rejected for hybrid models until sliding-state restoration is implemented. The verified hardware case is Gemma3-1B B1/context9216 with the upstream causal-prefill correction; this does not establish other hybrid model or multi-lane hardware support. diff --git a/opt_prime/docs/validation/batching/runtime207/.research/events.jsonl b/opt_prime/docs/validation/batching/runtime207/.research/events.jsonl new file mode 100644 index 00000000..dcce6224 --- /dev/null +++ b/opt_prime/docs/validation/batching/runtime207/.research/events.jsonl @@ -0,0 +1 @@ +{"id": "1b738ab428584b0f822c8c9392b3d203", "ts": "2026-10-23T17:38:12+09:00", "task": "Gemma hybrid KV tables (aicomp #207)", "result": "Actual AtomBatcher integration:16/16 Gemma outputs exact vs verified bridge through8K;20 regression tests pass,2 optional PD suites skipped. Hybrid unsupported cache modes now fail before SDK calls.", "next": "Review engine PR targeting optimus-rm; implement sliding-state restoration for hybrid prefix/PD; continue HTTP and cross-backend numerical validation.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {"model": "Gemma3-1B", "backend": "ATOM", "batch": "1"}, "artifacts": ["results/atom-gemma3-1b-runtime207.jsonl"], "experiments": ["experiments/2026-10-23-atom-hybrid-runtime.md"], "labels": ["bug"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "RNGD+ATOM host", "branch": "main", "commit": "351d39d-dirty"} diff --git a/opt_prime/docs/validation/batching/runtime207/STATUS.md b/opt_prime/docs/validation/batching/runtime207/STATUS.md new file mode 100644 index 00000000..f5865d2b --- /dev/null +++ b/opt_prime/docs/validation/batching/runtime207/STATUS.md @@ -0,0 +1,3 @@ +# ATOM hybrid runtime integration + +State: review. Issue #207. Real batcher16/16 Gemma outputs match the experimental bridge through8K;20 local tests pass,2 optional suites skip. Prefix/host cache and PD restoration, B>1 hardware and cross-backend numerical parity remain open. diff --git a/opt_prime/docs/validation/batching/runtime207/experiments/2026-10-23-atom-hybrid-runtime.md b/opt_prime/docs/validation/batching/runtime207/experiments/2026-10-23-atom-hybrid-runtime.md new file mode 100644 index 00000000..39f61759 --- /dev/null +++ b/opt_prime/docs/validation/batching/runtime207/experiments/2026-10-23-atom-hybrid-runtime.md @@ -0,0 +1,36 @@ +# Integrate hybrid cache ownership into AtomBatcher + +- **Date**: 2026-10-23 +- **Task**: aicomp #207; broader matrix #201 +- **Host / device**: RNGD+ATOM host, one allocated ATOM; rebel-compiler0.11.0, optimum-rbln0.11.0.post1, torch2.11.0+cpu, Transformers5.8.1 +- **Commit**: baseline optimus-rm0e69cc9d54d19e336d33920298c39507d1415811; candidate this PR's batching.py patch (uncommitted during measurement) + +## Question + +Does the verified experimental table/mask fix also work inside the actual engine batcher, without replacing its methods in a validation wrapper? + +## Setup + +Gemma3-1B revision5b11413a10db4e486ef16a20101fd028f8f2499c, B1/context9216/three KV blocks, no prefix/host cache or remote prefill. Reuse the artifact rebuilt with the causal-prefill correction from optimum-rbln PR827. Config SHA256204d77f67d2566125248a0cc3c74b6a40da7bf718632a14eb70726f4ad47b170. Four input lengths128/512/2048/8192, outputs32/128, two repeats. Fixed-output tests ignore EOS. + +Only the candidate batching module is loaded in the isolated validation process. validate_runtime207.py imports atom_batching207.py as the engine module and invokes the PR209 Gemma matrix harness. It does not alter the shared installed package. The harness has no experimental-hybrid-tables flag in this run; normal AtomBatcher methods execute. + +## Command + +Within the matrix allocation/bootstrap environment, with PR203/209 matrix helpers on the working directory/search path: + +```bash +cp "$ENGINE/opt_prime/opt_prime/accel/platforms/atom/batching.py" atom_batching207.py +python container_bootstrap.py validate_runtime207.py --model unsloth/gemma-3-1b-it --cache-dir atom-artifacts --output atom-gemma3-1b-runtime207.jsonl +PYTHONPATH=opt_prime python -m pytest -q -rs opt_prime/tests/test_atom_hybrid.py opt_prime/tests/test_atom_batching.py opt_prime/tests/test_pd_atom_reuse.py opt_prime/tests/test_pd_native_atom.py +``` + +## Results + +Before implementation, all4 initial regression cases fail: missing local table and absent early rejections for prefix/host/PD. After implementation and an additional staggered cancellation/refill case,20 tests pass. Two optional suites skip because hetero_p2p and rngd_pd packages are unavailable locally; they are not counted as passing hardware or PD coverage. + +Hardware:16/16 runs complete;8/8 repeat pairs identical;16/16 outputs equal the experimentally validated bridge from PR209. That bridge independently matched complete32-step logits against SDK-owned tables at all four lengths. This integration run compares output tokens, not every logit. Raw results: results/atom-gemma3-1b-runtime207.jsonl. + +## Conclusion + +The real engine can execute this standalone Gemma configuration without the experimental method replacement. Static-cache Qwen/Llama behavior is covered by the existing fake-runtime regression suite. Prefix/host caching and PD for hybrid models are rejected before SDK work because local sliding-state restoration remains unimplemented. B>1 is covered by ownership/cancellation unit tests, not Gemma hardware validation. GPU/ATOM numerical differences from PR209 remain; this is not a complete model-quality or serving-performance certification. diff --git a/opt_prime/docs/validation/batching/runtime207/results/atom-gemma3-1b-runtime207.jsonl b/opt_prime/docs/validation/batching/runtime207/results/atom-gemma3-1b-runtime207.jsonl new file mode 100644 index 00000000..b7793176 --- /dev/null +++ b/opt_prime/docs/validation/batching/runtime207/results/atom-gemma3-1b-runtime207.jsonl @@ -0,0 +1,16 @@ +{"input_tokens": 128, "output_tokens": 32, "repeat": 0, "batch": 1, "input_sha256": "557b0eedab01df046189400093ccf00995f3e74642d1159283992762e18d1e70", "ttft_s": 0.03371518000494689, "e2e_s": 0.40370554593391716, "tpot_s": 0.011932162417759818, "tokens": [9079, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669], "status": "passed"} +{"input_tokens": 128, "output_tokens": 32, "repeat": 1, "batch": 1, "input_sha256": "557b0eedab01df046189400093ccf00995f3e74642d1159283992762e18d1e70", "ttft_s": 0.02560331905260682, "e2e_s": 0.42961790901608765, "tpot_s": 0.013029861902337401, "tokens": [9079, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669], "status": "passed"} +{"input_tokens": 128, "output_tokens": 128, "repeat": 0, "batch": 1, "input_sha256": "557b0eedab01df046189400093ccf00995f3e74642d1159283992762e18d1e70", "ttft_s": 0.028461545007303357, "e2e_s": 1.9107618549605832, "tpot_s": 0.014820677000024597, "tokens": [9079, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107], "status": "passed"} +{"input_tokens": 128, "output_tokens": 128, "repeat": 1, "batch": 1, "input_sha256": "557b0eedab01df046189400093ccf00995f3e74642d1159283992762e18d1e70", "ttft_s": 0.02486087998840958, "e2e_s": 1.8448207559995353, "tpot_s": 0.01432968326844275, "tokens": [9079, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107], "status": "passed"} +{"input_tokens": 512, "output_tokens": 32, "repeat": 0, "batch": 1, "input_sha256": "70e4b24e2cf34281354882fbe1b899615c206153ac7f3be6cf01d6f2e11dee96", "ttft_s": 0.04974888893775642, "e2e_s": 0.5968867749907076, "tpot_s": 0.017646442742026862, "tokens": [9079, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669], "status": "passed"} +{"input_tokens": 512, "output_tokens": 32, "repeat": 1, "batch": 1, "input_sha256": "70e4b24e2cf34281354882fbe1b899615c206153ac7f3be6cf01d6f2e11dee96", "ttft_s": 0.048813628032803535, "e2e_s": 0.4256832260871306, "tpot_s": 0.012154390870202933, "tokens": [9079, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669], "status": "passed"} +{"input_tokens": 512, "output_tokens": 128, "repeat": 0, "batch": 1, "input_sha256": "70e4b24e2cf34281354882fbe1b899615c206153ac7f3be6cf01d6f2e11dee96", "ttft_s": 0.056926797959022224, "e2e_s": 1.986548307002522, "tpot_s": 0.015193108716744959, "tokens": [9079, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107], "status": "passed"} +{"input_tokens": 512, "output_tokens": 128, "repeat": 1, "batch": 1, "input_sha256": "70e4b24e2cf34281354882fbe1b899615c206153ac7f3be6cf01d6f2e11dee96", "ttft_s": 0.04907623992767185, "e2e_s": 1.5254827969474718, "tpot_s": 0.011624555629641405, "tokens": [9079, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107], "status": "passed"} +{"input_tokens": 2048, "output_tokens": 32, "repeat": 0, "batch": 1, "input_sha256": "2d38f6b9337dc2f56fff99b2995ef6cda76f7da1821ddd9bb7f9f4efc4e324e5", "ttft_s": 0.29576452693436295, "e2e_s": 0.7815623859642074, "tpot_s": 0.015668201549429325, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761], "status": "passed"} +{"input_tokens": 2048, "output_tokens": 32, "repeat": 1, "batch": 1, "input_sha256": "2d38f6b9337dc2f56fff99b2995ef6cda76f7da1821ddd9bb7f9f4efc4e324e5", "ttft_s": 0.19869860005564988, "e2e_s": 0.703318142099306, "tpot_s": 0.016275610031199553, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761], "status": "passed"} +{"input_tokens": 2048, "output_tokens": 128, "repeat": 0, "batch": 1, "input_sha256": "2d38f6b9337dc2f56fff99b2995ef6cda76f7da1821ddd9bb7f9f4efc4e324e5", "ttft_s": 0.20976138697005808, "e2e_s": 2.4802555789938197, "tpot_s": 0.01787703976357781, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761], "status": "passed"} +{"input_tokens": 2048, "output_tokens": 128, "repeat": 1, "batch": 1, "input_sha256": "2d38f6b9337dc2f56fff99b2995ef6cda76f7da1821ddd9bb7f9f4efc4e324e5", "ttft_s": 0.2570356719661504, "e2e_s": 2.3116221609525383, "tpot_s": 0.016177268724477903, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761], "status": "passed"} +{"input_tokens": 8192, "output_tokens": 32, "repeat": 0, "batch": 1, "input_sha256": "0a550211766d20a9650d69b1ffabfa3b26aca460e3cf07d640acbd722c18704a", "ttft_s": 0.9133721810067073, "e2e_s": 1.4257266820641235, "tpot_s": 0.01652411677529134, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761], "status": "passed"} +{"input_tokens": 8192, "output_tokens": 32, "repeat": 1, "batch": 1, "input_sha256": "0a550211766d20a9650d69b1ffabfa3b26aca460e3cf07d640acbd722c18704a", "ttft_s": 0.8985766950063407, "e2e_s": 1.2889886510092765, "tpot_s": 0.012590600804786288, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761], "status": "passed"} +{"input_tokens": 8192, "output_tokens": 128, "repeat": 0, "batch": 1, "input_sha256": "0a550211766d20a9650d69b1ffabfa3b26aca460e3cf07d640acbd722c18704a", "ttft_s": 0.8298595099477097, "e2e_s": 2.5776417369488627, "tpot_s": 0.013760936968981868, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761], "status": "passed"} +{"input_tokens": 8192, "output_tokens": 128, "repeat": 1, "batch": 1, "input_sha256": "0a550211766d20a9650d69b1ffabfa3b26aca460e3cf07d640acbd722c18704a", "ttft_s": 0.8298346350202337, "e2e_s": 2.563930473057553, "tpot_s": 0.013652600275579928, "tokens": [9079, 236761, 108, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761, 107, 2094, 563, 1695, 4403, 236761, 669, 2604, 3673, 6097, 36157, 236761, 669, 3730, 3673, 6097, 6249, 236761], "status": "passed"} diff --git a/opt_prime/docs/validation/batching/runtime207/validate_runtime207.py b/opt_prime/docs/validation/batching/runtime207/validate_runtime207.py new file mode 100644 index 00000000..b233b3c0 --- /dev/null +++ b/opt_prime/docs/validation/batching/runtime207/validate_runtime207.py @@ -0,0 +1,11 @@ +"""Load the candidate engine module only in this validation process.""" +import importlib.util +import runpy +import sys + +name='opt_prime.accel.platforms.atom.batching' +spec=importlib.util.spec_from_file_location(name,'atom_batching207.py') +module=importlib.util.module_from_spec(spec) +sys.modules[name]=module +spec.loader.exec_module(module) +runpy.run_module('atom_gemma_sweep',run_name='__main__') diff --git a/opt_prime/opt_prime/accel/platforms/atom/batching.py b/opt_prime/opt_prime/accel/platforms/atom/batching.py index d29064a8..2be40067 100644 --- a/opt_prime/opt_prime/accel/platforms/atom/batching.py +++ b/opt_prime/opt_prime/accel/platforms/atom/batching.py @@ -262,6 +262,9 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib= slow prefill device is therefore used only for overflow, and a fast one takes most prompts. 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") self.lanes = rc.batch_size decoders = getattr(model, "decoders", None) or {} if self.lanes not in decoders: @@ -313,9 +316,18 @@ def kv_usage(self): def _prefill(self, lane, prompt, start=0): kw = {} if self.blocks is None else dict(block_tables=torch.tensor([self.blocks.lane_block[lane]], dtype=torch.int16)) - return self.model.prefill_decoder(input_ids=torch.tensor([list(prompt[start:])], dtype=torch.int64), - cache_position=torch.arange(start, len(prompt), dtype=torch.int32)[None], - batch_idx=lane, **kw) + if self.hybrid_cache and self.blocks is not None: + if start: + raise ValueError("hybrid sliding cache requires a fresh prefill") + kw["local_block_tables"] = torch.tensor([lane], dtype=torch.int16) + out = self.model.prefill_decoder(input_ids=torch.tensor([list(prompt[start:])], dtype=torch.int64), + cache_position=torch.arange(start, len(prompt), dtype=torch.int32)[None], + batch_idx=lane, **kw) + if self.hybrid_cache and self.blocks is not None: + # External tables bypass the SDK's prompt mask reset; decode only sets new bits. + self.decoder.dec_attn_mask[lane].zero_() + self.decoder.dec_attn_mask[lane, :len(prompt)] = 1 + return out @torch.inference_mode() def step(self): @@ -367,6 +379,8 @@ def _decode_step(self): tokens[r.slot], positions[r.slot] = r.tokens[-1], self.kv.length[r.slot] kw = {} if self.blocks is None else dict(block_tables=self.blocks.table([r.slot for r in rows])) cache_position = torch.tensor(positions, dtype=torch.int32)[:, None].contiguous() + if self.hybrid_cache and self.blocks is not None: + kw["local_block_tables"] = torch.arange(self.lanes, dtype=torch.int16)[:, None] if getattr(self.model.rbln_config, "use_position_ids", False): # e.g. Gemma 3 (sliding attention) kw["position_ids"] = cache_position.clone() out = self.decoder(input_ids=torch.tensor(tokens, dtype=torch.int64)[:, None].contiguous(), diff --git a/opt_prime/tests/test_atom_hybrid.py b/opt_prime/tests/test_atom_hybrid.py new file mode 100644 index 00000000..0d3b32dc --- /dev/null +++ b/opt_prime/tests/test_atom_hybrid.py @@ -0,0 +1,79 @@ +"""Hybrid SDK contract: local cache belongs to lanes, global cache to blocks.""" + +import pytest +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 + + +class HybridFake(FakeBlocks): + def __init__(self): + super().__init__(_llama(), batch=2, blocks=4) + self.rbln_config.cache_impl = "hybrid" + self.rbln_config.use_position_ids = True + self.dec_attn_mask = torch.ones(2, 128) + self.decoders = {2: self} + self.local_calls = [] + + def prefill_decoder(self, *, local_block_tables=None, **kw): + assert local_block_tables is not None, "missing local cache ownership" + assert local_block_tables.tolist() == [kw["batch_idx"]] + self.local_calls.append((kw["batch_idx"], kw["block_tables"].tolist())) + return super().prefill_decoder(**kw) + + def __call__(self, *, local_block_tables=None, position_ids=None, **kw): + assert local_block_tables.tolist() == [[0], [1]] + assert torch.equal(position_ids, kw["cache_position"]) + for lane in range(2): + if int(kw["block_tables"][lane, 0]) == 3: + continue + pos = int(kw["cache_position"][lane, 0]) + self.dec_attn_mask[lane, pos] = 1 + expected = torch.arange(128) <= pos + assert torch.equal(self.dec_attn_mask[lane].bool(), expected), ( + "stale or missing prompt mask" + ) + return self._decode(**kw) + + +def test_hybrid_refill_reinitializes_only_request_lane(): + model = HybridFake() + b = AtomBatcher(model) + prompts = _prompts(4, 20, 32) + for i, p in enumerate(prompts): + f = b.submit(PDRequest(str(i), p, 3)) + b.run_until_idle() + assert f.result() == _greedy(model.m, p, 3) + assert any(lane != blocks[0] for lane, blocks in model.local_calls) + + +@pytest.mark.parametrize( + "options", + [{"prefix_cache": True}, {"host_kv_gib": 1}, {"prefill_worker": object()}], +) +def test_hybrid_unrestored_cache_modes_fail_before_sdk(options): + model = HybridFake() + with pytest.raises(ValueError, match="sliding"): + AtomBatcher(model, **options) + assert not model.local_calls + + +def test_hybrid_staggered_cancel_and_refill_preserves_other_lane(): + from opt_prime.accel.pd.runtime import RequestCancelled + + model = HybridFake() + b = AtomBatcher(model) + p, q, r = _prompts(3, 20, 35) + first = b.submit(PDRequest("a", p, 12)) + b.step() + second = b.submit(PDRequest("b", q, 8)) + b.step() + assert b.cancel("a") + third = b.submit(PDRequest("c", r, 4)) + b.run_until_idle() + with pytest.raises(RequestCancelled): + first.result() + assert second.result() == _greedy(model.m, q, 8) + assert third.result() == _greedy(model.m, r, 4) + assert sorted(b.free) == [0, 1]