Skip to content

cua_s1: accept image embeddings and 3D positions in native language model - #56

Draft
Levius-Fubuki wants to merge 2 commits into
ThinkFlowLab:mainfrom
Levius-Fubuki:codex/cua-native-multimodal-input
Draft

Levius-Fubuki wants to merge 2 commits into
ThinkFlowLab:mainfrom
Levius-Fubuki:codex/cua-native-multimodal-input

Conversation

@Levius-Fubuki

@Levius-Fubuki Levius-Fubuki commented Oct 1, 2026 •

Copy link
Copy Markdown
Collaborator

Purpose

Add Model::forward_multimodal to accept adapted BF16 image features and explicit int64 T/H/W positions ([3,1,S]). Validate the unpadded prompt, insert image features into ordered placeholders, and execute the language layers with interleaved MRoPE. Accept standalone merged-language weights and preserve the text CUDA Graph cache with separate immutable text position tables.

The diff contains only inputs.rs, lib.rs, and model.rs. CUDA ABI remains 3; the HTTP worker remains text-only.

Validation

After cleanup: workspace formatting, strict all-target Clippy, workspace tests and the locked release build passed. 19 tests passed; 3 existing GPU tests skipped. Runtime source was compared against the previous head: only new trailing test modules were removed; production code is unchanged. Existing baseline tests remain.

Feature-specific tests, examples, documentation and validation assets were removed from this PR diff to keep it focused on core implementation. They remain in the pre-cleanup commit and a local archive.

No new standalone GPU run was performed for this cleanup. The integrated language/Graph path was exercised as part of #64; that evidence applies to its combined runtime.

Self-review

Cleanup reviewed for unchanged runtime code, retained licensing and valid build targets. Draft status is retained for contributor/maintainer 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.

Copilot AI balanced review requested due to automatic review settings October 1, 2026 02:04

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@twu3202

twu3202 commented Oct 1, 2026

Copy link
Copy Markdown
Contributor

Reviewed 93a6cb0 on an RTX 6000 Ada (sm_89) with CUDA 13.2 and driver 595.91.07, against main 566dec1. Both builds used the same libqwen3_5_cuda.so, since the kernels are unchanged.

The HTTP text worker gives the same answers as before on this card:

  • I sent 567 requests to a main worker and a cua_s1: accept image embeddings and 3D positions in native language model #56 worker side by side: 560 text requests covering answers and validation errors, plus 7 HTTP edge cases. Status, content type and body are byte-identical on all 567, including the 134 answered ones.
  • 276 prompts (272 lengths from 173 to 2,024 tokens, plus 2,047, 2,048, 2,049 and 16,384), cold and warm, then 8 concurrent clients, four runs per build: every answer is byte-identical across builds and runs, and so is the 413 for 16,385 tokens.
  • Latency is the same to within about 1%. The spread between runs followed the card's temperature when a run started, not the build: runs starting at 55 to 58 C averaged 93.0 to 94.6 ms per warm request for both builds, and runs starting at 72 to 75 C averaged 100.4 to 102.5 ms.
  • The CPU tests and the three GPU tests pass here. For the multimodal GPU test I used our merged text weights with the base model's config.json, which adds image_token_id, because the native loader takes BF16 only and the base checkpoint stores A_log and the linear-attention norms in F32.

Today's worker never calls forward_multimodal, so the following only matters for a process that serves both. The first text call after a multimodal call rebuilds the rotary tables on the host for every scratch row, and the scratch keeps the size of the longest prompt so far. Medians of 30, with each timed text call preceded by a 2-token call (an image-free multimodal one, or a text one for the baseline):

Scratch rows Host rebuild alone Extra on a 139-token request Extra on a 712-token request
1,024 0.45 ms 0.51 ms 0.05 ms
2,048 0.94 ms 1.08 ms 0.11 ms
4,096 1.95 ms 2.19 ms 0.63 ms
8,192 3.96 ms 4.45 ms 1.73 ms
16,384 8.00 ms 8.72 ms 2.26 ms

