qlip-mlx

Z-Image-Turbo at 2048px: 37.9 s/step on stock MLX → 19.1 s/step — 2× faster, same quality.

INT8 diffusion inference for Apple Silicon.

qlip-mlx runs the two operations that dominate a diffusion transformer step — attention and linear layers — on the INT8 matrix hardware of the Apple M5 GPU, through MLX. The first release ships two custom Metal-4 primitives, drop-in usable from any MLX model:

• sage_attention — to our knowledge the first Sage-class attention kernel on Apple Silicon: INT8 QK^T on the M5 Neural Accelerators inside a flash-attention structure with exact online softmax. 10–20% faster than mx.fast.scaled_dot_product_attention at DiT-scale sequence lengths, at ~1% total quantization error vs fp32 attention.
• w8a8_linear / W8A8Linear — W8A8 GEMM with per-token activation quantization, INT8×INT8→INT32 matmul and dequant+bias fused into a single graph node. No calibration data needed — weights quantize offline in seconds, activations quantize dynamically inside the kernel.
• qlip_mlx.caching / qlip_mlx.progressive — two kernel-independent acceleration axes ported from our ComfyUI-Qlip stack: EasyCache-style step skipping (easycache/taylor/hermite predictors) and progressive resolution (early denoising steps on a downscaled latent, sigma-backbone ladder). Runtime install on mflux models, one line each; they compose with the kernels — on Krea-2 the full stack is ×3.0.
• w4a8_linear / W4A8Linear — the same INT8 compute with weights stored as packed int4 (0.53 bytes/weight): identical speed to W8A8 at DiT shapes, half the weight memory — plus a register-dequant GEMV path for small batches (M ≤ 16), so decode-style calls run at weight-bandwidth speed. On Wan 2.2 this removes the one-off first-step weight-materialization cost entirely; on Qwen-Image it is what lets the full stack fit a 20B model into 24 GB.

📝 Read the full write-up: SageAttention on Apple Silicon — an INT8 stack for diffusion models on M5 (link TBD).

Results

Full stack = both primitives applied to the model: attention routed to sage_attention and the DiT's linear layers swapped to W8A8Linear (one runtime conversion at load, see Quick start). Attention only = just the sage_attention hook, weights untouched.

Model	Setup	Baseline	qlip-mlx	Speedup
Z-Image-Turbo 6B @2048px	stock mflux, full stack	37.9 s/step	19.1 s/step	×2.0
Z-Image-Turbo 6B @1024px	stock mflux, full stack	4.6 s/step	2.5 s/step	×1.8
Qwen-Image 20B @1328px	stock mflux, attention only	38.9 s/step	30.8 s/step	×1.26
Qwen-Image 20B @1328px	stock mflux, full stack (sage + W4A8)	38.9 s/step	29.0 s/step	×1.34
Wan 2.2 I2V 14B @480×832×33f	video DiT pipeline, full stack	30.0 s/step	21.8 s/step	−27%
Krea-2 @1024px, 20 steps	stock mflux, W4A8 + cache + progressive	9.6 s/step	3.2 s/step	×3.0

All numbers: fresh-boot machine, warm (JIT amortized), steady-state step, same seed as baseline; output quality verified by PSNR and visual review on every model. MacBook M5 Pro, 24 GB. With the w4a8 weight format, full-stack step times hold at half the DiT weight memory: on Qwen-Image it is what lets the full stack fit a 20B model at all, and on Wan 2.2 it shrinks the checkpoint 14 → 7.4 GB and removes the one-off first-step weight-materialization cost (46 s → 27 s).

Key features

