mlx-kquant 0.1.0__tar.gz

mlx_kquant-0.1.0/CHANGELOG.md

@@ -0,0 +1,23 @@
+# Changelog
+
+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).
+
+## [0.1.0]
+
+First public release. A C++/Metal extension for a stock `mlx==0.31.2` wheel that
+adds the K-quant superblock and per-block integer codecs as native MLX ops
+(`kq.dequantize` / `quantized_matmul` / `gather_qmm` / `quantize`), with Metal and
+portable CPU paths for all ten codecs. On top of the ops, the `mlx-kquant` CLI
+quantizes an HF / mlx-lm model into a K-quant MLX safetensors checkpoint and runs,
+chats with, LoRA-fine-tunes, and fuses it, with importance-matrix calibration and
+per-tensor recipe inspection. A loader runs those checkpoints on stock mlx-lm.
+
+### Notes
+- `requires-python >= 3.10` (mlx 0.31.2 ships no cp39 wheel).
+- The GPU path is macOS 26 (Tahoe) or later on Apple Silicon (Metal); the NAX
+  matmul kernel needs the Metal 4 SDK (`MetalPerformancePrimitives`). Linux is
+  supported CPU-only - build against `mlx[cpu]==0.31.2`; model forwards there also
+  need `MLX_DISABLE_COMPILE=1` (an upstream MLX CPU-JIT limitation under GCC, not
+  mlx-kquant).

mlx_kquant-0.1.0/CMakeLists.txt

