Skip to content

open_jev: run shared prefixes once per request - #179

Merged
psv666 merged 8 commits into
InferMatrixLab:mainfrom
twu3202:open-jev-shared-prefix
Oct 27, 2026
Merged

psv666 merged 8 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 the default path (on 27B, Optimize Open-Jev-27B candidate execution #186's packing and sharing within one question; on 9B, one pass per candidate).
    • 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. metadata.prefix_cache is {"policy": "default" | "fixed" | "shared" | "auto", "enabled": <ran on shared prefixes>, "scope": "none" | "within_question" | "across_questions", "saved_tokens": <prefix tokens not recomputed>}, the fields agreed on Optimize Open-Jev-27B candidate execution #186. With both switches off it is what Optimize Open-Jev-27B candidate execution #186 reports: policy default, and scope within_question when its sharing ran. Usage still counts every candidate's tokens.
  • With both switches off, the executor takes Optimize Open-Jev-27B candidate execution #186's default path unchanged.
  • 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: 312e88d. 5235e50 merges main 28f4957 (#186) into the branch: requests without a switch, and those auto declines, take #186's default path. After the merge, 002f510 adds a GPU test with the 27B packing options selected and 312e88d records the checks after the merge in the validation notes. Before the merge, on main a41e379, 2a19ed2 adds the startup check, e02eca4 records the GPU tests below in the validation notes, and 97d6376 changes only the prefix_cache fields. 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 312e88d.

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/qwen3_5/prefix.rs, new packing_options_leave_fixed_rows_unchanged: with batched GDN and packed projections selected, as Open-Jev does at load for checkpoints that pack candidates (27B), shared prefixes and forward_fixed against forward_fixed without them, bit for bit, also after a packed batch on the same model.
  • 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. After the merge with #186 (5235e50), prefix and fixed_gemm_rows_do_not_depend_on_m pass again, the off, shared and auto responses to the 253 requests below keep the same answers as before it, and #179's startup refusal with Graphs holds. #186's new text_prefix_buffers_grow_and_shrink_without_reusing_prior_contents needs a checkpoint that packs candidates: on Open-Jev-9B it fails the same way on main 28f4957. At 002f510, packing_options_leave_fixed_rows_unchanged passes, with the other tests/qwen3_5/prefix.rs tests and fixed_gemm_rows_do_not_depend_on_m.

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.

Use the fields proposed for Open-Jev's prefix reuse on InferMatrixLab#186, so both
changes report reuse the same way: the configured policy (default with
no switch), whether the request ran on shared prefixes, whether that was
within one question or across questions, and the prefix tokens not
recomputed. The off case now carries the same fields instead of the
old mode name.
Without a switch, and when `auto` declines to share, requests take InferMatrixLab#186's
default path, including its sharing within one question. With a switch,
`prefix_cache` reports the switch as the policy and `across_questions`
when questions shared the request prefix, keeping the four fields agreed
on InferMatrixLab#186 and here.
@MrDongsls

Copy link
Copy Markdown
Collaborator

Re-reviewed at 5235e50. The code is sound, and the #186 merge resolved the conflict the right way: 5235e50 routes the two prefix-reuse mechanisms explicitly (switch on → forward_shared; off / auto fallback → #186's try_forward_shared_prefix + packed forward_batch), all eight conflicting files hand-resolved, the four-field prefix_cache contract unified with main's. Re-ran cargo test -p omni-open-jev-native on the merge tree: 0 failed. Two non-blocking follow-ups:

  • P3 — 27B with the switch on remains unmeasured: packed projections are enabled at load while forward_shared uses fixed-row semantics; only the startup check_fixed_rows gate covers that combination today.
  • P3 — the GPU records in validation-9b.md's Controls still point at 01c84c6/2a19ed2, but the merge brought in Optimize Open-Jev-27B candidate execution #186's model.rs (+887) and cuda.rs (+186) — the head binary is now well past the tested binary. One minimal GPU prefix re-run at head, plus updating the Controls pointer, closes the chain.

Open-Jev selects batched GDN and packed projections at load for
checkpoints that pack candidates. Shared prefixes and forward_fixed run
their projections with the fixed-algorithm handle either way; the test
checks that they still match forward_fixed without those options, bit
for bit, also after a packed batch on the same model.
The prefix and fixed-row tests and the 253 requests were re-run at
5235e50, and the new packing-options test at 002f510.
@twu3202

twu3202 commented Oct 26, 2026

Copy link
Copy Markdown
Contributor Author

Thanks. Both are in 312e88d.

  • 27B with the switch on: forward_shared and forward_fixed run every projection through the fixed-algorithm handle; the packed GEMM is only used by Optimize Open-Jev-27B candidate execution #186's forward_batch. On those paths the two options change only the conv kernel, and the vector conv does the same arithmetic as the scalar one. 002f510 adds a GPU test that selects batched GDN and packed projections on the 9B export and checks shared prefixes and forward_fixed against forward_fixed without them, bit for bit, also after a packed batch. It passes. I still haven't run 27B itself; its BF16 weights don't fit on my 48 GB card.
  • Controls: 5235e50 was re-run on the GPU before I pushed it (prefix, fixed_gemm_rows_do_not_depend_on_m, and the 253 requests with no switch, shared and auto, same answers). validation-9b.md now says so and records the new test at 002f510.

@psv666
psv666 merged commit 57bb016 into InferMatrixLab:main Oct 27, 2026
6 checks passed
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.

3 participants