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.

# 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.

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

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

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 isolates that separate precision choice. Cross-backend CUDA comparisons remain in the initial report; 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. 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 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.

# 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, compute-plan utilities, typed execution, Metal TensorOps guide, and M5 GPU tensor operations talk. The M5 GPU shader-core neural accelerators are distinct from the ANE selected by Core ML. The MLX implementation examined was 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.