Skip to content

cua_s1: cache each question's leading block across requests - #189

Merged
hsliuustc0106 merged 3 commits into
InferMatrixLab:mainfrom
twu3202:cua-s1-prefix-cache
Oct 30, 2026
Merged

hsliuustc0106 merged 3 commits into
InferMatrixLab:mainfrom
twu3202:cua-s1-prefix-cache

Conversation

@twu3202

@twu3202 twu3202 commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

Required human contributor answers

  1. 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.

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

    The latency of a repeated question on a cache hit compared with a miss.

  3. 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).

  4. 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.

  • Shared executor (qwen3_5): capture_fixed saves the state of a prompt's first tokens, and continue_fixed runs the rest of the prompt from it, both with the fixed GEMMs, as forward_shared does within one call. prefix_cache::PrefixCache keeps 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 one forward_fixed pass.
  • Cua-S1 text worker: chat_head returns the start of chat_text up to the state, and the processor marks the last multiple of 64 tokens inside it. With CUA_S1_PREFIX_CACHE_BYTES set, the executor runs every prompt through the cache. It is refused for unmerged exports and with CUA_S1_GRAPH=1, and /health reports the cache's counters. With it set, the worker also runs check_fixed_rows at startup, as Open-Jev does with prefix sharing, and refuses to start if it fails.
  • With the switch unset, preparation, execution and /health are unchanged.
  • With it set, every prompt uses the fixed GEMMs, cached or not, so an answer doesn't depend on what the cache holds. Answers differ from the default path by rounding.
  • Docs: a section in recipe/cua_s1/native.md and a paragraph in docs/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: 5a2dba2 on main 32346fc (#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 against forward_fixed on the whole prompt, bit for bit. fixed_gemm_rows_do_not_depend_on_m on 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.
  • Replay of six System1-Agents traces from benchmarks: record, measure and replay decision traces #185 through the worker, one episode after another with no think time, in three runs: the default path, the cache at 1 GiB, and the cache at 1 byte (every prompt one uncached forward_fixed pass). The latencies below come from the run at 3025b57 on main 28f4957, whose responses matched a run before Optimize Open-Jev-27B candidate execution #186. Rerun at 3950083, all three runs give the same responses apart from timing fields, and the same cache counters.
  • At 5a2dba2 (startup check): tests/cua_s1/prefix_cache_gpu.rs again, 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:

Trace Requests Hits / misses Mean: default / fixed / cache p50: default → cache p95: default → cache
Google Flights 93 175 / 89 451.1 / 455.5 / 358.9 ms 456.3 → 327.1 ms 678.4 → 481.4 ms
zurich-london 9 30 / 4 659.1 / 646.3 / 468.4 ms 728.1 → 523.6 ms 849.2 → 678.0 ms
Form fixtures 108 168 / 144 231.1 / 215.5 / 146.8 ms 311.1 → 122.9 ms 350.6 → 326.2 ms
2048 800 799 / 1 38.3 / 33.4 / 22.8 ms 38.4 → 22.7 ms 39.2 → 23.5 ms
Ticket router 30 29 / 1 41.1 / 38.3 / 22.6 ms 41.1 → 21.4 ms 44.2 → 22.9 ms
Blackjack 119 118 / 1 36.0 / 31.4 / 17.2 ms 35.9 → 16.9 ms 37.9 → 17.7 ms
  • Each episode's first step misses. The form episodes are three steps long, so their first steps set the p95.
  • Flights and the forms evicted 76 and 131 blocks at 1 GiB, as finished episodes made room.
  • The worker refuses the other requests before running them, in all three runs: 5 Flights and 3 zurich-london pages offer more than 26 click targets.
  • The worker took 10,360 MiB with the cache at 1 GiB against 9,208 MiB without it.

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

  • 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.

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.
@twu3202
twu3202 force-pushed the cua-s1-prefix-cache branch from 3025b57 to 3950083 Compare October 27, 2026 02:12
@twu3202 twu3202 mentioned this pull request Oct 27, 2026
1 of 10 tasks
@MrDongsls

Copy link
Copy Markdown
Collaborator

Reviewed at 3950083. The cache itself is sound: keys are exact prefix token sequences, LRU and budget accounting matches the real allocation, failure paths keep their invariants, and with the switch unset the default path and /health are unchanged. Findings:

P2: with CUA_S1_PREFIX_CACHE_BYTES set, startup only checks the merged export and CUA_S1_GRAPH, and nothing verifies on the actual device that fixed GEMM rows are independent of M, which is exactly what a cache hit matching a full forward_fixed pass relies on; the docs say it was checked on sm_89 only. Open-Jev already runs check_fixed_rows at startup when sharing is enabled, so wiring the same check into src/models/cua_s1/native/src/engine.rs:26 would turn that assumption into a startup rejection.

P3: a capture_fixed failure at src/models/qwen3_5/native/src/prefix_cache.rs:166 errors out while alloc_prefix failure falls back to one forward_fixed pass, so the documented out of device memory fallback covers allocation only; worth narrowing the wording.

P3: uncached counts only over budget and alloc failure prompts and skips prompts whose block is under 64 tokens, so it understates how much traffic bypasses the cache; a line in the field comment would fix it.

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.
@twu3202

twu3202 commented Oct 29, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the careful review. All three are in 5a2dba2:

  • P2: with CUA_S1_PREFIX_CACHE_BYTES set, the worker now runs check_fixed_rows right after loading the model (in Executor::load, since that's where the model exists) and refuses to start if it fails. On the RTX 6000 Ada it passes in about 0.3 s with the 4B export.
  • P3: the docs now say only a failed state allocation falls back to an uncached pass; any other error while saving a state fails the request.
  • P3: the uncached comment now says prompts with no block to cache aren't counted.

@hsliuustc0106
hsliuustc0106 merged commit e573260 into InferMatrixLab:main Oct 30, 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