mlx-diffusion-kit

Inference optimizations for diffusion and video super-resolution models on MLX / Apple Silicon. Training-free techniques that reduce compute by 2-5x without quality loss.

Current version: 0.2.1 — 23 optimization components, 284+ tests.

Foreword

This library was born from the same frustration that drove mlx-mfa: VSR inference on Apple Silicon is painfully slow. mlx-mfa tackles the attention kernel; mlx-diffusion-kit tackles everything else — step caching, token merging, cross-attention gating, VAE optimization, scheduling, and orchestration.

The two libraries are complementary:

mlx-mfa                              mlx-diffusion-kit
┌──────────────────────────┐        ┌──────────────────────────┐
│ Flash Attention kernels  │        │ Step-level caching       │
│ Sparse / GNA / Paged     │◄───────│ Token merging / pruning  │
│ KV cache management      │        │ Cross-attention gating   │
│ attn_bias (native Metal) │        │ VAE optimization         │
│ TurboQuant KV            │        │ Scheduling               │
│ SVDQuantLinear            │        │ Orchestrator (PISA)      │
└──────────────────────────┘        └──────────────────────────┘
        Kernel layer                     Optimization layer

Installation

pip install -e ".[dev]"

With mlx-mfa integration (for proportional attention via attn_bias):

pip install -e ".[mfa]"

Requirements: Python >= 3.10, MLX >= 0.25.0, Apple Silicon Mac.

Current Status

• 9 STABLE components — tested, integrated, production-ready API.
• 11 BETA components — functional and tested, API may evolve.
• 1 EXPERIMENTAL — functional, use with caution.
• 2 STUB — interface defined, implementation pending.
• 276+ tests pass, 0 failures.
• Primary validation hardware: Apple M1 Max.

Component Overview

Quick Start

import mlx_diffusion_kit as mdk

# 1. Cache text embeddings (all models)
emb_cache = mdk.TextEmbeddingCache()
embedding = emb_cache.get_or_compute("enhance 4x", my_t5_encoder, encoder_id="t5-xxl")

# 2. Step caching for multi-step models
from mlx_diffusion_kit.cache import TeaCacheConfig, load_coefficients
config = load_coefficients("cogvideox")  # Pre-calibrated

# 3. Token merging (all models)
merged, info = mdk.tome_merge(tokens, mdk.ToMeConfig(merge_ratio=0.5))
# ... run attention on merged tokens ...
output = mdk.tome_unmerge(merged_output, info)

# 4. Full orchestration
from mlx_diffusion_kit.orchestrator import DiffusionOptimizer, OrchestratorConfig
opt = DiffusionOptimizer(OrchestratorConfig(
    teacache=config,
    tome=mdk.ToMeConfig(merge_ratio=0.5),
    tgate=mdk.TGateConfig(gate_step=5),
    is_single_step=False,
))

See docs/API_MANUAL.md for complete API reference.

Target Models

Single-step (no inter-step caching)

Multi-step (step caching applicable)

Scripts

# Calibrate TeaCache coefficients for a new model
python scripts/calibrate_teacache.py --features-dir ./features/ --output coefficients.json

# Analyze layer redundancy for DeepCache
python scripts/analyze_layer_redundancy.py --weights model.npz --output scores.json

Documentation

• docs/API_MANUAL.md — Complete API reference for all 89 exports
• docs/ARCHITECTURE.md — Module structure and design principles
• CHANGELOG.md — Version history

License

MIT — see LICENSE for details.