feat(image): fused DiT kernels, a 3x3 VAE convolution and faster large-row NAX matmuls, on mlx 0.32.3

.github/workflows/ci.yml

@@ -62,7 +62,7 @@ jobs:
       - name: Build extension + test deps
         run: |
           python -m pip install --upgrade pip
-          pip install "setuptools>=77" wheel "mlx==0.32.1" "nanobind==2.13.0" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers<5.13"
+          pip install "setuptools>=77" wheel "mlx==0.32.3" "nanobind==3.0.1" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers!=5.13.0"
           pip install -e . --no-build-isolation
 
       - name: Codec lint (incl. runtime kq.codecs() leg)
@@ -155,7 +155,7 @@ jobs:
       - name: Build extension (CPU-only) + test deps
         run: |
           python -m pip install --upgrade pip
-          pip install "setuptools>=77" wheel "mlx[cpu]==0.32.1" "nanobind==2.13.0" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers<5.13"
+          pip install "setuptools>=77" wheel "mlx[cpu]==0.32.3" "nanobind==3.0.1" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers!=5.13.0"
           pip install -e . --no-build-isolation
 
       # A non-Metal build must report no metallib (and not throw doing so).

CHANGELOG.md

@@ -4,7 +4,101 @@ All notable changes to this project are documented here. The format follows
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project aims to
 adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
