[New Model]: Support Cua-S1 4B 0.2 — community help wanted

[hsliuustc0106]

The model to consider

Add support for Cua-S1 4B 0.2 in System1-Omni.

• Weights and model card: https://huggingface.co/cua-ai/cua-s1-4b-0.2
• Reference implementation: https://github.com/trycua/cua/tree/main/libs/cua-s1
• Base model: Qwen/Qwen3.5-4B

The release provides independently trained text/ and multimodal/ LoRA adapters for closed-option computer-use element/action decisions. The upstream model card loads them through cua_s1.four_b.FourBModel, selecting the adapter by modality. Version 0.2 is published separately from 0.1; this request specifically targets 0.2.

The closest planned or implemented model

System1-Omni is still in its design stage: no model engines or GPU backends are implemented, and LAYA is the first planned model. This is a request for a new integration, not a claim of existing support or a change to the first-model priority.

The broader model roadmap in #9 already lists Cua-S1 4B as a candidate. This issue tracks its implementation separately. Agent-side integration is tracked in ThinkFlowLab/system1-agents#14; that is distinct from engine support here.

What's needed to support it

Start by auditing and pinning the upstream reference implementation, then document the exact preprocessing, adapter loading, decision/scoring path, and output semantics for both modalities. Confirm the prefill-only execution requirements and compatibility with the frontend proposal in #1 before choosing the implementation approach.

Following the repository architecture, the model engine should own preprocessing, batching, execution, state, and postprocessing. Hardware operations belong in the CUDA or Metal backend. Potential work includes Qwen3.5 execution, LoRA loading or merging, image processing for the multimodal adapter, and matching upstream candidate scoring. The initial implementation may target one backend, with coverage and limitations stated explicitly.

Use case and motivation

Serve closed-option GUI element/action decisions through System1-Omni, including screenshot-conditioned decisions with the multimodal adapter. This would extend the planned model coverage to a Qwen3.5-based computer-use decision model.

Community help wanted

We welcome contributors who can help with any of the following:

• Audit the upstream inference path and propose a small, concrete integration plan.
• Implement model loading and text decision inference.
• Add multimodal preprocessing and screenshot-conditioned inference.
• Implement or integrate the required CUDA or Metal operations.
• Validate outputs against the upstream reference and contribute reproducible setup and usage instructions.

If you would like to help, please comment with the area you can take on and your available hardware. Upstream maintainers and users familiar with FourBModel are especially welcome to clarify inference requirements and useful validation cases.

Acceptance criteria

• Record the pinned upstream code and model revisions, inference contract, and initial backend scope.
• Load the 0.2 adapters and validate text and multimodal decisions against the upstream reference on fixed inputs, with comparison tolerances declared in advance. Track any deferred modality explicitly.
• Exercise the supported request path end to end and document unsupported input or output features.
• Provide reproducible setup, launch, and example-request instructions.
• Update the supported-model table with only the modalities and backends actually validated.

Performance claims, if contributed, should include the hardware, revisions, commands, raw results, and comparison methodology, with loading and warmup reported separately.

Before submitting

• Searched existing and past issues and checked the repository documentation. [RFC] Support emerging System 1 models, with CLM as the second model after LAYA #9 is the umbrella proposal; no dedicated Cua-S1 4B 0.2 issue existed at submission time.

[twu3202]

I'd like to take this, starting with the text adapter on CUDA. Hardware: one RTX 6000 Ada (48 GB, sm_89), and also rented NVIDIA GPU access.

Plan:

1. A docs PR first: pinned revisions (trycua/cua, the base model and the 0.2 adapters), the inference contract, and tolerances declared before any comparison. From reading FourBModel, a decision is one forward pass and a softmax over the answer-letter logits (A to Z) at the last position, so there is no decode and at most 26 options.
2. A text worker in src/models/cua_s1/ behind the Add Rust frontend for Jev workers #2 frontend, with setup and example requests in recipe/. It first calls upstream FourBModel at the pinned versions, so its outputs on a fixed input set become the reference for step 3. I'll report load, warmup and warm latency separately.
3. Then replace FourBModel inside the worker with a prefill-only Qwen3.5-4B, with the LoRA merged at load. The text adapter only touches the 8 full-attention layers and the MLPs, so the Gated DeltaNet layers run on base weights.