@@ -0,0 +1,209 @@
+cmake_minimum_required(VERSION 3.27)
+
+project(mlx_kquant LANGUAGES CXX C)
+
+# ----------------------------- Setup -----------------------------
+# C++20: MLX 0.31.2'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.
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
+
+# ----------------------------- Dependencies -----------------------------
+# Pick the right interpreter. mlx.extension.CMakeBuild invokes cmake WITHOUT
+# forwarding the building interpreter, so CMake's FindPython can otherwise latch
+# onto an unrelated system/uv Python that lacks mlx & nanobind (then
+# `python -m mlx --cmake-dir` returns empty and find_package(MLX) fails).
+# Precedence: an explicit -DPython_EXECUTABLE (e.g. via CMAKE_ARGS) wins; else an
+# active virtualenv ($VIRTUAL_ENV); else CMake's normal search.
+if(NOT Python_EXECUTABLE AND DEFINED ENV{VIRTUAL_ENV})
+  set(Python_EXECUTABLE "$ENV{VIRTUAL_ENV}/bin/python")
+endif()
+set(Python_FIND_VIRTUALENV FIRST)
+
+find_package(
+  Python 3.9
+  COMPONENTS Interpreter Development.Module
+  REQUIRED)
+
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE nanobind_ROOT)
+find_package(nanobind CONFIG REQUIRED)
+
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE MLX_ROOT)
+find_package(MLX CONFIG REQUIRED)
+
+# NOTE: do NOT splice ${MLX_CXX_FLAGS} into CMAKE_CXX_FLAGS - it is a CMake list
+# (-DACCELERATE_NEW_LAPACK;-D_METAL_), and embedding the ';' into the flags
+# *string* corrupts every compile command (`/bin/sh: -D_METAL_: command not
+# found`). MLXConfig already attaches these as INTERFACE_COMPILE_OPTIONS on the
+# imported `mlx` target, so -D_METAL_ propagates correctly to anything that
+# links `mlx` (below).
+
+# ----------------------------- gguflib (GGUF parser) -----------------------------
+# antirez/gguf-tools single-file GGUF parser, for the C++ kq.load_gguf loader.
+# PIN: a ref new enough to define GGUF_TYPE_Q5_0/Q5_1/Q2_K..Q6_K + correct
+# K-quant block sizes (older refs only know Q4_0/Q4_1/Q8_0).
+include(FetchContent)
+FetchContent_Declare(
+  gguflib
+  GIT_REPOSITORY https://github.com/antirez/gguf-tools/
+  GIT_TAG fdfafbed766db0a1e9019b07994cd88f133d1aab
+  # The pinned ref only knows GGML types up to BF16 (30); a tensor of any newer
+  # codec trips gguf_get_tensor's `>= GGUF_TYPE_COUNT` guard and truncates the
+  # tensor list. Teach the parser the post-BF16 type block-geometry (incl. the
+  # MXFP4/NVFP4 float codecs MLX has native kernels for). Idempotent; see script.
+  PATCH_COMMAND ${CMAKE_COMMAND} -DGGUFLIB_DIR=<SOURCE_DIR>
+                -P ${CMAKE_CURRENT_LIST_DIR}/cmake/patch-gguflib.cmake)
+FetchContent_MakeAvailable(gguflib)
+add_library(gguflib STATIC ${gguflib_SOURCE_DIR}/fp16.c
+                           ${gguflib_SOURCE_DIR}/gguflib.c)
+# Hidden visibility so the gguf_* symbols stay LOCAL to whichever image links
+# them whole-archive (below) and never interpose / collide at runtime with the
+# identically-named gguf_* symbols that libmlx also exports.
+set_target_properties(gguflib PROPERTIES POSITION_INDEPENDENT_CODE ON
+                                         C_VISIBILITY_PRESET hidden)
+# gguflib uses assert() to reject malformed tensor headers (e.g. ndim > 8).
+# -DNDEBUG (release builds) would compile those out, leaving out-of-bounds
+# reads/writes unguarded when loading untrusted GGUF files. Force NDEBUG off.
+target_compile_options(gguflib PRIVATE -UNDEBUG)
+
+# Whole-archive flag for the patched gguflib (used on both linking images below):
+# pull all its (hidden) gguf_* objects into the image so they win over libmlx's
+# exported copy. Apple ld spells this -force_load; GNU ld / lld use the
+# --whole-archive / --no-whole-archive bracket.
+if(APPLE)
+  set(KQ_WHOLE_ARCHIVE_GGUFLIB "-Wl,-force_load,$<TARGET_FILE:gguflib>")
+else()
+  set(KQ_WHOLE_ARCHIVE_GGUFLIB
+      "-Wl,--whole-archive,$<TARGET_FILE:gguflib>,--no-whole-archive")
+endif()
+
+# ----------------------------- C++ library -----------------------------
+add_library(mlx_kquant_ext)
+
+target_sources(
+  mlx_kquant_ext
+  PUBLIC ${CMAKE_CURRENT_LIST_DIR}/src/kquant_codec.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_cpu_decode.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_cpu_encode.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_ops.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_dequantize.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_matmul.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_gather.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_encode.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_gguf.cpp)
+
+target_include_directories(mlx_kquant_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR}/src)
+# gguflib include/link are PUBLIC: target_sources above lists the .cpp as PUBLIC,
+# so it is also compiled into the `_ext` nanobind module (consumer), which then
+# needs gguflib.h on its include path and the gguf_* symbols at link time.
+target_include_directories(mlx_kquant_ext PUBLIC ${gguflib_SOURCE_DIR})
+target_link_libraries(mlx_kquant_ext PUBLIC mlx)
+# Apple: route the CPU matmul's large-M GEMM through Accelerate (engages the
+# AMX/SME matrix units). PUBLIC so the nanobind module, which re-compiles the
+# PUBLIC sources above, sees the same define and link. Other platforms keep
+# the portable threaded-scalar fallback.
+if(APPLE)
+  target_compile_definitions(mlx_kquant_ext PUBLIC KQ_USE_ACCELERATE)
+  target_link_libraries(mlx_kquant_ext PUBLIC "-framework Accelerate")
+endif()
+# arm64: NEON-dotprod int8 GEMV kernels for the fused small-M CPU matmul.
+# The TU is compiled only on arm64/aarch64 hosts; execution is further gated
+# at runtime (KQ_CPU_NEON=0 kill switch; Linux aarch64 hwcap dotprod check).
+# Every other target keeps the portable scalar path (the header stubs
+# kq_neon_kernel to nullptr when KQ_CPU_NEON_TU is undefined).
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64)$")
+  target_sources(mlx_kquant_ext
+                 PUBLIC ${CMAKE_CURRENT_LIST_DIR}/src/kquant_cpu_neon.cpp)
+  target_compile_definitions(mlx_kquant_ext PUBLIC KQ_CPU_NEON_TU)
+  if(NOT APPLE)
+    # Linux toolchains may default to an armv8.0 baseline without dotprod;
+    # enable it for this one TU - dotprod instructions only execute after the
+    # runtime hwcap check passes.
+    set_source_files_properties(
+      ${CMAKE_CURRENT_LIST_DIR}/src/kquant_cpu_neon.cpp
+      PROPERTIES COMPILE_OPTIONS "-march=armv8.2-a+dotprod")
+  endif()
+endif()
+# libmlx bundles its OWN (older) gguflib and exports its gguf_* symbols. A plain
+# link can bind our gguf_get_tensor()/gguf_open() calls to MLX's copy, whose
+# GGUF_TYPE_COUNT predates the MXFP4/NVFP4 float codecs - so a tensor of type
+# >= 31 trips its `>= GGUF_TYPE_COUNT` guard and the tensor list silently
+# truncates at the first such tensor. Linking our patched gguflib whole-archive
+# pulls its objects into this image as hidden-visibility definitions, so every
+# gguf_* call resolves in-image to the patched parser instead of to libmlx.
+target_link_libraries(mlx_kquant_ext PUBLIC gguflib)
+target_link_options(mlx_kquant_ext PRIVATE ${KQ_WHOLE_ARCHIVE_GGUFLIB})
+
+# ----------------------------- 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)
+  mlx_build_metallib(
+    TARGET
+    mlx_kquant_metallib
+    TITLE
+    mlx_kquant
+    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_encode.metal
+    INCLUDE_DIRS
+    ${CMAKE_CURRENT_LIST_DIR}/metal
+    ${MLX_INCLUDE_DIRS}
+    OUTPUT_DIRECTORY
+    ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
+
+  add_dependencies(mlx_kquant_ext mlx_kquant_metallib)
+endif()
+
+# ----------------------------- Python bindings -----------------------------
+nanobind_add_module(
+  _ext
+  NB_STATIC
+  STABLE_ABI
+  LTO
+  NOMINSIZE
+  NB_DOMAIN
+  mlx
+  ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp)
+target_link_libraries(_ext PRIVATE mlx_kquant_ext)
+# kquant_gguf.cpp is a PUBLIC source (above), so it is also compiled into this
+# module and the LTO link keeps its own gguf_* references - force_load the
+# patched gguflib here too so they don't fall back to libmlx's exported copy
+# (see the mlx_kquant_ext force_load note for the full rationale).
+target_link_libraries(_ext PRIVATE gguflib)
+target_link_options(_ext PRIVATE ${KQ_WHOLE_ARCHIVE_GGUFLIB})
+
+if(BUILD_SHARED_LIBS)
+  if(APPLE)
+    # @loader_path finds the co-located libmlx_kquant_ext.dylib (same package
+    # dir); the second rpath resolves @rpath/libmlx.dylib against the *user's*
+    # 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.31.2. 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)
+  else()
+    # ELF equivalent of @loader_path: $ORIGIN is the dir of _ext.so. The
+    # co-located libmlx_kquant_ext.so sits beside it; ../mlx/lib reaches the
+    # user's installed mlx wheel's libmlx.so (auditwheel must *exclude* libmlx.so
+    # from vendoring, mirroring the macOS delocate exclude).
+    target_link_options(_ext PRIVATE "-Wl,-rpath,\$ORIGIN")
+    target_link_options(_ext PRIVATE "-Wl,-rpath,\$ORIGIN/../mlx/lib")
+  endif()
+endif()

mlx_kquant-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright © 2026 Asher Feldman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlx_kquant-0.1.0/MANIFEST.in

@@ -0,0 +1,23 @@
+# Build inputs the sdist needs to compile the extension from source.
+include CMakeLists.txt
+include setup.py
+include bindings.cpp
+recursive-include src *.cpp *.h
+recursive-include metal *.metal *.h
+recursive-include cmake *.cmake
+
+# Project metadata and tooling.
+include LICENSE
+include README.md
+include CHANGELOG.md
+recursive-include scripts *.py
+
+# Tests and fixtures, so an sdist build can run its own checks.
+recursive-include tests *.py *.npz
+include tests/fixtures/SHA256SUMS
+
+# Never ship machine-specific build artifacts in the sdist (the wheel carries
+# these via package-data; the sdist must not, or it would build them in stale).
+global-exclude *.so *.dylib *.metallib *.air *.pyc
+prune tests/fixtures/__pycache__
+global-exclude .DS_Store

mlx_kquant-0.1.0/PKG-INFO

@@ -0,0 +1,377 @@
+Metadata-Version: 2.4
+Name: mlx-kquant
+Version: 0.1.0
+Summary: GGUF K-quant dequantize / quantized-matmul / gather-qmm / quantize ops for MLX, via custom Metal kernels.
+Author: Asher Feldman
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/asher/mlx-kquant
+Project-URL: Repository, https://github.com/asher/mlx-kquant
+Project-URL: Issues, https://github.com/asher/mlx-kquant/issues
+Project-URL: Changelog, https://github.com/asher/mlx-kquant/blob/main/CHANGELOG.md
+Keywords: mlx,quantization,gguf,k-quant,metal,apple-silicon,llm
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Environment :: GPU
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mlx==0.31.2
+Provides-Extra: tools
+Requires-Dist: mlx-lm>=0.27; extra == "tools"
+Requires-Dist: transformers; extra == "tools"
+Requires-Dist: gguf; extra == "tools"
+Requires-Dist: numpy; extra == "tools"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: numpy; extra == "dev"
+Requires-Dist: gguf; extra == "dev"
+Dynamic: license-file
+
+# mlx-kquant
+
+[![ci](https://github.com/asher/mlx-kquant/actions/workflows/ci.yml/badge.svg)](https://github.com/asher/mlx-kquant/actions/workflows/ci.yml)
+
+Bring **K-quant precision to [MLX](https://github.com/ml-explore/mlx)** on Apple Silicon: a C++/Metal
+**extension** for a stock `mlx` wheel that adds the K-quant superblock and per-block integer codecs
+as native MLX ops, plus a toolchain that quantizes a model into a **K-quant MLX safetensors
+checkpoint** and runs, LoRA-trains, and fuses it.
+
+Two layers:
+
+- **Ops** (C++/Metal) - a `kq.*` namespace (`dequantize`, `quantized_matmul`, `gather_qmm`,
+  `quantize`) backed by Metal kernels compiled to a `.metallib` at build time (no runtime JIT). All
+  ten codecs: `q2_k, q3_k, q4_k, q5_k, q6_k` and `q4_0, q4_1, q5_0, q5_1, q8_0`.
+- **Tooling** (Python) - `mlx-kquant quantize / run / chat / lora / fuse` (plus `verify`, `inspect`,
+  `calibrate-imatrix`) and a `loader` that create and run K-quant checkpoints in **MLX-native
+  safetensors** format.
+
+## Why
+
+K-quant is the precision recipe behind the strongest small-footprint community quants; this makes it
+first-class on MLX. Quantize an HF / `mlx-lm` model to a uniform- or mixed-precision K-quant
+checkpoint, then load, generate, LoRA-train, and fuse it - all on a stock `mlx` wheel, all in MLX
+safetensors. The kernels are tuned for real models (matrix-contiguity handling for fused MoE experts,
+single-pass NAX matmul), see [Performance](#performance).
+
+## Install
+
+**macOS 26.2 (Tahoe) or later on Apple Silicon** (the GPU path), with the Metal toolchain
+(`xcrun metal`) and the exact pinned MLX wheel.
+
+```sh
+pip install "mlx==0.31.2"     # pinned, ABI-matched stock wheel (pulls the Metal backend)
+pip install -e .              # builds _ext + mlx_kquant.metallib
+pip install -e ".[tools]"     # + mlx-lm, for the CLI subcommands (quantize / run / chat / ...)
+```
+
+**Linux (CPU-only)** also builds, with no Metal toolchain. The ops run on their portable `eval_cpu`
+paths and no metallib is produced. (The tuned matmul/gather are Apple-Silicon-targeted: arm64 Linux
+picks up the NEON int8 GEMV when the CPU has dot-product, but the Accelerate GEMM is Apple-only, and
+x86_64 stays on the scalar/threaded path.) The base `mlx` wheel ships no backend on Linux, so
+install the CPU one explicitly:
+
+```sh
+pip install "mlx[cpu]==0.31.2"   # base frontend + libmlx CPU backend
+pip install -e . --no-build-isolation
+```
+
+CPU is for portability and CI, not throughput. Running a full model forward on Linux also needs
+`MLX_DISABLE_COMPILE=1`, see [Limitations](#limitations).
+
+Smoke-test the toolchain:
+
+```python
+import mlx_kquant as kq
+kq.codecs()          # -> ['q2_k', 'q3_k', ..., 'q8_0']
+kq.metallib_loads()  # -> True  (the bundled metallib opened on the Metal device)
+```
+
+> The extension links `libmlx` and its kernels `#include` MLX's steel-GEMM headers, so it is bound to
+> an exact MLX ABI **and** header API. The pin is intentionally `==`, never `>=`; moving to a newer
+> `mlx` may require updating the bundled headers and recompiling. See [Version pinning](#version-pinning).
+
+## Quickstart
+
+Quantize a checkpoint and run it, load it through mlx-lm, fine-tune it with LoRA, or build directly on
+the `kq.*` ops.
+
+### Create and run a checkpoint
+
+The CLI (the `[tools]` extra adds `mlx-lm`) quantizes an HF / `mlx-lm` model into a K-quant **MLX
+safetensors** checkpoint and runs it:
+
+```sh
+pip install "mlx-kquant[tools]"
+mlx-kquant quantize --model Qwen/Qwen3-0.6B --preset q4_k_m --mlx-path qwen3-q4
+mlx-kquant run --model qwen3-q4 --prompt "Explain entropy in one sentence."
+mlx-kquant chat --model qwen3-q4 --temp 0.7      # interactive REPL (mlx-lm chat)
+```
+
+`run` takes the usual sampling knobs (`--temp`, `--top-p`, `--top-k`, `--min-p`, `--seed`,
+`--repetition-penalty`, `--presence-penalty`, `--frequency-penalty`) and chat-template controls
+(`--system-prompt`, `--no-chat-template`, `--chat-template-config` for template kwargs such as
+`'{"enable_thinking": false}'`). The `chat` REPL has a line-editable prompt with persistent
+history (`--no-history` or in-chat `/history off|on|clear` to control it) and in-chat sampling
+control (`/temp`, `/top-p`, `/top-k`, `/min-p`, `/max-tokens`, and the three penalties;
+`/sampling` shows current values); `/load <file>` prefills the next prompt from a text file for
+editing; `/clear` resets the conversation and wipes the screen; Tab completes `/commands` and
+paths; Ctrl-C cancels the in-flight reply (at an idle
+prompt it exits, as does Ctrl-D). `--max-kv-size` bounds the KV cache for long sessions (a rotating
+window, set at start).
+
+The result is a standard MLX checkpoint (`config.json` + sharded safetensors, weights as K-quant wire
+bytes). Load it in code with the bundled loader:
+
+```python
+import mlx.core as mx
+from mlx_kquant.loader import load
+
+model, config = load("qwen3-q4")          # KQuant* layers swapped in, on a stock mlx-lm model
+mx.eval(model(mx.array([[1, 2, 3]])))
+```
+
+`mlx-kquant lora` (train an adapter) and `mlx-kquant fuse` (merge it back) round out the toolchain -
+see [LoRA fine-tuning](#lora-fine-tuning). Run `mlx-kquant --help` for every subcommand.
+
+### Using with mlx-lm
+
+In-process, a kquant checkpoint also loads through **stock mlx-lm**: one idempotent call installs the
+load shim, and from then on `mlx_lm.load` / `mlx_lm.generate` (and anything built on
+`mlx_lm.utils.load_model`, e.g. an eval harness or your own serving loop) open a kquant checkpoint
+transparently:
+
+```python
+from mlx_kquant.mlx_lm_patch import patch_mlx_lm_load
+patch_mlx_lm_load()                 # process-wide, idempotent; call once before mlx_lm.load
+
+from mlx_lm import load, generate
+model, tokenizer = load("qwen3-q4")
+print(generate(model, tokenizer, "Explain entropy.", max_tokens=64))
+```
+
+This is the load-only shim for inference / eval / serving; `patch_mlx_lm_lora()`
+([below](#lora-fine-tuning)) adds the train/merge shims on top. The bundled `mlx_kquant.loader.load`
+(above) is the standalone path when you don't need the rest of mlx-lm.
+
+### LoRA fine-tuning
+
+A kquant checkpoint is a frozen base you can adapt with LoRA. Attach an adapter for inference, train
+one (the matmul/gather ops define a gradient-through-the-base `vjp`, so the adapter is differentiable
+while the quantized weights stay frozen), and merge it back with `mlx-kquant fuse` (re-encode to
+kquant, or `--dequantize` to float). One call wires it into stock mlx-lm:
+
+```python
+from mlx_kquant.mlx_lm_patch import patch_mlx_lm_lora
+patch_mlx_lm_lora()   # before building LoRA layers / loading adapters; idempotent
+```
+
+See **[docs/lora.md](https://github.com/asher/mlx-kquant/blob/main/docs/lora.md)** for attach / train / merge workflows. (DoRA on a kquant base is
+not supported - use LoRA.)
+
+### Using K-quant ops directly
+
+Under the toolchain, the four `kq.*` ops operate on raw K-quant wire bytes. K-quant scales live
+*inside* the packed bytes, so the `scales` argument is a vestigial placeholder (the API keeps it for
+shape symmetry with MLX's affine quant); `kq.quantize` returns one for you.
+
+```python
+import mlx.core as mx
+import mlx_kquant as kq
+
+N, K = 256, 512                       # q4_k: K must be a multiple of 256
+w = mx.random.normal((N, K))
+
+# encode float -> K-quant wire bytes (CPU or Metal); optional imatrix steers the encoder
+wq, scales = kq.quantize(w, "q4_k")           # wq: uint8 [N, bytes_per_row]
+
+# dequantize back to float
+deq = kq.dequantize(wq, scales, "q4_k")       # float16 [N, K]
+
+# quantized matmul: x @ dequant(w).T   (transpose=True => w is [N, K])
+x = mx.random.normal((8, K))
+y = kq.quantized_matmul(x, wq, scales, "q4_k", transpose=True)   # [8, N]
+```
+
+Mixture-of-experts (gathered) matmul:
+
+```python
+E, N, K = 128, 704, 2816
+we = mx.random.normal((E, N, K))
+weq, sc = kq.quantize(we, "q4_k")             # per-expert wire bytes
+x = mx.random.normal((1, 8, K))               # (tokens, top_k, K)
+idx = mx.array([[0, 5, 9, 17, 33, 41, 88, 120]], dtype=mx.uint32)
+out = kq.gather_qmm(x, weq, sc, "q4_k", rhs_indices=idx, transpose=True)
+```
+
+Ready-made modules that store the wire bytes and dispatch the matching `kq.*` op ship in
+`mlx_kquant.nn`:
+
+```python
+from mlx_kquant.nn import KQuantLinear, KQuantEmbedding, KQuantSwitchLinear
+
+x = mx.random.normal((8, 512))   # a (tokens, in_dims) activation batch
+lin = KQuantLinear(in_dims=512, out_dims=256, bias=False, codec="q4_k")
+lin.weight = wq                  # the uint8 wire bytes from kq.quantize, above
+lin.scales = scales              # [1] placeholder (scales live in the bytes)
+y = lin(x)                       # kq.quantized_matmul under the hood
+```
+
+`KQuantEmbedding` (with a tied-`as_linear`), the `gather_qmm`-backed `KQuantSwitchLinear` for MoE
+experts, and `KQuantMultiLinear` (absorbed-MLA) are exported alongside it. To swap the quantizable
+leaves of a whole constructed mlx-lm model in one call, use
+`mlx_kquant.nn.install_kquant_modules(model, {"<path>.weight": "q4_k", ...})`.
+
+The `[tools]` layer is itself a worked reference for wiring `kq.*` into the MLX ecosystem: the loader,
+encoder, layer modules, and the mlx-lm monkeypatch are all small and self-contained. See
+**[docs/integration.md](https://github.com/asher/mlx-kquant/blob/main/docs/integration.md)** if you're building on the ops.
+
+## Performance
+
+The Metal kernels use a single-pass NAX matmul and matrix-contiguity handling for fused MoE expert
+weights. Measured on an M5 Max (128 GB):
+
+| Model | Codec | Decode (tok/s) | Prefill pp512 (tok/s) |
+|-------|-------|---------------:|----------------------:|
+| gemma-4-26B-A4B-it (MoE) | q4_k_xl | ~111 | ~2330 |
+| Qwen3.5-9B (dense)       | q5_k_xl | ~83  | ~2396 |
+
+Transposed matmuls with a small row count (the speculative-decode verify regime) automatically route
+through a weight-read-amortizing `verify_qmv` kernel; `KQ_DISABLE_VERIFY_QMV=1` forces the plain
+per-row `qmv` path (see [Environment variables](#environment-variables)).
+
+## How it works
+
+- **Own ops.** Four `Primitive` subclasses (`KQuantDequantize`, `KQuantMatmul`, `KQuantGatherQMM`,
+  `KQuantQuantize`) and their op functions live entirely in the extension.
+- **Precompiled metallib on stock headers.** The `kq_*` kernels are compiled against the stock
+  wheel's steel-GEMM headers into `mlx_kquant.metallib` at build time; host dispatch resolves them
+  through MLX's exported `Device::get_kernel`. No JIT, no steel host structs.
+- **Codec registry** derives `group_size`/`bits` from the codec name, so callers pass only
+  `kquant_type`.
+- **CPU and GPU execution.** Every op - the decode ops (`dequantize` / `quantized_matmul` /
+  `gather_qmm`) and `quantize` (encode) - runs on either stream, covering all 10 codecs, so the full
+  quantize/decode pipeline (and the op tests) runs in CI without a GPU. The per-block `dequantize` is
+  a scalar, bit-exact (per-codec, vs the `gguf.quants` reference quantizer) decoder. The CPU **matmul**
+  and **gather** are tuned for Apple Silicon: a shared worker pool over output rows, NEON int8
+  dot-product GEMV for the small-M (decode) shape, and an Accelerate (AMX/SME) GEMM for the large-M
+  (prefill) shape. The NEON path quantizes activations to int8 (lossy, as ggml does), so its matmul
+  matches at tolerance, not bit-exactly; `KQ_CPU_NEON=0` forces the scalar path for exact parity.
+
+## Environment variables
+
+All optional; the defaults are right for normal use.
+
+- `KQ_CPU_THREADS` - worker-pool size for the CPU ops (default: hardware concurrency; `1` runs them
+  inline). `KQ_CPU_SPIN_US` sets a spin-before-park window for the pool (default `0` = park).
+- `KQ_CPU_NEON=0` - disable the arm64 NEON int8 GEMV kernels and run the scalar decode-then-dot
+  matmul, which is bit-exact (the NEON path is tolerance-level; see [How it works](#how-it-works)).
+- `KQ_DISABLE_VERIFY_QMV=1` - on Metal, force the plain per-row `qmv` path instead of the
+  weight-read-amortizing `verify_qmv` kernel. An A/B debugging lever, not a tuning knob.
+
+## Quant recipes
+
+A **preset** is a named mixed-precision recipe. It classifies each tensor by *role* (attention
+q/k/v/o, embeddings, `lm_head`, MoE routed vs shared experts, the FFN down-projection) and maps each
+role to a codec - spending bits where they move the output most and staying frugal on the bulk
+feed-forward weights, to beat a uniform quant at the same byte budget.
+
+```sh
+mlx-kquant quantize --model <src> --preset q4_k_m   --mlx-path out   # a mixed recipe
+mlx-kquant quantize --model <src> --kquant-type q6_k --mlx-path out  # one codec, every tensor
+```
+
+Naming follows the ggml convention: the family (`q4_k`, `q5_k`, ...) sets the baseline codec and the
+suffix sets how much extra precision the recipe spends:
+
+- `_s` / `_m` / `_xl` - small / medium / extra: increasing bumps on the sensitive tensors (the value
+  and output projections, the down-projection on a subset of layers, the linear-attention
+  projections).
+- `_moe` - expert-aware: routed experts at the baseline, shared experts a step above.
+- bare `q6_k` / `q8` - uniform (every tensor at one codec), equivalent to passing `--kquant-type`.
+
+`mlx-kquant quantize --list-presets` prints the full, authoritative mapping for every preset; it is
+generated from the recipe tables, so it never drifts from what the encoder actually does. The recipes
+are informed by our analysis of the mixed-precision quants that [Unsloth][unsloth] and
+[bartowski][bartowski] publish on Hugging Face, together with llama.cpp's own per-layer
+"use more bits" schedule.
+
+[unsloth]: https://huggingface.co/unsloth
+[bartowski]: https://huggingface.co/bartowski
+
+## Codec reference
+
+| Codec | Block | Bits | Bytes/block | Notes |
+|-------|------:|-----:|------------:|-------|
+| q2_k  | 256 | 2 |  84 | K-quant superblock |
+| q3_k  | 256 | 3 | 110 | K-quant superblock |
+| q4_k  | 256 | 4 | 144 | K-quant superblock |
+| q5_k  | 256 | 5 | 176 | K-quant superblock |
+| q6_k  | 256 | 6 | 210 | K-quant superblock |
+| q4_0  |  32 | 4 |  18 | block scale |
+| q4_1  |  32 | 4 |  20 | block scale + min |
+| q5_0  |  32 | 5 |  22 | block scale |
+| q5_1  |  32 | 5 |  24 | block scale + min |
+| q8_0  |  32 | 8 |  34 | block scale |
+
+## Version pinning
+
+Pinned to `mlx==0.31.2`. 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.
+
+## Tests
+
+```sh
+python -m pytest tests/
+```
+
+## Requirements
+
+- **macOS 26.2 (Tahoe) or later on Apple Silicon** (M-series) with a working Metal toolchain
+  (`xcrun metal`) for the GPU build-from-source install.
+- **Linux** (x86_64 or aarch64) is supported CPU-only: build against `mlx[cpu]==0.31.2`, no Metal
+  toolchain required. See [Install](#install) and [Limitations](#limitations).
+- **Python >= 3.10** (the pinned `mlx==0.31.2` ships no cp39 wheel).
+- **`mlx==0.31.2`** 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
+
+- **GPU path is Apple-Silicon Metal only.** No ROCm or CUDA support. Every op also has a CPU path
+  (`stream=mx.cpu`) covering all 10 codecs, so the extension still builds and runs without Metal (see
+  [How it works](#how-it-works) and [Install](#install)).
+- **Linux model forwards need `MLX_DISABLE_COMPILE=1`.** Stock MLX's CPU compile JIT generates C++
+  that redeclares GCC's built-in `_Float32`/`_Float64`/`_Float128` types, which `g++` rejects, so any
+  model forward through MLX's compile path fails on Linux+GCC. Disabling the JIT runs those graphs
+  eagerly with identical numerics. This is an upstream MLX-on-Linux limitation independent of
+  mlx-kquant - the `kq.*` ops have their own `eval_cpu` and never touch the JIT.
+- **LoRA, not DoRA.** LoRA adapters train, attach, and fuse on a kquant base (see
+  [docs/lora.md](https://github.com/asher/mlx-kquant/blob/main/docs/lora.md)); DoRA is not yet supported. `fuse` re-encodes to kquant or, with
+  `--dequantize`, to float; both modes run on CPU or Metal.
+
+## License
+
+MIT - see [LICENSE](https://github.com/asher/mlx-kquant/blob/main/LICENSE).
+
+### Acknowledgements
+
+mlx-kquant builds on three MIT-licensed projects; their license texts ship in the wheel under
+[`mlx_kquant/licenses/`](https://github.com/asher/mlx-kquant/tree/main/mlx_kquant/licenses):
+
+- **[llama.cpp / ggml](https://github.com/ggml-org/llama.cpp)** - the K-quant and block codec formats
+  and the quantization / dequantization algorithms that encode and decode them are derived from
+  ggml's reference implementation.
+- **[gguf-tools](https://github.com/antirez/gguf-tools)** - used to implement a zero-copy GGUF loader
+  for downstream projects, statically linked into built wheels.
+- **[MLX](https://github.com/ml-explore/mlx)** - the extension links `libmlx`, the kernels compile
+  against MLX's bundled headers, and parts of the Metal kernels are adapted from MLX's quantized and
+  steel-GEMM kernels.