Limit Cua-S1 reused-request output projection to final token

[Levius-Fubuki]

Purpose

Project only the final language token for reused multi-question requests, preserving the final-token decision path and the single-question reference behavior.

Stacked on #17. 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: 378d101348644f855b193402cce56605c7fc9f9e.

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: 116 passed in 17.07s.
• Fresh full-checkpoint GPU checks confirm complete response equality between full and final-token projection on two cases, with output-head shapes confirming the restriction. Two paired runs of 10 timed iterations per case show a modest roughly 1–3% p50 reduction on this setup; this is not a general performance claim.
• Runtime source is unchanged from the maintainer-reviewed 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 c522e00 and the follow-up diff through 068eb52ddc4cefe1efefbcdbf67a6ce3020b9cdd. 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.72s. No accelerator execution was performed.

Validation and scope of the original snapshot review:

No new findings in final-token projection delta. 9 targeted tests passed including CPU tensor integration.

[hsliuustc0106]

cleanup it

[hsliuustc0106]

Review of 378d10134864 (incremental stack changes).

Reviewed the incremental change over #17. Restricting vocabulary projection to the final token with logits_to_keep=1 preserves the final-token decision path; the single-question path is unchanged. No actionable correctness defect found. I recommend accepting this after #17.

Validation: 116 archived CPU tests passed against this head's runtime source. No full-checkpoint GPU parity or new latency measurement was performed. Local environment used Torch 2.13 and Transformers 5.14.1, rather than the recipe's Torch 2.14 / Transformers 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 Following up on your recommendation to accept this after #17. Head 378d10134864 is unchanged from your reviewed version. Full-checkpoint RTX 4090 checks on the pinned stack confirm complete-response equality between full and final-token projection in both measured cases; output-head shapes confirm the restriction. The archived CPU suite reports 116 passes.

Please consider approval/merge after #17. The incremental change touches one runtime file. It merges cleanly with current main (566dec1), with the changed file matching the validated source.

Immutable evidence · raw #18 report. The evidence/source verifier was rerun successfully on October 1. The measured p50 gain is only roughly 1–3% on two synthetic cases; this is not a general speedup claim.