cua_s1: add opt-in exact-length CUDA Graph replay

[hsliuustc0106]

Purpose

Adds opt-in CUDA Graph replay to the native Cua-S1 text worker from merged #19. Rebased onto main at a6f7b23.

Repeated native Cua-S1 requests currently launch the entire prefill eagerly. Setting CUA_S1_GRAPH=1 warms the GEMM plans and captures the forward pass for each exact prompt length, then replays it with newly uploaded token ids. Up to eight captures are retained with FIFO eviction; scratch growth invalidates them before freeing buffers. Eager execution remains the default. The CUDA runtime wrapper owns graph cleanup, and the worker/library ABI moves to 3, requiring both to be rebuilt together.

Scope is CUDA Graph replay only. Kernel arithmetic and GEMM selection are unchanged. Benchmark evidence and reproduction commands are in docs/benchmarks/cua-s1-cuda-graphs/.

Test Plan

• cargo fmt --all --check
• cargo clippy --workspace --locked --all-targets -- -D warnings
• cargo test --workspace --locked
• cargo build --workspace --release --locked
• CUDA library build with CUDA 13.0, sm_89.
• Eager/Graph response comparison: 14 fixed fixtures plus 3 reuse probes; covers changed inputs, scratch growth, later shorter prompts reusing the allocation, and capture eviction/recapture.
• Sequential HTTP A/B on one reserved NVIDIA L20X (GPU 3), same executable/library/merged BF16 weights and inherited CPU affinity. One server per configuration; feasibility excluded, then two measured runs, each with 3 warmups and 20 measured requests per case. No cache clearing. Startup, first-use capture and warm measurements reported separately.
• GPU regression covers changed-input replay, ordinary recording failure, and actual capture invalidation by stream synchronization. It checks both a propagated closure error (900) and EndCapture's error (901), retained-graph replay, and successful embedding recapture on the same host thread, without clearing CUDA errors in the test.

System1-Omni Version / Commit: benchmark base 8367333; benchmark implementation 61238a0; current review fixes a55a111 on main at a6f7b23. The original latency measurements precede capture-failure cleanup and later test/documentation changes; the latency benchmark was not rerun for this recovery fix.

Test Result

Formatting, Clippy, release build and CUDA compilation passed. CI passed on the current head a55a111. CPU suite: 19 tests passed, 3 GPU tests ignored. Full-model comparison: all 17 response bodies matched exactly, including probabilities and choices.

Warm p50 HTTP latency in milliseconds, both runs:

Prompt tokens	Eager	Graph replay
139	6.31 / 6.24	5.73 / 5.74
154	6.60 / 6.60	5.82 / 5.87
218	7.25 / 7.27	6.63 / 6.93
292	8.58 / 8.69	7.99 / 8.01
712	16.17 / 16.11	16.04 / 16.07
15,446	355.10 / 355.05	355.50 / 355.73

The 139–292-token cases improved 6.6–11.4% in the mean of run medians; the long case had no meaningful gain. First post-readiness fixture: 25.23 ms eager versus 33.66 ms Graph, reflecting capture overhead. These findings apply to this fixed workload and L20X, not the B300/27B chart. Compact raw latency samples, warmups, the 17/17 correctness summary, controls and limitations are checked in; full response bodies are retained locally.

Review-fix validation on 2026-09-30: the stronger capture-invalidation regression fails against the original CUDA library with stale error 901 at embedding recapture, then passes with the cleanup fix. All three GPU tests pass on scheduler-reserved GPU 3 (NVIDIA L20X, sm_89), CUDA 13.0, driver 570.133.20, Rust 1.98.1. Formatting, Clippy, all 19 CPU tests, workspace release build and CUDA compilation also pass. The library clears already-reported CUDA last errors at graph begin/end/launch while preserving each original return code.

Minor-review validation on 2026-10-01: capture failure now returns the completed eager result, logs the error, releases cached graphs and disables graphs for that worker lifetime, avoiding retries. On the same reserved L20X, the expanded GPU regression covers begin and launch argument errors followed by successful recapture; it fails against the previous library at the begin-error recovery and passes with the fix. All three GPU tests and all 19 CPU tests pass; formatting, Clippy, release build and CUDA compilation pass. Separate task-owned fault-injection libraries force begin and end failures in the real worker: each produces 17 HTTP responses byte-for-byte equal to the eager control and exactly one failed capture attempt, with no retries across later requests. Fault copies and raw results are kept outside the PR source. The recipe and benchmark docs now distinguish scratch growth from shorter prompts reusing the allocation.

Unverified in this local validation: the normal successful-Graph full-model comparison and latency benchmark were not rerun for these recovery fixes; full fp32-reference accuracy, concurrency, other GPUs/CUDA versions, Compute Sanitizer and frontend overhead were not evaluated. The original successful-Graph 17/17 full-model comparisons and latency samples above remain historical evidence.

Self-review

• Reviewed the complete diff against cua_s1: add a native CUDA text worker #19's frozen head, including lifetime/cleanup and documentation.
• Scope follows model-owned native execution; no new Python runtime, dependencies or GEMM tuning.
• Ran appropriate available checks, including real GPU capture-invalidation recovery, and reported remaining validation limits.
• Performance claims are limited to preserved A/B evidence; eager remains default.

[Levius-Fubuki]

