[BobClawblaw]

Closes #1518.

MOSS-Image's transformer (OpenMOSS-Team/MOSS-Image-Base @ 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3, the released bf16 weights as published) on the CUDA runtime, beside the Metal family and driven by the same tf-moss-image case driver. The driver is now shared by both backends (zig/tests/moss_image_run.zig); on Metal it gives latents identical byte for byte to main's tool on the fox case.

Like the Metal family it is a denoiser, not a token family: the prompt encoder and image decoder stay in torch (tools/zig/moss_image_case.py), and nothing registers with the lane engine, the server or families.detect. Details: docs/recipes/moss-image-cuda.md.

What it adds

• zig/src/families/moss_image/cuda*.zig: the model, the weights (the released geometry only; any tensor left unread is refused), and the launchers. The config, rotary tables and the timestep's modulation are shared with Metal.
• zig/kernels/cuda/moss_image.cu: moss_image.metal's four kernels op for op (modulated RMSNorm, q/k norm and rotary into packed rows, situ, gated add), built with the glue flags (--fmad=false, IEEE division and square root).
• Projections use cuBLASLt bf16 GEMMs with fp32 sums (vit.Gemms / vit.linear). Attention uses prefill_attention.cu's full vision_kernel, which gets a new instance at head dim 128.
• Model.load refuses GPUs other than sm_120 and sm_121, the two it was checked on.
• Precision is bf16 activations and weights with fp32 sums, norms and softmax. There is no int8, fp8 or fast math.

Kernel contract

tf-cuda-test moss-contract registers the four kernels in each launched form, plus the head-128 attention. It passes on a GB10 (nvcc 13.0 and 13.2) and on an RTX 5090 (nvcc 13.3):

SUMMARY moss image: 8 kernels, 8 pass gates 1-4, 0 fail; 5 controls: 5 red as planted, 0 missed

Gate-5 records are in zig/src/core/contract/moss_image_cuda_records.zig (sm_80, 86, 89, 120, 121). kernel-contract-limits finds all 25 MOSS lines current on nvcc 13.0, 13.2 and 13.3, so no later release needs its own set. The audit's other problems are other families' records and are the same with and without this PR.

Shared code

• prefill_attention.cu: one new instance, vision_kernel<128, 8>. The six existing functions compile to identical SASS before and after on sm_80, 86, 89, 120 and 121, under nvcc 13.0, 13.2 and 13.3.
• vit.zig: fulls gains dim 128. As a result, every tower's Attention.load (SigLIP2 in gemma4 and lfm2, RADIO) also resolves the head-128 function and opts it into 69,632 B of dynamic shared memory. Dispatch at dims 72 and 80 is unchanged.
• A/B on a GB10, main vs this branch:
  • tf-cuda-test lfm2-vision and radio-kernels: the same output apart from timings.
  • radio-attention (head-dim-80 tower attention at five sizes): 0.008/0.012/0.152/2.975/12.884 ms on main vs 0.008/0.011/0.151/2.973/12.950 ms on the branch.
  • No full-model tower A/B was run: no Gemma 4, LFM2-VL or RADIO checkpoint with its references was at hand. The towers' kernels are byte-identical in SASS and their launches are unchanged.

Measured

Two cases with the Metal doc's prompts and seeds: a fox at 1024x1024 (seed 42) and the card's sticky note at 1200x880 (seed 7). The card's defaults apply (50 Euler steps, guidance 3). The reference is the checkpoint's pipeline_moss_image.py in bf16 (torch 2.11 on the GB10, torch 2.13 on the 5090; diffusers 0.38).

	GB10 (sm_121)	RTX 5090 (sm_120)
