devin-lai's picture
Release v0.4.0: selective MLX model downloads
3da88e3 verified
|
Raw History Blame Contribute Delete
19.9 kB
# Further Mac optimization experiments
These experiments run the original MiniMax-H3-Image checkpoint on the local
64 GB M5 Pro Mac mini. The preferred path remains MLX with BF16 weights,
FP32 precision islands, dense attention and all six original PDD heads.
The new exact optimizations are block compilation, compiled four-tile VAE
batches, cached refined conditioning, native lazy safetensors loading and
bounded overlap of weight reads with GPU execution. `packed12` adds an exact
weight-storage option when memory matters more than a small decoding overhead.
The measured optimum is specific to this Mac and these prompts. This is a
practical search across runtimes and operators, not proof of a hardware ceiling.
The final product samples measured **67.30 s median at 38.03 GiB peak MLX**
for exact optimized BF16 (two warm samples, 65.14–69.45 s), and **72.23 s at
29.90 GiB** for exact packed BF16 (one warm sample). Product and portrait pixels
matched their BF16 references. The packed 4K trial also matched every pixel,
at 31.20 GiB peak and 446.75 s including loading; its background load was not
controlled, so it is evidence of memory/equality rather than a 4K speed win.
## Run the measured configurations
The following memory settings were tested on **this 64 GB machine**. The runtime
rejects wired limits above Metal's recommended working set and caps block
residency below that set. `--wired-gb` affects this process; it does not change
the system's wired-memory setting.
```bash
# Full six-head BF16 generation, with exact optimizations.
h3-mlx generate --model /path/to/MiniMax-H3-Image \
--text-encoder /path/to/MiniMax-H3/text_encoder \
--compile-blocks --eval-every 0 --prefetch-blocks 4 --wired-gb 42 \
--compile-vae --vae-batch 4 \
"A detailed image description" --output outputs/optimized.png
# Same arithmetic, smaller resident weights. The first use builds a ~27 GiB cache.
h3-mlx generate --model /path/to/MiniMax-H3-Image \
--text-encoder /path/to/MiniMax-H3/text_encoder \
--weight-storage packed12 --compile-blocks --eval-every 0 \
--prefetch-blocks 4 --wired-gb 32 --compile-vae --vae-batch 4 \
"A detailed image description" --output outputs/compact.png
```
Keep one `Pipeline` alive for warm requests. The CLI starts a new process for
each invocation, so it includes loading. The refined-text cache is enabled by
default and is keyed by the exact BF16 conditioning bytes, checkpoint identity,
configuration, quantization/chunk policy, MLX version and TF32 policy.
`--no-refiner-cache` bypasses it. A cache hit avoids loading the two refiner
blocks and their projection, without approximating their output.
## Full-image measurements
The generated table comes from the raw ledger and complete PNG outputs. A warm
sample reuses resident DiT weights; it may still compile a new sequence shape.
Peak memory is the **MLX allocator peak**, not total system RAM. The images and
metric definitions are linked below. A single load-inclusive 4K sample is not
reported as a warm benchmark.
<!-- RESULTS_TABLE -->
| Case | Configuration | Warm samples | Warm median (s) | Peak MLX GiB | SSIM | Exact pixels |
|---|---|---:|---:|---:|---:|---|
| product | BF16 baseline (packing still active) | 1 | 80.67 | 38.29 | 1.000000 | yes |
| product | Exact optimized BF16, final quiet session | 2 | 67.30 | 38.03 | 1.000000 | yes |
| product | Exact optimized packed BF16, final quiet session | 1 | 72.23 | 29.90 | 1.000000 | yes |
| product | BF16 baseline | 2 | 78.21 | 38.31 | 1.000000 | yes |
| product | Compiled BF16, retained refiners | 2 | 65.34 | 39.52 | 1.000000 | yes |
| product | Lossless packed, retained refiners | 1 | 70.62 | 31.38 | 1.000000 | yes |
| product | Compiled BF16 + exact caches + prefetch | 1 | 68.04 | 38.08 | 1.000000 | yes |
| product | Lossless packed + exact caches + prefetch | 1 | 73.86 | 29.90 | 1.000000 | yes |
| product | FP16 projection failure (non-finite) | 0 | invalid | 37.58 | invalid state | non-finite |
| product | Five original passes (skip head 1) | 0 | β€” | 37.35 | 0.796451 | no |
| portrait | BF16 baseline | 3 | 72.06 | 38.30 | 1.000000 | yes |
| portrait | Compiled BF16, retained refiners | 3 | 67.38 | 39.59 | 1.000000 | yes |
| portrait | Lossless packed, retained refiners | 2 | 73.56 | 31.38 | 1.000000 | yes |
| portrait | Compiled BF16 + exact caches + prefetch | 2 | 69.54 | 38.22 | 1.000000 | yes |
| portrait | Lossless packed + exact caches + prefetch | 2 | 74.44 | 29.90 | 1.000000 | yes |
| portrait | FP16 projection failure (non-finite) | 1 | invalid | 38.33 | invalid state | non-finite |
| portrait | Five original passes (skip head 1) | 1 | 68.43 | 38.09 | 0.694111 | no |
| native-4k | Lossless packed + exact caches + prefetch | 0 | β€” | 31.20 | 1.000000 | yes |
<!-- END_RESULTS_TABLE -->
The original current-arithmetic product baseline was 72.78 s end to end,
including 67.87 s of denoising, at 38.31 GiB peak MLX memory. The retained-refiner
compiled trials subsequently reached 65.00–65.68 s on that prompt, with a
39.52 GiB peak. The later exact-cache/prefetch configuration frees the refiner
weights, so its speed and memory should be assessed from its own table row.
The baseline rows labeled with active packing are not clean speedup
denominators. Three background Git maintenance jobs were eventually identified;
initially only one had been paused. All three were paused for the final
`quiet-native` warm samples and `quiet-packed` samples, starting at 03:12:36 UTC,
and automatically resumed at 03:24:36 UTC. Other user applications remained
open. Our full-image GPU jobs ran sequentially. The 4K run and earlier trials
had mixed CPU background load; their memory and equality results remain valid,
but no controlled 4K speedup is claimed. The OS file cache was never flushed.
Hardware: Mac mini, M5 Pro, 15 CPU cores, 16 GPU cores, 64 GB; macOS 27.0.1,
MLX/Metal 0.32.4, Python 3.12.14, PyTorch 2.15.0 and Core ML Tools 9.0.
AC power mode remained Automatic. High Power Mode is supported, but changing it
required unavailable administrator credentials; no High Power measurements are
claimed. `MLX_ENABLE_TF32=0` and `MLX_CONV_WINOGRAD=0` retain the original
precision policy.
## Look at the actual losses
- [Product: complete-generation grid](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/product/grid.png)
and [native detail crops](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/product/details-grid.png).
- [Portrait: complete-generation grid](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/portrait/grid.png)
and [native detail crops](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/portrait/details-grid.png).
- [4K comparison](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/native-4k/index.html).
- [All Core ML devices on the same product latent](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/decoder-all-product/comparison/index.html).
- [Decoder alternatives on the same portrait latent](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/decoder-portrait/comparison/index.html).
Every gallery has side-by-side frames, identical native-resolution crops,
vertical stacks, RGB differences amplified eight times, full-resolution
PSNR/SSIM/MAE and source PNG hashes. The tool verifies prompt, seed, dimensions,
canvas, noise source, schedule and reference inputs. Fewer-head experiments are
explicitly labeled. Pixel similarity is not a general measure of visual quality
or prompt adherence. Decoder-only tests use the *identical saved final latent*.
No comparison imagery is synthesized or retouched.
Complete-generation equality is measured against the original **MLX BF16
recipe**, not byte identity with CUDA or the FP32 lossless decoder. The
[same-latent FP32 decoder comparison](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/fp32-reference/comparison/index.html)
isolates that separate precision choice. Cross-backend CUDA comparisons remain
in [the initial report](../RESULTS.md); different decoder quality modes must not be mistaken for
optimization-induced loss.
On the product latent, native BF16 versus the original FP32 decoder measured
SSIM 0.996629 and MAE 0.559/255. The FP32 compiled decoder took 5.21 s in its
single recorded sample. Use `--vae-dtype float32 --compile-vae --vae-batch 4`
to select that precision; weight packing and decoder precision are independent.
## Lossless 12-bit storage
The codec stores each BF16 low byte verbatim and uses a nibble for its sign and
one of eight adjacent high-byte exponent values. Rare values outside that
window use sorted uint32 indices plus their exact uint16 payload. The GPU
reconstructs the original BF16 bits immediately before an ordinary BF16 GEMM.
It performs no quantization, multiplication by a scale, pruning or low-rank fit.
Block-0 projection storage was 75.3–76.2% of its original size. Seven-repeat
operator medians, including per-call reconstruction, were:
| Projection, actual 4181-token shape | BF16 ms | Packed ms | Output error |
|---|---:|---:|---|
| QKV, 5376 β†’ 21504 | 39.11 | 41.56 | zero |
| Attention output, 7168 β†’ 5376 | 13.41 | 14.00 | zero |
| MLP input, 5376 β†’ 28672 | 52.17 | 54.94 | zero |
| MLP output, 14336 β†’ 5376 | 32.67 | 34.37 | zero |
The tests reconstruct all 65,536 BF16 bit patterns, including signed zero,
subnormals, infinities and NaN payloads. Real projection outputs are bit-exact.
Compiled blocks receive compressed arrays as explicit dynamic inputs so the
compiler cannot fold decoding into a permanently retained dense matrix. Cache
shards are written atomically and separated by source-checkpoint fingerprint.
The tradeoff is extra disk space and a small reconstruction cost, not accuracy.
## Compilation, scheduling and loading
The actual first DiT block, using saved H3 conditioning and noise, took a
193.87 ms median eagerly and 181.86 ms compiled, with identical BF16 output.
MLP chunks of 512/1024/2048 were slower on this machine. Compilation also
reduced the complete product VAE from about 2.87 s at batch one to about 1.56 s
at batch four. The batch-four image was identical; batch eight introduced small
rounding changes (product SSIM 0.998905, MAE 0.125/255).
Native lazy `mx.load` was 2.1–2.3Γ— faster than NumPy memmap/copy in three
page-cache-warm tensor-read tests. Full tensor SHA256s matched. `WeightStore`
discards each lazy shard dictionary immediately, so evaluated tensors cannot
bypass its explicit residency budget. The memmap reader remains available for
audits. These read tests do not measure an uncached whole-checkpoint load.
Prefetching asynchronously dispatches a block while loading the next weights,
synchronizing after a bounded number of blocks. Once all core blocks are
resident, it returns to the requested warm evaluation schedule. It preserves
all operations and dependency ordering. Load time may overlap GPU execution;
the raw records therefore expose block wall time and per-step wall time rather
than adding CPU load time to a supposed GPU compute time.
## Core ML / ML Package: CPU, GPU, ANE and all
The repository exports the **complete 24-block VAE tile decoder** as a fixed-batch
ML Package and supports using it behind the MLX denoiser. This is not a complete
33B Core ML denoiser port. Source weights and inputs are rounded to native BF16
values before conversion; Core ML's internal arithmetic is FP16 and changes
rounding. Package metadata binds the decoder to its original VAE and batch.
`MLComputePlan` was inspected for actual preferred devices. Selecting `ALL`
alone is not evidence of ANE execution, as described in Apple's
[compute-plan utilities](https://apple.github.io/coremltools/docs-guides/source/mlmodel-utilities.html).
In the eight-tile test:
| Compute units | Preferred placement | Median tile-batch ms | Same-latent product SSIM |
|---|---|---:|---:|
| CPU_ONLY | CPU | 655.40 | 0.995786 |
| CPU_AND_GPU | GPU | 85.27 | 0.998319 |
| CPU_AND_NE | ANE, with two CPU operations | 459.67 | **0.492298** |
| ALL | GPU | 85.20 | 0.998319 |
The ANE export produces conspicuous tile/reconstruction artifacts on real
latents despite finite outputs on random test data. Its cause is not established;
this is a failed correctness result, not proof that ANE inevitably loses this
much precision. It is slower as well. Treat this batch-eight ANE export as a
diagnostic, not a quality-preserving generation option. Its recorded model-load
wall time includes a deliberately paused client and is not an isolated compiler
benchmark. Prediction samples were taken after resuming.
The batch-one ANE export also completed, at 251.73 ms per tile. Its full-image
check reduced the batch-eight artifacts but still changed pixels: SSIM 0.988974,
MAE 1.179/255, max error 43/255, and 30.63 s decoder-only latency. Its
[image grid](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/contiguous-b1-CPU_AND_NE/comparison/index.html)
and raw NumPy-bridge audit are preserved. The output was contiguous and MLX
imported its values exactly, so that test did not reveal an output-import error.
A repeated batch-eight cold probe with explicit contiguous I/O did not complete
within 150 s. These are bounded feasibility probes, not comparable compilation
benchmarks. Smaller batches are not an ANE speed win here.
Core ML GPU decoded the full product in 1.56 s in the first three-repeat trial;
compiled MLX took 1.56 s with closer or identical pixels. A later run after CPU
tests was slower. The MLX decoder is preferred on this Mac. The package remains
available for reproducible device experiments and other hardware.
```bash
# Isolated conversion environment; no MLX dependency is required by the exporter.
uv venv --python 3.12 .venv-coreml
uv pip install --python .venv-coreml/bin/python coremltools==9.0 torch safetensors numpy
.venv-coreml/bin/python -m experiments.v2_coreml_vae \
--vae /path/to/MiniMax-H3-Image/vae --batch 8 \
--package cache/coreml-v2/vae-batch8.mlpackage
# Optional Core ML GPU decoder in the real MLX generation pipeline.
pip install -e '.[coreml]'
h3-mlx generate --model /path/to/MiniMax-H3-Image \
--text-encoder /path/to/MiniMax-H3/text_encoder \
--compile-blocks --eval-every 0 --prefetch-blocks 4 --wired-gb 42 \
--vae-backend coreml --vae-batch 8 \
--coreml-package cache/coreml-v2/vae-batch8.mlpackage \
--coreml-compute-units CPU_AND_GPU \
"A detailed image description" --output outputs/coreml-decoder.png
```
Core ML allocations are excluded from MLX allocator counters, so those counters
must not be used to claim a hybrid pipeline's total memory footprint. The
experimental converter completed successfully with the recorded PyTorch version;
Core ML Tools warns that this PyTorch release is outside its tested range.
## Other operator and runtime experiments
- **FP16 projection compute failed:** both prompts produced entirely non-finite
velocity/state tensors at the first head. The saved black images are invalid
diagnostic artifacts, not valid low-quality generations. This was missed by
normalized synthetic GEMM inputs. The runtime now rejects non-finite states;
the research profile explicitly opts into saving failure diagnostics.
- **Five-head sampling:** skip original head 1 and advance head 0 across the
larger interval. The five-head images are compared directly; this is a sampler
approximation, not a precision-preserving speedup. Product/portrait SSIM was
0.79645/0.69411, too large a change for the minimal-loss recommendation.
- **Affine Q8/Q6/Q4, MXFP8/MXFP4/NVFP4:** all four real block-0 projections were
measured at the actual 4181-token shape. Quantization did not provide a clear
overall compute advantage. Separate dequantize-then-GEMM tests occasionally
improved over fused quantized matmul but did not justify replacing the exact
path. Quantization remains useful for smaller weight storage; earlier full
Q8/Q4 images are included in the grids.
- **FP16 versus BF16, row padding, transposed weights:** padding and transposing
preserved tested outputs, but did not beat the native layout. FP16 changes
rounding and was not consistently faster.
- **Metal 4 TensorOps GEMM:** threadgroup shapes, SIMD counts, K barriers and
locality-preserving walk orders were swept on the actual QKV shape. The best
custom kernel was about 39.65 ms versus 39.12 ms native, with exact output.
It was not selected. MLX already uses M5 GPU neural acceleration here.
- **Custom online-softmax attention:** query/key tile sizes and SIMD counts were
varied. The correct prototype was much slower than native SDPA and introduced
small numeric differences. It remains an isolated operator experiment, with
no full-image quality or speed claim.
- **PyTorch MPS / Inductor and MPSMatrix:** real H3 matrix shapes were benchmarked
across CPU, MPS BF16/FP16/FP32 and raw MPSMatrix. MPS was competitive; CPU was
much slower. The compiled MPS MLP reached 83.18 ms and preserved its eager
output. This is an operator result, not a measured complete H3 MPS pipeline.
- **Memory residency:** default MLX wired memory is zero. An explicit, bounded
per-process limit prevented paging from becoming an uncontrolled variable in
resident-weight trials. Retaining refiners saves some load time for new
conditioning, but the exact refined-text cache is preferable for repeats.
The custom attention result is compared numerically to native SDPA. It was
rejected before full-image deployment because it was substantially slower;
operator-level error is not presented as an image-quality measurement.
## Reproduction and remaining opportunities
All scripts are under `experiments/v2_*.py`; raw operator samples are adjacent
JSON files. The image ledger and full traces are in the ignored local
`outputs/optimization-v2` directory. `v2_report.py` rebuilds the table, PNG grids,
hashes and chart from completed outputs. The normal test suite covers precision
boundaries, exact packed reconstruction, cached conditioning, conversion and
image comparison provenance.
Remaining opportunities include native mixed-format Metal tensor kernels,
vectorized lossless reconstruction, smaller ANE batches/layout rewrites with
intermediate tensor audits, and overlapping an ANE decoder with the next
request's GPU denoising for throughput. None should be advertised as a speed or
quality win without full-image checks. Further pruning, token merging, low-rank
compression or learned timestep reduction would need calibration or training;
they are not silently applied to the original checkpoint.
The runtime distinctions and operator design follow Apple's
[Core ML compute-unit documentation](https://apple.github.io/coremltools/docs-guides/source/model-prediction.html),
[compute-plan utilities](https://apple.github.io/coremltools/docs-guides/source/mlmodel-utilities.html),
[typed execution](https://apple.github.io/coremltools/docs-guides/source/typed-execution.html),
[Metal TensorOps guide](https://developer.apple.com/download/files/Metal-Performance-Primitives-Programming-Guide.pdf),
and [M5 GPU tensor operations talk](https://developer.apple.com/videos/play/wwdc2026/330/).
The M5 GPU shader-core neural accelerators are distinct from the ANE selected by
Core ML. The MLX implementation examined was
[v0.32.4](https://github.com/ml-explore/mlx/tree/v0.32.4), including its NAX GEMM
and SDPA routing. All performance and image-quality numbers above come from
local experiments, not vendor performance claims.