[BobClawblaw]

What this changes

Fixes #1459. MOSS-Image's prompt encoder on Metal, inside the moss_image family (#1457). The encoder is the checkpoint's text_encoder/, Qwen3.5's text model: 32 layers, three Gated DeltaNet layers to each gated full-attention one, hidden 4,096, bf16 as published. It runs as the pipeline's encode_prompt does: the prompt in the pipeline's chat format, one causal forward, and the final norm's rows. With --prompt, tf-moss-image makes both prompts' rows itself, so torch no longer loads the encoder.

• Family: encoder_config.zig (the released geometry, checked), encoder_weights.zig (the language_model. tensors; the vision tower is not read), encoder.zig (the chat format, the token ids, the forward), and the launchers in kernels.zig.
• Kernels: nine of the family's own in zig/kernels/metal/moss_image_encoder.metal: me_norm, me_head, me_attention, me_gate, me_silu_mul, me_conv, me_gdn_prep, me_gdn and me_gated_norm. The projections reuse the Qwen-VL tower's tfp::gemm_nax, as the transformer's do, and the residual adds reuse mi_gate_add. No shared kernel changes.
• Tokenizer: the native one, on the checkpoint's tokenizer/tokenizer.json. The engine module has no tokenizer, so the build gives it to the test programs (one import added to engine_programs in metal_targets.zig). No program's behaviour changes.
• Tools:
  • tf-moss-image --prompt FILE [--negative FILE] encodes both prompts natively and frees the encoder before the transformer loads.
  • tf-moss-image-encoder checks a case from moss_image_case.py encode. The ids must match. Each layer and the final norm, run alone (Encoder.steps) on the float32 run's input rounded to bf16, must be within twice torch bf16's distance from float32 on the same input, in cosine error and in its worst row. The whole prompt must repeat byte for byte. It exits 1 otherwise. reference.zig holds the gate, and its tests plant a mismatch that must fail with an identical repeat.
  • The case tool gains encode, start --no-encoder and exact --fp32-encoder. The last is the float32 ruler the docs quoted for moss_image: MOSS-Image-Base's transformer on the Metal runtime (new family, denoiser) #1457.
• Chips: Encoder.load refuses a GPU without tensor units and names it, like Model.load.

Receipt

• Environment:
  • Mac Studio M5 Ultra 256 GB, macOS 27; TensorFold f9651de30 (code) and 92c829887 (the gate) on main dd0a9cc5e.
  • Checkpoint OpenMOSS-Team/MOSS-Image-Base @ 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3.
  • Reference: the checkpoint's pipeline_moss_image.py (sha256-checked) on torch 2.14.1 (MPS, bf16), diffusers 0.42.0, transformers 5.21.0.
