Serve Cua-S1 multimodal requests with bounded segmented CUDA Graphs

[Levius-Fubuki]

Purpose

Add opt-in segmented CUDA Graph execution to the Python multimodal worker, with exact-logit admission, bounded resident ownership and eager fallback. Full-attention layers remain eager.

The latest follow-up addresses the maintainer’s P1/P2 findings directly in this PR: independently evictable shapes own separate capture streams; graphs reset before streams are released; resident bytes include reserved private-pool scratch plus external static buffers. Rejected/evicted candidates retire explicitly even if another reference remains. The safety implementation is moved forward from #33, so this PR is independently safe with respect to those findings. Capture-work admission remains in #33.

Stacked on #18. Incremental runtime comparison. The PR diff remains limited to runtime Python source; validation helpers, tests and generated evidence are preserved separately. This optimizes the Python multimodal worker, not the native Rust/CUDA worker.

Test Plan

System1-Omni Version / Commit: 552c02ce9ce48cfe1ada512dee3b9da55ce3b8ca.

Fresh validation started 2026-09-30 on RTX 4090 24 GiB, driver 595.71.05, Torch 2.14.0+cu130, Transformers 5.17.0 and PEFT 0.21.0; BF16 base with unmerged adapter, without FLA/causal-conv1d.

• Validation checkout and pinned setup: PYTHONPATH=src:recipe/cua_s1 python -m pytest tests/cua_s1 -q.
• Exact GPU commands and exit statuses, using upstream-verified weights and a fresh output directory.
• Runtime/validation source mapping: git diff <runtime> <validation> -- src is empty. Source hashes and numerical/performance reports are independently verified.

Test Result

• CPU suite: 135 passed in 17.68s.
• Real GPU stream/pool regressions pass: 40 unique live capture streams, 100 exact replays across retirement, 40 three-segment shape pools, reserved-byte-driven eviction and release checks.
• Full checkpoint: seven exact full-logit comparisons, including three changed-content replays of a surviving graph after another shape is evicted. Retained references confirm explicit cleanup; pool snapshots verify reserved scratch is included in the byte budget.
• Five full-model cases retain exact original/changed-image/changed-text responses; cache eviction and 1 MiB eager fallback pass.
• Negative performance result: under the correctly enforced 1 GiB resident budget, eight distinct questions cause repeated eviction/capture. Graph p50 is 4.57–4.65 s versus eager 0.84–0.85 s. The other four measured cases improve, but Graph is not a universal speedup. This matrix warms caches before timing; any recapture during predictions remains included.
• GitHub rust CI passed on the above runtime head.

Full evidence and tradeoffs · Raw primary report · Independent verification summary.

Timing is serial and synthetic on one GPU, after warmup. Two runs are not confidence intervals or proof of production workload behavior.

Self-review

Agent-assisted full-diff review and validation limits found no actionable correctness or architecture defect. Runtime scope, commands and claims were checked against the recorded sources/results. The contributor remains responsible for understanding the changes; this does not replace maintainer review.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[hsliuustc0106]

Reviewed original snapshot 9a818a7 and the follow-up diff through 8783e2160343557bc101fc41d893ab1092d11045. No new actionable findings.

The follow-up fixes the inherited extreme-aspect-ratio validation and lock-cleanup test race from #12. Reviewed the protocol/test changes and documentation updates. Protocol and HTTP tests on this head: 61 passed in 14.76s. No accelerator execution was performed.

Validation and scope of the original snapshot review:

No new findings in segmented graph runtime delta. 16 targeted CPU tests passed; archived report verifier confirmed 5 cases/400 timings and cache behavior. No GPU replay executed.

[hsliuustc0106]

Review of e869ef5f8bf3 (incremental stack changes).

I recommend changes before merging this independently. Two findings in src/models/cua_s1/multimodal/graph_runtime.py:

1. [P1] Give independently evictable shapes exclusively owned capture streams (lines 103–105). The dedicated stream above this block is only used for warmup. torch.cuda.graph(self.graph), without stream=, uses PyTorch's shared default capture stream. In pinned Torch 2.14, destroying/resetting a graph clears cuBLAS workspaces for its capture stream. When another cached graph recorded a workspace address owned by the retired graph's pool, eviction can invalidate the surviving graph's address and cause an illegal access or corruption on its next replay. Capture-time eager parity does not cover later eviction. Use exclusively owned streams and reset graphs before releasing their streams. #33 implements that ownership model. This is a source/lifetime finding, not a reproduced GPU fault. Pinned PyTorch source.

2. [P2] Count retained graph-pool scratch in the byte limit (line 106). The memory_allocated delta excludes inactive scratch retained in graph-private pools for replay. Consequently max_bytes can admit multiple shapes beyond the configured resident memory budget and cause avoidable OOM under shape churn. Account for each owned pool's reserved bytes plus external static buffers. #33 implements pool-specific accounting. PyTorch graph memory documentation.

Please fold the necessary #33 safety fixes into this PR, then validate capture, shape eviction and surviving-graph replay on the pinned GPU stack.

Validation: 127 archived CPU tests passed against this head's runtime source. CUDA tests use mocks; no real GPU safety or latency claim was verified. Local Torch 2.13 / Transformers 5.14.1 differs from the recipe's 2.14 / 5.17 pins.

The validation suites were retrieved from the immutable archived revisions linked in the PR descriptions; they are not retained in the current PR diffs.

[Levius-Fubuki]

@hsliuustc0106 The P1/P2 fixes requested in your review of e869ef5f8bf3 are now in #22 itself, at 552c02ce9ce4. Could you re-review the fix delta?

• P1: independently evictable shapes own separate capture streams/pools; retirement resets graphs before destroying the stream, including rejected candidates retained by another reference.
• P2: resident bytes include each private pool's reserved scratch plus external static buffers, rather than only an allocated-memory delta.
• Pinned-stack RTX 4090 validation covers 40 unique live capture streams, 100 exact replays across retirement, 40 three-segment shape pools, and reserved-byte-driven eviction/release. Full-checkpoint tests include three changed-content survivor replays after actual eviction; all seven full-vocabulary-logit comparisons are bitwise equal to eager.
• The archived CPU suite reports 135 passes; the rust CI check passed on this head. It remains conflict-free with current main (566dec1), with changed files identical to the validated source.

Immutable evidence and reproduction commands · full-model eviction/pool report. The evidence/source verifier was rerun successfully on October 1.

Graph remains opt-in: the eight-distinct-question case under the enforced 1 GiB cache budget regresses to p50 4.57–4.65 s versus eager 0.84–0.85 s because of recapture. Admission/capture-work controls remain the separate #33 delta. Proposed merge order: #17 → #18 → #22 → #33.