omnijev: reuse the prefix within a request

[twu3202]

Required human contributor answers

1. What outcome are you trying to achieve?

   Serve OmniJev-4B v1.1 on system1-omni's native Rust/CUDA path.

2. What is the highest-leverage variable for that outcome?

   answer agreement with the reference first, asn then per-request latency.

3. How will this work move that variable rather than create busywork?

   omnijev: add the native eager worker #129 tries to minimize the difference from the reference, omnijev: reuse the prefix within a request #134 makes it faster without changing any answer.

4. Who is the customer or user, and what do they actually need?

   Anyone who wants OmniJev decisions rom system1-omni's server, and they need correct answers with low latency.

After answering the four questions: is a PR needed at all to achieve this outcome?

yes, to answer #114.

For a bug fix: link actual observed failure evidence.

N/A,, not a bug fix.

Purpose

Third PR for #114, the "request-local reuse" step of the plan. It builds on #129 (the native eager worker) and its dependencies, #98 and #99 (the fixed-algorithm GEMMs and run_window) and #117 (the shared vision tower), and should be merged after them, in the order #98, #99, #117, #129, then this PR; until then the diff here includes them. Its own changes are two commits, ee1bf55 and b5fdc74 (a comment fix from review): 7 files, +595/−142.

The worker in #129 reruns the shared prefix for every (question, option) row. Here it runs:

• the request's prefix, up to its last multiple of 64 tokens, once;
• each question's text, up to the last multiple of 64 before its first option, once, keeping the state after it when it reaches one;
• each option's block, with the tokens that remain before it, continuing that state.

Every answer equals that of #129's per-row passes bit for bit: the splits fall on the Gated DeltaNet chunks and attention's query blocks, and the fixed-algorithm GEMMs keep a row's result independent of what else runs. zq, and each option's first-token log-probability, are read with the question's text when its first option starts a chunk, and otherwise in every row. Executor::execute_rows keeps the per-row passes, and Executor::head_outputs returns the head outputs of either path, for checking one against the other.

Shared crate: Model::readout_window(window, prefix, capture, rows, targets) runs one window of a prompt with the fixed GEMMs, continuing a captured prefix and optionally capturing the state after it. The window carries its own rows' token ids, image placeholders, image embeddings and positions, so a split may fall inside the image, and it reads hidden rows and option-text log-probabilities as forward_multimodal_readout, which is now the case without a prefix or capture. PrefixState::capacity lets the executor keep two states across requests.

Against the plan: the first two of its three additions (image rows and positions in any window, per-branch readouts) are here. The vocabulary pass stays #129's GEMM and CPU log-softmax, with no new kernel and no ABI change. Batching a question's option blocks and moving the heads to the GPU come next. The 262,144-token limit on row passes from #129 is gone: with reuse a request's passes stay under the 65,536 processed tokens plus 63 per question and per option.

Test Plan

System1-Omni Version / Commit: c2a5759, which merges #129's b220f47 (with #117 18094f1 and main 426e158) into b5fdc74 on #129 5cf7c4e. The runs below used ee1bf55's change on fc27b75; 5cf7c4e only changes the HTTP handler's admission and b5fdc74 only a comment. On c2a5759 the CPU checks pass, and so do the readout and worker GPU tests, including reuse against the per-row passes bit for bit.

• CPU: the five Rust steps in CI (format, Clippy, frontend tests, workspace tests, release build) and mkdocs build --strict.
• GPU, one RTX 6000 Ada (sm_89), CUDA 13.2, with the export of omnijev: add the request contract, preparation, heads and export #120:
  • tests/qwen3_5/readout.rs, new: a prompt run as three windows, the image crossing the first boundary, with two captures, against one pass: hidden rows in the captured and the last window, and log-probabilities;
  • tests/omnijev/worker.rs, new: head outputs with reuse against the per-row passes, on 64×48, 320×200 and 1600×900 images, for the gate questions plus questions the test makes split all three ways (text within one chunk, text crossing a boundary, first option on a boundary), which it asserts are present;
  • the other readout and worker tests from omnijev: add the native eager worker #129;
  • the worker's answers against omnijev: add the native eager worker #129's per-row worker on both question sets of omnijev: add the native eager worker #129 (140 answers), with the stills as PNG and as JPEG;
  • latency as in omnijev: add the native eager worker #129: one to 16 closed-loop clients on a 1600×900 still.

Test Result

CPU. All checks pass on macOS.

GPU tests. All pass: reuse and the per-row passes give the same head outputs bit for bit at all three sizes, and the windows equal one pass.

Answers. Equal to #129's per-row worker on all 280 answers (both sets, PNG and JPEG), so the comparison with the reference in #129 holds unchanged.

Latency (medians, one client):

Workload	This PR	#129	Reference
1 question	0.16 s	0.24 s	0.17 s
3 questions	0.21 s	0.45 s	0.18 s
6 questions	0.35 s	1.02 s	0.22 s
12 questions	0.66 s	2.01 s	0.36 s
32 options	0.70 s	1.90 s	0.25 s
255 options	5.1 s	14.9 s	1.10 s

The worker runs a question's option blocks one after another, which is most of what remains against the reference. Throughput stays about the same with 8 and 16 clients (1.5 requests per second for 12 questions), and the queue sets the latency: 5.0 s and 10.1 s for 12 questions, and 40.6 s for 255 options with 8 clients.