• Precision: bf16 weights as published and bf16 activations, rounded where the torch code rounds them. Every projection, norm and attention sums in fp32. The DeltaNet's q, k, decay and state are fp32; the state runs a column a thread over the rows in order, the recurrence the publisher's chunked algorithm computes. There is no int8, no cache and no lossy weight.
• Exactness, on the empty negative prompt, the fox prompt and the card's Chinese sticky-note prompt:

  • Tokens: the ids equal the publisher's tokenizer's on all three (10, 34 and 63 tokens).

  • Each step alone, on 9 prompts (10 to 75 tokens), every layer and the final norm on the float32 run's input in bf16: our distance from float32 is a mean of 0.50 of the gate's limit, which is torch bf16's own distance, and at most 0.65 (the fox, layer 13's worst row). All 297 gates pass. A row planted 5% off at layer 20 fails at layers 20 and 21, with an identical repeat.

  • Every layer of the whole prompt against torch bf16: cosine 0.99974 or better.

  • The final rows against a float32 run of the encoder, with torch bf16's own:

    prompt	ours: cosine, worst row	torch bf16: cosine, worst row
    empty	0.999826, 0.0261	0.999806, 0.0294
    fox	0.999931, 0.0261	0.999934, 0.0259
    sticky note	0.999926, 0.0454	0.999938, 0.0318

    These are not gated. Over 32 layers in bf16, each implementation's rounding is magnified differently, so two correct runs part by more than any one step's error. On the 75-token prompt, torch's own variants (MPS or CPU; chunks of 8, 16, 32 and 64, or the recurrent form; equal in exact arithmetic) end 0.86 to 4.97 times its default's cosine error from float32, and ours 1.82 times.

  • Pictures: with our rows, the transformer's latents through the pipeline's decoder are 47.07 dB PSNR from the pipeline's picture for the fox, and 28.71 dB for the sticky note. With torch's rows they are 48.52 and 28.93 dB. The ruler is a float32 run of the whole schedule, the encoder included, over 8 cases (both prompts, 4 seeds each):

    • torch bf16's picture averages 32.51 dB from it;
    • our transformer on torch's rows averages 31.20 dB;
    • our transformer on our rows averages 31.13 dB.

    Ours with our rows is the closer of the last two in 4 of the 8 cases.

  • Repeats: an encode run twice gives identical bytes.

  • Pixels, the fox case on this head:

    • from the case's precomputed prompt rows, the latents equal the pinned run's byte for byte (sha256 2e3100e8…);
    • with --prompt, the latents equal the first head's (6f05bcd94) byte for byte (sha256 9ba4fd4c…), so the decoded picture is the same.
• Speed, one session with no other GPU work:
  • Ours: 0.054 to 0.067 s a prompt (the same in three runs of each), and 0.6 s to load.
  • torch MPS: 0.158 to 0.165 s a prompt (the median of 5 warm), and 2.60 s to load.
  • The engine's process peaks at a 16.7 GB footprint (the encoder's 16 GB of bf16 weights), then frees it before the transformer loads.
• Prompt speed: n/a, not a token family.
• Kernel contract, 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
  

  • The new kernels: me_norm (the transformer's Norm test, generalized to plain, modulated and zero-centered), me_head, me_gate, me_silu_mul, me_conv, me_gated_norm, me_gdn_prep, me_gdn and me_attention.
  • How they are checked: each output against float64 at the torch code's bf16 rounding points, inside the range its fp32 steps allow. Eight new controls each drop or add one rounding point; the last adds bf16 rounding to the DeltaNet's state. Gate 3 catches every one.
  • me_gdn: checked against the recurrence in float64. Each output's allowance is its two sums' rounding over the state's magnitude, plus every earlier step's error, which a decay and a unit-key projection never grow.
  • Gate 4: for me_gdn, one-hot keys write j, 2^20 and -2^20 to three state rows, which its query sums back to j. me_conv uses taps of 256, 1, 1 and -256 over rows that cancel to 2. The attentions use equal scores over values whose sums reach about 2^22 and cancel.
  • Gate 5: 9 more launches, 17 for the family's source in all. As for moss_image: MOSS-Image-Base's transformer on the Metal runtime (new family, denoiser) #1457, --offline cannot compile the tensor-op source for M1/M2, 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 and dev-tools -Doptimize=Debug.
  • zig fmt --check build.zig zig.
  • tools/zig/lean_check.py: 425 problems, the same as main.
• 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).

atc overlap: every open PR it test-merged merges cleanly. Its conflicts are with local worktrees only: the shared metal_targets.zig and sources.zig lists, and older MOSS-Image branches.

A recipe branch that builds this PR's first head (6f05bcd94), whose latents this head reproduces byte for byte: https://github.com/BobClawblaw/moss-image-base-bf16-m5-ultra/tree/native-encoder

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (n/a: a one-shot prompt encoder for a denoiser, no decoding)
• 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 for the prompt encoder on Metal. Every new function has its kernel contract probes, and there are rounding and state controls for the whole family suite.

Two asks before a device run:

