MOSS-Image: the image decoder on Metal

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- MOSS-Image's image decoder (the checkpoint's FLUX.2 small VAE, in float32 as the pipeline runs it) runs on Metal inside the `moss_image` family, and `tf-moss-image` writes `picture.png` itself, so torch no longer decodes (#1687, #1504). Thanks to @BobClawblaw.
+
 - On CUDA, a driver older than CUDA 13.0 is refused at startup with a plain message naming the driver to install (R580 or newer), before any kernel loads; an R550 driver had ended the process with SIGILL inside libcuda on the first module load (#1791, #1724). Thanks to @rootsbymenda.
 
 - Spark 2.5 on CUDA: `TF_KV_SHARE=n` keeps prompts as the KV pool's pages in the shared prompt store, with at most n kept states, so a conversation resumes at the store's last 256-token mark from pages it already holds instead of a full copy. Unset or 0, nothing changes. The native CUDA adapter it uses is family-neutral; Kolibri-1 follows.

docs/recipes/moss-image.md

@@ -11,21 +11,23 @@ transformer on the Metal runtime from the released bf16 diffusers checkpoint:
 
 The card's guidance runs two passes a step: the prompt, and the empty negative prompt with block 22 skipped.
 
-The prompt encoder, the checkpoint's Qwen3.5 text model, runs on Metal too (below): `tf-moss-image --prompt` makes both
-prompts' rows itself. The image decoder (FLUX.2's small decoder) is not ported. `tools/zig/moss_image_case.py` runs it,
-and the reference pipeline, through the publisher's own `pipeline_moss_image.py` in torch.
+The prompt encoder, the checkpoint's Qwen3.5 text model, and the image decoder, FLUX.2's small VAE decoder, run on
+Metal too (below): `tf-moss-image --prompt` makes both prompts' rows itself, and `tf-moss-image` turns the final latents
+into the picture. `tools/zig/moss_image_case.py` runs the reference pipeline through the publisher's own
+`pipeline_moss_image.py` in torch, a development reference.
 
 It is a denoiser, not a token family. It does not register with the lane engine, the server or `families.detect`.
 
-It runs on the M5's tensor units alone, the only chip it was checked on: `Model.load` refuses a GPU without them.
+It runs on the M5 alone, the only chip it was checked on: `Model.load` and `Decoder.load` refuse a GPU without tensor
+units.
 
 The same transformer runs on CUDA too: [moss-image-cuda.md](moss-image-cuda.md).
 
 ## Checkpoint
 
 `OpenMOSS-Team/MOSS-Image-Base` at `85975793d8b3d432f1b6e27f352abaa4c2a4c0d3`, Apache 2.0. `transformer/config.json`
-must be `MOSSImageTransformer` at the released geometry; any other config is refused, and so is a checkpoint with
-tensors this family does not read.
+must be `MOSSImageTransformer` at the released geometry, and `vae/config.json` `AutoencoderKLFlux2` at its released
+decoder geometry. Any other config is refused, and so is a checkpoint with tensors this family does not read.
 
 ## Precision
 
@@ -37,6 +39,14 @@ tensors this family does not read.
 - The rotary tables in bf16, as the pipeline casts them; its frequencies are the checkpoint's `rope.bands`.
 - No int8.
 
+The decoder is float32 throughout, as the pipeline decodes: its weights as published (float32), fp32 sums in every
+convolution and in its attention, GroupNorm's mean and variance in two fp32 passes, then the pipeline's own 8-bit
+conversion (clamp, `* 0.5 + 0.5`, `* 255`, round to even). Its 3x3 convolutions run as Winograd F(2x2, 3x3)
+(`md_wino`): the weights' transform G g G^T is computed in float64 at load and rounded once, the input's and the
+output's transforms in fp32 in the kernel, 2.25x fewer products than a direct convolution. Its error stays the size
+of a direct convolution's (see Measured). Like any Winograd convolution, an infinite input can come out NaN where a
+direct one gives an infinity; the contract checks the classes against the Winograd formula's.
+
 The bf16 projections are the Qwen-VL tower's `tfp::gemm_nax`. Attention is the prefill `tfp::attention_nax` template's
 full (non-causal) form. The time MLP and both modulation projections run on the host, once a timestep.
 
@@ -118,10 +128,28 @@ Gate 4:
 
 Gate 5 holds the 17 launches of the family's source and the two split-K projections.
 
+`zig/tests/contract/metal_moss_image_decoder*.zig` registers every kernel of `moss_image_decoder.metal`:
+
+- `md_unpack` (the latents' batch norm undone, the patches unpacked);
+- `md_wino` (64 output channels a threadgroup) and `md_wino_32` (32, for 96 and 3 channels) in the forms the decoder
+  launches: 3x3, 3x3 read nearest from half size, and 3x3 with a residual, against the float64 convolution through the
+  exact transform, within a bound on the transformed terms' magnitude;
+- `md_conv`, the 1x1 convolutions, with and without a residual;
+- `md_norm`, with and without SiLU, and GroupNorm's statistics: `md_group_part` (plain and centered) and
+  `md_group_finish`;
+- `md_attention` and `md_pixels`.
+
+Each fp32 step's output must lie in the range float32 allows about float64. Gate 4 uses integer channels, pixels or
+values whose sums pass fp16's range and cancel. `md_pixels` is checked byte for byte against the same fp32 steps, with
+every tie between two bytes among its inputs. Nine controls each break 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 holds its nine launches.
+
 `tf-kernel-contract --only "moss image"` on the M5 Ultra:
 
 ```
-SUMMARY 22 kernels: 22 pass gates 1-4, 0 fail, 0 known findings; 16 controls: 16 red as planted, 0 missed; 0 limit problems
+SUMMARY 37 kernels: 37 pass gates 1-4, 0 fail, 0 known findings; 25 controls: 25 red as planted, 0 missed; 0 limit problems
 ```
 
 ## Run a case
@@ -133,9 +161,8 @@ only when the file's sha256 is the pinned revision's (`1dac1aca…8b31`) and ref
 hf download OpenMOSS-Team/MOSS-Image-Base --revision 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3 --local-dir MODEL_DIR
 python3 tools/zig/moss_image_case.py case MODEL_DIR case/ --prompt-file prompt.txt   # needs torch, diffusers, transformers
 zig build tf-moss-image -Dcpu=apple_m1
-zig-out/bin/tf-moss-image MODEL_DIR/transformer case/case.safetensors case/
-python3 tools/zig/moss_image_case.py decode MODEL_DIR case/latents.f32 case/zig.png
-python3 tools/zig/moss_image_case.py compare case/reference.png case/zig.png case/side.png
+zig-out/bin/tf-moss-image MODEL_DIR/transformer case/case.safetensors case/      # case/picture.png
+python3 tools/zig/moss_image_case.py compare case/reference.png case/picture.png case/side.png
 python3 tools/zig/moss_image_case.py exact MODEL_DIR case/      # the same schedule with a float32 transformer
 python3 tools/zig/moss_image_case.py exact MODEL_DIR case/ --fp32-encoder prompt.txt   # and a float32 encoder
 
@@ -155,10 +182,23 @@ The case holds:
 - with `--fp32-steps N`, a float32 run of the same transformer for the first N steps.
 
 `tf-moss-image` first runs both branches of each step on the pipeline's own input and compares the velocities. It then
-runs the whole guided schedule from the starting latents, times it, and writes the final latents. A case made by
-`start` holds only the prompts, the schedule and the noise; `tf-moss-image` then just makes the picture. With
-`--prompt` (and `--negative`, else the empty negative prompt), it encodes both prompts natively in place of the case's
-rows; `start --no-encoder` leaves them out of the case, and torch never loads the encoder.
+runs the whole guided schedule from the starting latents, times it, frees the transformer, and decodes the final
+latents with the checkpoint's `vae/` decoder into `picture.png`. A case made by `start` holds only the prompts, the
+schedule and the noise; `tf-moss-image` then just makes the picture. With `--prompt` (and `--negative`, else the empty
+negative prompt), it encodes both prompts natively in place of the case's rows; `start --no-encoder` leaves them out of
+the case, and torch never loads the encoder.
+
+To check the decoder alone against a float64 decode of the same latents:
+
+```sh
+# float64 on the CPU, and float32 on MPS as the pipeline decodes
+python3 tools/zig/moss_image_case.py vae MODEL_DIR case/latents.f32 case/decode.f64 case/torch.f32
+zig build tf-moss-image-decoder -Dcpu=apple_m1
+zig-out/bin/tf-moss-image-decoder MODEL_DIR case/latents.f32 1024 1024 case/decode.f64 case/torch.f32 case/decoded.png
+```
+
+Ours 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 fails otherwise.
 
 On MPS, which has no float64, the case tool computes the time sinusoid on the host in float64 and moves its float32
 result to the GPU; every other step of the publisher's code runs unchanged.
@@ -226,3 +266,12 @@ Results:
     pipeline's own picture: 47.07 dB.
   - Speed, one session with no other GPU work: 0.054 to 0.067 s a prompt against torch MPS's 0.158 to 0.165 s
     (warm). Loading takes 0.6 s against 2.60 s.
+- **Decoder**, on the transformer's final latents of both cases, against the publisher's decoder in float64 on the CPU:
+  - Ours: max |diff| 8.9e-6 (fox) and 5.2e-6 (sticky note) before the clamp; 162 of 3,145,728 and 45 of 3,168,000
+    8-bit values differ from float64's, each by 1.
+  - torch's own float32 decode on MPS, in the same session: 155 and 34 values differ from float64's, each by 1 (the
+    count moves between torch runs). Ours and torch's pictures differ by at most 1 (95.4 and 97.7 dB PSNR).
+  - Speed, back to back in one session with no other GPU work, the median of three: ours 0.181 s and 0.188 s on the
+    GPU; torch's warm float32 decode on MPS 0.186 s and 0.191 s. Ours loads in 0.3 to 0.4 s. Its scratch is five
+    float32 picture tensors at the widest stage: 4.7 GB at peak at 1024x1024, after the transformer is freed, so
+    `tf-moss-image` peaks at the transformer's 11.3 GB.

tools/zig/moss_image_case.py

@@ -20,6 +20,9 @@
                                            with --fp32-encoder F, F's rows from a float32 encoder too
                                            (CASE_DIR/exact_fp32_encoder.png)
   decode MODEL_DIR LATENTS.f32 PNG         raw float32 latents (rows x columns x 128) from tf-moss-image to a picture
+  vae MODEL_DIR LATENTS.f32 OUT.f64 TORCH.f32  the same decode in float64 on the CPU and in float32 as the
+                                           pipeline runs it, each output before the clamp [height, width, 3] raw,
+                                           for tf-moss-image-decoder
   encode MODEL_DIR OUT_DIR [--prompt-file F]  the prompt encoder alone: token ids, embedding rows, and every layer's
                                            output and the final rows in bf16 and in float32, and
                                            each step alone in bf16 on the float32 run's input
@@ -79,19 +82,24 @@ def rows(x):
     return x[0].float().permute(1, 2, 0).contiguous().cpu()
 
 
-def to_png(pipe, latents, path: str) -> None:
-    """Latents [1, 128, H, W] through the pipeline's own decode path (its batch-norm undo, unpack, VAE) to a PNG."""
+def decoded(vae, latents):
+    """Latents [1, 128, H, W] through the pipeline's own decode path (its batch-norm undo, unpack, VAE), unclamped."""
     import torch
 
     with torch.no_grad():
-        z = latents.to(device=pipe.vae.device, dtype=pipe.vae.dtype)
-        bn = pipe.vae.bn
-        eps = pipe.vae.config.batch_norm_eps
+        z = latents.to(device=vae.device, dtype=vae.dtype)
+        bn = vae.bn
+        eps = vae.config.batch_norm_eps
         z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1)
         b, c, h, w = z.shape
         z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
-        image = pipe.vae.decode(z, return_dict=False)[0].clamp(-1, 1)
-        pipe.image_processor.postprocess(image, output_type="pil")[0].save(path)
+        return vae.decode(z, return_dict=False)[0]
+
+
+def to_png(pipe, latents, path: str) -> None:
+    """Latents [1, 128, H, W] decoded as the pipeline does, clamped, to a PNG."""
+    image = decoded(pipe.vae, latents).clamp(-1, 1)
+    pipe.image_processor.postprocess(image, output_type="pil")[0].save(path)
 
 
 def case(args) -> int:
@@ -228,6 +236,28 @@ def decode(args) -> int:
     return 0
 
 
+def vae(args) -> int:
+    import numpy as np
+    import torch
+    from diffusers import AutoencoderKLFlux2
+
+    raw = np.fromfile(args.latents, dtype=np.float32)
+    h, w = args.height // 16, args.width // 16
+    if raw.size != h * w * 128:
+        sys.exit(f"{args.latents} holds {raw.size} values; {args.width}x{args.height} needs {h * w * 128}")
+    latents = torch.from_numpy(raw).reshape(1, h, w, 128).permute(0, 3, 1, 2)
+    model = AutoencoderKLFlux2.from_pretrained(Path(args.model_dir) / "vae", torch_dtype=torch.float32).eval()
+    torch_rows = decoded(model.to(args.device), latents)[0].permute(1, 2, 0).contiguous().cpu().numpy()
+    torch_rows.astype(np.float32).tofile(args.torch_out)
+    model = model.cpu().double()
+    mark = time.perf_counter()
+    out = decoded(model, latents)[0].permute(1, 2, 0).contiguous().numpy()
+    out.astype(np.float64).tofile(args.out)
+    print(json.dumps({"out": args.out, "torch_out": args.torch_out, "shape": list(out.shape),
+                      "decode_s": round(time.perf_counter() - mark, 1)}))
+    return 0
+
+
 def compare(args) -> int:
     import numpy as np
     from PIL import Image
@@ -399,14 +429,22 @@ def main() -> int:
     n.add_argument("out_dir")
     n.add_argument("--prompt-file", help="the prompt (default: the empty negative prompt)")
     n.add_argument("--device", default="mps")
+    v = sub.add_parser("vae")
+    v.add_argument("model_dir")
+    v.add_argument("latents")
+    v.add_argument("out")
+    v.add_argument("torch_out")
+    v.add_argument("--device", default="mps", help="where the float32 decode runs, as the pipeline would")
+    v.add_argument("--width", type=int, default=1024)
+    v.add_argument("--height", type=int, default=1024)
     p = sub.add_parser("compare")
     p.add_argument("a")
     p.add_argument("b")
     p.add_argument("side", nargs="?")
     n = sub.add_parser("noise")
     n.add_argument("out")
     args = parser.parse_args()
-    return {"case": case, "start": start, "exact": exact, "decode": decode, "encode": encode, "compare": compare, "noise": noise}[args.mode](args)
+    return {"case": case, "start": start, "exact": exact, "decode": decode, "encode": encode, "vae": vae, "compare": compare, "noise": noise}[args.mode](args)
 
 
 if __name__ == "__main__":

zig/build/metal_targets.zig

@@ -316,8 +316,9 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s
         .{ .name = "tf-glm-video", .path = "zig/tests/glm_video.zig", .about = "GLM-5.3-Flash's video preparation on a clip (libtfvideo), written for tools/zig/glm_video_fixtures.py" },
         .{ .name = "tf-qwen-vision", .path = "zig/tests/qwen_vision.zig", .about = "The Qwen-VL tower on picture files: patches, grids and merged rows written for a float32 reference" },
         .{ .name = "tf-qwen-image", .path = "zig/tests/qwen_image.zig", .about = "Qwen-Image-2.1's transformer on a reference case: each step compared, then the schedule timed" },
-        .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image's transformer on a reference case: each step's two branches compared, then the guided schedule timed" },
+        .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image on a reference case: each step's two branches compared, then the guided schedule timed and decoded" },
         .{ .name = "tf-moss-image-encoder", .path = "zig/tests/moss_image_encoder.zig", .about = "MOSS-Image's prompt encoder on a reference case: the ids checked, each step gated against float32, then timed" },
+        .{ .name = "tf-moss-image-decoder", .path = "zig/tests/moss_image_decoder.zig", .about = "MOSS-Image's image decoder on latents, gated against float64 beside torch float32, then timed" },
         .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" },
         .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" },
         .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" },