[drowzeys]

Thank you, Ash, for the invitation to send this for review, and for the Metal runtime and kernels it builds on.

This is the first image model on TensorFold's Zig runtime. Qwen-Image-2.1's 7B transformer now runs on TensorFold's Metal runtime from the released bf16 diffusers checkpoint, with no MLX and no Python in the transformer. Prompt rows go in, and each denoising step's velocity comes out. It reuses your kernels rather than adding a parallel stack: the projections are the Qwen-VL tower's tfp::gemm_nax, and attention is the prefill tfp::attention_nax template. This is the case for TensorFold beyond text: the same runtime, kernel contract and lean code carry a model whose step is a denoising pass over an image rather than a token.

On a Mac Studio M5 Ultra, Qwen/Qwen-Image-2.1-Turbo, 1344x768 (4,032 image rows), 8 steps, compared with mflux (MLX):

• Image parity: 43.35 dB PSNR between our final image and mflux's, from the same prompt rows and starting noise, both decoded by mflux's decoder. Side by side, they are visually identical.
• Velocity parity, on every one of the 8 real steps: cosine 0.99996 to 0.99999 against mflux bf16, run on the reference's own input.
• Against an fp32 run of the same transformer: our cosine is 0.99997 to 0.99998. mflux bf16's own cosine against that fp32 run is 0.99994 to 0.99999. At step 1 we are closer to fp32 than mflux bf16 is (0.999966 against 0.999944).
• Repeat runs are byte-identical.
• 15.1 GB peak process footprint, the bf16 weights (about 14 GB) included.
• Speed: 0.740 s a forward against mflux's 0.702 s, measured in the same session and alternated, so within 5% of MLX. The prompt pass, run once a prompt, takes 0.23 s.
• Precision is the checkpoint's own:
  • bf16 activations, rounded where the reference rounds them.
  • fp32 sums in every projection and in attention.
  • fp32 norms, rotary and softmax.
  • No int8, which addresses the precision point from qwen_image: Qwen-Image-2.1 text to image on MLX (draft proposal, non-lane family, stacked on #384) #393.
• Kernel contract: gate 5 holds all 15 new launches, and they pass on this GPU with 0 limit problems. The full run's SUMMARY is below.
• Lean check: 0 new findings.

What comes next, for the gate we agree on:

• The prompt encoder and the image decoder in Zig. Today, a development tool runs both through mflux: the Qwen3-VL language model and the VAE decoder. The prompt encoder is a Qwen3-VL text stack, close to what the engine already runs for the Qwen families, so it is the natural next step. With the decoder, text to picture would run with no Python at all.
• Closing the last 5% to MLX. By stage, the projections are about 80% of the GPU time, and they are already the MLX-identical gemm. The remainder is attention and the small row kernels, which are candidates for fusion.
• Gates 1–4 for the new row kernels. These are the modulated layer norm, the prompt RMSNorm, gelu, q/k norm and rotary, the prefix copy, the gated residual and SwiGLU. Today they are registered for gate 5, as the Qwen-VL tower's are. The non-causal attention entry points are the existing attention_nax template.
• CUDA. A CUDA backend for DGX Spark would follow as its own pull request.
• A tensorfold entry point. That means a generate-style command, or an endpoint, once we agree where a denoiser sits.

I did not open an issue first, since this came by invitation. Happy to move the gate discussion into one if you prefer.

Ours (left) and mflux (right), same prompt rows and starting noise, both decoded by mflux:

What it adds

These are new files, plus minimal wiring: 1 line in build.zig, 2 in zig/kernels/metal/sources.zig, 2 in zig/src/tensorfold.zig, 3 in the gate-5 runner and 15 records in registry.zig.

• zig/src/families/qwen_image/ (Zig over the Metal runtime):
  • config.zig checks transformer/config.json and refuses anything but QwenImage21Transformer2DModel at the released geometry.
  • weights.zig reads the diffusers shards into shared Metal buffers through core/checkpoint_metal.zig. It checks that the checkpoint holds 297 tensors, and the dtype and shape of each one it reads.
  • host.zig holds the rotary tables (three axes: prompt rows on the diagonal, the image centred at the prompt's frame), the timestep embedding, and the modulation every block shares.
  • kernels.zig holds the pipelines. The projections are qwv_gemm_bf16_nt_* (the tower's tfp::gemm_nax). Attention is tfp::attention_nax, causal for the prompt rows and full for the image rows.
  • model.zig holds the prompt pass, whose keys and values are kept for each block, and the per-step image pass.
• zig/kernels/metal/qwen_image.metal holds seven row kernels and eight attention_nax entry points: full and causal, by alignment.
• zig build tf-qwen-image runs the transformer on a reference case:
  • It runs every step on the reference's input and prints max |diff|, cosine and worst row against bf16, and against fp32 when the case has it.
  • It checks that a repeated step gives the same bytes.
  • It then runs the whole schedule, times it, and writes the final latents and run.json.
• tools/zig/qwen_image_case.py (development tool, using mflux and MLX):
  • case writes the case: prompt rows, noise levels, starting latents, mflux's per-step inputs and velocities, and optional fp32 velocities.
  • decode turns our latents into a picture.
  • compare gives PSNR and a side-by-side image.
• docs/recipes/qwen-image-2.1.md has how to run it and the receipts.

How it fits the engine, and how it does not

• What a step is. A step is one denoising pass, with no token decoded.
  • The prompt's encoder rows pass through the 32 blocks once, with causal attention and modulation at noise level 0. Each block's keys and values are kept.
  • Each step's image rows then attend over those kept prompt keys plus their own, with modulation at that step's noise level.
  • The tool takes an Euler step on the host in fp32, on the schedule saved with the checkpoint.
• It is a denoiser, not a token lane. There are no drafts, so the "drafted equals plain" contract does not apply. What it does keep is the same activation precision as the reference, fp32 sums and byte-identical repeats. It does not register with the lane engine, the server or families.detect, and touches nothing in the lane engine, the server or any existing family.
• Where I propose it sits. It sits beside the engine, as the Qwen-VL tower does: a family module in tensorfold.zig, driven today by a tf-* program. A served text-to-image endpoint would be a separate request type, not a lane. Image towers already run as one-shot passes outside the lane rounds, and a denoiser would sit in the same place.
• Shared code. No existing kernel or family is changed. The Qwen-VL tower's gemm and the prefill attention template are reused as they are; only new entry points are instantiated.

Receipts

Environment:

• Mac Studio M5 Ultra, 256 GB, macOS 27.
• Zig 0.17.0, Xcode Metal toolchain.
• This branch on upstream main 46a08c18.
• Checkpoint Qwen/Qwen-Image-2.1-Turbo.
• Reference: mflux add5164 on MLX 0.32.3, bf16.

Request:

• One public prompt: "A lighthouse on a rocky coast at dusk, waves breaking below, warm light in the lantern room, a small red boat moored in the harbour, painted in soft watercolour" (44 prompt rows).
• 1344x768, seed 42.
• The 8 saved sample_sigmas, given the resolution shift: 1.0, 0.989079, 0.976488, 0.961812, 0.944488, 0.915859, 0.826254, 0.585455, then 0.

Velocity, each step on mflux's own input (tf-qwen-image on the case):

Step	vs mflux bf16: max |diff|	cosine	worst row (rel. L2)	vs fp32: cosine	mflux bf16 vs fp32: cosine
1	0.122	0.9999637	0.0398	0.9999661	0.9999435
2	0.125	0.9999677	0.0274	0.9999661	0.9999729
3	0.109	0.9999775	0.0278	0.9999737	0.9999667
4	0.125	0.9999854	0.0296	0.9999745	0.9999732
5	0.172	0.9999823	0.0447	0.9999687	0.9999624
6	0.164	0.9999801	0.0464	0.9999748	0.9999788
7	0.156	0.9999898	0.0429	0.9999839	0.9999851
8	0.211	0.9999862	0.0688	0.9999810	0.9999825

The same step run twice gave identical bytes.

Image. Our free-running 8 steps, decoded by mflux's float32 decoder, against mflux's own image: PSNR 43.35 dB, mean |diff| 0.84 of 255. A side-by-side image is attached.

Speed, one session under a GPU lock shared with the machine's other jobs, run in the order mflux, ours, mflux:

	Per forward	On the GPU	Prompt pass	8 steps
mflux bf16 (add5164), run 1	0.702–0.703 s (step 1 0.875 s, prompt prefix included)			5.80 s
this branch	0.739–0.741 s	0.729–0.731 s	0.23 s (0.074 s on the GPU)	5.92 s, plus the prompt pass
mflux bf16 (add5164), run 2	0.702–0.703 s (step 1 0.85 s)			5.77 s

Where the GPU time goes, measured by switching off one stage class at a time in a local build. That session's full forward took 0.805 s on the GPU, so these are shares rather than absolute times:

Stage	GPU time
MLP projections	0.47 s
q/k/v/out projections	0.19 s
attention	0.12 s
norms, rotary, residuals and SwiGLU	about 0.03 s

I tried instantiating the gemm with K-alignment, a single change, and it did not move it: 0.729 s against 0.730 s.

Earlier alternated pairs ran while another job had a 35B model loaded in Ollama on the same GPU, which was not part of this work. Under that matched contention, mflux and this branch were within about 3% of each other: 2.13–2.48 s against 2.24–2.36 s. The table above is the uncontended session.

Memory:

• This branch: 15.1 GB peak process footprint (/usr/bin/time -l), with the bf16 weights in shared Metal buffers.
• mflux's case run: 44.2 GB, but that includes its prompt encoder, so it is not like for like.

Host tests, zig build test -Dcpu=apple_m1:

• The full suite gave 2,996 of 3,006 tests passed and 10 skipped, with 1 failing step. That step was this branch's own config test, which logged its refusal at error level.
• After the fix, the rerun gave 204/204 steps, with 560 of 561 tests passed and 1 skipped. Only the steps that rebuilt reran.

Kernel contract, zig build kernel-contract -Dcpu=apple_m1:

SUMMARY 26 kernels: 26 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems; gates 113995 ms, limits 151 ms

The 2 known findings are affine_mm's documented ones. With --only "qwen image", all 15 new launches pass the limits check on this GPU: 256 threads, or 128 for attention; 144 threadgroup bytes for the two norms and 0 for the rest.

Lean check, python3 tools/zig/lean_check.py: 517 problems on the branch, equal to main's baseline of 517. That is 0 new, and none of the 517 are in this branch's files.

Not tested, or not in this pull request

• Other machines and checkpoints:
  • M1–M4: the sources use nax.h's simdgroup fallback, as the Qwen-VL tower does, but I have not run them there.
  • Gate 5's offline M1/M2 compiles (--offline) were not run.
  • The base Qwen/Qwen-Image-2.1 checkpoint (40 steps) has the same transformer and is untested.
  • Other resolutions are untested.
• Features not here:
  • Image editing, reference images, adapters, and true CFG with a negative prompt.
  • Prompts with padding: the tool refuses them.
• Not in Zig yet: the prompt encoder and the image decoder (see above).
• CUDA: not in this pull request.

Licenses

• Code: Apache-2.0, as TensorFold.
• Qwen/Qwen-Image-2.1 and -Turbo weights: the Qwen Research License, which is non-commercial. They are not redistributed here, and no weight bytes or fixtures from them are in the tree.
• mflux (MIT) is only a development dependency of tools/zig/qwen_image_case.py, which imports it at run time. No mflux or diffusers code is copied.

Reproducing

python3 tools/zig/qwen_image_case.py case MODEL_DIR case/ --prompt-file prompt.txt --fp32-steps 8
zig build tf-qwen-image -Dcpu=apple_m1
zig-out/bin/tf-qwen-image MODEL_DIR/transformer case/case.safetensors case/
python3 tools/zig/qwen_image_case.py decode MODEL_DIR case/latents.f32 case/zig.png
python3 tools/zig/qwen_image_case.py compare case/reference.png case/zig.png case/side.png

The Studio that uses this model is #941.

🤖 Generated with Claude Code

https://claude.ai/code/session_01PfjmufXm4L7qsVNFyRw9ZK

[ashhart]

Thank you, @drowzeys, the native Qwen-Image denoiser and retained prompt KV are useful first components of the image family.

We checked fa73745. The config accepts an epsilon it never validates against the fixed kernel constant, and Scratch.of multiplies unbounded shapes. We can add semantic/bounds refusals and failure regressions before exposing requests.

We will land the qualified denoiser first, then native text encoding, noise/schedules, VAE and picture output under shared job admission/cancellation. The recipe currently leaves encoding and decoding to mflux. Please keep same-input step/pixel references and multiple-resolution controls for that completion, alongside numerical kernel adapters and exact-source memory measurements. No Python/framework service should remain required for the final native request.

Posted by Ash's AI Assistant Synthena

[drowzeys]

Rebased on main 6ee828c2; head fa737459 is now 0267610b. The pull request's diff against main is patch-identical to before (main's new commits only touch docs); zig fmt --check clean on the changed Zig files; python3 tools/zig/lean_check.py reports no new findings against main.

[drowzeys]

Thank you. Both points are fixed; the head is now b3cc38e9.

• Epsilon never validated: fixed in 3b09dc66. config.check now requires eps and compares it at f32 with the norm kernels' fixed cfg.eps (1e-6).
  • Another value, or a value that isn't a number, is refused with error.EpsilonMismatch.
  • A missing eps is error.BadConfig.
  • Test: a config whose eps is not the kernels' constant is refused. The released value passes written both as 1e-06 and as 0.000001. The values 1e-05, 0, 1, a string and null are refused, and so is a config with eps removed.
• Scratch.of unbounded: fixed in b3cc38e9. Every product and the 256-aligned total are now checked with std.math.mul/add.
  • An overflow, or a key count past the kernels' u32 row arguments, is refused with error.ScratchTooLarge.
  • The prompt store's bytes are computed and checked in Scratch.of too.
  • config.grid refuses a width × height whose latent rows overflow, with error.BadImageSize.
  • The existing 1344 × 768 size test still gives 533,978,112 bytes.
  • New test: scratch for shapes whose products overflow is refused by name.

Checks on the M5 Ultra:

Check	Result
zig fmt --check build.zig zig	clean
lean_check	517, the same as main
zig build test -Dcpu=apple_m1	204/204 steps, including the family's tests through tensorfold.zig
tf-qwen-image	builds

These are build and host tests only; I didn't do a device run. The same-input step and pixel references and the multiple-resolution controls are all kept.

[ashhart]

Thanks @drowzeys for the epsilon and overflow controls, including the unchanged 533,978,112-byte real-shape case. These are the semantic and allocation checks requested in the review; our integration carries those boundaries too.

The reviewed landing unit remains the native denoiser component. The host tests and build do not qualify its same-input device steps or make the external text encoder and VAE native. Keep the numerical adapters, resource checks, multiple-resolution controls and reference comparisons on the final integrated source. The new head's lack of a device rerun should stay explicit.

Posted by Ash's AI Assistant Synthena

[ashhart]

Thank you, @drowzeys, for Qwen-Image-2.1's Metal transformer. 1.1.0 includes #948 (33b46f7ba) as the tf-qwen-image denoiser test tool, using the shared runtime and family-owned kernels. External case preparation supplies encoded prompts and decodes the final image. Latent grids, scratch sizes and epsilon are checked before checkpoint/device work. Numerical/resource and matched denoising results: Host/tool builds passed; Qwen-Image checkpoint/mflux end-to-end image comparisons remain unrun.

Posted by Ash's AI Assistant Synthena

[ashhart]

The Qwen-Image Metal denoiser component shipped in 1.1.0 as 33b46f7ba; thanks @drowzeys. We're closing it as landed.

Posted by Ash's AI Assistant Synthena