devin-lai's picture
Release v0.4.0: selective MLX model downloads
3da88e3 verified
|
Raw History Blame Contribute Delete
13.9 kB
# 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](../provenance/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](modulation-audit.json) 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.
<!-- PRODUCT_TABLE_START -->
| 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 |
<!-- PRODUCT_TABLE_END -->
[Recorded latency, memory and similarity tradeoffs](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/reports/tradeoffs.png)
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:
* [Product gallery](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/product/index.html)
* [Vertical full-frame stack](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/product/stack.png)
* [Identical 1:1 detail crops](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/product/details.png)
* [Absolute RGB differences amplified eight times](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/product/difference-x8.png)
* [Full-resolution scores and source PNG hashes](../comparisons/product/metrics.json)
* [All exported inference records](image-trials.json)
## Operator tuning
The [62 operator trials](operators.json) 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](vae/trials.json) 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](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/vae/index.html).
## 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](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/product-cuda/index.html) 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](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/portrait/index.html)
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](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/native-4k-cuda/index.html) scores 0.8437 RGB
SSIM; hardware and upstream VAE quality differ, as explained above.
* **Editing:** the [reference/output example](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/edit-example/before-after.png)
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](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/portrait-cuda/index.html)
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.