Reviewed 6d2a96c against the frozen #19 base 8367333, and independently tested it on an RTX 4090 (sm_89), CUDA 13.0, driver 595.71.05 and Rust 1.98.1.

I found one failure-recovery issue worth fixing before merge:

[P2] Clear the returned CUDA error after an invalidated capture — runtime.cu:43–46.

I reproduced this by recording an embedding and then calling stream synchronization inside the capture closure. Synchronization returns 900, and EndCapture returns 901. The stream exits capture, but 901 remains in the calling thread's CUDA last-error state. A subsequent valid embedding on that same thread reports stale 901 through cudaGetLastError, so recapture fails despite the stream being usable. This reproduced normally and under Compute Sanitizer; even successfully replaying a retained graph and downloading its result between the failed and new captures does not consume the stale error.

Please preserve the original return code and clear the already-reported error inside this CUDA library, after graph cleanup:

if (graph) cudaGraphDestroy(graph);
if (e != cudaSuccess) (void)cudaGetLastError();
return e;

The exact failing probe passes with this isolated change, including when the recording closure propagates the synchronization error and when EndCapture supplies the error. NVIDIA documents last-error state per host thread and CUDA runtime instance; clearing through an unrelated dynamically loaded runtime would not be sufficient here.

The current regression's immediate Rust bail! does not invalidate CUDA capture. Please add actual invalidation followed by successful embedding recapture on the same thread, without manually clearing the error in the test. This finding is specific to injected recoverable capture failure; normal model replay did not fail in my checks.

Independent validation of the unchanged head:

• All three existing GPU tests passed, including the new updated-input replay regression.
• 70 eager/Graph last hidden states matched bitwise, covering changed IDs, FIFO eviction, scratch growth/shrink and recapture, with lengths through 15,446. Graph stress under Compute Sanitizer memcheck reported zero errors.
• 17 sequential HTTP requests plus 32 requests through eight concurrent clients matched in status, content type and response body bytes across modes; concurrent responses also matched serial ones.
• Formatting, Clippy, 19 CPU tests and workspace release build passed. I did not independently rerun the performance benchmark or FP32-reference accuracy.

Two small follow-ups: recipe/cua_s1/native.md still says every question runs eagerly, and now that #19 has merged this can follow the planned rebase/retarget to main. It would also help to agree the Graph implementation to land given the overlap with #50.

For my next CUDA contribution, I opened #53 to export multimodal reference tensors, including image features, language inputs and 3D position IDs. With Graph work underway in #52/#50, would you prefer a follow-up adding native language-model support for image embeddings and 3D rotary positions, using #53 as the correctness reference, or is there a specific CUDA kernel or validation gap you would like me to prioritize?

[twu3202]

Reviewed 30ae42a and tested it on an RTX 6000 Ada (sm_89) with CUDA 13.2 and driver 595.91.07, against eager mode from the same build. The prompts were 272 screen-like states from 173 to 2,024 tokens, each a different length, plus prompts of exactly 2,047, 2,048, 2,049 and 16,384 tokens, with two starts per mode.

Correctness looks good:

• Graph and eager responses are byte-identical on all 276 prompts in both starts, and the two starts match each other.
• The README float32 rule passes on all 276, with the cua_s1: add a text worker that loads Qwen3.5-4B through Transformers #13 worker in float32 and bfloat16 as the reference: the largest difference is 0.0172 against an allowance of 0.073.
• Lengths captured again after eviction, and 320 requests from 8 clients at once, match the first answers byte for byte. 16,385 tokens still gets a 413.

The cost is the first request at each length, which runs the forward pass twice (the eager warm-up, then the replay). Graph mode, milliseconds, mean of the two starts:

Prompt tokens	First request	Next three
173	31.6	16.7
530	113.8	54.1
1,062	193.4	96.6
1,962	343.8	172.6

On this card the replays ran at eager speed: over those 32 lengths the repeats averaged 100.1 and 102.9 ms in the two graph starts, against 102.5 ms eager (a first eager start ran about 11% faster throughout, so I left it out of this comparison). With the 240 other lengths sent once each, the mean was 209.6 ms with graphs against 99.9 ms eager, about twice, and the 8-client run took 65.5 s against 32.4 s. Cua-S1 prompts carry the current screen state, so I'd expect exact lengths to repeat rarely in real use.

A small change would remove most of this: on a miss, return the eager result and capture without replaying it, since capture only records. Here the first request cost two forward passes and nothing measurable beyond that (first minus twice the repeat time averaged under 0.5 ms in both starts), so capture itself is cheap. A length cap would also help: the 16,384-token prompt took 3.3 s the first time against 1.6 s eager, and your numbers show no gain for long prompts. Keeping the most recently used lengths instead of the first ones in would help when a few lengths do repeat.

Minor points:

• If capture fails, the request returns an error although the eager pass already produced the answer, and every later request of that length tries again and fails the same way.
• Only cs1_graph_end clears the last error. A failed cudaStreamBeginCapture or cudaGraphLaunch would show up in the next kernel call on that thread, the same way as the case Levius found.
• The benchmark README mentions scratch "growth/shrink", but the scratch only grows.

On the overlap with #50: let's go with this PR for the graphs. Once it lands, I'll rework #50 on top of it to keep only the GEMM timing, or close it if that isn't wanted.