[BobClawblaw]

What this changes

Fixes #1504. MOSS-Image's image decoder on Metal, inside the moss_image family (#1457, #1512). The decoder is the checkpoint's vae/, FLUX.2's small VAE decoder (AutoencoderKLFlux2), run as the pipeline's decode does, in float32: the latents' batch norm undone with the saved statistics, the 2x2 patches unpacked, post_quant_conv, conv_in, the middle block with its single-head attention, four upsampling stages, conv_out, then the pipeline's own 8-bit conversion. tf-moss-image now frees the transformer and writes picture.png itself, so torch no longer decodes.

Stacked on #1512: the first three commits are #1512's. This PR's own are the last two, 602714801 (code) and 7597295ff (docs). Once #1512 lands, I'll rebase this onto main.

• Family: decoder.zig (the released vae/config.json geometry, checked; the float32 weights as published; the forward and its scratch) and png.zig (the picture's PNG encoder).
• Kernels: nine of the family's own in zig/kernels/metal/moss_image_decoder.metal: md_unpack, md_wino and md_wino_32 (the 3x3 convolutions as Winograd F(2x2, 3x3)), md_conv (1x1), md_norm (GroupNorm, with or without SiLU), md_group_part and md_group_finish (its statistics), md_attention and md_pixels. No shared kernel changes.
• Winograd: the weights' transform G g G^T is computed in float64 at load and rounded once to float32; the input's and the output's transforms run in fp32 in the kernel, with 2.25x fewer products than a direct convolution. Its error stays the size of a direct convolution's (see Exactness). Like any Winograd convolution, an infinite input can come out NaN where a direct one gives an infinity. The contract checks the special-value classes against the Winograd formula's.
• Tools:
  • tf-moss-image decodes the final latents after the transformer is freed and writes picture.png beside latents.f32.
  • tf-moss-image-decoder MODEL_DIR LATENTS WIDTH HEIGHT REFERENCE TORCH OUT checks one decode. It must be within twice torch float32's distance from float64, in cosine error and in its worst row of pixels, and a second decode must repeat the first byte for byte. It exits 1 otherwise. The gate is reference.within, beside the encoder's, with a test that plants a row further off than the reference allows.
  • The case tool gains vae, which writes the publisher's decoder's output before the clamp in float64 on the CPU and in float32 on MPS, as the pipeline runs it.
• Chips: Decoder.load refuses a GPU without tensor units and names it, like Model.load and Encoder.load.

Receipt

• Environment:

  • Mac Studio M5 Ultra 256 GB, macOS 27; TensorFold 602714801 (code) on MOSS-Image: the prompt encoder on Metal #1512's 92c829887, on main dd0a9cc5e.
  • Checkpoint OpenMOSS-Team/MOSS-Image-Base @ 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3.
  • Reference: the checkpoint's pipeline_moss_image.py (sha256-checked) and diffusers 0.42.0's AutoencoderKLFlux2, on torch 2.14.1.

• Precision: float32 throughout, as the pipeline decodes. Weights as published (float32), fp32 sums in every convolution and in the attention, GroupNorm's mean and variance in two fp32 passes, and the pipeline's own 8-bit conversion (clamp, * 0.5 + 0.5, * 255, round half to even). No half, int8 or lossy weight.

• Exactness, on the transformer's final latents of the fox case (1024x1024) and the card's sticky-note case (1200x880), against the publisher's decoder in float64 on the CPU:

  case	ours: max |diff|, 1-cosine, worst row	ours: 8-bit values off float64	torch float32 on MPS: 1-cosine, worst row	torch: 8-bit values off
  fox	8.86e-6, 6.69e-13, 1.90e-6	162 of 3,145,728, each by 1	6.66e-13, 1.92e-6	155, each by 1
  sticky note	5.23e-6, 4.01e-13, 1.25e-6	45 of 3,168,000, each by 1	3.64e-13, 1.27e-6	34, each by 1

  • Gate: both pass at about half the limit (fox: 1-cosine 6.69e-13 of 1.33e-12, worst row 1.90e-6 of 3.83e-6).
  • Torch varies: torch's off-by-one counts move between its own runs (161 and 45 in an earlier session, against 155 and 34 in this one, on the same byte-identical float64 references), so the gate is on cosine error and the worst row, not on counts.
  • Ours against torch's picture: at most 1 apart, 95.4 and 97.7 dB PSNR.
  • Planted mismatch: one row of the fox's references 0.1% off fails the gate (worst row 9.98e-4 of 3.83e-6) and exits 1, with an identical repeat.
  • Repeats: three runs of each case, each decoding twice, all byte-identical.
  • End to end: with the case's precomputed prompt rows, tf-moss-image's latents equal the pinned run's byte for byte; with --prompt, they equal MOSS-Image: the prompt encoder on Metal #1512's byte for byte. The decoded picture is unchanged across rebases.

• Speed, back to back in one session with no other GPU work, the median of three:

  • Ours: 0.181 s (fox) and 0.188 s (sticky note) on the GPU, and 0.3 to 0.4 s to load.
  • torch's warm float32 decode on MPS: 0.186 s and 0.191 s.

• Memory: the decoder's scratch is five float32 picture tensors at the widest stage, 4.7 GB at 1024x1024, allocated after the transformer is freed. tf-moss-image peaks at an 11.3 GB footprint from a case's rows, and at 16.4 GB with --prompt, while the encoder is loaded.

• Prompt speed: n/a, not a token family.

• Kernel contract, tf-kernel-contract --only "moss image" on the M5 Ultra:

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

  • The new entries: 15, for the nine kernels in the forms the decoder launches. md_wino and md_wino_32 run as a plain 3x3, a 3x3 read nearest from half size, and a 3x3 with a residual. md_conv runs with and without a residual, md_norm with and without SiLU, and md_group_part plain and centered.
  • How they are checked: each fp32 step's output must lie in the range float32 allows about float64. The Winograd kernels are checked against the float64 convolution through the exact transform, within a bound on the transformed terms' magnitude. md_pixels is checked byte for byte against the same fp32 steps, with every tie between two bytes among its inputs.
  • Gate 4: integer channels, pixels or values whose sums pass fp16's range and cancel.
  • Controls: nine, each breaking one kernel: a swapped patch axis, a 1x1's inputs staged in half, a Winograd patch staged in half, a term dropped from its output transform, a SiLU in half, both statistics kernels summing in half, the attention's rescale dropped, and the pixels' halves rounded up. Each is caught at the gate it names.
  • Gate 5: nine more launch records, in moss_image_records.zig.

• Tests run: all passed.

  • zig build test, test-golden, test-kernel-contract, test-server-security test-pull-security, fuzz-cortheon, dev-tools -Doptimize=Debug and check-release-paths.
  • zig fmt --check build.zig zig.
  • tools/zig/lean_check.py: 426 problems, the same as MOSS-Image: the prompt encoder on Metal #1512's head.

• Not run, and why:

  • Chips: Metal M1 to M4 (no hardware; refused by the family). CUDA (a Metal-only family).
  • Image quality: no image-quality benchmark (GenEval and the like). The decoder's output is checked against float64 instead.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (n/a: an image decoder for a denoiser, no decoding of 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.

[BobClawblaw]

Rebased onto 20705d6 (with #1512 under it). With the case driver now shared by Metal and CUDA (#1528), the decoder is the Metal backend's decode hook: the transformer's run sits in its own scope, so the model is freed before the decoder loads, as before. On CUDA, tf-moss-image writes the latents and run.json as on main, and no picture.

Checked: fmt, lean (0 new) and zig build test on the pushed head; the kernel contract, golden, security, fuzz and dev-tools on the same commits one rebase earlier (787a9ef8e). On an M5 Ultra, the 1024 fox case with --prompt gives latents and picture.png byte-identical to the pre-rebase head 7597295 (2,097,152 and 2,272,723 bytes). On an RTX 5090, the 512 fox case's latents are byte-identical to main's (524,288 bytes, 5.72 s both), and --prompt exits 2. The CUDA build needs the one-brace fix in #1769, which main is missing. The rebase onto 20705d6 changed no file this PR touches; the hardware runs were on its ab4b5d8 rebase.

[ashhart]

Heads-up @BobClawblaw: #1512 is merged into main (1157ea5), so this one can be rebased whenever suits you. Once it's down to your two commits, we'll gate it and land it. Thanks for the encoder!

[BobClawblaw]

Rebased onto main bd10bd9: the PR is now its two commits (522108b, the decoder, then 02bf950, its docs); the decoder's files are unchanged. Checked on the same two commits over 1157ea5 (main since adds only a CHANGELOG line): zig build test, test-kernel-contract, test-golden, fmt and the lean check pass (424 lean problems, all on the base). On the M5 Ultra, tf-kernel-contract --only "moss image" gives SUMMARY 37 kernels: 37 pass gates 1-4, 0 fail, 0 known findings; 25 controls: 25 red as planted, 0 missed; 0 limit problems, and tf-moss-image-decoder on the fox latents passes its float64 gate (1−cosine 5.82e-13 of 1.45e-12, worst row 1.900e-6 of 3.918e-6), repeated byte for byte.

[ashhart]

Thank you @BobClawblaw. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

[ashhart]

Thanks @BobClawblaw! MOSS-Image's image decoder on Metal, with tf-moss-image writing picture.png itself, is in 1.2.1: https://github.com/ashhart/TensorFold/releases/tag/v1.2.1