Like the Cua-S1 Nano backend in system1-agents, the worker will accept choice questions only and return a clear error for score and noul. Multimodal (vision tower, image preprocessing, LoRA on the vision MLPs) stays open here as a follow-up.

[hsliuustc0106]

I'd like to take this, starting with the text adapter on CUDA. Hardware: one RTX 6000 Ada (48 GB, sm_89), and also rented NVIDIA GPU access.

Plan:

1. A docs PR first: pinned revisions (trycua/cua, the base model and the 0.2 adapters), the inference contract, and tolerances declared before any comparison. From reading FourBModel, a decision is one forward pass and a softmax over the answer-letter logits (A to Z) at the last position, so there is no decode and at most 26 options.
2. A text worker in src/models/cua_s1/ behind the Add Rust frontend for Jev workers #2 frontend, with setup and example requests in recipe/. It first calls upstream FourBModel at the pinned versions, so its outputs on a fixed input set become the reference for step 3. I'll report load, warmup and warm latency separately.
3. Then replace FourBModel inside the worker with a prefill-only Qwen3.5-4B, with the LoRA merged at load. The text adapter only touches the 8 full-attention layers and the MLPs, so the Gated DeltaNet layers run on base weights.

Like the Cua-S1 Nano backend in system1-agents, the worker will accept choice questions only and return a clear error for score and noul. Multimodal (vision tower, image preprocessing, LoRA on the vision MLPs) stays open here as a follow-up.

• directly import HF
• implement models natively

[Levius-Fubuki]

I'd like to take on the 0.2 multimodal adapter integration, using an RTX 4090 24GB.

@twu3202 I'll coordinate with your text/CUDA work and reuse the agreed revisions and worker interface. I'll start with upstream FourBModel, add screenshot-conditioned choice inference, and provide parity tests and reproducible examples.

Once that works, I'd also like to profile the pipeline and tackle a focused CUDA/Triton optimization, coordinating changes to shared execution code.

Does this split work for you?

[twu3202]

Works for me, the multimodal adapter is yours.

I'm putting the pinned revisions, the prompt format and the /v1/systemone mapping into a docs PR now and will link it here, so we both build on the same contract. Following the checklist above, the first worker loads the model through Transformers and PEFT directly, and upstream FourBModel is only the parity reference.

[twu3202]

Numbers from the text worker in #13 on one RTX 6000 Ada in bfloat16, for prompts of about 150 tokens: 46 ms with the reference path, and 21 ms with flash-linear-attention kernels, the LoRA merged, a float32 head over the 26 letter rows and CUDA graphs, within the tolerance from #11. What is left is mostly GEMMs at about 60% of memory bandwidth and about 2,000 small elementwise kernels per forward pass.

For the native step I'm assuming a Rust engine, per the repo description. Two questions before I start:

1. Should it build on an existing Rust stack or own everything down to the kernels? candle has Qwen3.5 work in progress, and mistral.rs has a Qwen3-Next implementation plus a fork that serves Kev, another Qwen3.5-based decision model from [RFC] Support emerging System 1 models, with CLM as the second model after LAYA #9, over /v1/systemone with prefill-only requests and Gated DeltaNet prefix-state reuse (Upstreaming Kev support: hidden-states prefill path + GDN prefix-state reuse EricLBuehler/mistral.rs#2443). A Qwen3.5 prefill engine here could serve both Cua-S1 and Kev.
2. For the kernels, CUDA C++ called from Rust, or is reusing Triton kernels such as flash-linear-attention's acceptable? xiaoyu-xyz asked a similar question for the CUDA backend in [Help wanted] Accelerated Laya serving on Apple Silicon #3.

[twu3202]

