Blame for docs/TECHNICAL.md

# MiniMax-H3-Image-mlx

Native MLX/Metal inference for MiniMax-H3-Image on Apple silicon. The default
uses the original BF16 transformer, all six distilled PDD steps, dense
attention, FP32 timestep/input/output projections, and the released overlapping
BF16 image VAE decoder. Q8, Q4, custom Metal normalization and fewer steps are
explicit experiments.

The port reads the original safetensors directly. It does not need CUDA,
SGLang, a remote API, or a second copy of the checkpoint. PyTorch is used only
for the original CPU random-number stream and independent numerical tests;
the model operators run in MLX on Metal.

The [latest operator-tuning report](../reports/optimization-r3/RESULTS.md) applies
the local `tmp/mlx-skills` workflow: native-order Q/K fusion, M5 FC2 tuning,
physical-memory accounting, real activation sweeps and complete trajectory
checks. [Matched image stacks](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-r3/index.html) include
the exact modes and links to earlier precision-loss examples.

On the tested Mac, paired warm latency improves **5.07%** with dense BF16
storage and **3.57%** with lossless packed storage, each against its own
current control. Every pixel and saved denoising tensor matches on four cases.
The largest sampled process footprint across candidate cases is **40.50 GiB
dense / 32.56 GiB packed**; see the report for timing ranges and memory scope.

The [further optimization report](../reports/optimization-v2/RESULTS.md) covers
compiled MLX, exact weight packing, asynchronous loading, PyTorch/MPS, custom
Metal kernels, and a complete Core ML VAE tested on CPU, GPU, ANE and all
devices. It includes real side-by-side image grids and failed experiments.

## Install and generate

An Apple silicon Mac and a compatible MLX Metal installation are required.
The tested machine is a 64 GB M5 Pro. BF16 inference uses a resident DiT core
of about 36 GiB, plus activations, decoder and allocator overhead. Smaller
memory budgets stream a fixed subset of blocks; Q8/Q4 reduce memory and change
the image. Disk space for the original image model and full encoder is roughly
134 GB, independent of this small code repository.

```bash
uv venv --python 3.12
source .venv/bin/activate
uv pip install -e '.[audit]'
h3-mlx doctor

h3-mlx generate \
  --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  "An editorial photograph of a cobalt perfume bottle on pale stone" \
  --width 2048 --height 2048 --seed 731 --output outputs/bottle.png
```

`requirements-macos.lock` records the exact tested environment, including audit
tools. Install it before the editable package when reproducing these measurements.

The neighboring checkpoint paths on the development machine are
`../minimax-h3-image-opt/models/MiniMax-H3-Image` and
`../minimax-h3-image-opt/models/MiniMax-H3/text_encoder`. Pass them explicitly;
the package has no dependency on that repository at runtime.

If weights are not already present, download the official repositories:

```bash
hf download MiniMaxAI/MiniMax-H3-Image \
  --revision 6534ea1c979a5f8cb00c2bd1025d5ca4ee8d8901 \
  --local-dir models/MiniMax-H3-Image
hf download MiniMaxAI/MiniMax-H3 --include 'text_encoder/*' \
  --revision 0919fb684102ef115126612f3fbb84d1027e7317 \
  --local-dir models/MiniMax-H3
```

The model expects detailed English prompts. Prompt expansion is not performed
silently. Requests smaller than 4 MP use the upstream larger generation canvas
and are resized with LANCZOS. `--min-generation-area 1024` disables that
enlargement for quick low-resolution development trials; these are outside the
model's usual 4–20 MP working range.

Every output PNG has a sibling JSON recording the prompt, seed, noise source,
canvas, schedule, precision options, per-step times, model-load times, memory
and software/hardware versions. `--save-trace` also saves conditioning, each
velocity/Euler state and the final latent for numerical investigations.

## Precision and performance controls

For the tested 64 GB M5 Pro, the recommended exact configuration is:

```bash
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 \
  --kernel metal-qk-exact --tune-matmul \
  "A detailed image description" --output outputs/optimized.png
```

This keeps all six heads and the original BF16/FP32 arithmetic. The M5 tuning
is scoped to MLX 0.32.4 and measured token ranges; other devices, versions and
shapes retain native execution. See the report for complete image/trajectory
checks and paired timing ranges. For lower
memory use, add `--weight-storage packed12` and use `--wired-gb 32` on this
machine. Packing is lossless, builds an approximately 27 GiB disk cache on first
use, and trades a small amount of compute for smaller resident weights. Wired
memory settings must fit the actual Mac; the runtime enforces Metal's limit.

