Repository navigation
cua_s1: cache each question's leading block across requests - #189
Conversation
4c04d9f to
3025b57
Compare
Add capture_fixed and continue_fixed, which save the state of a prompt's leading block and run the rest of the prompt from it with the fixed GEMM algorithms, and PrefixCache, which keeps one such state per block, keyed by the block's tokens and evicted least recently used within a byte budget. A hit, a miss and a prompt run uncached give the same result.
An agent asks the same questions at every step, and each question's prompt starts with the same system message, goal and header. With CUA_S1_PREFIX_CACHE_BYTES set, the text worker marks where that block ends and runs prompts through PrefixCache, so later steps continue from the saved block. Off by default; merged-adapter exports only, with CUA_S1_GRAPH off. /health reports the cache counters.
3025b57 to
3950083
Compare
|
Reviewed at P2: with P3: a P3: |
Cache hits equal one forward_fixed pass only where fixed GEMM rows do not depend on the row count, so run Open-Jev's check_fixed_rows at startup with CUA_S1_PREFIX_CACHE_BYTES set. Narrow the docs: only a failed state allocation falls back to an uncached pass, and the uncached counter skips prompts with no block to cache.
|
Thanks for the careful review. All three are in 5a2dba2:
|
Required human contributor answers
What outcome are you trying to achieve?
Give Cua-S1's text worker an optional cache, on the merged-adapter path, for questions that are asked again.
What is the highest-leverage variable for that outcome?
The latency of a repeated question on a cache hit compared with a miss.
How will this work move that variable rather than create busywork?
Keep the state at the end of the fixed leading block across requests and continue from it, which cuts much of the latency (on [Feature]: Cache Cua-S1's leading prompt block across requests #188's measurements, a hit took 30–65% less time than a full pass).
Who is the customer or user, and what do they actually need?
Agents that make multi-step decisions with Cua-S1's text layout and ask many of the same questions at every step. They need those repeated questions answered faster.
After answering the four questions: is a PR needed at all to achieve this outcome?
Yes. The cache needs code in the shared executor and the Cua-S1 worker; #188 describes it.
For a bug fix: link actual observed failure evidence.
N/A, not a bug fix.
Purpose
For #188, in the order agreed there: merged-adapter exports first, opt-in and eager only, with #143's unmerged profile unchanged.
Each Cua-S1 text prompt starts with the system message, the question's goal and a header, then the accessibility tree. An agent sends that block again at every step it asks the question.
qwen3_5):capture_fixedsaves the state of a prompt's first tokens, andcontinue_fixedruns the rest of the prompt from it, both with the fixed GEMMs, asforward_shareddoes within one call.prefix_cache::PrefixCachekeeps one state per block, keyed by the block's tokens, sized to it, and evicted least recently used within a byte budget. A hit continues from the saved state; a miss captures it first; a prompt without a block, or with one over the budget or whose state the device can't allocate, runs as oneforward_fixedpass.chat_headreturns the start ofchat_textup to the state, and the processor marks the last multiple of 64 tokens inside it. WithCUA_S1_PREFIX_CACHE_BYTESset, the executor runs every prompt through the cache. It is refused for unmerged exports and withCUA_S1_GRAPH=1, and/healthreports the cache's counters. With it set, the worker also runscheck_fixed_rowsat startup, as Open-Jev does with prefix sharing, and refuses to start if it fails./healthare unchanged.recipe/cua_s1/native.mdand a paragraph indocs/architecture.md.Not here: the unmerged profile, the screenshot worker, Graph mode, and cache hits in
metadata.prefix_cache(question 3 on #188).Test Plan
System1-Omni Version / Commit:
5a2dba2onmain32346fc(#179 and #184 in).CPU: the CI steps (format, Clippy, frontend tests, workspace tests, release build, Python checks) and
mkdocs build --strict. New tests cover the eviction order and budget, the switch and its refusals, and where checkpoints fall in each question's prompt.GPU, one RTX 6000 Ada (sm_89), Cua-S1 4B 0.2 merged-adapter text export:
tests/qwen3_5/prefix.rs: hits, misses, evictions, a hit after the scratch buffers grow, blocks that can't be cached, and a block over the budget, each againstforward_fixedon the whole prompt, bit for bit.fixed_gemm_rows_do_not_depend_on_mon the same build.tests/cua_s1/prefix_cache_gpu.rs: three agent steps with two questions each, through the worker's executor, with a warm cache, a cache that evicts at every step, and a 1-byte budget that caches nothing.forward_fixedpass). The latencies below come from the run at3025b57onmain28f4957, whose responses matched a run before Optimize Open-Jev-27B candidate execution #186. Rerun at3950083, all three runs give the same responses apart from timing fields, and the same cache counters.5a2dba2(startup check):tests/cua_s1/prefix_cache_gpu.rsagain, and worker startup with and without the cache, twice each. The check passes and takes about 0.3 s; startup goes from 2.3 s to 2.3–2.5 s, and peak device memory stays at what the worker holds once ready.Test Result
GPU tests. All pass.
Answers. With the cache at 1 GiB and at 1 byte, all 1,167 replayed responses are byte-identical. Against the default path, 27 of 1,559 choices differ (8 on Flights, 18 on 2048, 1 on blackjack): the fixed GEMMs round differently.
Latency per request (questions run one after another inside a request), and the cache's hits and misses:
zurich-londonzurich-londonpages offer more than 26 click targets.Demo / evidence
No demo. Raw responses and replay summaries are kept locally; the traces copy third-party pages and stay outside the repository, as #185 describes.
Self-review