Following up on my two questions above: unless you'd prefer another route, I'll start the native text engine as a Rust worker in src/models/cua_s1/ that owns the whole path from request to answer (prompt, tokenization, the Qwen3.5-4B forward pass with the LoRA merged at load, and the 26-row head). Its kernels will be CUDA C++ under src/backends/cuda/, with cuBLASLt for the GEMMs, and no mistral.rs, candle or Triton underneath.

Before that I checked how far a Rust worker gets on an existing stack. I built one on mistral.rs, using the prefill-only hidden-states request from a fork that serves Kev (EricLBuehler/mistral.rs#2443). It produces the same prompt token ids as the Python worker in #13, returns the same status and error bodies over 567 requests (the fixed inputs, error cases and random bodies), and stays within the #11 tolerance against float32 (largest difference 0.0065, allowance 0.039). Warm p50 on one RTX 6000 Ada in bfloat16, one request at a time, measured with recipe/cua_s1/bench_text.py:

Prompt tokens	Reference (#13)	Fast path (flash-linear-attention, CUDA graphs; not in a PR yet)	Rust on mistral.rs
139	46 ms	21 ms	20 ms
712	109 ms	74 ms	62 ms
15,446	3.1 s	1.7 s	1.5 s

Reading the weights of the 32 layers once takes about 7.4 ms at this card's peak bandwidth, so short prompts still have room. In a profile of the mistral.rs worker, GEMMs take about as long as in PyTorch, prefill runs without CUDA graphs, and Gated DeltaNet prefill on sm_89 goes through a cuTile kernel that takes 3.4 ms for the 24 layers at 148 tokens, against about 1 ms for flash-linear-attention. My first tensor-core version of that kernel in CUDA takes 1.3 ms, which beats cuTile but not yet flash-linear-attention. So the plan is CUDA graphs over length buckets as in the fast path, fused small kernels, and that kernel as the first in-repo piece.

Validation will be the same as for the prototype: the fixed inputs against the float32 reference under the #11 tolerance, and the HTTP diff against the #13 worker. Most of the language-model code could later be reused for the multimodal adapter, and I'll keep the kernel layout in line with the CUDA backend discussion in #3.

[Levius-Fubuki]

Following up on my earlier multimodal CUDA/Triton optimization plan: #12 provides the multimodal worker, parity checks, and an RTX 4090 baseline. I'd like to take on the following work in stages:

1. Profiling and reproducible benchmarks: break down preprocessing, transfers, vision encoding, language layers, and output projection across image sizes, prompt lengths, and question counts.
2. Same-image reuse within a request: reuse image preprocessing, then investigate sharing vision features across questions where they are independent of the question, with output-parity checks.
3. Focused CUDA/Triton optimization: select a measured hotspot, compare existing implementations, and implement a targeted kernel where justified. Report both kernel-level and end-to-end results.
4. Multimodal compilation and CUDA graphs: investigate graph breaks, shape bucketing, and graph capture for varying image and text lengths, including compilation cost, padding overhead, graph memory, and fallback behavior.
5. Controlled concurrency: evaluate bounded queueing, backpressure, and small dynamic batches, measuring throughput, p95/p99 latency, rejection rates, and memory under mixed workloads.
6. Cua-S1 multimodal on Apple Silicon: assess feasibility and correctness first, then propose a scoped Metal operator/kernel integration plan based on profiling and available hardware. Track reference-runtime results separately from native Metal support.

@twu3202 I saw your native Rust/CUDA C++ engine plan. I'll coordinate shared language-layer and execution changes with you, and CUDA kernel placement with @xiaoyu-xyz. My initial focus will be multimodal preprocessing and the vision path; any Triton production integration would be discussed first. For Apple Silicon, I'll coordinate with @cacheline999's LAYA/MPS work.

The first deliverables would be a stage-level profile and a request-local image-reuse PR. Subsequent work would follow the measurements and agreed interfaces, in separate focused PRs. Each optimization will use predefined numerical tolerances, fixed workloads, and reproducible measurements, with startup and warm performance reported separately.

Does this scope fit with your plans?

[twu3202]

Sounds good to me. The native text engine only takes token ids for now; when you get to the language layers, let's agree on how image embeddings and the 3D rotary positions go in.