| Option | Effect |
|---|---|
| Default | BF16 core, original schedule, dense Metal SDPA, FP32 precision islands |
| `--compile-blocks --eval-every 0` | Compile DiT blocks and evaluate the complete warm step |
| `--prefetch-blocks 4` | Bound asynchronous loading/execution overlap to four blocks when weights need loading |
| `--kernel metal-qk-exact` | Native-order Q/K RMSNorm + partial RoPE fusion; keeps native full-width normalization |
| `--tune-matmul` | Metal 4 FC2 tile selected on real activations; M5 Pro/MLX 0.32.4 and 4096–8192 joint tokens, native fallback otherwise |
| `--kernel metal-exact` | Also fuse full-width RMSNorm/AdaLN; exact on the recorded suite, with a measured physical-memory tradeoff |
| `--compile-vae --vae-batch 4` | Faster decoder; identical pixels on the tested product and portrait |
| `--vae-dtype float32` | Original FP32 decoder arithmetic; slower, with a [same-latent precision grid](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/fp32-reference/comparison/grid.png) |
| `--weight-storage packed12` | Lossless BF16 storage with exact GPU reconstruction before native GEMMs |
| `--wired-gb 42` | Process-local residency on the tested 64 GB Mac; does not modify system settings |
| `--no-refiner-cache` | Bypass exact refined-conditioning caching for audits |
| `--keep-refiners` | Keep refiner weights for new conditioning instead of releasing them |
| `--kernel metal --norm-threads 256` | Experimental fused RMSNorm/AdaLN; preserves intermediate BF16 rounding but changes reduction order |
| `--kernel metal-qk` | Also fuse per-head Q/K normalization and partial 3D RoPE |
| `--vae-batch 8` | Decode several complete native tiles together; GEMM rounding can change |
| `--mlp-chunk 1024` | Lower activation memory; `0` uses a full-sequence MLP |
| `--attention-chunk 512` | Bound attention query working size while keeping all keys; benchmark before using |
| `--dit-cache-gb 16` | Pin up to 16 GiB of blocks and stream the rest without LRU thrashing |
| `--quantize 8` / `--quantize 4` | Weight-only affine quantization of DiT attention/MLPs, group size 64; AdaLN and precision islands remain native |
| `--group-size 32` | Smaller quantization groups; separate quality/memory tradeoff |
| `--encoder-quantize 8` | Independently quantize the fifty language layers |
| `--experimental-steps 0,2,4` | Three original heads with larger Euler intervals; a preview approximation |
| `--offload-after-denoise` | Release DiT before decoding for tighter memory budgets |
| `--rng mlx` | Avoid the PyTorch dependency, using a different seed/noise stream |
| `--vae-backend coreml` | Experimental FP16 ML Package decoder; requires an exported package and matching tile batch |

The JSON memory report separates MLX allocation counters from macOS parent
physical footprint. Its process lifetime peak includes earlier requests;
the experiment harness also samples each request separately. These counters
overlap and must not be added. The text encoder now reads only requested
vocabulary rows, preserving original BF16 payloads and avoiding a full
vocabulary allocation on a conditioning-cache miss.

Core ML GPU and `ALL` had small image differences. The batch-eight ANE export
had severe reconstruction artifacts and is a diagnostic option. See the
[device comparison grid](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples/blob/v0.4.0/comparisons/optimization-v2/decoder-all-product/comparison/grid.png)
and report before selecting it. The full denoiser remains MLX in every runtime
configuration; only the VAE has a Core ML export.

For repeated quantized runs, preconvert attention/MLP matrices once. This
avoids online conversion at startup; the original checkpoint remains the source
of the exact modulation tables and FP32 layers:

```bash
python -m h3_mlx.convert --model /path/to/MiniMax-H3-Image --output models/core-q8 --bits 8
h3-mlx generate --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  --quantized-core models/core-q8 "A detailed image description" --output outputs/q8.png
```

Conversion writes one block per shard, verifies the source identity on load,
and records SHA256 hashes. It requires additional disk space (roughly half
the native core size for Q8). Quantization is a memory tradeoff; it does not
automatically improve image-generation latency.

