MLX implementation and experiment report

Measured on 11 November 2026: Apple M5 Pro, 16 GPU cores, 64 GB unified memory, macOS 27.0.1, MLX/MLX Metal 0.32.4. The port runs actual original MiniMax-H3-Image weights locally. No remote inference or substituted image model is used.

What the default preserves

The default keeps the DiT's BF16 matrices, the original six sigma intervals and six individual PDD output heads, dense attention, FP32 Euler state and FP32 input/timestep/output projections. It reads the first fifty Qwen3-VL language layers without the final normalization. The image VAE uses the native BF16 decoder, 256 px tiles, at least 64 px overlap and the released vertical-then-horizontal blending order. The checkpoint stores VAE weights in FP32; this decoder casting policy follows the inspected SGLang image path.

Specific precision fixes are covered by differential tests:

Original DiT QKV weights are head-interleaved, and must be reordered before a conventional Q/K/V split. The image VAE retains its own interleaved layout.
ATen DiT/VAE RMSNorm and LayerNorm apply their affine in FP32 before the BF16 cast. Qwen RMSNorm explicitly casts before its affine. Using one MLX normalization convention for both introduced many incorrect rounded values.
BF16 biased linear layers use an addmm epilogue. Separately rounding the matrix product and then adding bias changes results.
MLX_ENABLE_TF32=0 retains the FP32 precision islands on M5. MLX_CONV_WINOGRAD=0 also protects the FP32 reference-image encoder.
RoPE coordinates, partial rotation, BF16 products and SiLU rounding follow each component's original recipe. CPU torch noise follows the original generator family; --rng mlx explicitly uses a different stream.

This preserves storage precision and the model's mathematical operations. It does not promise CUDA-identical pixels. Metal/CPU/CUDA reductions and attention kernels can cross BF16 rounding boundaries. An evolving diffusion trajectory can amplify these differences.

The source checkpoint revisions and inspected implementation hashes are in sources.json. These are a port of the pinned SGLang inference contracts, not a benchmark of a pristine official MiniMax/Diffusers application.

Fitting the unquantized model

Each DiT block has a large modulation projection that depends on the sampling timestep, not the image. For the fixed schedule we compute its small outputs once, cache them, and release approximately 26 GB of original modulation weights. This is direct evaluation of the learned matrices: no curve fitting, low-rank approximation, pruning or interpolation is involved.

The remaining BF16 transformer core occupies about 36 GiB of active MLX memory. A persistent Pipeline retains it across all steps and requests. The conditioning encoder streams one layer at a time, and the VAE can coexist with the core on this Mac. A fixed resident subset avoids the pathological all-miss behavior of a rotating LRU when the whole core does not fit.

The actual-weight modulation audit checks layers 0, 25 and 49 against individual PyTorch CPU projections. Relative RMSE is 1.06e-4–1.63e-4, with 98.6–99.9% of BF16 values identical. Individual MLX GEMVs and batched MLX projection also round differently despite identical source weights. These results support retaining the original matrices and FP32 islands, while avoiding any claim of bit identity across backends.

Matched 2K product trials

All rows below use the same detailed product prompt, seed 731, CPU torch noise and 2048×2048 working/output size. Each profile has one initial request and one subsequent request reusing resident weights and cached conditioning. The table uses one recorded resident DiT sample per profile, not a statistically established latency median. A final decoder LayerNorm audit re-rendered the saved latents without rerunning diffusion. These decoder-only changes leave the recorded DiT timings applicable. Fresh end-to-end values are shown only for requests measured again with the final corrected decoder. Other rows report denoising time, with their separate decode records retained in JSON. Within each profile, its two original repetitions re-render to identical pixels.

Profile	Passes	DiT only	Fresh end-to-end	Recorded peak MLX	RGB SSIM	RGB MAE / 255
native	6	67.87 s	72.78 s	38.31 GiB	1.0000	0.000
tuned	6	65.02 s	—	38.33 GiB	0.9757	2.339
fused	6	63.47 s	66.35 s	38.31 GiB	0.9769	2.252
q8	6	72.28 s	—	21.48 GiB	0.9318	5.793
q4	6	67.71 s	—	12.51 GiB	0.7799	22.925
preview	3	32.30 s	—	38.33 GiB	0.6781	25.687

Recorded latency, memory and similarity tradeoffs

The final BF16 native and tuned requests took 72.78 s and 66.35 s, respectively: an observed 8.8% reduction. The final measurements appear above; there is only one resident timing sample per profile. The three-pass preview took about half the denoising time and changes the image much more. Q8 is useful for memory capacity here, not a speed improvement. Q4 substantially changes the trajectory and fine details. Neither is the default near-original-precision path.

In the final timing pass, a native request took 118.59 s with a fresh conditioning encode, a cached modulation table and model loading. Its subsequent resident request took 72.78 s. This distinction is critical: a new CLI process does not inherit the previous process's resident weights. Building the native modulation table also takes additional time on a fresh installation. Initial quantized-request totals are not a controlled cold-start comparison.

Memory figures are peak MLX allocations, not whole-system RAM, RSS or swap. The tests do not establish capacity/performance on every Mac.

Inspect the actual evidence:

Operator tuning

The 62 operator trials use actual H3 dimensions, synthetic BF16 inputs, two warmups and seven synchronized repetitions per setting. They measure Python/MLX dispatch plus GPU execution. They exclude weight loading and quantization, and are not image-generation latency measurements.