+## [0.5.0]
+
+### Added
+- `kq.qk_norm_rope` normalizes and rotates DiT attention q and k in one
+  dispatch, about 25x faster than the unfused ops at Z-Image shapes.
+- `kq.quantized_matmul_glu` computes a gated feed-forward product of two packed
+  weights. On NAX GPUs (M5) a call holds one output buffer, not two or three.
+- `kq.quantized_matmul_lora` adds a low-rank adapter term to a packed matmul in
+  one call. On NAX GPUs (M5) the call is 0 to 11% faster than the unfused ops on
+  the layer shapes of an image model, 4% at the median.
+- `kq.quantized_matmul_residual` adds a residual, gated or not, to a packed
+  matmul in one call. On NAX GPUs (M5) it costs what the matmul alone costs.
+- `kq.conv2d_3x3` runs the 3x3 convolutions of an image VAE decoder. On NAX
+  GPUs (M5) a bfloat16 call takes 0.25 to 0.34 of the time of MLX's convolution.
+- `KQ_TUNE_LIVE=1` makes the NAX tile levers take effect on the next call, so
+  tile A/B runs share one process. `KQ_QMM_LOG` now logs NAX tile dispatches.
+- `KQ_NAX_TILE` and `KQ_QMM_TILE` pin the quantized matmul tile, and raise
+  with the kernel name on a tile that is not built; `KQ_QMM_TILE` accepts
+  only `64,64,2,2` for now. `kq.metallib_kernels()` lists the built kernels.
+- `kq.max_threadgroup_memory()` reports the GPU's threadgroup memory limit.
+- `benchmarks/bench_dit_shapes.py` and `benchmarks/bench_roofline.py` time
+  matmul tiles against stock and MLX affine arms at diffusion model shapes.
+- `KQ_SPLITK_CAP` caps the K slices of the NAX and ALU split-K routes at 32 or
+  fewer rows by GPU core count, for A/B runs on wide ones such as a vocab head.
+- `KQ_DEQUANT_ROUTE` turns the dequantize route off or forces it on any GPU,
+  and `KQ_QMM_LOG` logs each call that takes it.
+- `KQ_SPLITK_F32=1` makes bfloat16 split-K matmuls, MTP verify among them,
+  round once instead of once per K slice. `0` stores float16 partials again,
+  which can return inf.
+
+### Changed
+- Requires mlx 0.32.3, which fixes wrong values from dequantizing a strided
+  array and a crash at process exit.
+- `mlx-kquant[tools]` now installs beside transformers 5.13.1 and later, as
+  mlx-vlm 0.7 requires. Only 5.13.0, which cannot import mlx-lm, is excluded.
+- `KQ_NAX_SWIZZLE=n` folds 2^n row tiles on the BM=64 and BM=128 NAX tiles.
+  `1` now folds two row tiles, where it folded four.
+- On NAX GPUs (M5), prefill and image steps of 4096 or more tokens no longer
+  change matmul kernels with the token count.
+- On NAX GPUs (M5), `q6_k` matmuls in prefill and image steps run 5 to 20%
+  faster, with the same output.
+- On NAX GPUs (M5), `q5_k` and `q3_k` matmuls in prefill and image steps run
+  up to 10% faster on most shapes, with the same output.
+- On NAX GPUs (M5), `q4_k` matmuls in image steps run up to 5% faster, with
+  the same output.
+- On NAX GPUs (M5), `q4_0`, `q4_1`, `q5_0`, `iq4_nl`, `iq3_xxs`, `iq3_s`,
+  `iq2_xs` and `iq2_s` matmuls in image steps run up to 8% faster, with the
+  same output.
+- On NAX GPUs (M5), matmuls of 4096 or more rows at a large K, in image steps
+  such as the Qwen-Image MLP down projection and in text prefill of 4096 or
+  more tokens, run 1 to 8% faster on most codecs.
+- On NAX GPUs (M5), output heads and `q4_0` and `q8_0` float16 matmuls at 9 to
+  32 rows, as in batched decode and speculative verify, run up to 30% faster.
+- On NAX GPUs (M5), batched gemma-4 decode of prompts of different lengths
+  runs 1 to 2% faster. `KQ_GQA_NAX=0` restores the previous kernel.
+- On M3 Max, `q3_k` to `q6_k`, `nvfp4` and `ptq1_0` matmuls from 2048 rows (`q4_k` 16384), K 3840-12288,
+  N 3840-24576, run up to 12% faster, `nvfp4` 39%; other pre-M5 GPUs without NAX take the rows unmeasured.
+- The bundled metallib drops the float32-activation matmul kernels, which no
+  call reached, about 22 MB.
+
+### Fixed
+- `quantized_matmul`, `quantized_matmul_qmv_bias` and `gather_qmm` return
+  correct values for a strided weight view that is not evaluated before the
+  call, such as `w[0::2]`.
+- `dequantize`, `quantize` and the bias of `quantized_matmul_qmv_bias` take a
+  strided view that is not evaluated before the call, such as `w[::-1]`.
+- `route_shed` returns correct slots for strided views of its indices,
+  scores or slot table.
+- The fused MoE ops and `moe_router_topk` take a strided activation or logits
+  view that is not evaluated before the call, and raise for such a stack.
+- `sdpa_prefill_block_sparse` raises for keys or values with a strided head
+  dim that are not evaluated before the call, instead of reading them as dense.
+- `sdpa_sparse_decode` and `sdpa_sparse_prefill` return correct values for a
+  window, pool or pool scales view with a strided or reversed last dim.
+- `dsa_indexer_score_decode` returns correct scores for one sequence whose
+  keys are a view with a strided or reversed last dim.
+- `sdpa_decode_gqa` and the cascade raise for quantized-cache scales and biases
+  that differ in layout or have a strided last dim, instead of misreading them.
+- The fused MoE ops with a limited activation (`silu_limit`, `swiglu_clamp`)
+  return NaN for a NaN gate or up value instead of the value of the limit.
+- `quantized_matmul` returns zeros for an inner dim of 0. On a CPU stream a
+  weight with no rows ended the process at more than 16 rows of x.
+- A float16 matmul of 128 or fewer rows split over K, MTP verify among them,
+  no longer returns inf where its output is finite.
+- Matmuls of 2 to 8 rows on q4_0, pq2_0, q2_0 and ptq1_0 weights, MTP verify among
+  them, no longer return inf for activations past 65504, and fp16 error drops 3-5x.
+- Decode and MTP verify steps on pq2_0 and q2_0 weights no longer return inf for
+  bfloat16 activations past 65504.
+- `KQ_MOE_HALF=1` no longer returns inf for activations of a few hundred and
+  up; only activations past 65504 still do.
+- Tuning levers that default on or to a per-codec rule, such as `KQ_QMV_SPLITK`
+  and `KQ_MOE_DEDUP`, keep that default when set empty instead of turning off.
+- After `patch_mlx_lm_load()`, `mlx_lm.utils.load_model` builds a kquant
+  checkpoint with the caller's `get_model_classes` instead of the stock class.
+  `mlx_kquant.loader.load` takes it too.
 
 ## [0.4.18]
 

CMakeLists.txt

@@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.27)
 project(mlx_kquant LANGUAGES CXX C)
 
 # ----------------------------- Setup -----------------------------
-# C++20: MLX 0.32.1's public headers (device.h, stream.h) use defaulted
+# C++20: MLX 0.32.3's public headers (device.h, stream.h) use defaulted
 # comparison operators, a C++20 feature. Apple Clang accepts them as an
 # extension under -std=c++17, but GCC (the Linux toolchain) rejects them, so a
 # C++17 standard breaks the Metal-free Linux build. 20 builds cleanly on both.
@@ -12,6 +12,7 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
 option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
+option(KQ_PROBE_KERNELS "Build the probe tiles of kernels under measurement" OFF)
 
 # ----------------------------- Dependencies -----------------------------
 # Pick the right interpreter. mlx.extension.CMakeBuild invokes cmake WITHOUT