The 139-token request itself takes about 12.8 ms. On the 712-token request the extra time was smaller and noisier, and I didn't look into why. Text outputs after a multimodal call matched the text-only ones every time. Keeping a device copy of the text tables and copying it back would remove nearly all of this.

#52 is on main now, so this needs a rebase, and the forward conflict has a trap. After this PR, run() no longer embeds, and the layers update s.res in place (add_rms_norm_kernel writes residual[i]). On a cache miss, main now runs run() eagerly and then launches the captured graph, which works only because main's run() embeds from s.ids first. If the conflict is resolved by calling embed_tokens before that block, the replay starts from the residual the eager pass already advanced, so with CUA_S1_GRAPH=1 every miss (a new length, or one evicted from the 8 entries) returns a wrong hidden state without an error. Using the eager result on a miss, or embedding again before the launch, avoids that. The numbers above are from 93a6cb0 against 566dec1, before #52 landed.

Minor:

  • The interleave unit test uses sections [12, 10, 10], so the real [11, 11, 10] layout, where index 31 belongs to H, is not covered.
  • The "Replay a reference boundary" subsection and the example depend on the bundle format from Add Cua-S1 multimodal reference tensor exports #53, which is still open.

@Levius-Fubuki
Levius-Fubuki marked this pull request as draft October 1, 2026 09:09
@Levius-Fubuki
Levius-Fubuki force-pushed the codex/cua-native-multimodal-input branch from 93a6cb0 to dea93f2 Compare October 1, 2026 09:09
@Levius-Fubuki
Levius-Fubuki force-pushed the codex/cua-native-multimodal-input branch from dea93f2 to 75a80cb Compare October 1, 2026 09:13
@Levius-Fubuki

Levius-Fubuki commented Oct 1, 2026 •

Copy link
Copy Markdown
Collaborator Author

@twu3202 Rebased onto main fd0d87e in 75a80cb and addressed the actionable items:

  • Graph misses return the warm-up eager result, then retain the captured layers for future hits. Hits first re-embed the uploaded token IDs. The captured graph excludes embedding, so the eager residual is never replayed on a miss. Existing failure fallback and cache invalidation before scratch growth/drop are retained.
  • Text rotary tables stay immutable on the device. Explicit-position calls use a separate pair of tables, so switching back to text does not rebuild host tables or copy tables back. This adds 2 MiB at 16,384 scratch rows; no CUDA backend/ABI change beyond main's ABI 3 is needed. I have not measured new latency.
  • The distinct-axis interleave oracle now covers actual [11,11,10], including H at frequency 31, in addition to the synthetic layout.
  • The recipe/example explicitly identify Add Cua-S1 multimodal reference tensor exports #53 as still open and pin its exporter revision. Only the optional replay fixture producer/verifier depend on that PR; the Rust model API can accept features/positions directly.

Added a full-model GPU regression for first-use misses, hits with changed IDs, eviction after eight cached lengths, mixed multimodal/text calls and scratch growth, comparing hidden states against eager execution and asserting capture succeeds. The existing feature-insertion regression remains.

Local format, Clippy, release build, all 25 Rust CPU tests, all 7 benchmark tests, smoke-manifest validation and strict Docs build pass. Five GPU tests are ignored by default. I could not execute the new/rebased GPU checks: the server used for the original evidence is shut down and SSH refuses connections. I have therefore changed the PR to draft pending fresh ABI-3 GPU parity/replay checks. The PR description distinguishes current CPU results from the historical 93a6cb0 GPU evidence; none of the original accuracy or latency numbers are relabelled as results for 75a80cb.

Update: GitHub Rust, benchmark and Docs CI all pass on 75a80cb. The five hardware tests remain unexecuted and the draft/GPU-validation status above is unchanged.

This branch has not been deployed

No deployments
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