Selected medians at 4096 tokens:

Operation	Native	Experiment	Result
RMSNorm + indexed AdaLN, width 5376	3.558 ms	0.543 ms, 256 threads	~6.6× operator speedup; relative RMSE ~1e-5
Q/K RMSNorm + partial RoPE, 56×128 heads	10.961 ms	2.057 ms	~5.3× operator speedup; relative RMSE ~9.6e-6
QKV GEMM, 5376→21504	37.589 ms	Q8/g64: 43.336 ms	Quantization is slower at this shape
Fused gate/up GEMM, 5376→28672	50.184 ms	Q8/g64: 57.969 ms	Same finding
Down GEMM, 14336→5376	29.900 ms	Q4/g64: 29.812 ms	Essentially unchanged; quantization changes values
Dense SDPA, 56 heads, head dim 128	24.146 ms	128-query chunks: 41.219 ms	Full memory-efficient SDPA is faster

The sweep also tests normalization threadgroups 128/256/512, Q4 and Q8 groups 32/64/128, and attention query chunks 0/128/512/1024. Query chunking retained identical outputs in these measured attention trials, but increased latency. Smaller quantization groups slightly reduced numerical error without making large-batch Q8 GEMM faster than BF16 on this device.

The custom kernels preserve explicit BF16 operation boundaries, but their parallel reduction order can change a few values. Their large isolated speedups translate to a much smaller image-generation gain. The earlier combined-kernel trials had very similar denoising times; these few samples do not establish a reliable end-to-end advantage over AdaLN fusion alone.

Decoder-only batching experiment

The VAE trials decode the same saved final latent, so diffusion drift cannot obscure decoder differences. Loaded weights and one warmup are excluded; each number is the median of two measured decodes.

Tile batch	Decode time	Pixel identity to batch 1	RGB SSIM	Max 8-bit error
1	2.972 s	Yes	1.0000	0
4	2.014 s	Yes	1.0000	0
8	1.931 s	No	0.99890	4
16	1.898 s	No	0.99889	4

Batch 4 is an attractive opt-in for this tested latent: 32.3% faster decode with identical pixels. This is a stage-specific, one-input result, not an end-to-end 32.3% speedup or a universal byte-identity guarantee. The default retains the native batch of one. Larger batches remain explicit tradeoffs. Inspect the decoder-only comparison stack.

Interpreting precision comparisons

Scores use full-resolution 8-bit RGB. SSIM is scikit-image's per-channel structural similarity with its default 7×7 window; PSNR uses RGB MSE and a 255 data range. These differ from the neighboring repository's luma/window metrics and should not be numerically conflated. Neither metric measures prompt adherence or overall visual quality. All crops use exactly the same source coordinates within a case.

The separate CUDA comparison uses the saved, hash-verified original BF16 RTX 5090 single-bf16-stream image at the same prompt/seed/resolution. Its upstream quality=lossless VAE path differs from this port's native tile recipe. The MLX default scores 0.9196 RGB SSIM, PSNR 21.15 dB, MAE 7.448/255 against that image. This is useful visual context, but does not isolate backend error or establish CUDA-identical output. Timings from different hardware/quality settings are not used as a port speedup claim.

Additional functional and quality checks

Portrait, 2048×2048: native BF16 and Q8 both complete. Q8 scores 0.8880 RGB SSIM and 10.525/255 MAE against native. The matched face crops make the changed skin detail visible. This is a second prompt/seed, not a broad perceptual-quality evaluation.
Native 4096×4096: all six native passes complete, producing the direct 16 MP image. Denoising took 397.76 s and recorded peak MLX allocation was 41.29 GiB. The separately revalidated final decoder took 12.11 s. The CUDA comparison scores 0.8437 RGB SSIM; hardware and upstream VAE quality differ, as explained above.
Editing: the reference/output example changes cobalt glass to emerald glass, retaining the broad scene and main label. Small lettering remains imperfect. The 1024×1024 output uses the native 2048×2048 minimum generation canvas. A fresh end-to-end run took 157.74 s. This is a one-reference functional check, not an editing-quality equivalence audit or validation of all sixteen-reference workloads.

The separate portrait CUDA comparison scores 0.9382 RGB SSIM. All input identities and source hashes remain in the individual gallery records.

Validation and next optimization opportunities

All 20 numerical tests pass on the development Mac, and the distributable wheel builds successfully. The suite checks native/Qwen normalization semantics, BF16 bias epilogues and SiLU, full/partial RoPE, dense causal GQA including query chunks, original QKV layout, BF16 binary loading, quantized checkpoint round trips, patch ordering, native tiling and blending, reflection padding, image output quantization and rejection of unmatched comparison metadata.

The most useful next work is shape-specific BF16 GEMM tuning and profiling of the M5 matrix-unit dispatch at the actual packed sequence lengths. Matrix projections dominate the remaining work; replacing them with slower Q8 operators does not help this 64 GB Mac. More efficient original-weight loading and persistent-service reuse can improve startup without changing arithmetic. The preconversion tool already removes online Q8/Q4 conversion from future quantized startups while retaining native modulation/FP32 islands.

Further precision work should expand same-prompt multi-seed tests, preserve exported original noise/conditioning tensors across backends, and compare against an identical upstream VAE quality setting. New custom attention or matrix kernels should be evaluated with both tensor errors and complete seeded image stacks. Sparse attention, residual reuse and timestep skipping would change the model and should remain separate experiments.