Skip to content

open_jev: run shared prefixes once per request - #179

Open
twu3202 wants to merge 4 commits into
InferMatrixLab:mainfrom
twu3202:open-jev-shared-prefix
Open

twu3202 wants to merge 4 commits into
InferMatrixLab:mainfrom
twu3202:open-jev-shared-prefix

Conversation

@twu3202

@twu3202 twu3202 commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Required human contributor answers

  1. What outcome are you trying to achieve?

    Run faster on the native Open-Jev worker for requests that ask several questions about the same thing.

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

    How many prefix tokens each request prefills.

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

    Make Open-Jev use forward_shared: when many questions or candidate answers are asked about the same page, the whole shared context is shared.

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

    agent using Open-Jev to ask several questions about one same thing such as a web page, and they need faster answers.

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

Yes, to answer #85.

For a bug fix: link actual observed failure evidence.

N/A, not a bug fix.

Purpose

For #85, as said on #150: Open-Jev moves onto forward_shared (#99), behind opt-in switches like Decider's (#113).

Open-Jev scores each candidate as its own prompt: chat prefix, state, question, then the candidate. So every candidate of a request repeats the state, and a question's candidates also repeat the question. Open-Jev-9B runs each candidate as its own pass; Open-Jev-27B-v1.1 packs them, but every token still runs.

  • Planner. Decider's PrefixPlan moves into the shared Qwen crate (src/models/qwen3_5/native/src/prefix.rs) behind a small PrefixRow trait: a row's tokens, and whether two rows are in the same group. Decider implements it for its rows and re-exports the planner; its planning, switches and health counters are unchanged. Open-Jev's rows are its candidates, grouped by question.
  • Switches, with the meaning of DECIDER_PREFIX and DECIDER_FIXED, both off by default:
    • OPEN_JEV_PREFIX=1: each shared prefix runs once through forward_shared, and each candidate continues from the state after it. A request with less than one reusable 64-token chunk runs its candidates one at a time with forward_fixed.
    • OPEN_JEV_PREFIX=auto: shares only when the plan saves at least 4,096 tokens and at least a third of the request's tokens, Decider's rule; other requests take today's path (packing on 27B, one pass per candidate on 9B).
    • OPEN_JEV_FIXED=1: every candidate with forward_fixed, the control for shared results.
    • The two can't be combined, and either is refused at startup with CUA_S1_GRAPH=1, since these paths run eagerly.
    • With OPEN_JEV_PREFIX=1 or auto, the worker first checks that the fixed GEMMs give each row the same result whatever the row count, for the model's own projection shapes (subsets of 4,096 random rows run alone against the full call, byte for byte), and refuses to start if not. Shared results equal fixed ones only where that holds.
  • Response. With a switch set, metadata.prefix_cache is {"enabled": <ran on shared prefixes>, "mode": "fixed" | "shared" | "auto", "saved_tokens": <prefix tokens not recomputed>}. With both off it stays {"enabled": false, "mode": "independent_candidates"}. Usage still counts every candidate's tokens.
  • With both switches off, the executor takes the same path as before.
  • Docs: a section in recipe/open_jev/native.md, the measurements in recipe/open_jev/validation-9b.md, and a sentence each in docs/architecture.md and both workers' READMEs.

Not here: running a prefix's branches in one batched pass (#150's next stage), and any run on 27B.

Test Plan