@@ -37,7 +38,7 @@ find_package(
 # "incompatible function arguments". Build isolation installs the pin;
 # non-isolated dev builds (build_ext --inplace, -e --no-build-isolation) use
 # whatever the venv has - guard that path here.
-set(KQ_NANOBIND_PIN "2.13.0")
+set(KQ_NANOBIND_PIN "3.0.1")
 execute_process(
   COMMAND "${Python_EXECUTABLE}" -c "import nanobind; print(nanobind.__version__)"
   OUTPUT_STRIP_TRAILING_WHITESPACE
@@ -153,6 +154,12 @@ target_sources(
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_kvarn.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_moe_glu.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_norm_fused.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_qk_norm_rope.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_matmul_fused.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_matmul_lora.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_matmul_resid.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_conv_nax.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_dense.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_hadamard.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_lora_epilogue.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_hc_glue.cpp
@@ -221,12 +228,27 @@ endif()
 target_link_libraries(mlx_kquant_ext PUBLIC gguflib)
 target_link_options(mlx_kquant_ext PRIVATE ${KQ_WHOLE_ARCHIVE_GGUFLIB})
 
+# A probe build carries the probe tiles and reads the levers that select
+# them. A release build compiles those levers out. PUBLIC, because the
+# sources above are PUBLIC and so are compiled into the Python module too,
+# and that copy is the one the bindings call.
+if(KQ_PROBE_KERNELS)
+  target_compile_definitions(mlx_kquant_ext PUBLIC KQ_PROBE_KERNELS)
+endif()
+
 # ----------------------------- Metal library -----------------------------
 # The kq_* kernels compiled into a single mlx_kquant.metallib.
 # The repo's metal/ dir is searched BEFORE the stock-wheel include tree so the
 # repo's kq_*.h + quantized_utils.h (which adds load_vector) win;
 # steel/gemm/*, utils.h, etc. fall through to the stock wheel headers.
 if(MLX_BUILD_METAL)
+  # Probe tiles are measured against the shipped ones and stay out of a
+  # release metallib.
+  set(KQ_PROBE_METAL_SOURCES "")
+  if(KQ_PROBE_KERNELS)
+    list(APPEND KQ_PROBE_METAL_SOURCES
+         ${CMAKE_CURRENT_LIST_DIR}/metal/kq_quantized_nax_epi_probe.metal)
+  endif()
   # The metallib has no depfile; header edits retrigger it only through DEPS.
   file(GLOB KQ_METAL_KERNEL_HEADERS CONFIGURE_DEPENDS
        ${CMAKE_CURRENT_LIST_DIR}/metal/mlx/backend/metal/kernels/*.h)
@@ -240,6 +262,10 @@ if(MLX_BUILD_METAL)
     SOURCES
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_quantized.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_quantized_nax.metal
+    ${CMAKE_CURRENT_LIST_DIR}/metal/kq_quantized_nax_epi.metal
+    ${CMAKE_CURRENT_LIST_DIR}/metal/kq_quantized_nax_lora.metal
+    ${CMAKE_CURRENT_LIST_DIR}/metal/kq_quantized_nax_resid.metal
+    ${CMAKE_CURRENT_LIST_DIR}/metal/kq_conv_nax.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_quantized_encode.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_sdpa.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_sdpa_nax.metal
@@ -256,12 +282,14 @@ if(MLX_BUILD_METAL)
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_moe_glu.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_moe_glu_kq.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_norm_fused.metal
+    ${CMAKE_CURRENT_LIST_DIR}/metal/kq_qk_norm_rope.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_hadamard.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_lora_epilogue.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_hc_glue.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_hc_lowrank.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_route_shed.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_skinny_mv.metal
+    ${KQ_PROBE_METAL_SOURCES}
     INCLUDE_DIRS
     ${CMAKE_CURRENT_LIST_DIR}/metal
     ${MLX_INCLUDE_DIRS}
@@ -296,7 +324,7 @@ if(BUILD_SHARED_LIBS)
     # installed mlx wheel at runtime. _ext.so lives in site-packages/mlx_kquant/,
     # libmlx in site-packages/mlx/lib/, so ../mlx/lib reaches it. (The absolute
     # build-tree rpath MLXConfig adds works only on the build machine; this is
-    # what makes a redistributed wheel find the ABI-pinned mlx==0.32.1. delocate
+    # what makes a redistributed wheel find the ABI-pinned mlx==0.32.3. delocate
     # must therefore *exclude* libmlx from vendoring - see [tool.cibuildwheel].)
     target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
     target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path/../mlx/lib)

CONTRIBUTING.md

@@ -12,7 +12,7 @@ front-end. A few notes to get productive quickly.
 
   ```sh
   python -m venv .venv && source .venv/bin/activate
-  pip install "mlx==0.32.1" "nanobind==2.13.0" "cmake>=3.27"
+  pip install "mlx==0.32.3" "nanobind==3.0.1" "cmake>=3.27"
   pip install -e ".[dev,tools]"   # builds _ext + mlx_kquant.metallib
   ```
 

README.md

@@ -52,14 +52,14 @@ pip install mlx-kquant            # the K-quant ops + precompiled metallib
 pip install "mlx-kquant[tools]"   # the CLI (quantize / run / chat / lora / fuse)
 ```
 
-Either pulls the ABI-matched `mlx==0.32.1` automatically.
+Either pulls the ABI-matched `mlx==0.32.3` automatically.
 
 **From source** (or to develop) needs the Metal toolchain (`xcrun metal`); the metallib compiles at
 install time, no runtime JIT:
 
 ```sh
 git clone https://github.com/asher/mlx-kquant && cd mlx-kquant
-pip install "mlx==0.32.1"         # pinned, ABI-matched stock wheel (pulls the Metal backend)
+pip install "mlx==0.32.3"         # pinned, ABI-matched stock wheel (pulls the Metal backend)
 pip install -e ".[tools]"         # builds _ext + mlx_kquant.metallib; adds mlx-lm for the CLI
 ```
 
@@ -70,7 +70,7 @@ x86_64 stays on the scalar/threaded path.) The base `mlx` wheel ships no backend
 install the CPU one explicitly:
 
 ```sh
-pip install "mlx[cpu]==0.32.1"   # base frontend + libmlx CPU backend
+pip install "mlx[cpu]==0.32.3"   # base frontend + libmlx CPU backend
 pip install -e . --no-build-isolation
 ```
 
@@ -299,9 +299,10 @@ All optional; the defaults are right for normal use.
   kernel. By default `sp` serves launches under 2048 threadgroups and `dd` serves wider ones at 2
   to 8 rows. Read per call. `loop` and `sp` give identical bytes. A/B lever.
 - `KQ_MOE_HALF=1` - on Metal, run the iq2_xs, iq2_xxs and iq3_xxs fused MoE decode gathers on the
-  half-dot kernels (half grid tables, half-staged activations, half chunk dots with a float sum per
-  chunk; shared-expert slots stay on the float path). Read per call. Default off: the outputs differ
-  from the float kernels at half rounding level. A/B lever.
+  half-dot kernels (half grid tables and half-staged activations, with each chunk dotted in float;
+  shared-expert slots stay on the float path). Read per call. Default off: the outputs differ from
+  the float kernels at half rounding level, and an activation past 65504 in magnitude turns to inf.
+  A/B lever.
 
 - `KQ_SDPA_IDX_NAX=0` - run `sdpa_fa_indexed` on its simdgroup kernel even where the NAX tile
   kernel is available. Read once per process. The simdgroup kernel matches `sdpa_fa_verify` over
@@ -376,7 +377,7 @@ are informed by our analysis of the mixed-precision quants that [Unsloth][unslot
 
 ## Version pinning
 
-Pinned to `mlx==0.32.1`. The kernels include MLX's steel headers and the extension links `libmlx`,
+Pinned to `mlx==0.32.3`. The kernels include MLX's steel headers and the extension links `libmlx`,
 binding it to that release's ABI and header API. To move to a newer MLX: update the bundled headers
 under `metal/mlx/backend/metal/kernels/` for that wheel, rebuild, and re-run the test suite.
 
@@ -390,10 +391,10 @@ python -m pytest tests/
 
 - **macOS 26.2 (Tahoe) or later on Apple Silicon** (M-series). Building from source needs the Metal
   toolchain (`xcrun metal`).
-- **Linux** (x86_64 or aarch64) is supported CPU-only. Build against `mlx[cpu]==0.32.1`, no Metal
+- **Linux** (x86_64 or aarch64) is supported CPU-only. Build against `mlx[cpu]==0.32.3`, no Metal
   toolchain required. See [Install](#install) and [Limitations](#limitations).
-- **Python >= 3.10** (the pinned `mlx==0.32.1` ships no cp39 wheel).
-- **`mlx==0.32.1`** exactly - the kernels include MLX's steel headers and the extension links
+- **Python >= 3.10** (the pinned `mlx==0.32.3` ships no cp39 wheel).
+- **`mlx==0.32.3`** exactly - the kernels include MLX's steel headers and the extension links
   `libmlx`, so the ABI is version-locked (see [Version pinning](#version-pinning)).
 
 ## Limitations