[BobClawblaw]

What this changes

Fixes #1456. Adds moss_image: MOSS-Image-Base's 5B joint transformer as a native Metal denoiser, from the released bf16 checkpoint as published. It is the same kind of component as Qwen-Image's (#948). The prompt encoder (Qwen3.5) and the image decoder are not ported: tools/zig/moss_image_case.py runs them through the publisher's own pipeline_moss_image.py to prepare cases and decode pictures. The family does not register with the lane engine, the server or families.detect.

• Family: zig/src/families/moss_image (config check, weights, host time MLP and modulation, launches, model), and its kernels in zig/kernels/metal/moss_image.metal.
• Reused kernels: the bf16 projections reuse the Qwen-VL tower's tfp::gemm_nax and the prefill split-K GEMM (fp32 sums for the biased projections). Attention is tfp::attention_nax in its full (non-causal) form. No shared kernel changes.
• Test tool: tf-moss-image checks both guidance branches of every step against the publisher's pipeline on the same inputs, then runs and times the whole guided schedule.
• Chips: it requires the M5's tensor units and refuses other GPUs, since nothing else was checked.

Receipt

• Environment:
  • Mac Studio M5 Ultra 256 GB, macOS 27; TensorFold 772283b6b (code) on main b30d5b9d3.
  • Checkpoint OpenMOSS-Team/MOSS-Image-Base @ 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3.
  • Reference: torch 2.14.1 (MPS, bf16), diffusers 0.42.0, transformers 5.21.0. MPS has no float64, so the case tool computes the publisher's float64 time sinusoid on the host; everything else runs unchanged.
• Precision: bf16 activations, rounded where the publisher's torch code rounds. Every sum is fp32. There is no int8, no cache and no lossy weight.
• Exactness: two cases at the card's defaults (50 steps, guidance 3, the empty negative prompt skipping block 22): a fox at 1024x1024, and the card's Chinese sticky-note prompt at 1200x880 (an odd 55 x 75 grid).
  • Per step, both branches on the pipeline's own input: cosine against torch bf16 is 0.99973 to 0.99998. Against a float32 run of the same transformer on the same inputs (sticky note, all 50 steps), our mean 1 - cosine is 1.602e-4 on the prompt branch and 1.735e-4 on the negative one. Torch bf16's own are 1.602e-4 and 1.715e-4.
  • Whole pictures: 48.52 dB PSNR from the pipeline's picture for the fox, 28.93 dB for the sticky note. Fifty guided steps magnify bf16 rounding into different fine detail, so the ruler is a float32 run of the whole schedule. Over 8 cases (both prompts, 4 seeds each), with the transformer in float32 but torch's bf16 encoder rows, ours is the closer to it in 6, at a mean of 32.3 dB against torch bf16's 31.4 dB. With the encoder in float32 too (measured with exact --fp32-encoder, a follow-up tool that this PR's case tool does not have; it comes with the prompt encoder, moss_image: the Qwen3.5 prompt encoder as a native Metal forward #1459), torch bf16 averages 32.5 dB and ours 31.2 dB: six cases are within 0.5 dB or favor ours, and two sticky-note seeds favor torch's by 4.6 and 8.7 dB. Those 8 ran on an earlier build of this family (0b69e127); this head gives byte-identical latents on both cases above.
  • Repeats: byte-identical across runs.
• Speed, one session with no other GPU work, the transformer only:
  • Ours: 0.489 s a forward at 1024x1024 and 0.486 s at 1200x880, the median of three 50-step runs.
  • Torch MPS bf16: 0.744 s and 0.760 s, the median of its 100 forwards. That makes ours 1.52x and 1.56x.
  • Peak process footprint 11.3 GB, of which the bf16 weights are 9.5 GB.