Demo / evidence

Test output and the answer files are kept locally; I can attach them.

Self-review

• I have reviewed the full diff and addressed the issues I found.
• I have checked that the change follows the project's architecture and stays focused on the stated purpose.
• I have run the checks appropriate to this change and reported commands, results, and anything I could not verify above.
• I have checked that the PR description, documentation, and any accuracy or performance claims match the implementation and available evidence.

[hsliuustc0106]

First review at head ee1bf55b09 (base 7cbce19f), full changed-file inventory considered. Stacked accounting: the GitHub diff (+5135/−1114 over 56 files) is mostly inherited — this PR's own contribution is exactly one commit (7 files, +594/−142; 736 changed lines) on top of #129's branch, plus #98/#99/#113/#117 stack commits; the body truthfully reports authored counts, stacking and split rationale, so the >3000-line CONTRIBUTING rule is not triggered by this PR's authored delta. Artifact hygiene over all 56 files: clean (maintained fixtures with provenance and consumers only).

The reuse mechanism itself checks out on inspection: split algebra verified (zq_at < q iff open % 64 == 0; per-option-independent row.zq justifies reading zq once per question), read/target/first positions match the kept per-row path exactly, lifetimes are safe (take/put-back before ?, initialized gating, freeing synchronizes, single in-order stream), reuse sits entirely inside the admitted request closure so admission-unit semantics are unchanged, MAX_PASS_TOKENS removal is justified by retained bounds, and no foreign model rules are imposed on omnijev.

One P3 finding — src/models/omnijev/native/src/executor.rs:29-31 (at head): the field doc says prefix states are "reused across requests", but the code recaptures content every request (prefix = None at the request-prefix capture, executor.rs:237-245) and only device-buffer capacity persists. A maintainer trusting that comment could later skip recapture when capacity suffices, silently corrupting answers with stale KV. One-line comment fix.

I ran mkdocs build --strict at the head: passes. Unverified: zero CI runs exist for this head; all bit-wise parity evidence (windows vs one pass; reuse vs per-row at 3 image sizes; 280-answer equality with #129) and the latency table are author-reported GPU output with no CI corroboration; the README memory figure and fixture regeneration unverified. Verdict: COMMENTED (draft; one P3).

[hsliuustc0106]

Incremental re-check at new head b5fdc7485c: one new commit since my review at ee1bf55b09 — "omnijev: say the prefix states keep only their buffers across requests".

This resolves the single P3 from that review, verified in the diff: the Executor field doc now states that each request or question recaptures the prefix/question state and only the allocations outlive a request — matching what the code does (per-request recapture, capacity-only persistence). Comment-only change; nothing else in the series moved (range-diff: prior commits patch-identical, one added).

Zero actionable findings at this head. Unverified: no CI has reported on this head (none existed for the prior one either); all parity evidence remains author-reported GPU output; no GPU/Rust toolchain on this host. Verdict: COMMENTED — still a draft, so no approval; it looks ready to mark ready for review.

[hsliuustc0106]

Incremental re-check at new head be68c68706 (merge of the omnijev-worker stack into omnijev-reuse): the branch absorbed main's advancement plus the same two reviewed commits as #129 — d8cd04f (manifest abi_version: 7, consistent with ops.h at this head; cuda-contract green in CI at the exact head) and a8f523a (omnijev native worker listed in the manifest's consumers; crate path verified). All previously reviewed commits are patch-identical (git range-diff against b5fdc7485c).

Zero actionable findings at this head; the doc-comment fix from my previous review remains in place. Still a draft, so no approval. Unverified (unchanged): all GPU parity evidence author-reported; no CI-reported GPU execution; no Rust toolchain on this host (fork CI rust/build/benchmarks/cuda-contract green at this exact head).

[hsliuustc0106]

Incremental re-check at new head 8cf558c360 after the rebase wave (prior review: zero findings at be68c68706). Wave-verification note: this pass reviewed the whole 10-branch rebase wave systematically — every range-diff shows the own-commit series patch-identical; all post-pin merge resolutions were checked file-by-file against both parents (source/test/config/docs) with zero dropped lines; approved-content markers verified present in every tree (jemm v2 export, decider fixture READMEs, serving-matrix docs, vision FIFO regressions). Fork/upstream CI states as noted per PR.

All commits patch-identical; the GEMM-region merge resolution was verified as a clean supersession (takes this branch's newer version, content confirmed identical to the reviewed state). Zero actionable findings; still a draft, so no approval. Template state (post-#138 answers/checklist) still to be completed before leaving draft.

[hsliuustc0106]

Incremental re-check at new head c2a5759e5f after the rebase/merge wave (prior review: zero findings at 8cf558c360). All commits patch-identical (git range-diff, named refs, cat-file verified). The three merge resolutions were checked: the backend manifest resolves to abi_version: 7 consistent with ops.h at this head, retaining the omnijev models entry plus main's new decider entry; the docs resolutions keep the supported-models table intact. Fork CI rust/build/benchmarks/cuda-contract green at this exact head. Zero actionable findings; still a draft, so no approval. Template state (post-#138 answers/checklist) still to be completed before leaving draft.