Velocity cosine vs torch bf16, every step and branch	0.99972-0.99998 / 0.99981-0.99998	0.99967-0.99998 / 0.99978-0.99998
Mean 1 - cos vs a float32 run, prompt (torch bf16's own)	1.598e-4 (1.596e-4) / 1.602e-4 (1.602e-4)	1.607e-4 / 1.612e-4
Picture PSNR vs the pipeline's own	45.93 / 32.76 dB	44.01 / 34.51 dB
Forward, median of 100 (torch bf16)	0.690 s (1.402) / 0.700 s (1.477)	0.243 s (0.307) / 0.245 s (0.309)
50 guided steps	69.6 s / 70.6 s	24.6 s / 24.8 s

Values are fox / sticky note. Against the float32 run, ours is the closer of the two bf16 pipelines at about half the steps on each case. The same step run twice, with another between, gives identical bytes, and repeated runs give identical latents. The sticky note's handwritten text comes out the same; fine detail differs, as on Metal. Scratch is 0.51 GB at 1024x1024 with a 34-row prompt, computed from the allocations. The 5090's latents differ from the GB10's in their low bits because cuBLASLt picks different GEMM kernels on each card.

Checks

• On the rebased branch on the GB10: test-cuda-host passes; moss-contract gives the summary above; the 512x512 fox case's latents are byte-identical to before the rebase; zig fmt is clean; lean_check.py reports 427 problems, the same 427 as main, with none in these files. There is no getenv in the diff.
• The 5090 numbers and the 1024 runs were taken before the rebase onto current main. The patch is identical (git range-diff reports =), and none of main's new commits touch the CUDA runtime, vit.zig or prefill_attention.cu.

Not covered

• GPUs other than sm_120 and sm_121 (refused at load).
• Ports of the prompt encoder and the VAE decoder (still torch).
• A full-model A/B of the vision towers (see Shared code).

Overlap

atc overlap lists #1439, #1508 and #1524 as conflicting. The conflicts are in shared registration lists: zig/build/cuda.zig, cuda_kernels.zig, kernels.zig, registry.zig, contract_controls.cu, tests/cuda/main.zig, kernel-contract.md, and prefill_attention.cu's instance list, where #1508 adds different pattn_kernel instances. Whichever PR lands second needs a trivial rebase.

[ashhart]

Hi @BobClawblaw, a status update on #1528.

One fix is needed before it can land, and we have it ready if you'd like us to carry it. The four MOSS records
name the image "moss_image", but the full kernel contract's image list in zig/tests/cuda/contract_limits.zig
doesn't have that name. So tf-cuda-test kernel-contract would stop with "unknown image name: moss_image" on
every GPU. moss-contract alone doesn't reach that list, which is why it passed. The fix adds the name. A host
test that every registered CUDA image resolves would catch the next one.

After that, what's owed is a run on one GB10: the build with the moss_image fatbin, moss-contract, the full
kernel contract and its limits, and the SASS identity of prefill_attention. There are no served cells, since
nothing in a serving path changes.

[BobClawblaw]

Thanks for catching that, please carry it. Adding moss_image to the contract's image list plus a host test that every registered CUDA image resolves sounds right; moss-contract passing alone hid it.

[BobClawblaw]

Thanks @ashhart. Rebased onto main (5ded79a), and no need to carry the fix: main's new host test that every registered CUDA launch resolves (from d704476) caught moss_image once rebased, so it's added to contract_limits.zig's image list as its own commit. Both conflicts in the rebase were in zig/build/cuda.zig beside main's qwen35-cuda step, which is kept as main has it.

On this head, zig fmt --check, zig build test (CUDA host tests included), test-kernel-contract, test-golden, the server and pull security tests, fuzz-cortheon and the Python suites pass. The GB10 run you listed (the build with the moss_image fatbin, moss-contract, the full kernel contract and limits, prefill_attention's SASS identity) is still to come.

[ashhart]

Thanks @BobClawblaw, and thanks for adding moss_image to the contract's image list yourself. That answers our last note, and the rebase onto current main is clean.

One more fix is needed before the GPU run. zig/build/cuda.zig installs tf-moss-image on Linux and on Windows, but seconds() in zig/src/families/moss_image/cuda_model.zig calls std.c.clock_gettime, which Zig doesn't provide for Windows. So a CUDA build on Windows (zig build) now stops at that tool:

compile exe tf-moss-image ... error: parameter of type 'void' not allowed in function with calling convention 'x86_64_win'
    extern "c" fn clock_gettime(clk_id: clockid_t, tp: *timespec) c_int;
referenced by: zig/tests/moss_image_run.zig:143, zig/tests/moss_image_cuda.zig:38

You can reproduce it from any host with zig build install -Dtarget=x86_64-windows-gnu. That tool is the only step that fails. The Windows CI job runs only the tests, which is why it passed.

The simplest fix is the platform shim the other CUDA code uses. The module already imports sys, so in seconds():

const c = @import("sys").c;
var ts: c.timespec = undefined;
_ = c.clock_gettime(c.CLOCK.MONOTONIC, &ts);

(as zig/src/cuda/memory_receipt.zig does). Skipping the install on Windows, the way h3 is skipped, would also work, but the shim keeps the tool buildable there.

After that, what's owed is the same as before, and we'll run it ourselves on one GB10: the build with the moss_image fatbin, moss-contract, the full kernel contract and its limits, and prefill_attention's SASS identity for the existing instances.

[BobClawblaw]

Fixed in 667839f: seconds() now goes through the sys shim, as in memory_receipt.zig. zig build install -Dtarget=x86_64-windows-gnu failed before with your error and now builds cleanly, and a native Windows CUDA build (sm_89) installs tf-moss-image, which runs. Thanks for running the GB10 checks.

[BobClawblaw]

An RTX 4090 (sm_89) run of this family, for reference. Model.load refuses sm_89 here, so it ran with the gate widened to admit it; nothing else changed. That one commit, with the doc lines, is 6d2b69de on BobClawblaw:moss-image-cuda-sm89 (on top of this PR's head). Happy to add it here or open it after this lands, or leave it out, whichever you prefer.

Setup: RTX 4090 24 GB, Windows 11, driver 616.56, nvcc 13.3, -Dsm=89, ReleaseFast. The cases were made on the same card by tools/zig/moss_image_case.py with the Metal doc's fox and sticky-note prompts and seeds (torch 2.12.0 CUDA bf16, diffusers 0.42.1, transformers 5.22.1). The tool's pipeline holds the 16 GB encoder and the transformer at once, which doesn't fit in 24 GB, so a wrapper had the two take turns on the card. The tool itself was unchanged.

• tf-cuda-test moss-contract: SUMMARY moss image: 8 kernels, 8 pass gates 1-4, 0 fail; 5 controls: 5 red as planted, 0 missed

	RTX 4090 (sm_89)
Velocity cosine vs torch bf16, every step and branch	0.99972-0.99998 / 0.99978-0.99998
Mean 1 - cos vs a float32 run, prompt (torch bf16's own)	1.600e-4 (1.609e-4) / 1.652e-4 (1.638e-4)
Picture PSNR vs the pipeline's own	46.54 / 29.85 dB
Forward, median of 100 (torch bf16)	0.352 s (0.490) / 0.358 s (0.499)
50 guided steps	35.2 s / 35.8 s

Values are fox / sticky note. Against the float32 runs, ours is the closer at 25 and 20 of the fox's 50 steps (prompt, negative) and 22 and 23 of the sticky note's. The same step run twice, with another between, gives identical bytes, and repeated runs give identical latents. The sticky note's handwritten text comes out the same.

Against a float32 run of the whole schedule, over 8 cases (both prompts, seeds 42/7, 1, 2, 3), ours is the closer in 5, at a mean of 31.62 dB PSNR against torch bf16's 31.35 dB (table).

Fox, 1024x1024, seed 42: the pipeline (left), ours (right)

Sticky note, 1200x880, seed 7: the pipeline (left), ours (right)

Every case's run.json, engine log, float32 picture and side-by-side are in the evidence folder, with four more prompts in the gallery.