• Prompt speed: n/a, not a token family.
• Kernel contract, tf-kernel-contract --only "moss image" on the M5 Ultra:

  SUMMARY 11 kernels: 11 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 red as planted, 0 missed; 0 limit problems
  

  • What is registered: mi_norm (plain and modulated), mi_qk (separate and packed q k v rows), mi_situ (both layouts), mi_gate_add (in place), mi_bias, and the attention's three forms.
  • How the steps are checked: each output against float64 at the publisher's bf16 rounding points, inside the range its fp32 steps allow. Eight controls each drop or add one rounding point, and gate 3 catches every one.
  • Attention: checked against attention_ref's float64 formula and bf16-P bound. Its gate 4 uses equal scores over values whose sums reach about 2^22 and cancel.
  • Gate 5: holds the 8 launches and the two split-K projections. --offline cannot compile the tensor-op source for M1/M2 (N/A, as for FSR's tensor kernels), and the family refuses those chips.
• Tests run: all passed.
  • zig build test, test-golden, test-kernel-contract, test-server-security test-pull-security, fuzz-cortheon, dev-tools -Doptimize=Debug.
  • zig fmt --check build.zig zig.
  • tools/zig/lean_check.py: 428 problems, the same as main.
• Not run, and why:
  • Chips: Metal M1 to M4 (no hardware; refused by the family). CUDA (Metal-only family).
  • Full contract: the unfiltered tf-kernel-contract needs a tf-rotary-sweep receipt for this macOS compiler (FlashNext's rotary probes), unrelated to this change.
  • Image quality: no image-quality benchmark (GenEval and the like).

atc overlap: four LFM2 pull requests (#1377, #1391, #1393, #1394) touch the same registration lists (sources.zig, registry.zig, metal_main.zig). Each conflict is one appended line, for whichever lands second.

A recipe that builds this head and makes pictures with the publisher's encoder and decoder: https://github.com/BobClawblaw/moss-image-base-bf16-m5-ultra

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (n/a: a denoiser, no tokens)
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (n/a: no drafts, prompt state or concurrency)
• No precision traded for speed. If the bits change, the description says which sums change and why.
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
• Prompt processing is no slower than the last release. (no shared code changed)
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks @BobClawblaw, this is a tidy component. It changes no shared kernel source, its eleven kernels are in the contract with planted controls, and it reads the released bf16 weights as published. It's also checked step by step against the publisher's own pipeline. On main b30d5b9d3, zig build test passes, fmt is clean and lean adds nothing.

A few small things:

• tools/zig/moss_image_case.py imports and runs the checkpoint's pipeline_moss_image.py. Please say so in the recipe, add the hf download --revision 85975793… line, and check that file's sha256 before running it.
• tests/moss_image.zig:151 leaks the first Prompt if the second prompt() fails, and model.zig:116 refuses with a bare error.TensorUnitsRequired. A short log line would help someone on an M4.
• The reused bf16 GEMMs (qwv_gemm_bf16_nt_*, split-K) still have no gates 1-4. That gap predates this PR, so it isn't a blocker here.

We'll run the contract and one tf-moss-image case on our M5 Ultra.

[BobClawblaw]

Thanks, Ash. All three points are fixed in 88626ebdb, a new commit on top. The branch is also rebased onto current main (b50d4e239); nothing in it conflicted.

• The publisher's Python: moss_image_case.py now hashes pipeline_moss_image.py before importing it. It runs only revision 85975793's file (sha256 1dac1aca…8b31, the same blob as on the Hub) and exits on any other. The docs and the recipe both say that the torch steps run that file, and both carry the hf download --revision 85975793… line. A copy of the file with one line added is refused.
• tests/moss_image.zig: each prompt is now released on its own, so a failed second prompt() frees the first.
• Model.load: logs the GPU's name and why it is refused before returning error.TensorUnitsRequired. I could only check that it builds; no pre-M5 Mac was available to trigger it.

On the M5 Ultra, the rebased 88626ebdb gives byte-identical latents to the measured 772283b6b on the fox case. zig build test, fmt and lean pass, and the contract's line is unchanged:

SUMMARY 11 kernels: 11 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 red as planted, 0 missed; 0 limit problems

A correction to one claim in the description. The "closer to float32 in 6 of 8" ruler ran the transformer in float32 but kept torch's bf16 encoder rows. With the encoder in float32 too (exact --fp32-encoder, added on a follow-up branch), torch bf16's pictures average 32.5 dB PSNR and ours 31.2 dB over the same 8 cases. Six cases are within 0.5 dB of each other or favor ours. Two sticky-note seeds favor torch's by 4.6 and 8.7 dB, where the schedule settled on different fine detail. The per-step velocity numbers are unchanged.

[BobClawblaw]

A small update: the head is now 6aa593720, rebased onto main 9ed939bb2 after #1493 changed the contract's offline limits. It adds one docs commit, which states the float32 ruler both ways (as in my comment above); the code is unchanged.

On the M5 Ultra, the rebased head still gives the fox case's byte-identical latents, and zig build test, fmt and lean pass. The contract's line is the same:

SUMMARY 11 kernels: 11 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 red as planted, 0 missed; 0 limit problems

[ashhart]

Thanks @BobClawblaw, all three are in.

• The recipe now says the case tool imports and runs the checkpoint's pipeline_moss_image.py. It gives the hf download --revision line (docs/recipes/moss-image.md:65-69), and moss_image_case.py:30-39 checks the file's sha256 before exec_module.
• tests/moss_image.zig:151-155 releases each prompt on its own.
• model.zig:117 names the GPU before refusing it.

We merged it onto current main. fmt and lean are clean, test-kernel-contract passes (120 tests), and tf-kernel-contract and tf-moss-image build.

One small thing for the recipe: moss-image.md:115 gives the float32-encoder numbers (32.5 dB against 31.2 dB), but this tree's exact uses the case's bf16 encoder rows and has no --fp32-encoder. Please say they came from a follow-up tool, or add that option here.

Next we'll run tf-kernel-contract --only "moss image" and the fox case through tf-moss-image on our M5 Ultra.

[BobClawblaw]

Thanks, Ash. You're right: this tree's exact has no --fp32-encoder. I didn't add it here, since it belongs with the prompt encoder (#1459). 0ebed8b8b is docs only: docs/recipes/moss-image.md now says the float32-encoder ruler (32.5 dB against 31.2 dB) came from that follow-up tool, which comes with #1459. The PR description says the same. The code is unchanged since 6aa593720.

[ashhart]

Landed on main with your authorship (e82837273 onwards, 2551c57 with the changelog). On our M5 Ultra: tf-kernel-contract --only "moss image" passes 11 of 11 kernels with all 8 controls red, and a 1024x1024 case through tf-moss-image stays at cosine 0.9997982 or better against the publisher's torch bf16 pipeline at every step, repeats byte for byte, runs 50 steps in 47.3 s, and decodes to 31.9 dB PSNR against torch's picture. Thanks @BobClawblaw, a tidy component.