System1-Omni Version / Commit: e02eca4 on main a41e379; 2a19ed2 adds the startup check, and e02eca4 only records the tests below in the validation notes. The five-way measurement ran on an earlier state of this branch on main 7b4c194, before the CUA_S1_GRAPH check, the self-review fixes to tests and docs, and the rebases. With the same code on main 51f80ad (#134 in), the GPU tests pass and the no-switch, shared and auto responses equal that run's; #165, the only later commit on main, adds a Python worker and docs. The interleaved timing check and the memory check ran with the same worker code on main 27176cb (#129). The CPU checks pass on e02eca4.

CPU: the CI steps (format, Clippy, frontend tests, workspace tests, release build, Python checks) and mkdocs build --strict. New tests cover the switches and the graph refusal, candidates grouped by question with the planner rebuilding every prompt in order, and prefix_cache in each mode. The planner's own tests stay in tests/decider/prefix.rs.

GPU, one RTX 6000 Ada (sm_89), Open-Jev-9B merged export, eager:

  • tests/open_jev/prefix.rs (new): on a two-question request whose Choice candidates share a question prefix past the next 64-token boundary, shared scores equal fixed scores bit for bit, twice, and save more than the request prefix alone would; under auto, a small request isn't shared and equals the default path. Also perf(open-jev): pack candidate projections within requests #102's prefill_batch test, on main and here.
  • The 253 requests from [RFC]: Reuse shared observations across decision questions — community help wanted #85's plan (M1 to M3, 22, and the 231 public JevBench tasks, Public231), one at a time, each sent twice, through main and through this worker with no switch, fixed, shared and auto. Responses compared with each other, and JevBench answers with the reference from open_jev: support the Open-Jev-9B checkpoint #92's validation and with the labels.
  • The no-switch path against main with requests interleaved A B B A, so slow clock and temperature drift hits both.
  • Device memory with no switch and with OPEN_JEV_PREFIX=1.

Test Result

CPU. All pass on macOS.

GPU tests. prefix passes; prefill_batch passes on main and here. With OPEN_JEV_PREFIX=auto and CUA_S1_GRAPH=1 the worker refuses to start. Both tests and fixed_gemm_rows_do_not_depend_on_m pass again at 01c84c6 and at 2a19ed2; on the 9B export the startup check takes 0.37 s.

Responses.

  • No switch: all 253 equal main's apart from timing.
  • OPEN_JEV_PREFIX=1: the 95 requests with a reusable chunk ran on shared prefixes, and all 253 responses equal OPEN_JEV_FIXED=1's apart from timing and prefix_cache.
  • auto: the 46 requests it shared equal the fixed responses; the other 207 equal the no-switch ones.

Fixed against default GEMMs. Shared results equal fixed ones, and fixed ones differ from today's by rounding, about as far as today's path is from the reference:

Requests Questions Fixed vs default: changed decisions largest difference Against the reference, default / fixed: changed decisions largest difference
M1 to M3 (22) 87 0 0.035 0 / 0 0.041 / 0.042
JevBench noul (74) 74 0 0.098 0 / 0 0.068 / 0.030
JevBench, 3 to 6 candidates (157) 157 1 0.064 2 / 1 0.104 / 0.141

The changed decision is a near tie that #92's validation already lists: 0.52/0.48 with today's GEMMs, 0.46/0.54 fixed, 0.44/0.56 in the reference. Against JevBench's labels: reference 165 of 231, today's path 163, fixed and shared 164, auto 163.

Time, sums of per-request client times (the faster of two sends), in seconds:

Requests Input tokens Not recomputed when shared main No switch Fixed Shared auto
M1 to M3 (22) 1,966,702 95% 305.3 308.7 371.1 35.5 37.9
JevBench noul (74) 31,802 0 5.76 5.74 6.44 6.16 5.69
JevBench, 3 to 6 candidates (157) 464,274 69% 80.0 79.8 92.0 34.3 42.0
of which over 8,192 input tokens (30) 347,670 75% 54.9 54.8 65.2 17.5 17.5
  • 255 candidates over a 3,072-token state: 6.7 s shared instead of 122.5 s; 32 over 1,024 tokens: 0.86 s instead of 5.3 s; support_28: 1.06 s instead of 1.88 s.
  • The fixed GEMMs alone add 12 to 20%. Under OPEN_JEV_PREFIX=1, requests with nothing to share pay that (the noul tasks: 6.16 s instead of 5.74 s); auto keeps them on today's path.
  • auto's threshold is Decider's and leaves gains: M2's requests of 2,600 to 5,500 tokens run 2 to 3 times faster shared but save fewer than 4,096 tokens. I kept Decider's rule rather than tune it on these requests.
  • With no switch the worker takes as long as main: sent alternately to both (A B B A), the 253 requests took 393.5 s against 392.8 s, median per-request ratio 1.001 (0.99 to 1.02 from the 5th to the 95th percentile), responses identical. The 1% gap in the table above, up to 11% on the first requests with many short candidates, is drift between the two consecutive runs.

Memory. After warmup, and after the 22 M1 to M3 requests and the five longest JevBench tasks: 15,874 and 16,674 MiB with no switch, 15,906 and 16,884 MiB with OPEN_JEV_PREFIX=1. The difference is the fixed-algorithm GEMM workspace (32 MiB, allocated at warmup) and the prefix states.

Demo / evidence

No demo. Raw responses, manifests and scripts are kept locally; JevBench's license allows publishing aggregate results only.

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.

Move PrefixPlan and PrefixMode from the Decider worker into the shared
crate, behind a small PrefixRow trait (a row's tokens and whether two
rows belong to the same group), so other Qwen3.5 workers can plan
forward_shared the same way. Decider re-exports them and implements
PrefixRow for its rows; its planning, switches and counters are
unchanged. One test now names the row type for an empty slice.
OPEN_JEV_PREFIX=1 plans the candidates' shared token prefixes with the
shared Qwen planner and runs each prefix once through forward_shared.
OPEN_JEV_PREFIX=auto shares only when Decider's auto rule holds, and
OPEN_JEV_FIXED=1 runs every candidate as its own fixed-algorithm pass.
The switches default off and are refused with CUA_S1_GRAPH=1. With a
switch set, the response's prefix_cache metadata reports how the
request ran.
@twu3202

twu3202 commented Oct 26, 2026

Copy link
Copy Markdown
Contributor Author

I replayed recorded System1-Agents decisions against this branch, on the same RTX 6000 Ada: 28 Google Flights episodes (98 decisions) and 36 runs of its local form fixtures (108). They were recorded with Open-Jev-9B deciding and a local chat model planning. Each request is resent as recorded, one agent at a time, with its think time.

No switch: p50 / p90 Over 5 s OPEN_JEV_PREFIX=auto: p50 / p90 Over 5 s
Flights (93 answered) 6.55 / 8.02 s 78 1.07 / 1.21 s 0
Forms (108) 1.13 / 1.32 s 0 0.59 / 1.14 s 0

auto shared all 93 Flights requests and 48 of the form ones. 3 of 264 Flights answers and 4 of 312 form answers changed, with largest probability differences of 0.054 and 0.070: the GEMM rounding described above. 5 s is System1-Agents' default decision deadline. The other 5 Flights requests were results pages over the 9B export's 4,096-token limit, so the worker refused them in both modes.

@MrDongsls

Copy link
Copy Markdown
Collaborator

Thanks for this. Just two small non-blocking follow-ups:

  1. Re-run the bit-equality gate at head. The shared-vs-fixed bit-identical test was last exercised at the 7b4c194 baseline, and validation-9b.md still points there. One run of cargo test --release --locked -p omni-open-jev-native --test prefix -- --ignored at head, plus a one-line pointer in the Controls section, closes the chain.

  2. Surface the fixed-GEMM assumption. The replay shows the rounding path is real (3/264 + 4/312 answers changed, deltas up to 0.070). The docs already require Decider's fixed-algorithm check on new GPUs; optionally tagging the assumption in metadata.prefix_cache, or a one-token M-independence self-check at startup, would make non-conforming hardware fail loudly instead of silently flipping near-ties.

Shared prefixes reproduce OPEN_JEV_FIXED=1 only on GPUs where the fixed
GEMM algorithms keep each row independent of the row count. Check that
for the loaded model's projection shapes when OPEN_JEV_PREFIX is set and
refuse to start otherwise, instead of relying on the documented manual
kernel test alone.
@twu3202

twu3202 commented Oct 26, 2026

Copy link
Copy Markdown
Contributor Author

Thanks, both done.

  1. At 2a19ed2 the prefix test passes on the RTX 6000 Ada, along with prefill_batch and fixed_gemm_rows_do_not_depend_on_m; all three also passed at 01c84c6. The Controls section in validation-9b.md now points there.
  2. With OPEN_JEV_PREFIX=1 or auto, the worker now checks at startup that the fixed GEMMs give each row the same result whatever M is, for the model's own projection shapes: subsets of 4,096 random rows run alone against the full call, byte for byte, with row counts like the kernel test's. If any differ, it refuses to start. On the 9B export this takes 0.37 s. Decider makes the same assumption and could call the same check_fixed_rows; I left that out of this PR.

One note on the replay numbers: those answer changes come from the fixed algorithms rounding differently from the default ones, which happens on any GPU. The startup check covers the other risk, shared results not matching fixed ones.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants