[jayleaton]

Stacked on #1766 (Idol's denoiser on Metal). This branch holds #1766's commits too; its own change is the last two commits. Independent of #1767 and #1773.

What this changes

Idol's FLUX.2 VAE image encoder on Metal: the 1,024 x 1,024 picture to the denoiser's feature2 [4101, 128], as idol-sdk's ImageEncoder._flux2_tokens builds it, with tf-idol-vae and an idol_case.py vae case. For #1757. Details: the "VAE encoder" part of docs/recipes/idol.md.

• The encoder: conv_in, four DownEncoderBlock2D stages (GroupNorm 32 and SiLU, stride-2 downsamples padded right and below), the mid block with its single-head attention over 16,384 positions, conv_norm_out, conv_out, quant_conv. Then the posterior's sample (the SDK samples, it does not take the mean), the 2 x 2 patches as 128 channels, the batch norm's saved statistics, and five zero rows ahead of the 4,096 patches.
• The SDK draws the posterior's noise from torch's MPS generator, which a native encoder cannot reproduce, so encode takes the noise [128, 128, 32] as an input and the case records the SDK's.
• zig/kernels/metal/idol_vae.metal, eight kernels of Idol's own. No shared kernel changes. MOSS-Image: the image decoder on Metal #1687 (MOSS-Image's FLUX.2 VAE decoder) is the mirror image of this; the two share no code yet, and once both land the convolutions and GroupNorm could move to one place.

Weights and precision

black-forest-labs/FLUX.2-klein-4B vae/ at e7b7dc27f91deacad38e78976d1f2b499d76a294 (Apache 2.0, ungated, bf16). Klein's file is FLUX.2-dev's fp32 VAE rounded to bf16; the loader takes either and rounds an fp32 file to bf16 as torch does, and FLUX.2-dev's file gives the same rows bit for bit.

The SDK runs the encoder in bf16 on the tensor units. This port computes in fp32 on the same bf16 weights: Winograd F(2x2, 3x3) with G g G^T made in float64 and rounded once, implicit-GEMM stride-2 and 1x1 convolutions, two-pass GroupNorm, fp32 attention, precise::exp. None of it uses the tensor units, whose fp32 GEMM truncates its operands (#1772). iv_tokens keeps a NaN logvar as torch.clamp does.

Kernel contract

zig/tests/contract/metal_idol_vae*.zig: all eight kernels in 13 forms against float64, with twelve controls (half-staged inputs, symmetric padding, a dropped Winograd term, half SiLU, half sums, attention without its rescale, swapped patch axes, a NaN logvar lost to clamp).

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

Three controls go red by small margins: the half-staged 1x1 at 1.25 times its bound, Winograd's half patch at 1.05 times, attention without rescale at 1.15 times. These kernels use no tensor ops; their records reserve their thread counts in the source, and the --offline M1/M2 compile was not run.

Receipt

• Environment: Mac Studio M5 Ultra 256 GB, macOS 27; this branch on idol: Idol's denoiser as a native Metal family #1766's 94a039249; idol-sdk 4702a7f0 on torch 2.13.0 (MPS) with diffusers 0.42.1 (sha256-pinned files); the SDK's preprocessed avocado example, seed 42.

• Exactness, against float64 on the CPU with the SDK's own noise:

  	worst row	1 - cosine
  ours (fp32)	5.1e-5	5.0e-11
  the SDK (bf16, MPS)	2.73e-2	5.05e-5

  The gate (twice the SDK's distance, worst row and cosine, zero padding rows) passes by about 500 times on worst row. Against the SDK: cosine 0.999949, worst row 0.0275; the case's SDK rows equal the SDK's own recorded feature2 bit for bit. Each down block, the mid block and the moments sit within 4e-5 (worst pixel) of float64. Five more encodes repeat byte for byte. vae.zig's test plants a 2% row error and a non-zero padding value, and the gate fails each.

• Speed: 0.185 s an encode against the SDK's 0.095 s (medians of five): this port is half the SDK's speed, because it runs fp32 off the tensor units where the SDK runs bf16 on them. The attention (about 40 ms) reads K and V once every 8 queries; wider query tiles, or the hi/lo split onto the tensor units, are the obvious next steps. It runs once an image, against about 28 s of denoising.

• Memory: 3.28 GB peak footprint (/usr/bin/time -l, the whole run with its case and stage copies), 2 GB of it four 512 MB tensors at the 1,024 x 1,024 x 128 stage.

• Tests: zig build test test-golden -Dcpu=apple_m1: 585 of 585 steps, 0 fail (before the rebase onto 94a039249, a docs-only change; after it, plain zig build 248 of 248). zig fmt --check build.zig zig: clean. lean_check.py: 424, the same as main.

• Not run: CUDA (none in this PR); Metal M1 to M4.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (No token path.)
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (n/a.)
• 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. (Not touched.)
• 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 @jayleaton. Computing in fp32 on the bf16 weights, well inside the float64 gate, is the right trade for an encode that runs once an image. It waits on #1766; please rebase after it lands. Alongside the contract we'll also run the --offline M1/M2 compile you couldn't run. Two follow-ups for later, neither blocking: the encode is half the SDK's speed (0.185 s against 0.095 s), and now that #1687 is merged, the convolutions and GroupNorm could be shared with MOSS-Image's decoder.