• Real speedups where diffusion actually spends time. Weight-only quantization (int4/nvfp4/mxfp8) compresses storage but still computes in bf16 — at DiT batch sizes it is slower than bf16. qlip-mlx changes the arithmetic itself: INT8 compute on the M5's matrix units, ×1.6 over a plain bf16 matmul at DiT layer shapes.
• Exact where it matters. No softmax approximations, no calibration, no tuning knobs: quantization touches only QK^T (per-token / per-block scales); softmax and P·V stay in floating point. Kernel output is bitwise deterministic and matches a bit-faithful reference emulation of the SageAttention algorithm at fp16 rounding level.
• Drop-in integration. sage_attention takes the same [B, heads, L, D] tensors as mx.fast.sdpa; W8A8Linear is a standard mlx.nn module with .stub() for load_weights(). Accelerating a stock mflux model requires no fork — see examples/zimage_turbo.py (runtime monkeypatch + in-place layer conversion, ~8 s one-time at load).
• Built for diffusion transformers. Covers the standard DiT attention configuration out of the box — head dim 128, bidirectional (mask-free) attention; trivial all-ones masks are detected and handled automatically. On M4 and older the package installs as a clean no-op (is_available() → False).

Quick start

Requirements: Apple M5 or later, macOS 26+ (Metal 4), Python ≥ 3.12, MLX ≥ 0.31.

git clone <repo-url> qlip-mlx
cd qlip-mlx

python3.12 -m venv .venv && source .venv/bin/activate   # NB: plain `python3` on macOS is often 3.9
pip install -e .

Two-minute sanity check on bare tensors (no model download):

python examples/quickstart.py
# sdpa:  ~227 ms   sage: ~194 ms   speedup: x1.17     (40 heads x 14040 tokens)
# bf16:  ~33 ms    w8a8: ~20 ms    speedup: x1.6      (14040x5120 @ 5120x5120)
# w4a8:  ~21 ms    == w8a8 speed at half the weight bytes

Full stack on a public model through stock mflux (pip install mflux):

python examples/zimage_turbo.py "a red fox in the snow, cinematic" --res 2048
python examples/zimage_turbo.py "..." --res 2048 --weights w4a8   # half the weight memory
python examples/krea2_fullstack.py "a red fox in the snow" --steps 20   # W4A8+cache+progressive, ×3.0

One rule for w4a8: int4-grid weights must never touch modulation, embedding or output-projection layers (the example's converter skips them by name automatically — converting scale/shift/gate producers to int4 destroys a model).

Use it from Python:

import mlx.core as mx
from qlip_mlx import sage_attention, w8a8_linear, quantize_weight_int8

# attention: q/k/v [B, heads, L, 128] (or tokens-first [B, L, heads, 128])
out = sage_attention(q, k, v, scale=128 ** -0.5)          # fp16 out, same layout

# linear: quantize weights once, run W8A8 forwards
w_i8, sw = quantize_weight_int8(w_np)                     # [N, K] -> int8 + f32 scales
y = w8a8_linear(x, mx.array(w_i8), mx.array(sw), bias=b)  # per-token act quant inside

More examples:

python examples/gemm_formats_bench.py    # bf16 / int4 / nvfp4 / w8a8 at real DiT shapes
python tests/test_sage_reference.py      # correspondence vs the reference Sage algorithm

Build notes

The nanobind version must match the ABI of the mlx wheel: nanobind==2.12.0 for mlx 0.31.x (pinned — the mflux ecosystem currently sits on 0.31), nanobind==2.13.0 for mlx 0.32. A plain pip install -e . handles all of this via the pinned build requirements. If you build with --no-build-isolation instead, make sure the venv has cmake>=3.27 and nanobind==2.12.0 and that the venv's cmake shadows any older system one (PATH="$PWD/.venv/bin:$PATH"). First call in a process pays a one-off Metal JIT (~seconds); a metallib disk cache is on the roadmap, together with a D=64 head variant and long-context (32k+) tile configurations. w4a8_linear unpacks weights per call into a shared device scratch (one per layer shape, ≤70 MB) — in-kernel threadgroup staging measured 2.7× slower on M5.

License

Released under the MIT License, © 2026 thestage.ai labs.