MLX's M5 defaults can reduce the precision of FP32 matrix operations. The
package sets `MLX_ENABLE_TF32=0` and `MLX_CONV_WINOGRAD=0` **before importing
MLX**, retaining full precision for the model's sensitive layers. Explicit
environment overrides are respected and recorded. When embedding this package,
import `h3_mlx` before using MLX, or set those variables before starting Python.
See [MLX's precision documentation](https://ml-explore.github.io/mlx/build/html/usage/precision.html).

Cross-backend byte identity with CUDA is not promised. Reduction order,
attention arithmetic and GEMM algorithms can change a seeded diffusion
trajectory. The tests distinguish ATen RMSNorm from Qwen's different BF16
rounding order and allow only narrowly bounded cross-device errors.

## Reuse resident models

A fresh CLI process includes checkpoint loading. Use one `Pipeline` for
multiple requests to retain the DiT and VAE in unified memory. Conditioning
and native modulation tables are also cached on disk.

```python
from h3_mlx.runtime import Pipeline

pipe = Pipeline("/path/to/MiniMax-H3-Image", "/path/to/MiniMax-H3/text_encoder")
try:
    for seed in (731, 732, 733):
        image, record = pipe.generate(
            "A detailed editorial photograph of a blue perfume bottle on pale stone",
            seed=seed, output=f"outputs/bottle-{seed}.png",
        )
finally:
    pipe.clear()
```

The package does not change system wired-memory settings. It caps block
residency below Metal's recommended working set and keeps a fixed resident
subset when the full core does not fit. Use a smaller block cache and MLP
chunks for larger canvases or lower-memory Macs.

## Image editing

The MLX implementation includes the Qwen3-VL vision tower, all three DeepStack
features, multimodal RoPE, a single-frame reduction of the causal CNN VAE
encoder, posterior sampling, and fixed noised reference anchors. References
are processed in input order, up to sixteen:

```bash
h3-mlx generate \
  --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  "Change the bottle in <Picture 1> to emerald glass while preserving the label" \
  --image bottle.png --seed 42 --output outputs/edit.png
```

Qwen receives a 1024 px short edge. VAE references retain short edges in
1024–2048 px, with the native 16 px crop and even latent grid. Target noise,
anchor noise and posterior samples consume the CPU generator in upstream
order. This path requires more conditioning work than text-only generation.

## Measurements and image comparisons

[The experiment report](../reports/RESULTS.md) describes measured latency,
precision, failed performance hypotheses and remaining work. Raw operator
trials are in [operators.json](../reports/operators.json). Measurements describe
this Mac and the recorded prompts; they are not a general speed or image
quality guarantee.

Generate a vertical stack, identical 1:1 crops and amplified RGB differences:

```bash
python -m h3_mlx.compare \
  --reference outputs/suite/product/native-1.png \
  --candidate outputs/suite/product/tuned-1.png \
  --candidate outputs/suite/product/q8-1.png \
  --label 'BF16 native' --label 'BF16 tuned' --label 'Q8' \
  --output comparisons/product
```

The comparison tool verifies prompt, seed, dimensions, generation canvas,
noise generator and reference inputs from the inference records. It scores
full-resolution RGB, includes the exact crop coordinates and hashes every
source PNG. PSNR/SSIM quantify similarity to the reference trajectory; visual
inspection remains necessary. The tool never fabricates or retouches evidence.
It also writes `grid.png` and `details-grid.png` for direct side-by-side review.

## Reproduce experiments

```bash
python -m pytest -q
python -m h3_mlx.benchmark --tokens 256 4096 --repeats 7
python experiments/run_suite.py \
  --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  --profiles native tuned q8 q4 preview --cases product portrait --repeat 3
```

The suite labels requests that include model loading separately from requests
that reuse resident models. Conditioning-cache hits and all per-stage times
are retained in the raw records. Do not run another GPU workload concurrently
when measuring latency.

The central memory optimization computes the exact schedule's small AdaLN
outputs once and releases roughly 26 GB of modulation-only DiT weights.
It preserves every learned matrix and uses no curve fitting, pruning,
interpolation or timestep skipping in the default path. The fifty remaining
transformer blocks can then reside in this 64 GB Mac's unified memory.

Code is Apache-2.0; model weights retain their separate MiniMax licenses.
[NOTICE](../NOTICE) and [source provenance](../provenance/sources.json) identify the
SGLang and Transformers implementations used for the port.