1. The reference check can pass on a mismatch. zig/tests/moss_image_encoder.zig prints every layer's error (:96) and computes the final errors against the bf16 and fp32 references (:103-104), but returns success on repeatability alone (:114). Please make the layer and final comparisons enforce the reference gate, and add a planted mismatch that must fail even when repeated output is identical.
2. Precision on the worst row. In the recipe (docs/recipes/moss-image.md:198), the sticky-note case's worst relative row error against the independent fp32 encoder is 0.0454, against 0.0318 for the bf16 reference; the fox is 0.0261 against 0.0259. Please find what makes the sticky-note row worse, rather than letting an average hide it, and attach a run whose final pixels equal the pinned reference, with the existing precomputed-prompt path held byte-identical.

[BobClawblaw]

Thanks. Pushed 198ed90 (one commit on the two, rebased onto main 1b9a987):

1. The check now fails on a mismatch. The gate is per step, not on the whole prompt; point 2 is why.

• Encoder.steps runs each layer and the final norm alone on given rows. For each step, tf-moss-image-encoder gives it the float32 run's input rounded to bf16.
• Each step's rows against float32 must be within twice torch bf16's distance on the same input (or plain bf16 rounding's, when larger), in cosine error and in its worst row. The whole prompt must still repeat byte for byte. It exits 1 otherwise.
• moss_image_case.py encode now records torch bf16's rows for each step alone (alone.{i}, alone.final), beside every layer's float32 rows.
• zig/src/families/moss_image/reference.zig holds the comparison and the gate. Its tests plant a row 5% off that must fail with an identical repeat, and a passing gate with a different repeat that must fail.
• On the device: 9 prompts (10 to 75 tokens), 33 steps each. All 297 pass, at a mean of 0.50 of the limit (that is, torch bf16's own distance) and at most 0.65. Planting row 7 of layer 20 5% off, in the float32 rows and in torch's, fails at layers 20 and 21 and exits 1, with an identical repeat.

2. The sticky note's row 58. Row 58 is <|im_start|> opening the assistant turn. Our kernels add no error there; the whole prompt in bf16 magnifies small rounding differences:

• Each layer alone matches torch. Given the same bf16 input, every layer is torch bf16's distance from float32. Over the 297 steps, our worst row is 1.00 times torch's on average (0.93 to 1.31), and it is the same row in 288. On the sticky note, ours drifts from torch's at row 58 from about layer 11. That is the end-to-end chain, not one step.
• torch's own correct variants spread as widely. These are the same code with the delta rule chunked at 8, 16, 32 or 64, or recurrent, on MPS or CPU, all equal in exact arithmetic. Against float64, row 58 is 0.0153 to 0.0320 over eight of them, and at chunk 16 it is torch's worst row too. On a 75-token prompt, torch's final cosine error from float32 ranges over 0.86 to 4.97 times its default; ours is 1.82 times. On that prompt, torch on CPU at the default chunk of 64 reaches a worst row of 0.24.
• The ruler holds. Float32 chunked and float32 recurrent both match float64 to four places.
• A gate on the whole prompt would therefore measure which rounding a run happened to draw, and a longer prompt would fail a correct build. The whole prompt's numbers are still printed, but not gated.

Pixels, the fox on this head:

• from the case's precomputed prompt rows, the latents equal the pinned run's byte for byte (sha256 2e3100e8776f10f0…);
• with --prompt, the latents equal the first head's (6f05bcd94) byte for byte (sha256 9ba4fd4c8f79471e…), so the picture is the same. Against the pipeline's own picture it is 47.07 dB.

Final pixels can't equal the torch pipeline's bit for bit: the encoder's rows differ from torch's by bf16 rounding.

On this head: zig build test, test-golden, test-kernel-contract, test-server-security test-pull-security, fuzz-cortheon, dev-tools and dev-tools -Doptimize=Debug all pass. zig fmt --check build.zig zig is clean, and lean is 426, as on main. The contract is unchanged:

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

The description is updated with the gate, the numbers and the head.