diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 24f691e5..6dbf798d 100644 --- a/.github/workflows/ci.yml +++ b/.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). diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c814047..0259f217 100644 --- a/CHANGELOG.md +++ b/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] diff --git a/CMakeLists.txt b/CMakeLists.txt index 3922fe00..e32d8444 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b89cb23a..3a8411e4 100644 --- a/CONTRIBUTING.md +++ b/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 ``` diff --git a/README.md b/README.md index abdccbbf..48e9f284 100644 --- a/README.md +++ b/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 diff --git a/benchmarks/bench_dit_first_dispatch.py b/benchmarks/bench_dit_first_dispatch.py new file mode 100644 index 00000000..b8c1cf19 --- /dev/null +++ b/benchmarks/bench_dit_first_dispatch.py @@ -0,0 +1,225 @@ +"""First-dispatch cost of the DiT kernel list, in fresh processes. + +The parent starts --repeats child processes. Each child times the metallib +load on its own, then walks the DiT cells (shape matrix x rows x batches x +codecs at one dtype). Per cell it builds weights and input, primes the +allocator with a buffer of the output's size, and times the first and the +second call of the kq arm at the current policy. The first call of a new +kernel builds its pipeline state; a cell whose kernel an earlier cell built +is counted once. The report sums first minus second over the distinct +kernels, as the median over children. + +--stock adds stock MLX ops a DiT block runs (dense matmul, attention at +head dim 128, RMS norm, RoPE, SiLU gate) for comparison. + +The OS keeps compiled shaders in a cache that outlives the process, so a +child measures a first image with that cache warm; clearing it is outside +this script. Markdown and JSON out. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import sys +import time + +import dit_common as dc + +CHILD = "--child" + + +def child(args): + t0 = time.perf_counter() + import mlx.core as mx + + t_import = time.perf_counter() - t0 + import mlx_kquant as kq + + t1 = time.perf_counter() + kq.metallib_loads() + t_metallib = time.perf_counter() - t1 + dt = dc.mx_dtype(args.dtype) + store = dc.Store(0) + seen = {} + rows = [] + + def timed(label, fn, out_shape, log_kernel): + mx.eval(mx.zeros(out_shape, dtype=dt)) + if log_kernel: + with dc.KernelLog() as log: + ta = time.perf_counter() + mx.eval(fn()) + first = time.perf_counter() - ta + kname = log.kernels[-1] if log.kernels else label + else: + ta = time.perf_counter() + mx.eval(fn()) + first = time.perf_counter() - ta + kname = label + if kname in seen: + # An earlier cell built this kernel; one call is enough. + return + seen[kname] = label + ta = time.perf_counter() + mx.eval(fn()) + second = time.perf_counter() - ta + rows.append( + { + "cell": label, + "kernel": kname, + "first_ms": first * 1e3, + "second_ms": second * 1e3, + } + ) + + for fam in args.families: + for op in dc.FAMILIES[fam]["ops"]: + for L in dc.op_rows(fam, op, args.rows): + for B in args.batches: + g = dc.cell(fam, op, L, B, args.shrink) + for codec in args.codecs: + fn = dc.make_fn("kq", codec, g, dt, store) + mx.eval(store.act(g, dt), *store.kq(g, codec)) + label = f"{fam}.{op} {L}x{B} {codec}" + timed(label, fn, (B, g["L"], g["N"]), True) + store.x.clear() + if args.stock: + g = dc.cell("zimage", "w13", 4608, 1, args.shrink) + fn = dc.make_fn("dense", None, g, dt, store) + mx.eval(store.act(g, dt), store.dense_w(g, dt)) + timed("matmul", fn, (1, g["L"], g["N"]), False) + H, D, L = 4, 128, max(64, 4608 // args.shrink) + q = (mx.random.normal((1, H, L, D)) * 0.1).astype(dt) + w = mx.ones((D,), dtype=dt) + mx.eval(q) + timed( + "sdpa", + lambda: mx.fast.scaled_dot_product_attention(q, q, q, scale=D**-0.5), + (1, H, L, D), + False, + ) + timed("rms_norm", lambda: mx.fast.rms_norm(q, w, 1e-6), q.shape, False) + timed( + "rope", + lambda: mx.fast.rope( + q, D, traditional=False, base=10000.0, scale=1.0, offset=0 + ), + q.shape, + False, + ) + timed("silu_gate", lambda: mx.sigmoid(q) * q * q, q.shape, False) + print( + json.dumps( + {"import_s": t_import, "metallib_s": t_metallib, "rows": rows}, + ) + ) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--families", default=",".join(dc.FAMILIES)) + ap.add_argument("--rows", default="") + ap.add_argument("--batches", default="1,2") + ap.add_argument("--codecs", default=",".join(dc.MAIN_CODECS)) + ap.add_argument("--dtype", default="bfloat16", choices=list(dc.DTYPES)) + ap.add_argument("--repeats", type=int, default=3) + ap.add_argument("--stock", action="store_true") + ap.add_argument("--shrink", type=int, default=1) + ap.add_argument(CHILD, action="store_true", help=argparse.SUPPRESS) + ap.add_argument("--json-out") + ap.add_argument("--md-out") + args = ap.parse_args() + args.families = [v for v in args.families.split(",") if v] + args.rows = [int(v) for v in args.rows.split(",") if v] + args.batches = [int(v) for v in args.batches.split(",") if v] + args.codecs = [v for v in args.codecs.split(",") if v] + if args.child: + child(args) + return + + argv = [a for a in sys.argv[1:] if a != CHILD] + env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + runs = [] + for _ in range(args.repeats): + r = subprocess.run( + [sys.executable, os.path.abspath(__file__), CHILD, *argv], + env=env, + capture_output=True, + text=True, + ) + if r.returncode != 0: + sys.exit(r.stderr[-4000:]) + runs.append(json.loads(r.stdout.strip().splitlines()[-1])) + desc = dc.machine_descriptor() + + per_kernel = {} + for run in runs: + for row in run["rows"]: + per_kernel.setdefault(row["kernel"], []).append( + (row["cell"], row["first_ms"] - row["second_ms"], row["second_ms"]) + ) + kernels = [] + for k, v in per_kernel.items(): + kernels.append( + { + "kernel": k, + "cell": v[0][0], + "tile": dc.tile_of(k), + "extra_ms": statistics.median(x[1] for x in v), + "second_ms": statistics.median(x[2] for x in v), + } + ) + sums = [sum(r["first_ms"] - r["second_ms"] for r in run["rows"]) for run in runs] + kq_sums = [ + sum( + r["first_ms"] - r["second_ms"] + for r in run["rows"] + if r["kernel"].startswith("kquant_") + ) + for run in runs + ] + result = { + "machine": desc, + "args": vars(args), + "metallib_ms": statistics.median(r["metallib_s"] * 1e3 for r in runs), + "import_ms": statistics.median(r["import_s"] * 1e3 for r in runs), + "sum_extra_ms": statistics.median(sums), + "kq_sum_extra_ms": statistics.median(kq_sums), + "kernels": kernels, + "runs": runs, + } + lines = ["# DiT first dispatch", "", dc.descriptor_md(desc), ""] + lines.append( + f"{args.repeats} fresh processes, {args.dtype}. mlx import " + f"{result['import_ms']:.0f} ms, metallib load {result['metallib_ms']:.1f} ms, " + f"first-dispatch extra summed over {len(kernels)} kernels " + f"{result['sum_extra_ms']:.1f} ms " + f"(kq kernels {result['kq_sum_extra_ms']:.1f} ms)." + ) + lines.append("") + lines.append("| kernel | tile | first cell | extra ms | second ms |") + lines.append("|---|---|---|---|---|") + for k in kernels: + lines.append( + f"| {k['kernel']} | {k['tile'] or '-'} | {k['cell']} " + f"| {k['extra_ms']:.2f} | {k['second_ms']:.3f} |" + ) + md = "\n".join(lines) + if args.json_out: + with open(args.json_out, "w") as f: + json.dump(result, f, indent=1) + if args.md_out: + with open(args.md_out, "w") as f: + f.write(md + "\n") + else: + print(md) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/bench_dit_shapes.py b/benchmarks/bench_dit_shapes.py new file mode 100644 index 00000000..14261a13 --- /dev/null +++ b/benchmarks/bench_dit_shapes.py @@ -0,0 +1,1072 @@ +"""Quantized matmul arms at DiT linear shapes, paired in one process. + +Cells are (op, rows per item L, batch B, codec). A batch of two runs as +[2, L, K], which the op flattens to 2L rows, so the tile policy can pick a +different tile at B 2 than at B 1. The modulation linears run at one row +per item, where the tile arms do not apply. Arms, switched by live levers +(KQ_TUNE_LIVE=1 is set on import of dit_common): + + dense stock matmul on a dense weight in the activation dtype + kq kq.quantized_matmul at the current policy + kq_bm64 kq with KQ_NAX_TILE=64,64,2,2 (KQ_QMM_TILE without NAX) + kq_bm128 kq with KQ_NAX_TILE=128,64,2,2 + deq kq.dequantize, then stock matmul, every call + affine mx.quantized_matmul, MLX affine, group 64, 4 bits (8 for q8_0) + kq_aa the kq arm a second time, for the A versus A noise floor + +dense and affine do not depend on the codec, so they run once per shape +visit and pair with every codec in it. + +Order: one discarded warmup pass, which also records the kernel each kq +arm ran, then rounds of Thue-Morse slots. Each slot visits every shape in +a fresh interleaved order. A visit opens with a discarded sample of the +first --visit-warmup arms, since the first call after a shape switch runs +slow; --visit-warmup-calls sets its call count. Inside a visit +the arm list reverses on odd slots, so the two calls of every pair swap +order with the Thue-Morse bit and a round of 8 slots splits each pair 4/4. +A paired ratio compares calls made seconds apart. Round 0 gets its own +column and later rounds the settled column. Round 0 is not a cold figure: +the GPU slows within seconds of load, and bench_dit_thermal.py gives the +cold rate per machine. + +A sample is a chain of --inner calls in one evaluation, divided by the +count. Each call is fed a tie of the previous output (the activation +again, selected on the output's first element), so the GPU runs the calls +one after another, as in a model forward; calls evaluated together without +the tie overlap, which slows some kernels and speeds others. Every visit +also times the tie chain alone ("glue"), and each sample of the visit has +it subtracted. Shapes whose fastest arm runs under --min-sample-ms per call +take more calls, up to --max-inner, and every arm of a shape uses the same +count. Reports and raw logs carry "sample_mode": "chain"; a log without it +is from the earlier independent-call bench, and its figures do not compare +with chained ones. + +The later call of a visit runs slower than the earlier one (up to 20 +percent at compute-bound shapes), so a cell's paired ratios form two +modes, one per call order, split 4/4 by the Thue-Morse bit. Per cell, the +ratios of each order are scaled to meet at the geometric mean of the two +order medians, and the factor between them is reported as "order"; the +raw pairs stay in the JSON as "pairs_raw". + +A sample is a time less its glue time, so at a tiny shape it lands on +either side of 0. A visit whose reference sample is at or below 0 has no +ratio and is left out, and each ratio carries the count as "dropped". A +ratio whose reference sits at the timer floor (a reference sample at or +below 0, or a reference median under twice the p10 to p90 width of the +glue times) is marked with * in the Markdown and "ref_floor" in the JSON: +its value says little. + +Per cell, ratios carry a 95 percent CI of the paired median, from order +statistics; its half-width is the cell's noise floor. Per op, the pooled +median of all settled pairs is the aggregate, reported beside the spread +of the per-cell medians. Markdown and JSON out; the JSON keeps every +sample and pair. + +Every sample also goes to a JSON Lines raw log (--raw-out, by default +beside --json-out) as each slot ends, so a run that stops or fails keeps +every finished slot. --reaggregate rebuilds the JSON and Markdown from a +raw log without running anything, and marks a partial run. --resume +continues a stopped run from its raw log: it keeps the finished slots and +the warmup's call counts, and runs the slots left, in the same order. +--time-budget-min splits a long run into processes that each stop after +the slot that passes the budget. + +Memory: MLX keeps freed buffers in its cache for reuse, and every shape +leaves its own sizes there, so a sweep of many shapes would otherwise hold +the outputs of all of them. --cache-limit-gib bounds the cache, and the +visit warmup runs a full sample by default so the timed calls reuse +cached buffers. Each slot record carries the active, cache +and process footprint bytes after every visit; --mem-trace prints them +as they are taken. + +A versus A: --arms kq,kq_aa. Plan only: --plan-only. + +Lever arms: --env-arm NAME:VAR=VALUE[,VAR=VALUE] defines a kq arm that +runs with those environment levers set (read live), adds it to --arms and +pairs it with kq as NAME/kq, for example --env-arm sw1:KQ_NAX_SWIZZLE=1. +A value may hold commas: --env-arm t:KQ_NAX_TILE=128,128,2,4,KQ_NAX_Q4K_BK128=0. +""" + +from __future__ import annotations + +import argparse +import ctypes +import ctypes.util +import gc +import json +import math +import os +import random +import re +import sys +import time + +import bench_stats as bs +import dit_common as dc + + +def parse_list(s, conv=str): + return [conv(v) for v in s.split(",") if v != ""] if s else [] + + +def build_groups(args, heavy_codecs): + groups = [] + for fam in args.families: + spec = dc.FAMILIES[fam] + ops = [o for o in spec["ops"] if not args.ops or o in args.ops] + for op in ops: + for L in dc.op_rows(fam, op, args.rows): + for B in args.batches: + g = dc.cell(fam, op, L, B, args.shrink) + g["heavy"] = (fam, op, L, B) in dc.HEAVY + g["codecs"] = list(args.codecs) + if g["heavy"]: + g["codecs"] += [c for c in heavy_codecs if c not in args.codecs] + groups.append(g) + return groups + + +CODEC_ARMS = ("kq", "kq_bm64", "kq_bm128", "deq", "kq_aa") +SHARED_ARMS = ("dense", "affine") + + +def lever_arms(arms): + """Arms defined by --env-arm, in --arms order.""" + return [a for a in arms if a not in CODEC_ARMS and a not in SHARED_ARMS] + + +def parse_env_arms(specs): + """{name: {VAR: VALUE}} from --env-arm NAME:VAR=VALUE[,VAR=VALUE].""" + out = {} + for spec in specs: + name, sep, body = spec.partition(":") + # A comma starts the next pair only before VAR=, so a value may hold + # commas, as KQ_NAX_TILE=bm,bn,wm,wn does. + pairs = re.split(r",(?=[A-Za-z_][A-Za-z0-9_]*=)", body) + env = dict(kv.split("=", 1) for kv in pairs if kv) + if not sep or not name or not env or name in CODEC_ARMS + SHARED_ARMS: + raise SystemExit(f"--env-arm {spec!r}: want NAME:VAR=VALUE[,VAR=VALUE]") + out[name] = env + return out + + +def set_arm_env(arm, env_arms): + """Sets the levers of a lever arm and clears those of the others.""" + for env in env_arms.values(): + for k in env: + os.environ.pop(k, None) + os.environ.update(env_arms.get(arm, {})) + + +def arm_instances(group, arms): + """(arm, codec) pairs one visit runs; codec is None for shared arms.""" + out = [] + if "dense" in arms: + out.append(("dense", None)) + if "affine" in arms: + for bits in sorted({dc.affine_bits(c) for c in group["codecs"]}): + out.append((f"affine{bits}", None)) + for c in group["codecs"]: + for a in CODEC_ARMS + tuple(lever_arms(arms)): + # Tile and lever arms do not apply at one row per item, where the + # mat-vec routes run. + forced = a in dc.ARM_TILE_FORCED or a not in CODEC_ARMS + if a in arms and not (group["vector"] and forced): + out.append((a, c)) + return out + + +def plan_summary(groups, arms, rounds, slots, est_tflops, inner): + calls = 0 + flops = 0.0 + for g in groups: + n = len(arm_instances(g, arms)) * inner + calls += n + flops += n * 2.0 * g["M"] * g["K"] * g["N"] + per_round = {"calls": calls * slots, "tflop": flops * slots / 1e12} + # Visit warmup calls and the extra calls of fast shapes are not counted. + per_round["est_min"] = per_round["tflop"] / est_tflops / 60 + total_min = per_round["est_min"] * (rounds + 1.0 / slots) + return per_round, total_min + + +def cell_key(arm, codec): + return f"{arm}/{codec}" if codec else arm + + +def paired(samples, a, b): + """Per-visit ratios a/b over the visits that ran both arms, each with + True when a ran first in the visit, False when b did, and None when the + times do not tell. Returns the ratios and the number of visits left + out. A visit whose b sample is at or below 0 has no ratio: a sample is + a time less its glue time, so at tiny shapes it lands on either side + of 0, and a ratio to such a reference measures nothing.""" + sb = {v: (t, ms) for (v, _r, _s, t, ms) in samples.get(b, [])} + out = [] + dropped = 0 + for v, r, _s, t, ms in samples.get(a, []): + if v in sb: + tb, mb = sb[v] + if mb <= 0: + dropped += 1 + continue + out.append((r, ms / mb, None if t == tb else t < tb)) + return out, dropped + + +def at_timer_floor(samples, b): + """True when the samples of reference arm b sit at the timer floor: one + of them is at or below 0, or their median is under twice the p10 to + p90 width of the glue times that were subtracted from them.""" + ref = [ms for (_v, _r, _s, _t, ms) in samples.get(b, [])] + if not ref: + return False + if min(ref) <= 0: + return True + glue = bs.summarize([ms for (_v, _r, _s, _t, ms) in samples.get("glue", [])]) + if not glue: + return False + return bs.summarize(ref)["median"] < 2 * (glue["p90"] - glue["p10"]) + + +def order_balance(pairs): + """The later call of a visit runs slower, so the ratios of the visits + where a ran first and of those where b ran first form two modes, and + their pooled median falls anywhere between them. Scale each order's + ratios by the square root of the order factor (median of b first over + median of a first) so both meet at the geometric mean of the two + medians. Returns the scaled (round, ratio) pairs and the factor, which + is None unless both orders ran with positive medians (a sample under its + glue time gives a ratio at or below 0).""" + af = [v for _r, v, f in pairs if f is True] + bf = [v for _r, v, f in pairs if f is False] + ma = bs.summarize(af)["median"] if af else 0 + mb = bs.summarize(bf)["median"] if bf else 0 + if ma <= 0 or mb <= 0: + return [(r, v) for r, v, _f in pairs], None + order = mb / ma + k = math.sqrt(order) + scaled = [] + for r, v, f in pairs: + scaled.append((r, v * k if f is True else v / k if f is False else v)) + return scaled, order + + +def split_stats(rows, rounds, round_at, value_at): + """Round 0, settled (later rounds) and all-round summaries.""" + round0 = [p[value_at] for p in rows if p[round_at] == 0] + settled = [p[value_at] for p in rows if p[round_at] > 0] + allv = [p[value_at] for p in rows] + return { + "round0": bs.summarize(round0), + "settled": bs.summarize(settled) if rounds > 1 else None, + "all": bs.summarize(allv), + } + + +_LIBC = None + + +def footprint_bytes(): + """The process physical footprint on macOS, which memory limits + count, cache included; None elsewhere.""" + global _LIBC + if sys.platform != "darwin": + return None + if _LIBC is None: + _LIBC = ctypes.CDLL(ctypes.util.find_library("c")) + buf = ctypes.create_string_buffer(512) + # RUSAGE_INFO_V4; ri_phys_footprint is at byte 72. + if _LIBC.proc_pid_rusage(os.getpid(), 4, buf) != 0: + return None + return int.from_bytes(buf.raw[72:80], "little") + + +def mem_row(mx, gi, trace=False): + row = [gi, mx.get_active_memory(), mx.get_cache_memory(), footprint_bytes()] + if trace: + gib = [None if v is None else round(v / 2**30, 3) for v in row[1:]] + print("mem", gi, *gib, flush=True) + return row + + +def complete_lines(path): + """Drops a last line cut short, so records can be appended.""" + with open(path) as f: + text = f.read() + if text and not text.endswith("\n"): + text = text[: text.rfind("\n") + 1] + with open(path, "w") as f: + f.write(text) + + +class RawLog: + """JSON Lines log of a run: a meta record, then per dtype a warmup + record and one record per finished slot, then an end record. Each + record is flushed and synced as it is written.""" + + def __init__(self, path, mode="w"): + self.f = open(path, mode) if path else None + + def write(self, rec): + if self.f is None: + return + self.f.write(json.dumps(rec) + "\n") + self.f.flush() + os.fsync(self.f.fileno()) + + def close(self): + if self.f is not None: + self.f.close() + self.f = None + + +def load_raw(path): + """Meta record, per-dtype measurements and end record (empty when the + run stopped early) from a raw log. A last line cut short is skipped.""" + with open(path) as f: + lines = [line for line in f.read().split("\n") if line.strip()] + recs = [] + for i, line in enumerate(lines): + try: + recs.append(json.loads(line)) + except json.JSONDecodeError: + if i < len(lines) - 1: + raise + if not recs or recs[0].get("kind") != "meta": + raise ValueError(f"{path}: no meta record") + meta, measured, end = recs[0], {}, {} + ngroups = len(meta["groups"]) + for rec in recs[1:]: + if rec["kind"] == "warmup": + measured[rec["dtype"]] = { + "inner": rec["inner"], + "kernels": {(gi, k): v for gi, k, v in rec["kernels"]}, + "errors": {(gi, k): v for gi, k, v in rec["errors"]}, + "samples": [dict() for _ in range(ngroups)], + "slots": 0, + } + elif rec["kind"] == "slot": + m = measured[rec["dtype"]] + for gi, key, visit, r, s, t, ms in rec["samples"]: + m["samples"][gi].setdefault(key, []).append((visit, r, s, t, ms)) + m["slots"] += 1 + elif rec["kind"] == "end": + end = rec + return meta, measured, end + + +def measure_pass(dtype, groups, args, store, desc, t_start, raw, prior=None): + """Warmup and timed rounds for one dtype. Samples go to `raw` as each + slot ends; no statistics here. `prior` is the part of the pass a raw + log holds: its slots are skipped, and its warmup's call counts and + kernels are kept.""" + import mlx.core as mx + + dt = dc.mx_dtype(dtype) + nax = desc["nax"] + env_arms = parse_env_arms(args.env_arm) + kernels = {} + errors = {} + calib = {} + samples = [dict() for _ in groups] + mem = [] + + # Warmup pass, discarded; records the kernel each kq arm dispatches. + for gi, g in enumerate(groups): + for arm, codec in arm_instances(g, args.arms): + key = cell_key(arm, codec) + x0, f = dc.make_call(arm, codec, g, dt, store) + dc.set_tile(dc.ARM_TILE.get(arm), nax) + set_arm_env(arm, env_arms) + try: + with dc.KernelLog() as log: + mx.eval(f(x0)) + kernels[(gi, key)] = log.kernels[-1] if log.kernels else None + calib.setdefault(gi, []).append(dc.chain_ms(f, x0, 1)) + except ValueError as e: + errors[(gi, key)] = str(e) + dc.set_tile(None, nax) + set_arm_env(None, env_arms) + mem.append(mem_row(mx, gi, args.mem_trace)) + inner = [] + for gi in range(len(groups)): + fastest = min(calib.get(gi, [args.min_sample_ms])) + n = math.ceil(args.min_sample_ms / max(fastest, 1e-3)) + inner.append(min(args.max_inner, max(args.inner, n))) + done = set() + visit = 0 + slots = 0 + if prior is None: + raw.write( + { + "kind": "warmup", + "dtype": dtype, + "inner": inner, + "kernels": [[gi, k, v] for (gi, k), v in kernels.items()], + "errors": [[gi, k, v] for (gi, k), v in errors.items()], + "mem": mem, + } + ) + else: + inner = prior["inner"] + kernels = prior["kernels"] + errors = prior["errors"] + samples = prior["samples"] + slots = prior["slots"] + for sm in samples: + for rows in sm.values(): + for v, r, s, _, _ in rows: + done.add((r, s)) + visit = max(visit, v + 1) + print(f"[{dtype}] warmup done at {time.perf_counter() - t_start:.0f} s", flush=True) + + for r in range(args.rounds): + for s in range(args.slots): + if (r, s) in done: + continue + rows = [] + mem = [] + tm = dc.THUE_MORSE[s % 8] + warm = args.visit_warmup_calls[s % len(args.visit_warmup_calls)] + order = list(range(len(groups))) + random.Random(args.seed * 7919 + r * 131 + s).shuffle(order) + for gi in order: + g = groups[gi] + inst = [ + (a, c) + for a, c in arm_instances(g, args.arms) + if (gi, cell_key(a, c)) not in errors + ] + if tm: + inst.reverse() + wcalls = inner[gi] if warm == "full" else int(warm) + for arm, codec in inst[: args.visit_warmup]: + x0, f = dc.make_call(arm, codec, g, dt, store) + dc.set_tile(dc.ARM_TILE.get(arm), nax) + set_arm_env(arm, env_arms) + dc.chain_ms(f, x0, wcalls) + x0 = store.act(g, dt) + glue = dc.chain_ms(lambda a: a, x0, inner[gi]) + t = round(time.perf_counter() - t_start, 4) + samples[gi].setdefault("glue", []).append((visit, r, s, t, glue)) + rows.append([gi, "glue", visit, r, s, t, glue]) + for arm, codec in inst: + x0, f = dc.make_call(arm, codec, g, dt, store) + dc.set_tile(dc.ARM_TILE.get(arm), nax) + set_arm_env(arm, env_arms) + ms = dc.chain_ms(f, x0, inner[gi]) - glue + t = round(time.perf_counter() - t_start, 4) + key = cell_key(arm, codec) + samples[gi].setdefault(key, []).append((visit, r, s, t, ms)) + rows.append([gi, key, visit, r, s, t, ms]) + dc.set_tile(None, nax) + set_arm_env(None, env_arms) + mem.append(mem_row(mx, gi, args.mem_trace)) + visit += 1 + raw.write( + { + "kind": "slot", + "dtype": dtype, + "round": r, + "slot": s, + "visit_warmup_calls": warm, + "samples": rows, + "mem": mem, + } + ) + slots += 1 + fp = [m[3] for m in mem if m[3] is not None] + fp = f", footprint up to {max(fp) / 2**30:.1f} GiB" if fp else "" + print( + f"[{dtype}] round {r} slot {s} done at " + f"{time.perf_counter() - t_start:.0f} s{fp}", + flush=True, + ) + budget = args.time_budget_min + left = slots < args.rounds * args.slots + if budget and left and time.perf_counter() - t_start >= budget * 60: + return { + "inner": inner, + "kernels": kernels, + "errors": errors, + "samples": samples, + "slots": slots, + "stopped": True, + } + return { + "inner": inner, + "kernels": kernels, + "errors": errors, + "samples": samples, + "slots": slots, + } + + +def summarize_pass(groups, m, rounds, arm_names, desc): + """Per-cell records and paired ratios from one dtype's measurements.""" + nax = desc["nax"] + arch = desc["architecture"] + kernels, errors, inner = m["kernels"], m["errors"], m["inner"] + out_groups = [] + for gi, g in enumerate(groups): + sm = m["samples"][gi] + M, N, K = g["M"], g["N"], g["K"] + + def arm_record(arm, codec, gi=gi, sm=sm, M=M, N=N, K=K): + key = cell_key(arm, codec) + if arm == "dense": + tile = dc.mlx_dense_tile(M, N, K, arch, nax) + elif arm.startswith("affine"): + tile = dc.mlx_affine_tile(nax) + elif arm == "deq": + tile = dc.mlx_dense_tile(M, N, K, arch, nax) + elif arm == "glue": + tile = None + else: + tile = dc.tile_of(kernels.get((gi, key))) + rec = { + "kernel": kernels.get((gi, key)), + "tile": tile, + "error": errors.get((gi, key)), + "samples": [[r, s, t, ms] for (_v, r, s, t, ms) in sm.get(key, [])], + } + rec.update(split_stats(sm.get(key, []), rounds, 1, 4)) + return rec + + shared = {} + for arm, codec in arm_instances(g, arm_names): + if codec is None: + shared[arm] = arm_record(arm, None) + if "glue" in sm: + shared["glue"] = arm_record("glue", None) + cells = [] + for codec in g["codecs"]: + arms = {} + for arm, c in arm_instances(g, arm_names): + if c == codec: + arms[arm] = arm_record(arm, codec) + ratios = {} + aff = f"affine{dc.affine_bits(codec)}" + for name, a, b in ( + ("kq/dense", cell_key("kq", codec), "dense"), + ("kq/affine", cell_key("kq", codec), aff), + ("bm64/kq", cell_key("kq_bm64", codec), cell_key("kq", codec)), + ("bm128/kq", cell_key("kq_bm128", codec), cell_key("kq", codec)), + ( + "bm128/bm64", + cell_key("kq_bm128", codec), + cell_key("kq_bm64", codec), + ), + ("deq/kq", cell_key("deq", codec), cell_key("kq", codec)), + ("aa/kq", cell_key("kq_aa", codec), cell_key("kq", codec)), + *( + (f"{a}/kq", cell_key(a, codec), cell_key("kq", codec)) + for a in lever_arms(arm_names) + ), + ): + raw_pairs, dropped = paired(sm, a, b) + if raw_pairs or dropped: + pr, order = order_balance(raw_pairs) + ratios[name] = split_stats(pr, rounds, 0, 1) + ratios[name]["order"] = order + ratios[name]["pairs"] = [[r, v] for r, v in pr] + ratios[name]["pairs_raw"] = [list(p) for p in raw_pairs] + ratios[name]["dropped"] = dropped + ratios[name]["ref_floor"] = at_timer_floor(sm, b) + cells.append({"codec": codec, "arms": arms, "ratios": ratios}) + out_groups.append(dict(g, inner=inner[gi], shared=shared, cells=cells)) + return out_groups + + +def aggregates(groups, rounds): + """Per op and ratio: the pooled median of every settled pair across the + op's cells, with its CI, beside the spread of per-cell medians and of + per-cell CI half-widths.""" + by = {} + for g in groups: + op = f"{g['family']}.{g['op']}" + for c in g["cells"]: + for name, rt in c["ratios"].items(): + e = by.setdefault( + (op, name), {"pairs": [], "cells": [], "half": [], "order": []} + ) + e["pairs"] += [v for r, v in rt["pairs"] if r > 0 or rounds == 1] + if rt.get("order") is not None: + e["order"].append(rt["order"]) + st = rt.get("settled") or rt.get("all") + if st: + e["cells"].append(st["median"]) + if st["ci95_half"] is not None: + e["half"].append(st["ci95_half"]) + out = [] + for (op, name), e in by.items(): + out.append( + { + "op": op, + "ratio": name, + "pooled": bs.summarize(e["pairs"]), + "cell_medians": bs.summarize(e["cells"]), + "cell_ci95_half": bs.summarize(e["half"]), + "cell_ci95_half_max": max(e["half"]) if e["half"] else None, + "cell_order": bs.summarize(e["order"]), + "cells": len(e["cells"]), + } + ) + return out + + +def build_result(meta, measured, end, raw_path): + """The report from a meta record and per-dtype measurements, live or + read back from a raw log.""" + a = meta["args"] + groups = meta["groups"] + planned = a["rounds"] * a["slots"] + passes = [] + last_t = None + for dtype, m in measured.items(): + pg = summarize_pass(groups, m, a["rounds"], a["arms"], meta["machine"]) + passes.append( + { + "dtype": dtype, + "slots_done": m["slots"], + "slots_planned": planned, + "groups": pg, + "aggregate": aggregates(pg, a["rounds"]), + } + ) + for sm in m["samples"]: + for rows in sm.values(): + for row in rows: + last_t = row[3] if last_t is None else max(last_t, row[3]) + return { + "machine": meta["machine"], + "args": a, + "heavy_codecs": meta["heavy_codecs"], + "wall_s": end.get("wall_s", last_t), + "peak_memory_bytes": end.get("peak_memory_bytes"), + "mlx_kquant_file": meta["mlx_kquant_file"], + "raw": raw_path, + "sample_mode": meta.get("sample_mode", "independent"), + "complete": bool(end), + "passes": passes, + } + + +def write_outputs(result, json_out, md_out): + if json_out: + with open(json_out, "w") as f: + json.dump(result, f) + md = markdown(result) + if md_out: + with open(md_out, "w") as f: + f.write(md + "\n") + else: + print(md) + + +def fmt(v, spec=".3f"): + return "-" if v is None else format(v, spec) + + +def pick(stats, key="median"): + if not stats: + return None + s = stats.get("settled") or stats.get("all") + return s[key] if s else None + + +def floor_mark(stats): + return "*" if stats and stats.get("ref_floor") else "" + + +def pick_round0(stats): + s = (stats or {}).get("round0") + return s["median"] if s else None + + +def markdown(result): + lines = ["# DiT shape arms", "", dc.descriptor_md(result["machine"]), ""] + a = result["args"] + if result.get("sample_mode", "independent") == "chain": + lines.append( + "Samples are dependent call chains with the tie chain (glue) subtracted." + ) + else: + lines.append( + "Samples are independent calls evaluated together (earlier " + "bench): do not compare with chained runs." + ) + lines.append("") + lines.append( + f"Rounds {a['rounds']} x {a['slots']} Thue-Morse slots, inner {a['inner']} " + f"(up to {a['max_inner']} under {a['min_sample_ms']} ms per call), " + f"shrink {a['shrink']}. Times are settled medians in ms (all rounds when " + "there is one); ratios are medians of per-visit paired ratios, " + "balanced over call order: the later call of a visit runs slower, " + "so per cell the pairs in which each arm ran first are scaled to " + "meet at the geometric mean of their two medians, and order is the " + "factor between them (later over earlier); round 0 = median of " + "round 0, which is not a cold figure (see bench_dit_thermal.py for " + "the cold rate)." + ) + lines.append("") + if not result.get("complete", True): + lines.append( + "Partial run: the raw log has no end record, so the figures cover " + "the slots it holds." + ) + lines.append("") + for p in result["passes"]: + lines.append(f"## {p['dtype']}") + lines.append("") + if p.get("slots_done", 0) < p.get("slots_planned", 0): + lines.append(f"{p['slots_done']} of {p['slots_planned']} slots ran.") + lines.append("") + lines.append( + "| op | L x B | M | codec | kq tile | kq ms | kq IQR | dense | affine " + "| bm64 | bm128 | deq |" + ) + lines.append("|---|---|---|---|---|---|---|---|---|---|---|---|") + for g in p["groups"]: + for c in g["cells"]: + ar = c["arms"] + aff = g["shared"].get(f"affine{dc.affine_bits(c['codec'])}") + kqa = ar.get("kq", {}) + lines.append( + f"| {g['family']}.{g['op']} | {g['L']} x {g['B']} | {g['M']} " + f"| {c['codec']} | {kqa.get('tile') or '-'} " + f"| {fmt(pick(kqa))} | {fmt(pick(kqa, 'rel_iqr'), '.1%')} " + f"| {fmt(pick(g['shared'].get('dense')))} | {fmt(pick(aff))} " + f"| {fmt(pick(ar.get('kq_bm64')))} " + f"| {fmt(pick(ar.get('kq_bm128')))} | {fmt(pick(ar.get('deq')))} |" + ) + lines.append("") + names = ["kq/dense", "kq/affine", "bm128/bm64", "deq/kq", "aa/kq"] + names += [f"{arm}/kq" for arm in lever_arms(a["arms"])] + lines.append( + "| op | L x B | codec | " + + " | ".join(names) + + " | round 0 kq/affine | round 0 aa/kq |" + ) + lines.append("|---|---|---|" + "---|" * (len(names) + 2)) + floored, left_out = False, [] + for g in p["groups"]: + for c in g["cells"]: + rt = c["ratios"] + row = ( + f"| {g['family']}.{g['op']} | {g['L']} x {g['B']} | {c['codec']} |" + ) + row += "".join( + f" {fmt(pick(rt.get(n)))}{floor_mark(rt.get(n))} |" for n in names + ) + row += f" {fmt(pick_round0(rt.get('kq/affine')))} |" + row += f" {fmt(pick_round0(rt.get('aa/kq')))} |" + lines.append(row) + for n, r in rt.items(): + if r.get("ref_floor"): + floored = True + if r.get("dropped"): + total = r["dropped"] + len(r["pairs_raw"]) + left_out.append( + f"{g['family']}.{g['op']} {g['L']} x {g['B']} " + f"{c['codec']} {n}: {r['dropped']} of {total}" + ) + lines.append("") + if floored: + lines.append( + "A ratio marked * has its reference at the timer floor: a " + "reference sample at or below 0, or a reference median under " + "twice the p10 to p90 width of the glue times." + ) + lines.append("") + if left_out: + lines.append( + "Visits left out of a ratio, their reference sample at or " + "below 0: " + "; ".join(left_out) + "." + ) + lines.append("") + lines.append( + "Aggregate per op: pooled median of settled pairs with its 95 percent " + "CI; per-cell medians p10 to p90; per-cell CI half-width median and " + "max; per-cell order factor median." + ) + lines.append("") + lines.append( + "| op | ratio | cells | pairs | pooled | CI | cell p10 | cell p90 " + "| cell CI half med | cell CI half max | order |" + ) + lines.append("|---|---|---|---|---|---|---|---|---|---|---|") + for ag in p["aggregate"]: + po, cm, ch = ag["pooled"], ag["cell_medians"], ag["cell_ci95_half"] + co = ag.get("cell_order") + ci = po["ci95"] if po else None + lines.append( + f"| {ag['op']} | {ag['ratio']} | {ag['cells']} " + f"| {po['n'] if po else 0} | {fmt(po and po['median'], '.4f')} " + f"| {f'{ci[0]:.4f} to {ci[1]:.4f}' if ci else '-'} " + f"| {fmt(cm and cm['p10'])} | {fmt(cm and cm['p90'])} " + f"| {fmt(ch and ch['median'], '.1%')} " + f"| {fmt(ag['cell_ci95_half_max'], '.1%')} " + f"| {fmt(co and co['median'])} |" + ) + lines.append("") + errs = [ + (g, c["codec"], arm, rec["error"]) + for g in p["groups"] + for c in g["cells"] + for arm, rec in c["arms"].items() + if rec["error"] + ] + if errs: + lines.append("Arms not run:") + lines.append("") + for g, codec, arm, e in errs: + lines.append( + f"- {g['family']}.{g['op']} {g['L']} x {g['B']} {codec} {arm}: {e}" + ) + lines.append("") + return "\n".join(lines) + + +def main(argv=None): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--families", default=",".join(dc.FAMILIES)) + ap.add_argument("--ops", default="", help="subset of op names, comma separated") + ap.add_argument( + "--rows", + default="", + help="rows per item; default per family (modulation ops keep one row)", + ) + ap.add_argument("--batches", default="1,2") + ap.add_argument("--codecs", default=",".join(dc.MAIN_CODECS)) + ap.add_argument( + "--heavy-codecs", + default="nax", + help="codecs added at the two heaviest cells: 'nax' (every codec " + "with NAX kernels), 'none', or a list", + ) + ap.add_argument( + "--arms", + default="dense,kq,kq_bm64,kq_bm128,deq,affine", + help="from dense,kq,kq_bm64,kq_bm128,deq,affine,kq_aa", + ) + ap.add_argument("--dtypes", default="bfloat16,float16") + ap.add_argument("--rounds", type=int, default=3) + ap.add_argument("--slots", type=int, default=8) + ap.add_argument("--inner", type=int, default=3, help="calls per sample") + ap.add_argument( + "--min-sample-ms", + type=float, + default=10.0, + help="shapes faster than this per call take more calls per sample", + ) + ap.add_argument("--max-inner", type=int, default=64) + ap.add_argument( + "--visit-warmup", + type=int, + default=1, + help="arms that run a discarded sample at the start of each shape visit", + ) + ap.add_argument( + "--env-arm", + action="append", + default=[], + metavar="NAME:VAR=VALUE[,VAR=VALUE]", + help="a kq arm run with these environment levers, paired with kq as " + "NAME/kq (repeatable)", + ) + ap.add_argument( + "--visit-warmup-calls", + default="full", + help="calls in that discarded sample: 'full' (the sample's count) or " + "a number; a comma list cycles over the slots (A versus B of the " + "warmup in one process)", + ) + ap.add_argument("--shrink", type=int, default=1, help="divide rows, K, N") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument( + "--est-tflops", type=float, default=25.0, help="rate for the time estimate" + ) + ap.add_argument("--plan-only", action="store_true") + ap.add_argument("--json-out") + ap.add_argument("--md-out") + ap.add_argument( + "--raw-out", + help="JSON Lines log of every sample, written as each slot ends " + "(default: --json-out plus .raw.jsonl)", + ) + ap.add_argument( + "--cache-limit-gib", + type=float, + default=8.0, + help="MLX buffer cache limit; it must hold the outputs of one sample " + "of the largest shape, or the timed calls allocate. Negative keeps " + "the MLX default", + ) + ap.add_argument( + "--mem-trace", + action="store_true", + help="print active, cache and footprint GiB after every visit", + ) + ap.add_argument( + "--time-budget-min", + type=float, + default=0, + help="stop after the slot that ends past this many minutes, without " + "an end record, so the run continues under --resume (0: no budget)", + ) + ap.add_argument( + "--resume", + action="store_true", + help="continue the run in the raw log (same arguments): keep its " + "finished slots and run the rest", + ) + ap.add_argument( + "--reaggregate", + metavar="RAW", + help="rebuild the JSON and Markdown from a raw log without running " + "anything; only --json-out and --md-out apply", + ) + args = ap.parse_args(argv) + if args.reaggregate: + meta, measured, end = load_raw(args.reaggregate) + result = build_result(meta, measured, end, args.reaggregate) + write_outputs(result, args.json_out, args.md_out) + return 0 + args.families = parse_list(args.families) + args.ops = parse_list(args.ops) + args.rows = parse_list(args.rows, int) + args.batches = parse_list(args.batches, int) + args.codecs = parse_list(args.codecs) + args.arms = parse_list(args.arms) + for name in parse_env_arms(args.env_arm): + if name not in args.arms: + args.arms.append(name) + unknown = set(lever_arms(args.arms)) - set(parse_env_arms(args.env_arm)) + if unknown: + ap.error(f"arms without a definition: {sorted(unknown)}") + args.dtypes = parse_list(args.dtypes) + args.visit_warmup_calls = parse_list(args.visit_warmup_calls) + + import mlx.core as mx + + import mlx_kquant as kq + + if args.heavy_codecs == "nax": + heavy_codecs = dc.nax_codecs() + elif args.heavy_codecs == "none": + heavy_codecs = [] + else: + heavy_codecs = parse_list(args.heavy_codecs) + groups = build_groups(args, heavy_codecs) + per_round, total_min = plan_summary( + groups, args.arms, args.rounds, args.slots, args.est_tflops, args.inner + ) + cells = sum(len(g["codecs"]) for g in groups) + print( + f"{len(groups)} shapes, {cells} codec cells per dtype; per round " + f"{per_round['calls']} calls, {per_round['tflop']:.1f} TFLOP, " + f"about {per_round['est_min']:.1f} min at {args.est_tflops} TFLOPS; " + f"about {total_min * len(args.dtypes):.0f} min for {len(args.dtypes)} " + f"dtype(s) x {args.rounds} rounds plus warmup", + flush=True, + ) + if args.plan_only: + return + + for k in ("KQ_NAX_TILE", "KQ_QMM_TILE", "KQ_QMM_LOG"): + os.environ.pop(k, None) + if args.cache_limit_gib >= 0: + mx.set_cache_limit(int(args.cache_limit_gib * 2**30)) + desc = dc.machine_descriptor() + store = dc.Store(args.seed) + raw_out = args.raw_out + if raw_out is None and args.json_out: + raw_out = args.json_out + ".raw.jsonl" + meta = { + "kind": "meta", + "machine": desc, + "args": vars(args), + "heavy_codecs": heavy_codecs, + "groups": groups, + "mlx_kquant_file": os.path.dirname(kq.__file__), + "sample_mode": "chain", + } + prior = {} + if args.resume and raw_out and os.path.exists(raw_out): + complete_lines(raw_out) + meta0, prior, end0 = load_raw(raw_out) + if ( + meta0.get("sample_mode") != "chain" + or meta0["groups"] != groups + or any( + meta0["args"].get(k) != getattr(args, k) + for k in ( + "arms", + "env_arm", + "dtypes", + "rounds", + "slots", + "seed", + "shrink", + ) + ) + ): + sys.exit(f"{raw_out}: another run's raw log; --resume needs the same plan") + if end0: + result = build_result(meta0, prior, end0, raw_out) + write_outputs(result, args.json_out, args.md_out) + return 0 + meta = meta0 + raw = RawLog(raw_out, "a") + else: + raw = RawLog(raw_out) + raw.write(meta) + t_start = time.perf_counter() + measured = {} + for dtype in args.dtypes: + store.clear_dtype() + gc.collect() + mx.clear_cache() + p = prior.get(dtype) + if p is not None and p["slots"] == args.rounds * args.slots: + measured[dtype] = p + continue + measured[dtype] = measure_pass( + dtype, groups, args, store, desc, t_start, raw, p + ) + if measured[dtype].get("stopped"): + break + if any(m.get("stopped") for m in measured.values()): + raw.close() + print("stopped at the time budget; --resume continues the run", flush=True) + result = build_result(meta, {**prior, **measured}, {}, raw_out) + write_outputs(result, args.json_out, args.md_out) + return 0 + end = { + "kind": "end", + "wall_s": time.perf_counter() - t_start, + "peak_memory_bytes": mx.get_peak_memory(), + } + raw.write(end) + raw.close() + try: + result = build_result(meta, measured, end, raw_out) + write_outputs(result, args.json_out, args.md_out) + except Exception: + if raw_out: + print( + f"summaries failed; the samples are in {raw_out}, and " + f"--reaggregate {raw_out} rebuilds the reports", + file=sys.stderr, + flush=True, + ) + raise + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/bench_dit_thermal.py b/benchmarks/bench_dit_thermal.py new file mode 100644 index 00000000..69ccb4b7 --- /dev/null +++ b/benchmarks/bench_dit_thermal.py @@ -0,0 +1,183 @@ +"""Thermal curve of one DiT cell: a timed loop, then idle recovery. + +After --cold-wait seconds idle and one discarded call, the cell loops for +--seconds and every call is recorded. Reported: + + cold the fastest median of --cold-calls consecutive calls in the + first --cold-span seconds, which skips the clock ramp of + the first calls after idle + first slowdown seconds until the rolling median of --window calls, after + the cold window, first exceeds cold by --slowdown (default + 3 percent) + settled ratio median of the last --tail fraction of calls over cold + idle gap after the loop the GPU idles along --idle-schedule + (cumulative seconds); at each point --probe-calls calls run, + and the gap is the first point whose fastest median of 3 + consecutive calls is within --recover of cold + +These define one thermal round for bench_dit_shapes.py. Markdown and JSON +out; the JSON keeps every call. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import time + +import dit_common as dc + + +def cold_window(ts, ms, n, span): + """(median, end index) of the fastest n consecutive calls that end + within `span` seconds.""" + best = (float("inf"), min(n, len(ms))) + for i in range(min(n, len(ms)), len(ms) + 1): + if ts[i - 1] > span and i > n: + break + best = min(best, (statistics.median(ms[i - n : i]), i)) + return best + + +def rolling_first_over(ts, ms, start, window, limit): + for i in range(max(window, start + 1), len(ms) + 1): + if statistics.median(ms[i - window : i]) > limit: + return ts[i - 1] + return None + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--family", default="qwen_image", choices=list(dc.FAMILIES)) + ap.add_argument("--op", default="gate_up") + ap.add_argument("--rows", type=int, default=4096) + ap.add_argument("--batch", type=int, default=1) + ap.add_argument("--codec", default="q4_k") + ap.add_argument("--arm", default="kq", choices=["kq", "dense", "affine", "deq"]) + ap.add_argument("--dtype", default="bfloat16", choices=list(dc.DTYPES)) + ap.add_argument("--seconds", type=float, default=120.0) + ap.add_argument("--cold-wait", type=float, default=0.0) + ap.add_argument("--cold-calls", type=int, default=5) + ap.add_argument("--cold-span", type=float, default=10.0) + ap.add_argument("--window", type=int, default=9) + ap.add_argument("--slowdown", type=float, default=0.03) + ap.add_argument("--tail", type=float, default=0.2) + ap.add_argument("--idle-schedule", default="5,10,20,40,80,160,320") + ap.add_argument("--probe-calls", type=int, default=8) + ap.add_argument("--recover", type=float, default=0.01) + ap.add_argument("--shrink", type=int, default=1) + ap.add_argument("--json-out") + ap.add_argument("--md-out") + args = ap.parse_args() + + import mlx.core as mx + + g = dc.cell(args.family, args.op, args.rows, args.batch, args.shrink) + dt = dc.mx_dtype(args.dtype) + store = dc.Store(0) + arm = f"affine{dc.affine_bits(args.codec)}" if args.arm == "affine" else args.arm + fn = dc.make_fn(arm, args.codec, g, dt, store) + desc = dc.machine_descriptor() + with dc.KernelLog() as log: + mx.eval(fn()) + kernel = log.kernels[-1] if log.kernels else None + + if args.cold_wait > 0: + time.sleep(args.cold_wait) + dc.time_fn(fn, 1) + t0 = time.perf_counter() + ts, ms = [], [] + while time.perf_counter() - t0 < args.seconds: + v = dc.time_fn(fn, 1) + ts.append(time.perf_counter() - t0) + ms.append(v) + loop_end = time.perf_counter() + cold, cold_end = cold_window(ts, ms, args.cold_calls, args.cold_span) + limit = cold * (1 + args.slowdown) + first = rolling_first_over(ts, ms, cold_end, args.window, limit) + tail = ms[int(len(ms) * (1 - args.tail)) :] or ms[-1:] + settled_ratio = statistics.median(tail) / cold + + probes = [] + gap = None + for point in (float(v) for v in args.idle_schedule.split(",") if v): + wait = loop_end + point - time.perf_counter() + if wait > 0: + time.sleep(wait) + pms = [dc.time_fn(fn, 1) for _ in range(args.probe_calls)] + pm = cold_window(list(range(len(pms))), pms, 3, float("inf"))[0] + probes.append({"idle_s": point, "ms": pm, "ratio": pm / cold, "calls": pms}) + # Probes add a little heat, so the clock restarts after each one. + loop_end = time.perf_counter() - point + if pm <= cold * (1 + args.recover): + gap = point + break + + result = { + "machine": desc, + "args": vars(args), + "cell": g, + "arm": arm, + "kernel": kernel, + "tile": dc.tile_of(kernel), + "calls": len(ms), + "cold_ms": cold, + "cold_window_end_s": ts[cold_end - 1], + "first_slowdown_s": first, + "settled_ratio": settled_ratio, + "idle_gap_s": gap, + "probes": probes, + "series": [[round(t, 4), m] for t, m in zip(ts, ms, strict=True)], + } + lines = ["# DiT thermal curve", "", dc.descriptor_md(desc), ""] + lines.append( + f"Cell {g['family']}.{g['op']} {g['L']} x {g['B']} (M {g['M']}, K {g['K']}, " + f"N {g['N']}), {arm} {args.codec if arm.startswith(('kq', 'deq')) else ''} " + f"{args.dtype}, tile {result['tile'] or '-'}, {len(ms)} calls in " + f"{args.seconds:.0f} s." + ) + lines.append("") + lines.append("| cold ms | first slowdown s | settled ratio | idle gap s |") + lines.append("|---|---|---|---|") + lines.append( + f"| {cold:.3f} | {'-' if first is None else f'{first:.1f}'} " + f"| {settled_ratio:.3f} | {'-' if gap is None else f'{gap:g}'} |" + ) + lines.append("") + lines.append("| t s (bin end) | median ms | over cold |") + lines.append("|---|---|---|") + step = max(args.seconds / 24, 1e-9) + edge, bucket = step, [] + for t, m in zip(ts, ms, strict=True): + if t > edge and bucket: + med = statistics.median(bucket) + lines.append(f"| {edge:.1f} | {med:.3f} | {med / cold:.3f} |") + bucket = [] + while t > edge: + edge += step + bucket.append(m) + if bucket: + med = statistics.median(bucket) + lines.append(f"| {ts[-1]:.1f} | {med:.3f} | {med / cold:.3f} |") + lines.append("") + lines.append("| idle s | probe ms | over cold |") + lines.append("|---|---|---|") + for p in probes: + lines.append(f"| {p['idle_s']:g} | {p['ms']:.3f} | {p['ratio']:.3f} |") + md = "\n".join(lines) + if args.json_out: + with open(args.json_out, "w") as f: + json.dump(result, f) + if args.md_out: + with open(args.md_out, "w") as f: + f.write(md + "\n") + else: + print(md) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/bench_glu.py b/benchmarks/bench_glu.py new file mode 100644 index 00000000..caadd91f --- /dev/null +++ b/benchmarks/bench_glu.py @@ -0,0 +1,380 @@ +"""kq.quantized_matmul_glu against the unfused ops at DiT feed-forward shapes. + +One call computes silu(x @ Wg^T) * (x @ Wu^T) for x [1, M, K] and two +packed weights [N, K] of one codec. Arms, all in one process: + + composed two kq.quantized_matmul calls, then stock silu and multiply + composed_aa the same again (A versus A) + two_mm the two matmuls alone, both evaluated, with no activation + and no product: the part a fused kernel cannot remove + glu kq.quantized_matmul_glu with no lever set + glu:BMxBK kq.quantized_matmul_glu on one tile, in a build with + KQ_PROBE_KERNELS=ON only (--tiles). BK takes a suffix: w for + the two-walk kernel, d for the one-walk kernel on two + unpadded buffers + +Where the op has no fused kernel for the codec, the dtype, the shape or +the GPU, glu is the composition, and composed against two_mm prices what a +fused kernel could save there. + +--cells is a list of DTYPE:K:N:ROWS sets separated by semicolons, ROWS a +comma list of M. The default is the sweep that decides whether a codec's +fused kernel is the op's default: it is, where glu is at or below composed +in every cell. + +Order: every arm of a cell runs once to build its pipelines and is compared +with composed. Then the MLX peak of one call of each arm is taken, over the +memory that is live once the buffers of the calls before it are released. +Then each round visits every arm once, in an order that rotates with the +round and reverses on odd cycles. A visit is one discarded call and --calls +timed calls, each evaluated alone; the sample is their median. Ratios pair +the arms inside a round. The JSON is rewritten after every round, with the +process footprint of that round. +""" + +from __future__ import annotations + +import argparse +import ctypes +import json +import os +import statistics +import sys +import time + +# The tile levers are read per call. +os.environ["KQ_TUNE_LIVE"] = "1" + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import dit_common as dc # noqa: E402 + +DEFAULT_CELLS = ( + "bfloat16:3840:10240:1056,2080,4128,4608,8208,16416;" + "float16:3840:10240:1056,4128,16416;" + "bfloat16:4096:12288:4128,16416" +) +LEVERS = ("KQ_GLU_BM", "KQ_GLU_BK", "KQ_GLU_DUAL128", "KQ_GLU_WALKS") + + +# struct rusage_info_v2 of , whole: the call fills it all. +class _RUsageV2(ctypes.Structure): + _fields_ = [("ri_uuid", ctypes.c_uint8 * 16)] + [ + (n, ctypes.c_uint64) + for n in ( + "ri_user_time", + "ri_system_time", + "ri_pkg_idle_wkups", + "ri_interrupt_wkups", + "ri_pageins", + "ri_wired_size", + "ri_resident_size", + "ri_phys_footprint", + "ri_proc_start_abstime", + "ri_proc_exit_abstime", + "ri_child_user_time", + "ri_child_system_time", + "ri_child_pkg_idle_wkups", + "ri_child_interrupt_wkups", + "ri_child_pageins", + "ri_child_elapsed_abstime", + "ri_diskio_bytesread", + "ri_diskio_byteswritten", + ) + ] + + +def phys_footprint(): + """Physical footprint of this process in bytes, 0 where unknown.""" + try: + lib = ctypes.CDLL("/usr/lib/libproc.dylib") + r = _RUsageV2() + if lib.proc_pid_rusage(os.getpid(), 2, ctypes.byref(r)) != 0: + return 0 + return int(r.ri_phys_footprint) + except Exception: + return 0 + + +def med(v): + return statistics.median(v) + + +def settle(mx): + """Active MLX memory once the buffers of finished calls are released. + A buffer is returned when its command buffer completes, which can be + after the call that dropped the array, so read until five readings + 10 ms apart agree.""" + mx.synchronize() + last, same = -1, 0 + for _ in range(500): + mx.clear_cache() + cur = mx.get_active_memory() + same = same + 1 if cur == last else 0 + last = cur + if same >= 5: + break + time.sleep(0.01) + return last + + +def tile_env(tile): + bm, bk = tile.split("x") + return { + "KQ_GLU_BM": bm, + "KQ_GLU_BK": bk.rstrip("dw"), + "KQ_GLU_DUAL128": "1" if bk.endswith("d") else "0", + "KQ_GLU_WALKS": "1" if bk.endswith("w") else "0", + } + + +def build_arms(a, dt, K, N): + import mlx.core as mx + import mlx.nn as nn + + import mlx_kquant as kq + + s0 = mx.zeros((1,), dtype=mx.uint8) + ws = [] + for i in range(2): + wf = mx.random.normal((N, K), key=mx.random.key(11 + i)) * K**-0.5 + w, _ = kq.quantize(wf, a.codec) + mx.eval(w) + ws.append(w) + wg, wu = ws + + def mm(x, w): + return kq.quantized_matmul(x, w, s0, a.codec, True) + + def composed(x): + y = nn.silu(mm(x, wg)) * mm(x, wu) + mx.eval(y) + return y + + def two_mm(x): + g, u = mm(x, wg), mm(x, wu) + mx.eval(g, u) + return g + + def glu(env): + def f(x): + for name in LEVERS: + os.environ.pop(name, None) + os.environ.update(env) + y = kq.quantized_matmul_glu(x, wg, wu, a.codec) + mx.eval(y) + return y + + return f + + arms = { + "composed": composed, + "composed_aa": composed, + "two_mm": two_mm, + "glu": glu({}), + } + for tile in a.tiles: + arms["glu:" + tile] = glu(tile_env(tile)) + for name in a.skip: + arms.pop(name, None) + return arms + + +def run(a): + import mlx.core as mx + + import mlx_kquant as kq + + mx.set_cache_limit(a.cache_gib << 30) + res = { + "machine": dc.machine_descriptor(), + "args": {k: v for k, v in vars(a).items() if k != "out"}, + "kq_file": kq.__file__, + "cells": [], + } + print(dc.descriptor_md(res["machine"]), flush=True) + for spec in a.cells: + dtype, K, N, rows = spec.split(":") + K, N = int(K), int(N) + dt = dc.mx_dtype(dtype) + mx.clear_cache() + arms = build_arms(a, dt, K, N) + names = list(arms) + for M in (int(r) for r in rows.split(",")): + x = mx.random.normal((1, M, K), key=mx.random.key(M)) * 0.5 + x = x.astype(dt) + mx.eval(x) + cell = { + "dtype": dtype, + "K": K, + "N": N, + "M": M, + "arms": names, + "check": {}, + "peak_outputs": {}, + "samples": [], + "rounds": [], + } + res["cells"].append(cell) + ref = None + for n in names: + y = arms[n](x) + if n == "composed": + ref = y.astype(mx.float32) + elif n != "two_mm" and ref is not None: + d = mx.linalg.norm(y.astype(mx.float32) - ref) + cell["check"][n] = (d / mx.linalg.norm(ref)).item() + del y + ref = None + # MLX peak of one call over the memory live before it, in + # outputs. The baseline waits for the arm before to let go. + for n in names: + act0 = settle(mx) + mx.reset_peak_memory() + y = arms[n](x) + pk = mx.get_peak_memory() - act0 + cell["peak_outputs"][n] = pk / (M * N * 2) + del y + mx.clear_cache() + for rnd in range(a.rounds): + r = rnd % len(names) + order = names[r:] + names[:r] + if (rnd // len(names)) % 2: + order = order[::-1] + for pos, n in enumerate(order): + f = arms[n] + f(x) + mx.synchronize() + ts = [] + for _ in range(a.calls): + t0 = time.perf_counter() + f(x) + ts.append(time.perf_counter() - t0) + cell["samples"].append( + {"round": rnd, "arm": n, "pos": pos, "s": med(ts)} + ) + fp = phys_footprint() / 2**30 + cell["rounds"].append({"round": rnd, "footprint_gib": fp}) + with open(a.out + ".json", "w") as fh: + json.dump(res, fh, indent=1) + print( + f"[glu] {a.codec} {dtype} K {K} N {N} M {M} round {rnd} " + f"footprint {fp:.2f} GiB, MLX cache " + f"{mx.get_cache_memory() / 2**30:.2f}", + flush=True, + ) + del x + mx.clear_cache() + del arms + report(res, a.out + ".md") + + +def _ratios(cell, arm, base): + b = {s["round"]: s["s"] for s in cell["samples"] if s["arm"] == base} + return sorted( + s["s"] / b[s["round"]] + for s in cell["samples"] + if s["arm"] == arm and s["round"] in b + ) + + +def report(res, path): + a = res["args"] + lines = [ + f"# kq.quantized_matmul_glu against the unfused ops, {a['codec']}", + "", + dc.descriptor_md(res["machine"]), + "", + f"{a['rounds']} rounds, {a['calls']} calls per visit. Ratios are " + "medians of per-round pairs against composed; the range is the " + "smallest and largest round. Peak is the MLX peak of one call over " + "the memory live before it, in output buffers.", + "", + "| dtype | K | N | M | Arm | ms per call | Ratio to composed | Range | " + "Rounds below 1 | Peak, outputs | Difference to composed |", + "|---|---|---|---|---|---|---|---|---|---|---|", + ] + verdict = [] + for cell in res["cells"]: + head = f"| {cell['dtype']} | {cell['K']} | {cell['N']} | {cell['M']} | " + for n in cell["arms"]: + t = [s["s"] for s in cell["samples"] if s["arm"] == n] + if not t: + continue + row = head + f"{n} | {1e3 * med(t):.3f} | " + r = _ratios(cell, n, "composed") + if r and n != "composed": + below = sum(x < 1 for x in r) + row += f"{med(r):.3f} | {r[0]:.3f} to {r[-1]:.3f} | " + row += f"{below}/{len(r)} | " + if n == "glu": + verdict.append(med(r)) + else: + row += " | | | " + row += f"{cell['peak_outputs'].get(n, 0):.2f} | " + d = cell["check"].get(n) + row += (f"{d:.2e}" if d is not None else "") + " |" + lines.append(row) + lines.append("") + aa = [med(r) for c in res["cells"] if (r := _ratios(c, "composed_aa", "composed"))] + if verdict: + lines.append( + f"- glu over composed, by cell: smallest {min(verdict):.3f}, " + f"largest {max(verdict):.3f}, {sum(v <= 1 for v in verdict)} of " + f"{len(verdict)} cells at or below 1." + ) + if aa: + lines.append( + f"- A versus A (composed_aa over composed), by cell: " + f"{min(aa):.3f} to {max(aa):.3f}." + ) + fp = max( + (r["footprint_gib"] for c in res["cells"] for r in c["rounds"]), + default=0.0, + ) + lines.append(f"- Largest process footprint {fp:.2f} GiB.") + text = "\n".join(lines) + "\n" + with open(path, "w") as fh: + fh.write(text) + print(text) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--out", required=True, help="path stem of the .json and .md") + ap.add_argument("--codec", default="q4_k") + ap.add_argument("--cells", default=DEFAULT_CELLS) + ap.add_argument("--rounds", type=int, default=16) + ap.add_argument("--calls", type=int, default=6) + ap.add_argument( + "--tiles", + default="", + help="comma list of BMxBK tiles, such as 64x64,128x128d; each adds an " + "arm glu:BMxBK", + ) + ap.add_argument("--skip", default="", help="comma list of arms to leave out") + ap.add_argument("--cache-gib", type=int, default=2, help="MLX cache limit") + ap.add_argument("--report", action="store_true", help="rebuild the .md") + a = ap.parse_args() + if a.report: + with open(a.out + ".json") as fh: + report(json.load(fh), a.out + ".md") + return + a.cells = [c for c in a.cells.split(";") if c] + a.tiles = [t for t in a.tiles.split(",") if t] + if a.tiles: + from mlx_kquant._ext import _probe_kernels + + if not _probe_kernels(): + ap.error( + "--tiles needs a build with KQ_PROBE_KERNELS=ON: this build " + "has one tile per codec and reads no tile lever" + ) + a.skip = [s for s in a.skip.split(",") if s] + run(a) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_qk_norm_rope.py b/benchmarks/bench_qk_norm_rope.py new file mode 100644 index 00000000..4e017733 --- /dev/null +++ b/benchmarks/bench_qk_norm_rope.py @@ -0,0 +1,298 @@ +"""kq.qk_norm_rope against the unfused ops at DiT attention shapes. + +One call normalizes and rotates q and k of one attention layer, q and k +sliced from one fused qkv output as a DiT block produces them. Arms, all in +one process: + + composed mx.fast.rms_norm, then the rotation of interleaved pairs from + primitive ops in f32, then the transpose to [B, H, L, D] + stock_fused mx.fast.rms_norm, then mx.fast.rope(traditional=True). Its + angles follow the sequence index, not per-row tables, so it + computes a different rotation; it prices stock fused ops + kq kq.qk_norm_rope + copy a transposed copy of q and k to [B, H, L, D], the same bytes + as kq with no math, the practical bandwidth floor + +Order: one discarded warmup call per arm and shape, then rounds of 8 +Thue-Morse slots. Each slot visits every shape; a visit opens with one +discarded call per arm, then times each arm, in reversed order on odd +slots. Round 0 gets its own column and later rounds the settled column. + +A sample is --inner calls queued back to back with one sync at the end; +shapes whose fastest arm runs under --min-sample-ms per sample take more +calls, up to --max-inner, and every arm of a shape uses the same count. + +The bandwidth floor of a shape is the bytes kq must move, q and k read +once and written once plus the tables, at --gbs GB/s. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import bench_stats as bs # noqa: E402 +import dit_common as dc # noqa: E402 + +ARMS = ("composed", "stock_fused", "kq", "copy") +# (rows per item L, batch B): 1024-pixel Z-Image rows (4096 image rows plus +# 32 caption rows), 1024x1152-class rows, and 2048-pixel rows. +SHAPES = [(4128, 1), (4608, 1), (4128, 2), (4608, 2), (16416, 1)] +EPS = 1e-6 + + +def make_inputs(L, B, H, D, dt, seed=0): + import mlx.core as mx + import numpy as np + + rng = np.random.default_rng(seed) + key = mx.random.key(seed) + qkv = mx.random.normal((B, L, 3 * H * D), key=key).astype(dt) + qkv = qkv.reshape(B, L, 3, H, D) + ang = rng.uniform(0.0, 64.0, (L, D // 2)).astype(np.float32) + cos, sin = mx.array(np.cos(ang)), mx.array(np.sin(ang)) + w = (1.0 + 0.1 * rng.standard_normal((2, D))).astype(np.float32) + q_w, k_w = mx.array(w[0]).astype(dt), mx.array(w[1]).astype(dt) + mx.eval(qkv, cos, sin, q_w, k_w) + return { + "q": qkv[:, :, 0], + "k": qkv[:, :, 1], + "q_w": q_w, + "k_w": k_w, + "cos": cos, + "sin": sin, + } + + +def composed(q, k, q_w, k_w, cos, sin): + import mlx.core as mx + + def rope(x): + b, n, h, d = x.shape + x = x.astype(mx.float32).reshape(b, n, h, d // 2, 2) + x0, x1 = x[..., 0], x[..., 1] + c = cos[None, :, None, :] + s = sin[None, :, None, :] + out = mx.stack([x0 * c - x1 * s, x0 * s + x1 * c], axis=-1) + return out.reshape(b, n, h, d) + + dt = q.dtype + qo = rope(mx.fast.rms_norm(q, q_w, EPS)).astype(dt) + ko = rope(mx.fast.rms_norm(k, k_w, EPS)).astype(dt) + return qo.transpose(0, 2, 1, 3), ko.transpose(0, 2, 1, 3) + + +def make_fn(arm, x): + import mlx.core as mx + + import mlx_kquant as kq + + q, k = x["q"], x["k"] + if arm == "composed": + return lambda: composed(q, k, x["q_w"], x["k_w"], x["cos"], x["sin"]) + if arm == "stock_fused": + d = q.shape[-1] + + def f(): + qn = mx.fast.rms_norm(q, x["q_w"], EPS).transpose(0, 2, 1, 3) + kn = mx.fast.rms_norm(k, x["k_w"], EPS).transpose(0, 2, 1, 3) + return ( + mx.fast.rope(qn, d, traditional=True, base=256.0, scale=1.0, offset=0), + mx.fast.rope(kn, d, traditional=True, base=256.0, scale=1.0, offset=0), + ) + + return f + if arm == "kq": + return lambda: kq.qk_norm_rope( + q, k, x["q_w"], x["k_w"], x["cos"], x["sin"], EPS + ) + if arm == "copy": + return lambda: ( + mx.contiguous(q.transpose(0, 2, 1, 3)), + mx.contiguous(k.transpose(0, 2, 1, 3)), + ) + raise ValueError(arm) + + +def time_fn(fn, inner): + """ms per call over `inner` calls queued back to back, one sync at the + end. Each call's outputs are released once it is queued, so memory + stays at the calls in flight.""" + import mlx.core as mx + + mx.synchronize() + t0 = time.perf_counter() + for _ in range(inner): + mx.async_eval(fn()) + mx.synchronize() + return (time.perf_counter() - t0) * 1e3 / inner + + +def floor_bytes(L, B, H, D, itemsize): + return 4 * B * L * H * D * itemsize + 2 * L * (D // 2) * 4 + + +def run_dtype(dtype, args): + import mlx.core as mx + + dt = dc.mx_dtype(dtype) + want = args.shapes + shapes = [(L, B) for (L, B) in SHAPES if not want or f"{L}x{B}" in want] + inputs = {s: make_inputs(s[0], s[1], args.heads, args.head_dim, dt) for s in shapes} + fns = {(s, a): make_fn(a, inputs[s]) for s in shapes for a in args.arms} + + inner = {} + for s in shapes: + fastest = None + for a in args.arms: + time_fn(fns[(s, a)], 1) + ms = time_fn(fns[(s, a)], 3) + fastest = ms if fastest is None else min(fastest, ms) + n = int(args.min_sample_ms / max(fastest, 1e-3)) + 1 + inner[s] = max(1, min(args.max_inner, n)) + counts = {f"{L}x{B}": n for (L, B), n in inner.items()} + print(f"[{dtype}] inner {counts}", flush=True) + + samples = {(s, a): [] for s in shapes for a in args.arms} + for r in range(args.rounds): + for slot in range(8): + arms = list(args.arms) + if dc.THUE_MORSE[slot]: + arms.reverse() + for s in shapes: + for a in arms: + time_fn(fns[(s, a)], 1) + for a in arms: + samples[(s, a)].append((r, slot, time_fn(fns[(s, a)], inner[s]))) + print(f"[{dtype}] round {r} done", flush=True) + + cells = [] + for s in shapes: + L, B = s + fb = floor_bytes(L, B, args.heads, args.head_dim, mx.array(0, dt).itemsize) + floor_ms = fb / (args.gbs * 1e9) * 1e3 + rec = { + "L": L, + "B": B, + "rows": L * B, + "inner": inner[s], + "floor_ms": floor_ms, + "bytes": fb, + "arms": {}, + "ratios": {}, + } + for a in args.arms: + xs = samples[(s, a)] + rec["arms"][a] = { + "round0": bs.summarize([ms for r, _, ms in xs if r == 0]), + "settled": bs.summarize([ms for r, _, ms in xs if r > 0]), + "samples": xs, + } + if "kq" in args.arms: + for base in args.arms: + if base == "kq": + continue + pairs = [ + k[2] / b[2] + for k, b in zip(samples[(s, "kq")], samples[(s, base)], strict=True) + if k[0] > 0 + ] + rec["ratios"][f"kq/{base}"] = bs.summarize(pairs) + cells.append(rec) + return cells + + +def fmt(v, nd=3): + return "n/a" if v is None else f"{v:.{nd}f}" + + +def report_md(res): + a = res["args"] + lines = [ + "# qk_norm_rope at DiT attention shapes", + "", + dc.descriptor_md(res["machine"]), + "", + f"{a['heads']} heads of {a['head_dim']}, q and k sliced from one fused qkv " + f"output. {a['rounds']} rounds of 8 Thue-Morse slots. Round 0 is the first " + f"column and later rounds are settled. Times are median ms per call, and " + f"spread is the relative IQR of the settled samples. The floor is q and k " + f"read and written once plus the tables, at {a['gbs']} GB/s. The kq ratios " + f"to other arms are medians of settled pairs.", + "", + ] + for dtype, cells in res["dtypes"].items(): + lines += [f"## {dtype}", ""] + hdr = "| L x B | floor |" + sep = "|---|---|" + for arm in a["arms"]: + hdr += f" {arm} first | {arm} settled | {arm} spread |" + sep += "---|---|---|" + hdr += " kq / floor | kq / composed | kq / copy |" + sep += "---|---|---|" + lines += [hdr, sep] + for c in cells: + row = f"| {c['L']} x {c['B']} | {fmt(c['floor_ms'])} |" + for arm in a["arms"]: + st = c["arms"][arm] + r0 = st["round0"]["median"] if st["round0"] else None + s = st["settled"] or {} + iqr = s.get("rel_iqr") + spread = "n/a" if iqr is None else fmt(100 * iqr, 1) + "%" + row += f" {fmt(r0)} | {fmt(s.get('median'))} | {spread} |" + kq_s = (c["arms"].get("kq", {}).get("settled") or {}).get("median") + row += f" {fmt(kq_s / c['floor_ms'] if kq_s else None, 2)} |" + for base in ("composed", "copy"): + r = c["ratios"].get(f"kq/{base}") + row += f" {fmt(r['median'] if r else None, 3)} |" + lines.append(row) + lines.append("") + return "\n".join(lines) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--dtypes", default="bfloat16,float16") + ap.add_argument("--arms", default=",".join(ARMS)) + ap.add_argument("--shapes", default="", help="L x B cells, e.g. 4128x1,16416x1") + ap.add_argument("--heads", type=int, default=30) + ap.add_argument("--head-dim", type=int, default=128) + ap.add_argument("--rounds", type=int, default=3) + ap.add_argument("--min-sample-ms", type=float, default=10.0) + ap.add_argument("--max-inner", type=int, default=64) + ap.add_argument("--gbs", type=float, default=614.0, help="bandwidth for the floor") + ap.add_argument("--json-out", default="bench_qk_norm_rope.json") + ap.add_argument("--md-out", default="bench_qk_norm_rope.md") + args = ap.parse_args() + args.arms = [x for x in args.arms.split(",") if x] + args.shapes = [x for x in args.shapes.split(",") if x] + for arm in args.arms: + if arm not in ARMS: + ap.error(f"unknown arm {arm}") + + res = { + "args": {k: v for k, v in vars(args).items()}, + "machine": dc.machine_descriptor(), + "dtypes": {}, + } + t0 = time.perf_counter() + for dtype in args.dtypes.split(","): + res["dtypes"][dtype] = run_dtype(dtype, args) + res["seconds"] = time.perf_counter() - t0 + with open(args.json_out, "w") as f: + json.dump(res, f, indent=1) + md = report_md(res) + with open(args.md_out, "w") as f: + f.write(md) + print(md) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_roofline.py b/benchmarks/bench_roofline.py new file mode 100644 index 00000000..109f9a22 --- /dev/null +++ b/benchmarks/bench_roofline.py @@ -0,0 +1,331 @@ +"""Three floors per DiT cell from stock-op rates, against measured times. + +Reads bench_dit_shapes.py JSON and, per cell and arm, computes: + + NAX floor 2*M*K*N over the stock matmul rate at a cube of --cube + (per activation dtype) + ALU floor decoded weights N*K*ceil(M/BM), times the codec's ALU ops + per weight, over the rate of an fp32 multiply-add chain + written with mx.fast.metal_kernel + bytes floor compulsory traffic: x, w and out each moved once, over the + copy bandwidth of a large elementwise op + +and reports measured time over the largest floor, and which floor binds. +Two columns sit beside the floors without entering them: + + loads scalar byte loads per decoded weight (Q6_K: its 210-byte + blocks rule out vector loads), counted as issue slots at the + ALU rate + no reuse x_bytes*ceil(N/BN) + w_bytes*ceil(M/BM) + out_bytes, the + traffic with no cache reuse across tiles, an upper bound + +BM and BN come from the kernel each kq arm logged; the dense and affine +arms use MLX's dispatch rule. Ops per weight are hand counts of the NAX +loaders (kq_quantized_nax.h); where a count is a range the floor takes the +low end. The chain rate counts a multiply-add as one op, as the counts do, +so the ALU floor is an estimate; --ops-per-weight overrides a count. Rates +are measured here unless --rates names a JSON a previous run wrote. +""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import sys +import time + +import dit_common as dc + +# Per decoded weight in load_unsafe of each NAX loader: integer ops, +# conversions, float ops, table reads, plus per-call scale setup spread +# over the 32 weights a thread decodes. +OPS_PER_WEIGHT = { + "q8_0": 4.09, + "q5_1": 7.16, + "q4_0": 5.09, + "q4_1": 4.13, + "q5_0": 8.19, + "q4_k": 5.22, + "q5_k": 8.3, + "q6_k": 8.4, + "q3_k": 8.44, + "q2_k": 4.81, + "iq4_nl": 5.09, + "iq4_xs": 5.63, + "iq3_xxs": 5.76, + "iq3_s": 6.41, + "iq2_xxs": 7.44, + "iq2_xs": 5.97, + "iq2_s": 5.85, + "iq1_s": 6.01, + "iq1_m": 6.95, + "stq1_0": 9.21, + "pq2_0": 5.25, + "ptq1_0": 9.50, +} +# High end where a count is a range. +OPS_PER_WEIGHT_HIGH = {"q5_k": 9.3} +# Scalar byte loads per decoded weight, outside the ALU count. +SCALAR_LOADS_PER_WEIGHT = {"q6_k": 2.0} +# MLX affine: nibble or byte extract, two conversions, one multiply-add. +AFFINE_OPS = 4.0 + + +def _median_ms(fn, reps, warmup=3): + import mlx.core as mx + + for _ in range(warmup): + mx.eval(fn()) + ts = [] + for _ in range(reps): + t0 = time.perf_counter() + mx.eval(fn()) + ts.append((time.perf_counter() - t0) * 1e3) + return statistics.median(ts) + + +def measure_rates(dtypes, cube, alu_elems, alu_chain, bw_mib, reps): + import mlx.core as mx + + rates = {"matmul_flops": {}, "cube": cube} + for name in dtypes: + dt = dc.mx_dtype(name) + a = mx.random.normal((cube, cube)).astype(dt) + b = mx.random.normal((cube, cube)).astype(dt) + mx.eval(a, b) + ms = _median_ms(lambda a=a, b=b: a @ b, reps) + rates["matmul_flops"][name] = 2.0 * cube**3 / (ms * 1e-3) + + # Four independent multiply-add chains per thread in registers; the + # rate counts one op per multiply-add. + kernel = mx.fast.metal_kernel( + name="roofline_alu_chain", + input_names=["inp"], + output_names=["out"], + source=""" + uint i = thread_position_in_grid.x; + float a = inp[i]; + float b = a * 0.5f; + float c = a + 1.0f; + float d = a - 1.0f; + for (int k = 0; k < CHAIN; ++k) { + a = fma(a, 0.9990234375f, 0.0009765625f); + b = fma(b, 0.9990234375f, 0.0009765625f); + c = fma(c, 0.9990234375f, 0.0009765625f); + d = fma(d, 0.9990234375f, 0.0009765625f); + } + out[i] = a + b + c + d; + """, + ) + x = mx.random.uniform(shape=(alu_elems,)) + mx.eval(x) + + def run(): + return kernel( + inputs=[x], + template=[("CHAIN", alu_chain)], + grid=(alu_elems, 1, 1), + threadgroup=(256, 1, 1), + output_shapes=[x.shape], + output_dtypes=[mx.float32], + )[0] + + ms = _median_ms(run, reps) + alu_chain = 4 * alu_chain + rates["alu_ops"] = alu_elems * alu_chain / (ms * 1e-3) + rates["alu_chain"] = alu_chain + rates["alu_elems"] = alu_elems + + n = bw_mib * (1 << 20) // 2 + y = mx.random.normal((n,)).astype(mx.bfloat16) + mx.eval(y) + ms = _median_ms(lambda: mx.abs(y), reps) + rates["copy_bytes"] = 2.0 * y.nbytes / (ms * 1e-3) + rates["bw_mib"] = bw_mib + return rates + + +def floors(M, K, N, wbits, bm, bn, ops, loads, dtype, rates): + """Floors in ms, the binding one, and the two side columns.""" + nax = 2.0 * M * K * N / rates["matmul_flops"][dtype] + decoded = N * K * math.ceil(M / bm) + alu = decoded * ops / rates["alu_ops"] + x_bytes, w_bytes, out_bytes = M * K * 2, N * K * wbits / 8, M * N * 2 + compulsory = (x_bytes + w_bytes + out_bytes) / rates["copy_bytes"] + no_reuse = x_bytes * math.ceil(N / bn) + w_bytes * math.ceil(M / bm) + out_bytes + fl = {"nax": nax * 1e3, "alu": alu * 1e3, "bytes": compulsory * 1e3} + bind = max(fl, key=fl.get) + side = { + "loads": decoded * loads / rates["alu_ops"] * 1e3, + "no_reuse_bytes": no_reuse / rates["copy_bytes"] * 1e3, + } + return fl, bind, side + + +def settled(rec): + s = rec.get("settled") or rec.get("all") + return s["median"] if s else None + + +def cell_rows(result, rates, ops_over, arms): + out = [] + for p in result["passes"]: + dtype = p["dtype"] + for g in p["groups"]: + M, K, N = g["M"], g["K"], g["N"] + ref = [] + if "dense" in arms and "dense" in g["shared"]: + ref.append(("dense", None, g["shared"]["dense"], 16.0, 0.0, 0.0)) + for c in g["cells"]: + codec = c["codec"] + bits = dc.affine_bits(codec) + aff = g["shared"].get(f"affine{bits}") + if ( + "affine" in arms + and aff + and not any(r[0] == f"affine{bits}" for r in ref) + ): + ref.append( + (f"affine{bits}", None, aff, bits + 0.5, AFFINE_OPS, 0.0) + ) + ops = ops_over.get(codec, OPS_PER_WEIGHT.get(codec)) + loads = SCALAR_LOADS_PER_WEIGHT.get(codec, 0.0) + for arm in ("kq", "kq_bm64", "kq_bm128"): + rec = c["arms"].get(arm) + if arm in arms and rec and not rec.get("error"): + wb = dc.weight_bits(codec) + ref.append((arm, codec, rec, wb, ops, loads)) + for arm, codec, rec, wbits, ops, loads in ref: + ms = settled(rec) + bm, bn = dc.tile_bm_bn(rec.get("tile")) + if ms is None or bm is None: + continue + fl, bind, side = floors( + M, K, N, wbits, bm, bn, ops or 0.0, loads, dtype, rates + ) + out.append( + { + "dtype": dtype, + "family": g["family"], + "op": g["op"], + "L": g["L"], + "B": g["B"], + "M": M, + "K": K, + "N": N, + "arm": arm, + "codec": codec, + "tile": rec.get("tile"), + "ms": ms, + "floors_ms": fl, + "side_ms": side, + "binding": bind, + "over_floor": ms / fl[bind], + "ops_per_weight": ops, + "ops_per_weight_high": OPS_PER_WEIGHT_HIGH.get(codec), + "scalar_loads_per_weight": loads, + } + ) + return out + + +def markdown(rows, rates, desc): + lines = ["# DiT roofline", "", dc.descriptor_md(desc), ""] + mm = ", ".join( + f"{k} {v / 1e12:.1f} TFLOPS" for k, v in rates["matmul_flops"].items() + ) + lines.append( + f"Rates: matmul at {rates['cube']} cube {mm}; ALU chain " + f"{rates['alu_ops'] / 1e12:.2f} Tops/s; copy " + f"{rates['copy_bytes'] / 1e9:.0f} GB/s." + ) + lines.append("") + for dtype in dict.fromkeys(r["dtype"] for r in rows): + lines.append(f"## {dtype}") + lines.append("") + lines.append( + "| op | L x B | M | arm | codec | tile | ms | NAX | ALU | bytes " + "| binds | ms / floor | loads | no reuse |" + ) + lines.append("|---|---|---|---|---|---|---|---|---|---|---|---|---|---|") + for r in rows: + if r["dtype"] != dtype: + continue + f, sd = r["floors_ms"], r["side_ms"] + lines.append( + f"| {r['family']}.{r['op']} | {r['L']} x {r['B']} | {r['M']} " + f"| {r['arm']} | {r['codec'] or '-'} | {r['tile']} | {r['ms']:.3f} " + f"| {f['nax']:.3f} | {f['alu']:.3f} | {f['bytes']:.3f} " + f"| {r['binding']} | {r['over_floor']:.2f} " + f"| {sd['loads']:.3f} | {sd['no_reuse_bytes']:.3f} |" + ) + lines.append("") + return "\n".join(lines) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("inputs", nargs="+", help="bench_dit_shapes.py JSON files") + ap.add_argument("--rates", help="rates JSON from an earlier run") + ap.add_argument("--rates-out") + ap.add_argument("--cube", type=int, default=4096) + ap.add_argument("--alu-elems", type=int, default=1 << 24) + ap.add_argument("--alu-chain", type=int, default=1024) + ap.add_argument("--bw-mib", type=int, default=1024) + ap.add_argument("--reps", type=int, default=20) + ap.add_argument("--arms", default="kq,kq_bm64,kq_bm128,affine,dense") + ap.add_argument( + "--ops-per-weight", default="", help="overrides, e.g. q6_k=6,q5_k=7.5" + ) + ap.add_argument("--json-out") + ap.add_argument("--md-out") + args = ap.parse_args() + + results = [] + for path in args.inputs: + with open(path) as f: + results.append(json.load(f)) + ops_over = { + k: float(v) + for k, v in (kv.split("=") for kv in args.ops_per_weight.split(",") if kv) + } + dtypes = sorted({p["dtype"] for r in results for p in r["passes"]}) + desc = dc.machine_descriptor() + if args.rates: + with open(args.rates) as f: + rates = json.load(f)["rates"] + else: + rates = measure_rates( + dtypes, args.cube, args.alu_elems, args.alu_chain, args.bw_mib, args.reps + ) + if args.rates_out: + with open(args.rates_out, "w") as f: + json.dump({"machine": desc, "rates": rates}, f, indent=1) + arms = set(args.arms.split(",")) + rows = [] + for r in results: + if r["machine"].get("metallib_sha256") != desc["metallib_sha256"]: + print("warning: input was measured on another build", file=sys.stderr) + rows += cell_rows(r, rates, ops_over, arms) + out = {"machine": desc, "rates": rates, "ops_per_weight": OPS_PER_WEIGHT} + out["ops_per_weight_high"] = OPS_PER_WEIGHT_HIGH + out["scalar_loads_per_weight"] = SCALAR_LOADS_PER_WEIGHT + out["ops_overrides"] = ops_over + out["rows"] = rows + if args.json_out: + with open(args.json_out, "w") as f: + json.dump(out, f, indent=1) + md = markdown(rows, rates, desc) + if args.md_out: + with open(args.md_out, "w") as f: + f.write(md + "\n") + else: + print(md) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/bench_stats.py b/benchmarks/bench_stats.py new file mode 100644 index 00000000..5fcb1d8e --- /dev/null +++ b/benchmarks/bench_stats.py @@ -0,0 +1,59 @@ +"""Summary statistics for the benchmarks: medians, spreads and the +distribution-free CI of the median. + +Standard library only, with no side effects on import, so tests can load it +directly. +""" + +from __future__ import annotations + +import statistics +from fractions import Fraction + + +def median_ci(s, conf=0.95): + """Distribution-free CI of the median of sorted values `s` from order + statistics: [s(k), s(n+1-k)] with the largest k whose binomial tail + stays within (1 - conf) / 2. None when n is too small. The tail is + summed in integers, so it stays exact and finite for any n.""" + n = len(s) + alpha = Fraction(1.0 - conf).limit_denominator(10**6) / 2 + bound = alpha.numerator << n + k, cum, term = 0, 0, 1 + while k < n: + nxt = cum + term + if nxt * alpha.denominator > bound: + break + cum, term, k = nxt, term * (n - k) // (k + 1), k + 1 + if k < 1: + return None + return s[k - 1], s[n - k] + + +def summarize(samples): + """Median, relative IQR, p10, p90 and the 95 percent CI of the median, + with its half-width relative to the median.""" + if not samples: + return None + s = sorted(samples) + n = len(s) + + def q(p): + if n == 1: + return s[0] + i = p * (n - 1) + lo = int(i) + hi = min(lo + 1, n - 1) + return s[lo] + (s[hi] - s[lo]) * (i - lo) + + med = statistics.median(s) + ci = median_ci(s) + return { + "n": n, + "median": med, + "rel_iqr": (q(0.75) - q(0.25)) / med if med else None, + "p10": q(0.10), + "p90": q(0.90), + "ci95": list(ci) if ci else None, + "ci95_half": (ci[1] - ci[0]) / 2 / med if ci and med else None, + } diff --git a/benchmarks/dit_common.py b/benchmarks/dit_common.py new file mode 100644 index 00000000..0d0db267 --- /dev/null +++ b/benchmarks/dit_common.py @@ -0,0 +1,429 @@ +"""Shared pieces of the DiT benchmarks: shape matrix, weights, arms, kernel +log capture, Thue-Morse order and the machine descriptor. + +Imported by the DiT benchmarks in this directory, which run from the +repository root or from here. +""" + +from __future__ import annotations + +import hashlib +import os +import platform +import re +import sys +import tempfile +import time +import zlib + +# Tile levers must be live before mlx_kquant reads them on its first call. +os.environ["KQ_TUNE_LIVE"] = "1" + +TESTS_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tests" +) + +# (K, N) per linear. Rows per item are image tokens plus padded caption +# tokens at 1024 and 2048 pixels, each also run as a batch of two; Qwen-Image +# adds two text lengths that are not multiples of 32. The modulation linears +# take one row per image (the timestep embedding), so they run at M 1 and 2. +IMAGE_ROWS = [4096, 4128, 4192, 4256, 4608, 4640, 16384, 16896] +FAMILIES = { + "zimage": { + "ops": { + "qkv": (3840, 11520), + "out": (3840, 3840), + "w13": (3840, 10240), + "w2": (10240, 3840), + "adaln": (256, 15360), + }, + "rows": IMAGE_ROWS, + "op_rows": {"adaln": [1]}, + }, + "qwen_image": { + "ops": { + "qkvo": (4096, 4096), + "gate_up": (4096, 24576), + "mlp_out": (12288, 4096), + "mod": (4096, 16384), + }, + "rows": IMAGE_ROWS + [4096 + 77, 4096 + 213], + "op_rows": {"mod": [1]}, + }, +} +# The two heaviest cells by FLOPs per step, where every NAX codec runs. +HEAVY = [("zimage", "w13", 4608, 1), ("qwen_image", "gate_up", 4096, 1)] +MAIN_CODECS = ["q4_k", "q5_k", "q6_k", "q8_0"] + +# t[i] = parity of popcount(i): balances linear drift without the quadratic +# aliasing of ABBA blocks. +THUE_MORSE = [bin(i).count("1") & 1 for i in range(8)] + +TILE_BM64 = "64,64,2,2" +TILE_BM128 = "128,64,2,2" +ARM_TILE_FORCED = ("kq_bm64", "kq_bm128") + +DTYPES = ("bfloat16", "float16") + + +def mx_dtype(name): + import mlx.core as mx + + return {"bfloat16": mx.bfloat16, "float16": mx.float16, "float32": mx.float32}[name] + + +def shrink_dim(v, f, mult): + """v / f rounded up to a multiple of `mult`, at least `mult`.""" + if f <= 1: + return v + return max(mult, -(-(v // f) // mult) * mult) + + +def nax_codecs(): + """Codecs with NAX qmm_t kernels in the bundled metallib, independent of + whether this GPU has NAX, so every machine runs the same list.""" + import mlx_kquant as kq + from mlx_kquant.codec_geometry import CODEC_GEOMETRY + + pat = re.compile(r"^kquant_(.+?)_qmm_t_nax_") + found = {m[1] for n in kq.metallib_kernels() if (m := pat.match(n))} + return sorted(found & set(CODEC_GEOMETRY)) + + +def kq_weight(codec, n, k, seed): + """Quantized (N, K) weight. Codecs the tests cannot encode get random + structurally valid wire with a sane scale; the rest quantize a normal + matrix, so the decode sees realistic scales.""" + import mlx.core as mx + import numpy as np + + import mlx_kquant as kq + from mlx_kquant.codec_geometry import CODEC_GEOMETRY, DECODE_ONLY_CODECS + + if TESTS_DIR not in sys.path: + sys.path.insert(0, TESTS_DIR) + from kqref import is_synth, synth_wire + + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + if is_synth(codec) or codec in DECODE_ONLY_CODECS: + rng = np.random.default_rng(seed) + wire = synth_wire(rng, codec, bpb, n * (k // wpb)) + w = mx.array(wire.reshape(n, (k // wpb) * bpb)) + s = mx.zeros((1,), dtype=mx.uint8) + else: + wf = mx.random.normal((n, k), key=mx.random.key(seed)) * 0.02 + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + return w, s + + +def weight_bits(codec): + """Stored bits per weight from the codec geometry.""" + from mlx_kquant.codec_geometry import CODEC_GEOMETRY + + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + return 8.0 * bpb / wpb + + +def affine_bits(codec): + return 8 if codec == "q8_0" else 4 + + +class KernelLog: + """Capture the KQ_QMM_LOG lines (kernel name, M, N, K, grid) printed + while a block runs, by pointing fd 2 at a temporary file.""" + + LINE = re.compile( + r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+)$" + ) + + def __enter__(self): + sys.stderr.flush() + self._tmp = tempfile.TemporaryFile(mode="w+b") + self._saved = os.dup(2) + os.dup2(self._tmp.fileno(), 2) + self._prev = os.environ.get("KQ_QMM_LOG") + os.environ["KQ_QMM_LOG"] = "1" + self.kernels = [] + return self + + def __exit__(self, *exc): + if self._prev is None: + os.environ.pop("KQ_QMM_LOG", None) + else: + os.environ["KQ_QMM_LOG"] = self._prev + sys.stderr.flush() + os.dup2(self._saved, 2) + os.close(self._saved) + self._tmp.seek(0) + other = [] + for line in self._tmp.read().decode(errors="replace").splitlines(): + m = self.LINE.match(line) + if m: + self.kernels.append(m[1]) + elif line.strip(): + other.append(line) + self._tmp.close() + if other: + sys.stderr.write("\n".join(other) + "\n") + return False + + +def tile_of(kname): + """'bmX bnY wmA wnB' from a kq qmm kernel name; the default ALU tile has + no tag in its name.""" + if kname is None: + return None + m = re.search(r"_bm(\d+)_bn(\d+)(?:_bk\d+)?_wm(\d+)_wn(\d+)", kname) + db = "_db" if kname.endswith("_db") else "" + if m: + return f"bm{m[1]} bn{m[2]} wm{m[3]} wn{m[4]}{db}" + if "_qmm_t_" in kname: + return "bm64 bn64 wm2 wn2" + return None + + +def tile_bm_bn(tile): + m = re.match(r"bm(\d+) bn(\d+)", tile or "") + return (int(m[1]), int(m[2])) if m else (None, None) + + +def mlx_dense_tile(M, N, K, arch, nax): + """Tile of MLX's GEMM for a half or bfloat x @ w.T, derived from its + dispatch rule; MLX does not log it.""" + letter = arch[-1:] if arch else "" + if nax: + if letter in ("s", "c", "d"): + bk = 64 if (K >= 8192 and K > M + N) else 256 + return f"bm64 bn128 bk{bk} wm2 wn4 (rule)" + return "bm128 bn128 bk512 wm4 wn4 (rule)" + if letter in ("g", "p"): + return "bm64 bn32 bk32 wm2 wn2 (rule)" + if letter == "d": + if M * N >= 1 << 20 and 2 * max(M, N) > K: + return "bm64 bn64 bk16 wm1 wn2 (rule)" + return "bm64 bn32 bk32 wm2 wn2 (rule)" + return "bm64 bn64 bk16 wm2 wn2 (rule)" + + +def mlx_affine_tile(nax): + """Tile of MLX's affine quantized matmul, from its dispatch rule.""" + if nax: + return "bm64 bn64 bk64 wm2 wn2 (rule)" + return "bm32 bn32 bk32 wm2 wn2 (rule)" + + +ARM_TILE = {"kq": None, "kq_aa": None, "kq_bm64": TILE_BM64} +ARM_TILE["kq_bm128"] = TILE_BM128 + + +def set_tile(tile, nax): + os.environ.pop("KQ_NAX_TILE", None) + os.environ.pop("KQ_QMM_TILE", None) + if tile: + os.environ["KQ_NAX_TILE" if nax else "KQ_QMM_TILE"] = tile + + +class Store: + """Weights and activations, built on first use.""" + + def __init__(self, seed): + self.seed = seed + self.kqw = {} + self.dense = {} + self.affine = {} + self.x = {} + + def clear_dtype(self): + self.dense.clear() + self.affine.clear() + self.x.clear() + + def kq(self, g, codec): + key = (g["family"], g["op"], codec) + if key not in self.kqw: + h = zlib.crc32("/".join(key).encode()) % 100000 + self.kqw[key] = kq_weight(codec, g["N"], g["K"], self.seed + h) + return self.kqw[key] + + def dense_w(self, g, dt): + import mlx.core as mx + + key = (g["family"], g["op"]) + if key not in self.dense: + w = mx.random.normal((g["N"], g["K"]), key=mx.random.key(self.seed)) + self.dense[key] = (w * 0.02).astype(dt) + mx.eval(self.dense[key]) + return self.dense[key] + + def affine_w(self, g, dt, bits): + import mlx.core as mx + + key = (g["family"], g["op"], bits) + if key not in self.affine: + q = mx.quantize(self.dense_w(g, dt), group_size=64, bits=bits) + mx.eval(*q) + self.affine[key] = q + return self.affine[key] + + def act(self, g, dt): + import mlx.core as mx + + key = (g["K"], g["L"], g["B"]) + if key not in self.x: + shape = (g["B"], g["L"], g["K"]) + x = mx.random.normal(shape, key=mx.random.key(self.seed + 1)) * 0.5 + self.x[key] = x.astype(dt) + mx.eval(self.x[key]) + return self.x[key] + + +def make_call(arm, codec, g, dt, store): + """(x0, f): the arm's op as a function of its activation, and the + activation of the shape.""" + import mlx.core as mx + + import mlx_kquant as kq + + x = store.act(g, dt) + if arm == "dense": + wd = store.dense_w(g, dt) + return x, lambda a: a @ wd.T + if arm.startswith("affine"): + bits = int(arm[len("affine") :]) + wq, sc, bi = store.affine_w(g, dt, bits) + return x, lambda a: mx.quantized_matmul( + a, wq, sc, bi, transpose=True, group_size=64, bits=bits + ) + w, s = store.kq(g, codec) + if arm == "deq": + return x, lambda a: a @ kq.dequantize(w, s, codec, dtype=dt).T + return x, lambda a: kq.quantized_matmul(a, w, s, codec, transpose=True) + + +def make_fn(arm, codec, g, dt, store): + x, f = make_call(arm, codec, g, dt, store) + return lambda: f(x) + + +def tie(o, x0): + """x0 again, computed from the first element of `o`, so a call fed the + result starts after the call that made `o` has finished.""" + import mlx.core as mx + + return mx.where(o.reshape(-1)[:1] > 1e30, 0.0, x0) + + +def chain(f, x0, n): + """n calls of f, each fed the tie of the previous output: one lazy + dependent chain.""" + x = x0 + for _ in range(n): + x = tie(f(x), x0) + return x + + +def chain_ms(f, x0, n): + """ms per call over a chain of n calls in one evaluation. The GPU runs + the calls one after another, as in a model forward; independent calls in + one evaluation would overlap.""" + import mlx.core as mx + + x = chain(f, x0, n) + t0 = time.perf_counter() + mx.eval(x) + return (time.perf_counter() - t0) * 1e3 / n + + +def time_fn(fn, inner): + """ms per call over `inner` calls evaluated together, so one sync (about + 0.17 ms) is spread over all of them.""" + import mlx.core as mx + + t0 = time.perf_counter() + ys = [fn() for _ in range(inner)] + mx.eval(*ys) + ms = (time.perf_counter() - t0) * 1e3 / inner + del ys + return ms + + +def op_rows(family, op, rows=None): + """Rows per item for an op: its fixed rows if it has them, else `rows` + or the family's image rows.""" + spec = FAMILIES[family] + return spec["op_rows"].get(op) or rows or spec["rows"] + + +def cell(family, op, rows, batch, shrink=1): + """One shape visit: family op at `rows` per item and `batch` items.""" + K, N = FAMILIES[family]["ops"][op] + L = max(1, rows // shrink) if shrink > 1 else rows + return { + "family": family, + "op": op, + "K": shrink_dim(K, shrink, 256), + "N": shrink_dim(N, shrink, 64), + "L": L, + "B": batch, + "M": batch * L, + "L_model": rows, + "vector": op in FAMILIES[family]["op_rows"], + } + + +def metallib_sha256(): + import mlx_kquant as kq + + path = os.path.join(kq.metallib_dir(), "mlx_kquant.metallib") + if not os.path.exists(path): + return None + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def machine_descriptor(): + """Chip, arch generation and letter, GPU cores, NAX, threadgroup memory, + OS, versions, metallib hash and the tuning environment.""" + import mlx.core as mx + + import mlx_kquant as kq + + info = mx.device_info() + arch = info.get("architecture", "") + m = re.match(r"applegpu_g(\d+)(\w)", arch) + env = { + k: v + for k, v in sorted(os.environ.items()) + if k.startswith(("KQ_", "MLX_", "KQUANT_")) + } + return { + "device": info.get("device_name"), + "architecture": arch, + "arch_gen": int(m[1]) if m else None, + "arch_letter": m[2] if m else None, + "gpu_cores": kq.gpu_core_count(), + "nax": kq.nax_available(), + "max_threadgroup_memory": kq.max_threadgroup_memory(), + "memory_bytes": info.get("memory_size"), + "os": platform.mac_ver()[0] or platform.platform(), + "python": platform.python_version(), + "mlx": mx.__version__, + "mlx_kquant": kq.__version__, + "metallib_sha256": metallib_sha256(), + "tf32": os.environ.get("MLX_ENABLE_TF32", "default"), + "env": env, + } + + +def descriptor_md(d): + return ( + f"Device: {d['device']} ({d['architecture']}, gen {d['arch_gen']}, " + f"letter {d['arch_letter']}), {d['gpu_cores']} GPU cores, " + f"NAX {d['nax']}, threadgroup memory {d['max_threadgroup_memory']} B. " + f"macOS {d['os']}, mlx {d['mlx']}, mlx-kquant {d['mlx_kquant']}, " + f"metallib {str(d['metallib_sha256'])[:12]}, TF32 {d['tf32']}." + ) diff --git a/benchmarks/probe_dit_numerics.py b/benchmarks/probe_dit_numerics.py new file mode 100644 index 00000000..1097123e --- /dev/null +++ b/benchmarks/probe_dit_numerics.py @@ -0,0 +1,269 @@ +"""Numerics probes against float64, with MLX_ENABLE_TF32 at 0 and at 1. + +MLX reads MLX_ENABLE_TF32 once per process, so the parent runs one child +per setting and compares them. + + reduction K 12288: each row of x holds one large term, then small + terms of either sign. Arms: stock matmul in f32, bf16 and + fp16; MLX affine 4-bit in f32 and bf16; kq q8_0 and q4_k in + bf16 and fp16 (kq casts f32 activations to bf16, so it has + no f32 arm). + same sign as reduction, with every small term and weight positive, so + the partial sums only grow and accumulation error shows. + f32 rows f32 matmul at M 1 and M 2 over --f32-shapes (timestep and + modulation widths). TF32 applies to products of more than one + row, so bits should match across the settings at M 1 only. + attention 16896 keys at head dim 128 whose scores are near uniform. + Arms: mx.fast.scaled_dot_product_attention and the composed + softmax(q k^T) v, in f32, bf16 and fp16. + calibration kq q4_k, q5_k, q6_k and q8_0 at M 256, K 3840, N 512 in bf16 + and fp16, the observed R1 and R2 against the R1 bounds the + tests use. + +References use float64 over the inputs as the arm sees them. R1 uses the +exact dequantized weights (kq.dequantize, mx.dequantize); R2 uses those +weights rounded to the activation dtype. Each arm reports relative +Frobenius error for both, max abs error, the error of rounding the exact +result to the output dtype (the floor), and whether the output bits match +between the two settings. Markdown and JSON out. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import sys + +import dit_common as dc + +CHILD = "--child" +# The R1 bounds tests/test_tile_override.py uses. +R1_BOUND = {"bfloat16": 6e-3, "float16": 1e-3} + + +def child(args): + import mlx.core as mx + import numpy as np + + import mlx_kquant as kq + + out = [] + + def f64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + def record(probe, arm, dtype, y, ref, ref2=None, m=None): + mx.eval(y) + got = f64(y) + err = got - ref + rounded = f64(mx.array(ref.astype(np.float32)).astype(y.dtype)) + nref = np.linalg.norm(ref) + rel2 = None if ref2 is None else float(np.linalg.norm(got - ref2) / nref) + out.append( + { + "probe": probe, + "arm": arm, + "dtype": dtype, + "m": m if m is not None else int(np.prod(y.shape[:-1])), + "rel": float(np.linalg.norm(err) / nref), + "rel_r2": rel2, + "max_abs": float(np.abs(err).max()), + "floor": float(np.linalg.norm(rounded - ref) / nref), + "sha": hashlib.sha256(np.array(y.view(mx.uint8)).tobytes()).hexdigest(), + } + ) + + def reduction(probe, x, w): + for name in ("float32", "bfloat16", "float16"): + dt = dc.mx_dtype(name) + xd, wd = x.astype(dt), w.astype(dt) + ref = f64(xd) @ f64(wd).T + record(probe, "matmul", name, xd @ wd.T, ref, ref) + for name in ("float32", "bfloat16"): + dt = dc.mx_dtype(name) + xd, wd = x.astype(dt), w.astype(dt) + wq, sc, bi = mx.quantize(wd, group_size=64, bits=4) + deq = mx.dequantize(wq, sc, bi, group_size=64, bits=4) + y = mx.quantized_matmul( + xd, wq, sc, bi, transpose=True, group_size=64, bits=4 + ) + ref = f64(xd) @ f64(deq.astype(mx.float32)).T + ref2 = f64(xd) @ f64(deq.astype(dt)).T + record(probe, "affine4", name, y, ref, ref2) + for codec in ("q8_0", "q4_k"): + wq, s = kq.quantize(w, codec) + deq = kq.dequantize(wq, s, codec, dtype=mx.float32) + for name in ("bfloat16", "float16"): + dt = dc.mx_dtype(name) + xd = x.astype(dt) + y = kq.quantized_matmul(xd, wq, s, codec, transpose=True) + ref = f64(xd) @ f64(deq).T + ref2 = f64(xd) @ f64(deq.astype(dt)).T + record(probe, f"kq {codec}", name, y, ref, ref2) + + # Reduction over K with one large leading term. + M, K, N = args.m, args.k, args.n + x = mx.random.normal((M, K), key=mx.random.key(1)) * 0.01 + x[:, 0] = 1024.0 + w = mx.random.normal((N, K), key=mx.random.key(2)) * 0.05 + mx.eval(x, w) + reduction("reduction", x, w) + reduction("same sign", mx.abs(x), mx.abs(w)) + + # f32 products at one and two rows. + for shape in args.f32_shapes.split(","): + k, n = (int(v) for v in shape.split("x")) + wf = mx.random.normal((n, k), key=mx.random.key(7)) * 0.05 + for m in (1, 2): + xf = mx.random.normal((m, k), key=mx.random.key(8)) + ref = f64(xf) @ f64(wf).T + record(f"f32 rows {k}x{n}", "matmul", "float32", xf @ wf.T, ref, ref, m) + + # Attention over near-uniform keys. + H, Lq, Lk, D = args.heads, args.lq, args.lk, 128 + q = mx.random.normal((1, H, Lq, D), key=mx.random.key(3)) + base = mx.random.normal((1, H, 1, D), key=mx.random.key(4)) + k = base + 1e-3 * mx.random.normal((1, H, Lk, D), key=mx.random.key(5)) + v = mx.random.normal((1, H, Lk, D), key=mx.random.key(6)) + mx.eval(q, k, v) + scale = D**-0.5 + for name in ("float32", "bfloat16", "float16"): + dt = dc.mx_dtype(name) + qd, kd, vd = q.astype(dt), k.astype(dt), v.astype(dt) + qn, kn, vn = f64(qd), f64(kd), f64(vd) + sc = (qn @ np.swapaxes(kn, -1, -2)) * scale + sc -= sc.max(axis=-1, keepdims=True) + p = np.exp(sc) + p /= p.sum(axis=-1, keepdims=True) + ref = p @ vn + y = mx.fast.scaled_dot_product_attention(qd, kd, vd, scale=scale) + record("attention", "sdpa", name, y, ref) + s2 = (qd * scale) @ mx.swapaxes(kd, -1, -2) + y2 = mx.softmax(s2, axis=-1, precise=True) @ vd + record("attention", "composed", name, y2, ref) + + # R1 and R2 calibration at a DiT-like shape. + xc = mx.random.normal((256, 3840), key=mx.random.key(9)) * 0.5 + wc = mx.random.normal((512, 3840), key=mx.random.key(10)) * 0.02 + for codec in ("q4_k", "q5_k", "q6_k", "q8_0"): + wq, s = kq.quantize(wc, codec) + deq = kq.dequantize(wq, s, codec, dtype=mx.float32) + for name in ("bfloat16", "float16"): + dt = dc.mx_dtype(name) + xd = xc.astype(dt) + y = kq.quantized_matmul(xd, wq, s, codec, transpose=True) + ref = f64(xd) @ f64(deq).T + ref2 = f64(xd) @ f64(deq.astype(dt)).T + record("calibration", f"kq {codec}", name, y, ref, ref2) + print(json.dumps(out)) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--m", type=int, default=64) + ap.add_argument("--k", type=int, default=12288) + ap.add_argument("--n", type=int, default=256) + ap.add_argument("--heads", type=int, default=2) + ap.add_argument("--lq", type=int, default=128) + ap.add_argument("--lk", type=int, default=16896) + ap.add_argument( + "--f32-shapes", + default="256x1024,256x15360,4096x16384", + help="KxN list for the f32 rows probe", + ) + ap.add_argument(CHILD, action="store_true", help=argparse.SUPPRESS) + ap.add_argument("--json-out") + ap.add_argument("--md-out") + args = ap.parse_args() + if args.child: + child(args) + return + + argv = [a for a in sys.argv[1:] if a != CHILD] + runs = {} + for tf32 in ("0", "1"): + env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + env["MLX_ENABLE_TF32"] = tf32 + r = subprocess.run( + [sys.executable, os.path.abspath(__file__), CHILD, *argv], + env=env, + capture_output=True, + text=True, + ) + if r.returncode != 0: + sys.exit(r.stderr[-4000:]) + runs[tf32] = json.loads(r.stdout.strip().splitlines()[-1]) + desc = dc.machine_descriptor() + rows = [] + for a, b in zip(runs["0"], runs["1"], strict=True): + rows.append( + { + "probe": a["probe"], + "arm": a["arm"], + "dtype": a["dtype"], + "m": a["m"], + "tf32_0": {k: a[k] for k in ("rel", "rel_r2", "max_abs")}, + "tf32_1": {k: b[k] for k in ("rel", "rel_r2", "max_abs")}, + "floor": a["floor"], + "bits_equal": a["sha"] == b["sha"], + } + ) + result = {"machine": desc, "args": vars(args), "rows": rows} + lines = ["# DiT numerics probes", "", dc.descriptor_md(desc), ""] + lines.append( + f"Reduction M {args.m}, K {args.k}, N {args.n}; attention {args.heads} " + f"heads, {args.lq} queries, {args.lk} keys, head dim 128. R1 and R2 are " + "relative Frobenius errors against float64 with exact and with " + "dtype-rounded weights; floor is the error of rounding the exact " + "result to the output dtype." + ) + lines.append("") + lines.append( + "| probe | arm | dtype | M | R1 TF32 0 | R1 TF32 1 | R2 TF32 0 | R2 TF32 1 " + "| max abs TF32 0 | max abs TF32 1 | floor | bits equal |" + ) + lines.append("|---|---|---|---|---|---|---|---|---|---|---|---|") + + def e(v): + return "-" if v is None else f"{v:.2e}" + + for r in rows: + t0, t1 = r["tf32_0"], r["tf32_1"] + lines.append( + f"| {r['probe']} | {r['arm']} | {r['dtype']} | {r['m']} " + f"| {e(t0['rel'])} | {e(t1['rel'])} | {e(t0['rel_r2'])} " + f"| {e(t1['rel_r2'])} | {e(t0['max_abs'])} | {e(t1['max_abs'])} " + f"| {e(r['floor'])} | {'yes' if r['bits_equal'] else 'no'} |" + ) + lines.append("") + lines.append("Calibration against the R1 bounds of the tests:") + lines.append("") + lines.append("| codec | dtype | R1 | R1 bound | R1 / bound | R2 |") + lines.append("|---|---|---|---|---|---|") + for r in rows: + if r["probe"] != "calibration": + continue + bound = R1_BOUND[r["dtype"]] + rel = r["tf32_1"]["rel"] + lines.append( + f"| {r['arm'].split()[-1]} | {r['dtype']} | {e(rel)} | {bound:.0e} " + f"| {rel / bound:.2f} | {e(r['tf32_1']['rel_r2'])} |" + ) + md = "\n".join(lines) + if args.json_out: + with open(args.json_out, "w") as f: + json.dump(result, f, indent=1) + if args.md_out: + with open(args.md_out, "w") as f: + f.write(md + "\n") + else: + print(md) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/probe_route_numerics.py b/benchmarks/probe_route_numerics.py new file mode 100644 index 00000000..9ecadbf9 --- /dev/null +++ b/benchmarks/probe_route_numerics.py @@ -0,0 +1,223 @@ +"""Where the quantized matmul routes round, against float64, per M band. + +Every route can be modeled one of two ways: + + float weights decode to f32 in registers, accumulate in f32, round + the output once (qmv, mv_ext). Its reference R1 uses + the exact dequantized weights. + dtype weights decode into the activation dtype for the matrix units, + accumulate in f32, round the output once (the qmm + tiles). Its reference R2 uses the dequantized weights + rounded to the activation dtype. + +A route that matches one model lands on that reference's floor, the error +of rounding the reference itself to the output dtype, and its excess over +the floor is near 1e-7. A positive excess on both models means a rounding +the models do not have, such as a K slice stored in the activation dtype +before the fold. The excess is sqrt(R^2 - floor^2). + +Inputs are nonnegative, so partial sums only grow. In the lead design +each row of x opens with one large term (1024), so the first K slice holds +most of the output and a rounding of that slice shows fully; the flat +design has no large term. Arms per cell: stock matmul and MLX affine +4-bit in the activation dtype, kq per codec on the default route, and kq +per forced route (KQ_QMM_ROUTE through M 32, the plain tile through +KQ_QMM_MIDM=0 from M 33 to 128). A forced route that cannot serve the cell +is left out. Markdown and JSON out. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +import dit_common as dc + +ROUTES = [ + "qmv", + "verify_qmv", + "mv_ext", + "verify_mma", + "verify_nax", + "splitk", + "nax_splitk", + "nax", +] +DEFAULT_CELLS = ( + # K sweep at a mid-M split band row count and at a DiT prefill row count. + "64:1024:256,64:3072:256,64:12288:256," + "4096:1024:256,4096:3072:256,4096:12288:256," + # The small-M bands at one K. + "2:12288:256,8:12288:256,16:12288:256,32:12288:256,48:12288:256," + "96:12288:256,128:12288:256,129:12288:256," + # DiT modulation widths at one and two rows. + "1:256:15360,2:256:15360,1:4096:16384,2:4096:16384" +) + + +def probe(args): + import mlx.core as mx + import numpy as np + + import mlx_kquant as kq + + dt = dc.mx_dtype(args.dtype) + os.environ["KQ_QMM_ROUTE_STRICT"] = "1" + + def f64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + def rel(a, b): + return float(np.linalg.norm(a - b) / np.linalg.norm(b)) + + def rounded(ref): + return f64(mx.array(ref.astype(np.float32)).astype(dt)) + + def run(fn, env): + saved = {k: os.environ.get(k) for k in env} + os.environ.update(env) + try: + with dc.KernelLog() as log: + y = fn() + mx.eval(y) + finally: + for k, v in saved.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + return y, (log.kernels[-1] if log.kernels else None) + + out = [] + for design in args.designs.split(","): + for spec in args.cells.split(","): + M, K, N = (int(v) for v in spec.split(":")) + w = mx.abs(mx.random.normal((N, K), key=mx.random.key(2))) * 0.05 + x = mx.abs(mx.random.normal((M, K), key=mx.random.key(1))) * 0.01 + if design == "lead": + x[:, 0] = 1024.0 + xd = x.astype(dt) + mx.eval(xd, w) + xn = f64(xd) + wd = w.astype(dt) + arms = [ + ("matmul", "-", lambda xd=xd, wd=wd: xd @ wd.T, f64(wd), f64(wd), {}) + ] + wq, sc, bi = mx.quantize(wd, group_size=64, bits=4) + deq = mx.dequantize(wq, sc, bi, group_size=64, bits=4) + arms.append( + ( + "affine4", + "-", + lambda xd=xd, wq=wq, sc=sc, bi=bi: mx.quantized_matmul( + xd, wq, sc, bi, transpose=True, group_size=64, bits=4 + ), + f64(deq.astype(mx.float32)), + f64(deq.astype(dt)), + {}, + ) + ) + for codec in args.codecs.split(","): + kw, ks = kq.quantize(w, codec) + kdeq = kq.dequantize(kw, ks, codec, dtype=mx.float32) + w1, w2 = f64(kdeq), f64(kdeq.astype(dt)) + + def fn(xd=xd, kw=kw, ks=ks, codec=codec): + return kq.quantized_matmul(xd, kw, ks, codec, transpose=True) + + arms.append((f"kq {codec}", "default", fn, w1, w2, {})) + if M <= 32: + for route in ROUTES: + arms.append( + (f"kq {codec}", route, fn, w1, w2, {"KQ_QMM_ROUTE": route}) + ) + elif M <= 128: + arms.append( + (f"kq {codec}", "tile", fn, w1, w2, {"KQ_QMM_MIDM": "0"}) + ) + refs = {} + for arm, route, fn, w1, w2, env in arms: + try: + y, kname = run(fn, env) + except RuntimeError as e: + # A forced route that cannot serve the cell raises under + # KQ_QMM_ROUTE_STRICT; any other error, such as a + # missing kernel, is a failure. + if "does not serve" not in str(e): + raise + continue + got = f64(y) + row = { + "design": design, + "M": M, + "K": K, + "N": N, + "arm": arm, + "route": route, + "kernel": kname, + } + for tag, wref in (("1", w1), ("2", w2)): + key = (id(wref), tag) + if key not in refs: + ref = xn @ wref.T + refs[key] = (ref, rel(rounded(ref), ref)) + ref, floor = refs[key] + r = rel(got, ref) + row[f"r{tag}"] = r + row[f"floor{tag}"] = floor + row[f"excess{tag}"] = float(np.sqrt(max(0.0, r * r - floor**2))) + out.append(row) + return out + + +def markdown(rows, args): + lines = [ + "# Route numerics", + "", + f"Activation dtype {args.dtype}. R1 and R2 are relative Frobenius errors " + "against float64 with exact and with dtype-rounded dequantized weights; " + "floor is the error of rounding that reference to the output dtype; " + "excess is sqrt(R^2 - floor^2). The route column is `-` for the stock " + "arms and names the forced route for kq.", + "", + "| design | M | K | N | arm | route | R1 | excess 1 | R2 | excess 2 | kernel |", + "|---|---|---|---|---|---|---|---|---|---|---|", + ] + for r in rows: + kname = r["kernel"] or "-" + lines.append( + f"| {r['design']} | {r['M']} | {r['K']} | {r['N']} | {r['arm']} " + f"| {r['route']} | {r['r1']:.2e} | {r['excess1']:.1e} " + f"| {r['r2']:.2e} | {r['excess2']:.1e} | {kname} |" + ) + return "\n".join(lines) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--cells", default=DEFAULT_CELLS, help="M:K:N list") + ap.add_argument("--designs", default="lead,flat") + ap.add_argument("--codecs", default="q8_0,q4_k") + ap.add_argument("--dtype", default="bfloat16") + ap.add_argument("--json-out") + ap.add_argument("--md-out") + args = ap.parse_args() + rows = probe(args) + if args.json_out: + with open(args.json_out, "w") as f: + json.dump({"machine": dc.machine_descriptor(), "rows": rows}, f) + md = dc.descriptor_md(dc.machine_descriptor()) + "\n\n" + markdown(rows, args) + if args.md_out: + with open(args.md_out, "w") as f: + f.write(md + "\n") + else: + print(md) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/bindings.cpp b/bindings.cpp index 7573bb8e..260a009e 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -66,6 +66,12 @@ NB_MODULE(_ext, m) { &mlx_kquant::metallib_loads, "Load the bundled metallib via the Metal device (toolchain self-check)."); + m.def( + "metallib_kernels", + &mlx_kquant::metallib_kernels, + "Sorted names of the kernels in the bundled metallib (empty on builds " + "without Metal). Tile overrides select from these names."); + m.def( "cpu_neon_available", &mlx_kquant::kq_cpu_neon_available, @@ -82,6 +88,61 @@ NB_MODULE(_ext, m) { &mlx_kquant::gpu_core_count, "GPU core count from IOKit (KQ_GPU_CORES overrides it); 0 when unknown."); + m.def( + "max_threadgroup_memory", + &mlx_kquant::max_threadgroup_memory, + "The GPU's threadgroup memory limit in bytes; 0 without Metal."); + + m.def( + "_probe_kernels", + &mlx_kquant::probe_kernels, + "True in a build with the CMake option KQ_PROBE_KERNELS=ON, for the " + "kernel name tests."); + + m.def( + "_split_kernel_names", + &mlx_kquant::split_kernel_names, + "Every kernel name the split-K dispatchers can build, for the kernel " + "name tests (empty without Metal)."); + + m.def( + "_nax_large_row", + &mlx_kquant::nax_large_row, + "kquant_type"_a, + "The codec's NAX qmm_t large-M row as [bm, bn, wm, wn, swizzle_log, " + "min_m, min_n, max_k, inherited], for the policy tests (empty without " + "Metal)."); + + m.def( + "_alu_dequantize_row", + &mlx_kquant::alu_dequantize_row, + "arch"_a, + "nax"_a, + "kquant_type"_a, + "M"_a, + "N"_a, + "K"_a, + "The row that runs a plain transposed quantized_matmul as dequantize " + "then the stock matmul by default, for an architecture string, NAX, " + "codec, rows of x, N and K, as [min_m, min_k, max_k, min_n, max_n]; " + "empty where the qmm_t kernel runs."); + + m.def( + "_dequantize_route_takes", + &mlx_kquant::dequantize_route_takes, + "arch"_a, + "nax"_a, + "mode"_a, + "pinned"_a, + "kquant_type"_a, + "M"_a, + "N"_a, + "K"_a, + "Whether a plain transposed quantized_matmul on a 2-D weight on a GPU " + "stream takes the dequantize route, for an architecture string, NAX, " + "the KQ_DEQUANT_ROUTE mode, a pinned KQ_QMM_TILE, codec, rows of x, N " + "and K."); + m.def( "qmv_fine_max_n", &mlx_kquant::qmv_fine_max_n, @@ -564,6 +625,265 @@ NB_MODULE(_ext, m) { array: same shape and dtype as x. )"); + m.def( + "qk_norm_rope", + &mlx_kquant::qk_norm_rope, + "q"_a, + "k"_a, + "q_w"_a, + "k_w"_a, + "cos"_a, + "sin"_a, + "eps"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + Per-head RMSNorm and RoPE of attention q and k in one dispatch: + y = rms_norm(x, w, eps) over the head dim, then each interleaved + pair (y[2i], y[2i + 1]) rotates by the angle with cosine cos[l, i] + and sine sin[l, i]. The output is written in the attention layout, + so no transpose copy follows. f32 math with one round at the write. + + q and k are read through their strides, so slices of one fused qkv + output need no copy. On the GPU head dims 64, 96, 128 and 256 run + the fused kernel; other even head dims and CPU streams run the same + math as stock ops. The op is for inference. On the GPU the fused + kernel has no gradient or vmap, and mx.grad or mx.vmap through it + raises. + + Args: + q (array): [B, L, H, D], float16/bfloat16, D even. + k (array): same shape and dtype as q. + q_w (array): [D] q norm weight, any float dtype. + k_w (array): [D] k norm weight, any float dtype. + cos (array): [L, D / 2] per-row cosines, any float dtype, shared + by the batch. + sin (array): [L, D / 2] per-row sines. + eps (float): norm epsilon. + + Returns: + tuple(array, array): q and k as row-contiguous [B, H, L, D] in + the input dtype. + )"); + + m.def( + "quantized_matmul_glu", + &mlx_kquant::quantized_matmul_glu, + "x"_a, + "w_gate"_a, + "w_up"_a, + "kquant_type"_a, + "activation"_a = "silu", + nb::kw_only(), + "stream"_a = nb::none(), + R"( + GLU product of two quantized matmuls that share x: + y = act(x @ Wg^T) * (x @ Wu^T). + + On a NAX GPU, for q4_k, q5_k and q6_k weights and for q8_0 weights + on bfloat16, with K a multiple of the codec's K tile (64 for q8_0, + 128 for the others) and at least 1024 rows, one dispatch computes + the product with y as its only buffer. It walks K for the gate + weight and stores the activation of that sum in y, then walks K for + the up weight and multiplies the stored activation into that sum at + the final store. The sums, the activation and the product run in + f32, with one rounding of the activation to the dtype of y and one + of the product. A float16 store keeps a finite value in the finite + half range, and a NaN or an infinity passes. Every other case + builds the same product from quantized_matmul and stock ops, which + round each matmul to the activation dtype first. mx.grad gives the + gradient with respect to x at every row count: the gradient of the + fused kernel runs the unfused ops. The weights are frozen, as in + quantized_matmul. mx.jvp and mx.vmap raise at every row count. + + Args: + x (array): [..., K], float16/bfloat16. float32 runs in bfloat16. + w_gate (array): [N, K bytes] uint8 wire bytes, the weight under + the activation. + w_up (array): same shape and codec as w_gate. The two may be + row slices of one wider weight. + kquant_type (str): the codec of both weights. + activation (str): "silu" or "gelu_tanh". + + Returns: + array: [..., N] in the dtype quantized_matmul gives for x. + )"); + + m.def( + "quantized_matmul_lora", + [](mx::array x, + mx::array w, + const std::string& kquant_type, + mx::array lora_a, + mx::array lora_b, + const std::optional& strength, + const std::optional>>& bands, + mx::StreamOrDevice s) { + std::optional> flat; + if (bands.has_value()) { + flat.emplace(); + for (const auto& row : *bands) { + if (row.size() != 4) { + throw std::invalid_argument( + "[mlx_kquant.quantized_matmul_lora] bands must hold rows " + "of (col_start, col_end, z_start, rank)."); + } + flat->insert(flat->end(), row.begin(), row.end()); + } + } + return mlx_kquant::quantized_matmul_lora( + std::move(x), + std::move(w), + kquant_type, + std::move(lora_a), + std::move(lora_b), + strength, + flat, + s); + }, + "x"_a, + "w"_a, + "kquant_type"_a, + "lora_a"_a, + "lora_b"_a, + nb::kw_only(), + "strength"_a = nb::none(), + "bands"_a = nb::none(), + "stream"_a = nb::none(), + R"( + Quantized matmul with a low-rank term: + y = x @ W^T + (strength * (x @ A)) @ B. + + On a NAX GPU, for q4_k, q5_k, q6_k and q8_0 weights on bfloat16 + with K up to 4096 and N of at least K, with at least 4096 rows + and at most 128 ranks per column (from 1024 rows and at every rank + where N is at least 8192), every band starting on a multiple of 64 + columns and ending on a multiple of 32 columns or at N, at most 16 + bands and a rank of at most 256 per band (without bands, a rank of + at most 256), stock ops compute z = strength * (x @ A) and one + dispatch walks the rank for z and B and then K for the weight in + the same f32 sums. y is the only + output buffer and is rounded once, at the store. KQ_QMM_LORA=2 + runs the kernel from 1024 rows at every K and N, and 0 turns the + kernel off. Every other case builds the same sum from + quantized_matmul and stock ops, which round the matmul to the + activation dtype before the add. mx.grad gives the gradient with + respect to x, lora_a, lora_b and strength at every row count: the + gradient of the kernel runs the unfused ops. The weight is frozen, + as in quantized_matmul. mx.jvp and mx.vmap run with respect to + lora_a, lora_b and strength at every row count. With respect to x + both raise "Not implemented for KQuantMatmul" at every row count, + since quantized_matmul has neither. + + Several adapters are one call: join their A columns and their B + rows, and give one strength per column. + + A function compiled with shapeless=True keeps the path of its + trace at every later shape: a trace that ran the kernel runs it + below 1024 rows too, and the result is right there. A shapeless + function cannot hold a call with bands that builds the unfused + ops: that call raises ValueError. + + Args: + x (array): [..., K], float16/bfloat16. float32 runs in bfloat16. + Fewer than 2^31 rows. + w (array): [N, K bytes] uint8 wire bytes. K and N below 2^25. + kquant_type (str): the codec of w. + lora_a (array): [K, R], any float dtype. R below 2^25. + lora_b (array): [R, N], any float dtype. + strength (array, optional): the scale of each column of + x @ A, in any float dtype: a 0-D array or [1] for one + scale, or [R] for one per column. Any other shape raises + ValueError. A new value needs no new trace. + bands (sequence, optional): rows of (col_start, col_end, + z_start, rank), in column order and disjoint. The columns + [col_start, col_end) of y take the ranks [z_start, z_start + + rank) and that block of lora_b. A column outside every + band takes no low-rank term. None (the default): every + column takes every rank. An empty sequence: no column + takes a rank, and y is quantized_matmul(x, w), bit for + bit. + + Returns: + array: [..., N] in the dtype quantized_matmul gives for x. + )"); + + m.def( + "conv2d_3x3", + &mlx_kquant::conv2d_3x3, + "x"_a, + "w"_a, + nb::kw_only(), + "bias"_a = nb::none(), + "stream"_a = nb::none(), + R"( + 3x3 convolution, stride 1, padding 1, channels last: + y = conv(x, w) + bias, as mx.conv2d(x, w, padding=1) + bias. + + On a NAX GPU, for bfloat16 with C a multiple of 16, from 1024 + output pixels, one dispatch computes y as an implicit GEMM: f32 + sums over the 9 taps and C channels, the bias added in f32, one + rounding. KQ_CONV_NAX=2 runs the kernel at every pixel count, and + 0 turns it off. Every other case builds MLX's convolution on the + flipped kernel, which takes its implicit GEMM with f32 sums and not + its Winograd transform, then a stock add of the bias. mx.grad, + mx.jvp and mx.vmap run on that composition. + + Args: + x (array): [B, H, W, C], a float dtype. B H W C and B H W O + below 2^31 for the kernel. + w (array): [O, 3, 3, C] in any float dtype. It runs in the + dtype of x. + bias (array, optional): [O] in any float dtype. + + Returns: + array: [B, H, W, O] in the dtype of x. + )"); + + m.def( + "quantized_matmul_residual", + &mlx_kquant::quantized_matmul_residual, + "x"_a, + "w"_a, + "kquant_type"_a, + "residual"_a, + nb::kw_only(), + "gate"_a = nb::none(), + "stream"_a = nb::none(), + R"( + Quantized matmul with a residual epilogue: + y = residual + gate * (x @ W^T). + + The residual add of a transformer block after its output + projection, with the gate of an adaptive norm where the block has + one. On a NAX GPU, for q4_k, q5_k, q6_k and q8_0 weights on + bfloat16, from 1024 rows, one dispatch walks K into f32 sums and + stores residual + gate * sum, rounded once: the product is never + written. KQ_QMM_RESID=2 runs the kernel at every row count, and 0 + turns it off. Every other case builds the same sum from + quantized_matmul and stock ops, which round the product and the + gated product to the output dtype before the add. mx.grad gives + the gradient with respect to x, residual and gate: the gradient of + the kernel runs the unfused ops. The weight is frozen, as in + quantized_matmul. mx.jvp and mx.vmap run with respect to residual + and gate. With respect to x both raise "Not implemented for + KQuantMatmul", since quantized_matmul has neither. + + Args: + x (array): [..., K], float16/bfloat16. float32 runs in bfloat16. + Fewer than 2^31 rows. + w (array): [N, K bytes] uint8 wire bytes. K and N below 2^25. + kquant_type (str): the codec of w. + residual (array): the output shape, x's shape with N last, in + any float dtype. It runs in the output dtype. + gate (array, optional): [N], or x's batch dims, then 1, then N + (one row per batch), in any float dtype. It runs in the + output dtype. None (the default): y = residual + x @ W^T. + + Returns: + array: [..., N] in the dtype quantized_matmul gives for x. + )"); + m.def( "kda_conv", &mlx_kquant::kda_conv, diff --git a/docs/integration.md b/docs/integration.md index 59bbcd4b..77476d4c 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -10,7 +10,7 @@ Everything here runs against a **stock `mlx` wheel** - importing `mlx_kquant` ad namespace; it does **not** register `mode="kquant"` onto `mx.*`, so integrations call `kq.*` explicitly. -**Compatibility.** The ops require stock `mlx==0.32.1` (the C++ extension is ABI-pinned to it). The +**Compatibility.** The ops require stock `mlx==0.32.3` (the C++ extension is ABI-pinned to it). The `[tools]` integration points below (loader, recipes, the mlx-lm patch) additionally need `mlx-lm>=0.27`. The on-disk checkpoint format is stable and mirrors what a future in-core kquant mode would read. @@ -210,7 +210,8 @@ points, each a narrow shim over a stable mlx-lm function: `patch_mlx_lm_load()` is the **load-only** shim - all an inference / eval / serving consumer needs. It wraps `mlx_lm.utils.load_model` so a kquant config routes through `loader.load`, and every non-kquant -checkpoint falls through unchanged. Because everything mlx-lm exposes is built on `load_model`, +checkpoint falls through unchanged. A `get_model_classes` the caller passes reaches `loader.load` and +picks the model class there too. Because everything mlx-lm exposes is built on `load_model`, `mlx_lm.load`, `mlx_lm.generate`, and the CLIs all then open a kquant checkpoint transparently: ```python diff --git a/docs/kernels.md b/docs/kernels.md index f805184d..af3452c4 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -63,13 +63,64 @@ output width N. Per-row qmv serves M 2 up to a per-codec limit, the mat-vec path from a per-codec entry through M 32. From M 33 to 128, split-K also takes the calls where the BM=64 grid wastes the GPU, as described below. Elsewhere above M 32 the BM=64 tile runs, with a double-buffered `_db` variant on the M 33-64 band at large N, and a BM=128 tile from M 193 when -ceil(M/64) is even. Above +ceil(M/64) is even. The k-quants `q3_k` to `q6_k`, the flat block codecs `q4_0`, `q4_1` and `q5_0`, +and `iq4_nl`, `iq3_xxs`, `iq3_s`, `iq2_xs` and `iq2_s` run the BM=64 and BM=128 tiles on a K tile +of 128 (the `_BK128` levers). Above the qmv limit, the verify kernels take the widths through M 8 on the codecs that have them. The NAX verify kernel serves `pq2_0`, `q4_0`, `q8_0` and `q2_k` to `q6_k` from a per-codec entry, and the register-resident MMA verify kernel serves `ptq1_0`, as described at the end of this section. Every limit and floor is a measured per-codec policy (`kq_nax_small_m` and `kq_smallbm_policy` in `src/kquant_matmul.cpp`). +On GPUs without NAX, the `qmm_t` kernel decodes each weight tile inside its K loop, once for every +64-row tile of the activations. From a floor in rows it is faster to decode the whole weight once +with `kq.dequantize` and run the stock matmul. A transposed `quantized_matmul` on a 2-D weight takes +that route on the GPUs without NAX before M5 when its codec has a row below, its rows of x reach the +floor (a batch counts all of its rows), K is from 3840 to 12288 and N from 3840 to 24576 +(`kDequantizeRows` in `src/kquant_ops.cpp`). Those GPUs are the ones whose architecture string is +`applegpu_g` with a generation from 13 to 16 (M1 to M4) and a suffix other than `p` +(phones), without NAX. Any other string keeps the kernel, as do NAX GPUs, CPU streams, calls with the +LoRA epilogue arguments, a 3-D weight, `transpose=False` and every call while `KQ_QMM_TILE` is set. +`gather_qmm` without indices on a 2-D weight is the plain call and takes the route. +`quantized_matmul_glu`, `quantized_matmul_lora` and `quantized_matmul_residual` build their +products from `quantized_matmul` on these GPUs, so they take it too. + +| codecs | rows from | K | N | +|---|---|---|---| +| `q3_k`, `q5_k`, `q6_k`, `nvfp4`, `ptq1_0` | 2048 | 3840 to 12288 | 3840 to 24576 | +| `q4_k` | 16384 | 3840 to 12288 | 3840 to 24576 | + +The route holds a copy of the weight in the activation dtype while the call runs, N x K x 2 bytes +per call, at most 604 MB at 24576 x 12288; the range keeps a vocabulary head on the kernel. Each +routed call holds its copy until its command buffer completes, so consecutive routed layers can each +hold one at the same time. +`quantized_matmul_glu` builds its product from two calls on these GPUs, so it holds two copies, up +to about 1.2 GB. The gradient wrt x runs the quantized matmul the other way, as on the kernel path, +so a backward pass holds no dense copy, and a gradient wrt the weight or the scales raises as it +does there. The choice is made when the op builds its graph, so a +function compiled with `shapeless=True` keeps the choice of its trace at other row counts. On the +output elements above 1% of the root mean square of a float64 reference, the route is within 1 ulp +of the kernel in bfloat16 and within 2 ulp in float16, and its error against that reference is +within 1% of the kernel's. Nearer zero, where the terms of a sum cancel, the two summation orders +differ by more ulp. + +The rows and the range were measured on M3 Max (`applegpu_g15s`, 40 GPU cores), and the other GPUs +of the class take them unmeasured. Measured with `benchmarks/bench_dit_shapes.py` over the seven +Z-Image and Qwen-Image linears at 2048, 4128 and 16416 rows, batch 1 and 2, bfloat16 and float16, +against an A versus A arm that stays within 0.98-1.02 per step: the kernel takes 1.49-1.63x the +route's time per step on `nvfp4`, 1.06-1.14x on `ptq1_0`, 1.03-1.12x on `q3_k` and `q6_k`, +1.01-1.10x on `q5_k` (1.01-1.03x at 2048 rows) and 1.02-1.05x on `q4_k` from its floor. No linear +runs slower outside the A versus A range, and the route takes 0.96-1.07x the time of the stock +matmul on unquantized weights. + +The decode pass runs at 1.0-1.15x the time of a copy of its output and costs the same at any row +count, so the route loses below its floor: at 128 rows it takes 1.4-1.9x the kernel's time. On the +other codecs the route measured inside the noise of an A versus A arm at 2048 to 16416 rows, and on +`q8_0` it gains 1-4% of a step, also inside it; they keep the kernel. A larger kernel tile does not +help on this chip: five tiles of 128 rows or columns ran 0.98-1.27x the time of bm64 bn64 per cell +at 4128 rows, slower on 51 of 56 cells of the two closest, against an A versus A range of +0.98-1.03. + Per-row qmv reads the weight once per activation row, the later rows from the cache, and its grid grows with M. It holds widest where the mat-vec grids are too small to fill the GPU, so its limit falls with N. At N 512 and below, 13 of the 22 NAX codecs keep it through M 7 or 8. At N 1024 most @@ -108,12 +159,34 @@ than the BM=64 tile when each shape is weighted by its calls in a forward of a 2 and no cell runs more than 2% slower. Whole forwards of that model run 19-22% faster at 33 to 64 rows and 17% faster at 80 and 96. `KQ_QMM_MIDM=0` keeps the BM=64 tile. +With float partials (the float16 default, and `KQ_SPLITK_F32=1`), `q4_0` and `q8_0` target 10 slices +from M 33 to 64, and from M 65 to 96 they target 5 up to K 5120 and 8 above. They cut K in units of +64 weights, so at K 5120 the targets of 16 and 8 write 16 or 8 partials where the 256-weight codecs +write 10 or 5, and a float partial is twice the size of a bfloat16 one. On M5 Max over the eight +projection shapes on float16 activations, each cell of shape and M takes 0.79-1.01x the time of a +target of 16 from M 33 to 64, and 0.67-1.02x the time of 16 or 8 from M 65 to 96 at K 5120. Above +K 5120 a target of 5 leaves 4 slices of the 96 or 272 units, so the target there stays 8. Those +cells take 0.88-0.97x the time of 4 slices, and 0.99-1.02x at K 6144 and M 96. Weighted by calls per +forward, M 65 to 96 run 0.96-0.99x the time of a target of 5 at every K. The K threshold was +measured at K 5120, 6144 and 17408 only. +The rule keys on the partials, so bfloat16 with `KQ_SPLITK_F32=1` takes these targets too, though +they were not measured there. bfloat16 on its default bfloat16 partials keeps 16 and 8. + The NAX split-K route cuts K into units of max(block, 64) weights and targets 16 slices of whole units. It takes the largest count at or under the target that divides the units evenly, unless equal slices with a shorter last one give more than twice as many. This covers inner dimensions whose unit count has no divisor near 16, such as 11008, 17408 and 18944 on the 256-weight codecs. `KQ_SPLITK_RAGGED=2` takes the ragged count wherever it is larger. +At M 32 and below, a call whose BM=32 grid already holds 16 threadgroups per GPU core, such as a +vocab head, takes one slice where the call without the split runs the same BM=32 tile. `q6_k` and +`iq2_xxs` below M 32 and `q8_0` below M 13 keep the target of 16 there, as do calls that would +leave the BM=32 tile for an mv route or the BM=64 tile. Below a full grid, `q4_0` and `q8_0` on +float partials hold the split grid to 16 threadgroups per core, except `q8_0` below M 13. On +M5 Max, against a target of 16 at N 49152, 151936 and 248320 and M 9 to 32, the single slice runs +0.68-1.05x the time, and summed over a forward pass, weighted by the calls per shape, no codec +runs slower. The float-partial cap runs 0.75-1.02x at N 5120 to 17408. + Measured on M5 Max at M 8 to 24 with the weights streamed from DRAM, the shorter last slice runs `q4_k` 1.1x faster at 5120x17408, 1.5x at 3584x18944, 2x at 4096x11008 and 4-5x at 896x4864. The other codecs range from no change (`pq2_0` at 3584x18944) to 5.7x (`iq1_s` at 896x4864). Where the @@ -143,17 +216,97 @@ the two kernels are within 3% of each other on every shape measured. Tuning levers (defaults are right for normal use): +- `KQ_TUNE_LIVE` - `1` makes the NAX tile levers (`KQ_NAX_SMALL_BM`, `KQ_NAX_BM128`, + `KQ_NAX_BM128_ODD_M`, `KQ_NAX_DB64`, the `_BK128` levers, `KQ_NAX_SWIZZLE`), the tile overrides + (`KQ_NAX_TILE`, `KQ_QMM_TILE`, `KQ_CONV_TILE`), `KQ_SPLITK_F32` and `KQ_QMM_LOG` read their + variables on every dispatch, so one process can switch tiles between calls on one resident copy of + the weights; `benchmarks/bench_dit_shapes.py` runs its tile arms this way. It reads the + environment on every dispatch, so it is for single-threaded A/B runs that set these variables + between calls. The `quantized_matmul_glu` levers (`KQ_GLU`, `KQ_GLU_MIN_ROWS`), the + `quantized_matmul_lora` levers (`KQ_QMM_LORA`, `KQ_QMM_LORA_MIN_ROWS`), the + `quantized_matmul_residual` lever (`KQ_QMM_RESID`), the dequantize route lever + (`KQ_DEQUANT_ROUTE`) and the `conv2d_3x3` lever (`KQ_CONV_NAX`) are + read when the op is called and builds its graph, not when a kernel is dispatched, so a compiled + function keeps the kernel choice of its trace. Unset (off) by default, and then each of them keeps + its first read for the life of the process. Read once per process. +- `KQ_NAX_TILE` - `bm,bn,wm,wn` pins the NAX `qmm_t` tile of every transpose call that reaches it, + in place of the BM=32, BM=64 and BM=128 policy. BK stays 64, except where a `_BK128` lever sets + it. The `_db` variant still follows `KQ_NAX_DB64` on the bm64 bn64 wm2 wn2 tile. A tile the + metallib does not carry raises a `ValueError` that names the missing kernel, and + `kq.metallib_kernels()` lists the ones it carries. Unset (off) by default. Read once per process + unless `KQ_TUNE_LIVE=1`. +- `KQ_QMM_TILE` - the same override for the `qmm_t` tile on the path without NAX (GPUs without + NAX, `KQ_DISABLE_NAX=1`, or a codec or shape the NAX tile does not serve). The metallib + carries only the shipped tile, so `64,64,2,2` is the one value that runs today; it keeps its + plain kernel name. Any other tile is looked up under a name that carries + `_bm_bn_wm_wn` before `_alN` and raises until such a kernel is built. Set, it + also keeps calls on the kernel where the dequantize route would take them (see above), and that + choice is read when the op builds its graph. Unset (off) by default. Read once per process + unless `KQ_TUNE_LIVE=1`. +- `KQ_DEQUANT_ROUTE` - `auto` runs the dequantize route described above on the GPUs it is for, + `off` keeps the `qmm_t` kernel on every GPU, and `force` runs the route on any GPU, NAX GPUs + included, for the codecs, rows and range of its table, so that tests run its code on any + machine. Any other value warns once and reads as `auto`. `auto` by default. Read when the op + builds its graph, once per process unless `KQ_TUNE_LIVE=1`. +- `KQ_QMM_LOG` - set to any value, logs the kernel name, shape and grid of every NAX and ALU qmm + tile dispatch, every split-K dispatch, every `verify_mma` and `verify_nax` dispatch, split or not, + and every fused `quantized_matmul_glu`, `quantized_matmul_lora` and `quantized_matmul_residual` + dispatch, and every `conv2d_3x3` dispatch to stderr, one `[kq qmm]` line each. The grid depth is + the K slice count. A call that takes the dequantize route logs one + `[kq dequant] M= N= K=` line when the op builds its graph. The mat-vec + routes (`qmv`, `verify_qmv`, `mv_ext`) do not log. Unset (off) by default. Read once per process + unless `KQ_TUNE_LIVE=1`. +- `KQ_SPLITK_F32` - partial dtype of the split-K routes (`splitk`, `nax_splitk`, the + `KQ_QMM_MIDM` band, and `verify_mma` and `verify_nax` when they split K). Float partials let the + fold round the result once, as the unsplit kernels do. Partials in the activation dtype round + each slice once more before the fold, and where one K slice carries most of an output, that puts + the split result more than one output rounding from the unsplit one. On float16 they also stop at + 65504, so a slice sum past it turns to inf even where the output is finite. Unset, float16 + activations store float partials and bfloat16 activations keep the previous bfloat16 partials. + `1` stores float partials on both, and `0` stores partials in the activation dtype on both. + `true`, `on` and `yes` read as `1`, and `false`, `off` and `no` as `0`. An empty value reads as + unset, and any other value reads as unset with a warning on stderr. Read once per process unless + `KQ_TUNE_LIVE=1`. - `KQ_NAX_SMALL_BM` - small-M routing. `0` restores the old routing (mat-vec paths below M 13 and no BM=32 tile), `2` forces BM=32 for policy-excluded codecs, unset or `1` follows the per-codec - policy. + policy. Read once per process unless `KQ_TUNE_LIVE=1`. - `KQ_NAX_BM128` - BM=128 band. `0` off, `1` forces the floor to M 193 for every codec, `2` drops the floor entirely (any even ceil(M/64), probing the M65-128 wash band), unset follows the per-codec entry floors (193/449/961 tiers, measured on M5 Max by - `benchmarks/bench_qmm_bm128_ab.py`; re-run it before trusting them on new silicon). + `benchmarks/bench_qmm_bm128_ab.py`; re-run it before trusting them on new silicon). Read once + per process unless `KQ_TUNE_LIVE=1`. +- `KQ_NAX_BM128_ODD_M` - the M from which the BM=128 band also takes an odd ceil(M/64), whose last + row tile holds 64 rows or fewer. There the simdgroup row with no row below M skips its + activation loads, matrix work and store, so that tile does the loads and matrix work of one + BM=64 tile. Unset, the codec's large-M row sets it (see below). `0` keeps the even ceil(M/64) + rule at every M, the policy before this lever, and turns off the large-M rows with their K + ceiling, so BM=128 also runs above the ceiling at an even ceil(M/64). Read once per process + unless `KQ_TUNE_LIVE=1`. - `KQ_NAX_DB64` - double-buffered M 33-64 band. `0` off, `1` drops the N floor, unset follows the per-codec N floors. Only the five policy-enabled codecs (q6_k, q8_0, q4_1, q5_1, q5_0) carry `_db` instantiations, so `1` is bounded by availability; probing another codec needs its - instantiation restored and a metallib rebuild. + instantiation restored and a metallib rebuild. Read once per process unless `KQ_TUNE_LIVE=1`. +- `KQ_NAX__BK128` (`KQ_NAX_Q6K_BK128`, `KQ_NAX_Q5K_BK128`, `KQ_NAX_Q4K_BK128`, + `KQ_NAX_Q3K_BK128`, `KQ_NAX_Q4_0_BK128`, `KQ_NAX_Q4_1_BK128`, `KQ_NAX_Q5_0_BK128`, + `KQ_NAX_IQ4_NL_BK128`, `KQ_NAX_IQ3_XXS_BK128`, `KQ_NAX_IQ3_S_BK128`, `KQ_NAX_IQ2_XS_BK128`, + `KQ_NAX_IQ2_S_BK128`) - K tile of the codec on the BM=64 and BM=128 tiles. Unset or `1`, the tile + walks K 128 weights per step. On the k-quants it takes one half of a super-block per step, and + each thread reads its code bytes once, where BK=64 reads them on two or four steps. On the other + codecs two BK=64 loaders fill the tile, which halves the barriers per K, and a 32-weight block + codec needs K % 128 == 0. The output is bit-identical. On an M5 Max, per cell against BK=64 on + bfloat16 and float16 activations, on the BM=128 tile at diffusion model shapes (M 4096 and 4608, N + and K 3840 to 24576): `q6_k` 0.87-0.96x the time, `q5_k` 0.93-1.00x, `q3_k` 0.92-1.03x. On the + BM=64 tile from M 128 to 2112 at N and K 3840 to 17408: `q6_k` 0.79-1.01x, `q5_k` 0.89-1.00x, + `q3_k` 0.88-1.00x, except on grids of 240 threadgroups (N 5120 at M 160 and 192), where `q5_k` + takes 1.06-1.17x and `q3_k` 1.01-1.06x, and of 480 (N 10240 at M 160 and 192), where `q5_k` takes + 1.02-1.06x. No threshold on threadgroups per core, K or M removes these at a lower total time. At + the diffusion model rows (M 4128 to 16896 at N and K 3840 to 12288, bfloat16 and float16, with the + K ceilings below), against an A/A spread of 0.98-1.03x: `q4_k` 0.95-1.00x, `q4_0` 0.97-1.01x, + `q4_1` 0.97-1.03x, `q5_0` 0.96-1.01x, `iq4_nl` 0.93-1.00x, `iq3_xxs` 0.93-1.02x, `iq3_s` + 0.92-1.03x, `iq2_xs` and `iq2_s` 0.95-1.01x. `q8_0`, `q5_1`, `iq4_xs`, `iq1_m` and `iq1_s` took + 1.01-1.06x at some of those cells and `iq2_xxs` 1.04-1.07x at all of them, so they keep BK=64 and + carry no `_bk128` kernels. `0` keeps BK=64. The BM=32 tile and the `_db` variant keep BK=64. Read + once per process unless `KQ_TUNE_LIVE=1`. - `KQ_FORCE_QMM_MIN_M` - probe lever: routes transpose shapes with M at or above the value straight to the NAX qmm, bypassing the mat-vec route claims, for crossover measurement below M 13. Unset (off) by default. @@ -165,9 +318,10 @@ Tuning levers (defaults are right for normal use): - `KQ_QMM_ROUTE_STRICT` - `1` makes a transpose call that the forced `KQ_QMM_ROUTE` cannot serve raise instead of taking the default routing, so a sweep never times the default under a forced route's name. Off by default. -- `KQ_NAX_SWIZZLE` - `1` enables the row-tile traversal swizzle (folds row-tiles into grid.x for - SLC reuse of the weight band). Falsified on M5 Max, where the M>64 band is per-threadgroup-bound - rather than DRAM-bound; kept as a probe for future silicon. Default off. +- `KQ_NAX_SWIZZLE` - row-tile traversal swizzle on the BM=64 and BM=128 tiles: `n` folds 2^n row + tiles into grid.x, so consecutive threadgroups share one weight column band, capped at the + row-tile count. `0` is off. Unset, the codec's large-M row sets the fold on that row's tile and + nothing else folds. Read once per process unless `KQ_TUNE_LIVE=1`. - `KQ_MV_EXT_NR` - `2` selects the two-rows-per-thread `mv_ext` variant (q6_k, M 5-12), which halves activation cache traffic but measured no faster than the shipped kernels. Kept as a probe for future silicon. Default `1` (shipped behavior). @@ -196,16 +350,28 @@ Tuning levers (defaults are right for normal use): - `KQ_QMM_SPLITK` - the same lever for the plain small-M qmm, used when NAX is absent or disabled. Entry points come from a per-device table. K-quants, legacy quants and the IQ codecs, M <= 32. `0` also turns off the `KQ_QMM_MIDM` route, and a larger value sets its slice target. -- `KQ_QMM_MIDM` - split-K for M 33-128. `0` keeps the BM=64 qmm and `1` forces split-K at every M - in the band. Unset takes split-K where the qmm grid, ceil(N/64) x ceil(M/64) threadgroups, is - under three per GPU core (`KQ_GPU_CORES` sets the count), such as the narrow projections of a - short prompt. On NAX GPUs it also takes M 65-96 up to 16 threadgroups per core and runs the NAX - split-K tile on every codec with NAX kernels. When NAX is absent or disabled it runs the BM=32 - tile on the codecs of `KQ_QMM_SPLITK`. Read live per call. +- `KQ_QMM_MIDM` - split-K for M 33-128. `0` keeps the BM=64 qmm and `1` forces split-K at every M in + the band. Unset takes split-K where the qmm grid, ceil(N/64) x ceil(M/64) threadgroups, is under + three per GPU core (`KQ_GPU_CORES` sets the count), such as the narrow projections of a short + prompt. On NAX GPUs it also takes M 65-96 up to 16 threadgroups per core and runs the NAX split-K + tile on every codec with NAX kernels, and there `q4_0` and `q8_0` on float partials target 10 + slices through M 64, and from M 65 to 96 5 up to K 5120 and 8 above. When NAX is absent or + disabled it runs the BM=32 tile on the codecs of `KQ_QMM_SPLITK`. Read live per call. - `KQ_SPLITK_RAGGED` - slice count of the NAX split-K route. `0` keeps equal slices, the largest count at or under the target that divides the slice units, and `2` takes the ragged count whenever it is larger. Unset takes it only where it more than doubles the equal count. Read live per call. +- `KQ_SPLITK_CAP` - probe lever: caps the default slice target of the small-M split-K routes + (`nax_splitk` and the ALU `splitk` route, M <= 32) so the split grid holds at most that many + threadgroups per GPU core (`KQ_GPU_CORES` sets the count). A call whose grid already holds that + many takes one slice and runs what it runs with the split route off. From M 13 that is usually + the BM=32 tile, or per-row qmv below the qmv batch limit, which is 18 on M5 Max at K and N up to + 2048 and up to 32 on Ultra parts. At M 9 to 12 it is the small-M qmm route, `mv_ext` or per-row + qmv. The ALU route falls through the same way, to the ALU qmm tile in place of the NAX one. + `verify_mma` and `verify_nax` are not capped. `0` keeps the target of 16 at every N on both + routes, so a wide N such as a vocab head writes up to 16 partials. Unset keeps it on the ALU route + and takes the NAX default described with the NAX split-K route. A target forced through + `KQ_QMM_SPLITK_NAX` or `KQ_QMM_SPLITK` is not capped. Read live per call. - `KQ_VERIFY_MMA` - the register-resident MMA verify route (below). A value of `2` or more forces it at that M and above on any GPU, `0` or `1` disables it. Unset takes the per-codec entry from the device's table (`kq_verify_mma_min_m` without NAX, `kq_verify_mma_min_m_nax` with it); the @@ -221,6 +387,14 @@ Tuning levers (defaults are right for normal use): - `KQ_VERIFY_NAX_SPLITS` - probe lever: the split-K target for `verify_nax` in place of the per-codec rule (the largest divisor of the K steps at or under the value is used). Unset (off) by default. +- `KQ_VERIFY_MMA_SPLITS` - probe lever: the same split-K target for `verify_mma`; `1` runs it + unsplit. Unset (off) by default. Read live per call. +- `KQ_VERIFY_MMA_HALF` - probe lever: `1` runs the `verify_mma` half kernels (name suffix `_h`) of + pq2_0, q2_0, q4_0 and q8_0, which stage the activations as half and sum each block in half, so an + activation or a block sum past 65504 turns to inf. On pq2_0 the codes enter that sum unscaled, so + a 128-weight block can pass 65504 with activations of magnitude 256 when the codes and activations + agree in sign. ptq1_0 has no half kernels and stays on its float kernels. For timing against the + float kernels only. Unset (off) by default. Read live per call. - `KQ_VERIFY_NAX_Q4_0_SB` - probe lever: `0` keeps `q4_0` on the two-block `verify_nax` kernel at a K that is a multiple of 256, where the eight-block kernel runs by default. Read live per call. - `KQ_MV_EXT_SB` / `KQ_MV_EXT_NX` / `KQ_MV_EXT_HD` - `mv_ext` activation-traffic experiments: @@ -228,6 +402,42 @@ Tuning levers (defaults are right for normal use): only. `HD` measured +4-5% at M 8; the rest flat to negative on M5 Max. Kept as probes. Default off. +Large-M rows. On NAX GPUs a transpose call at M 4096 or more takes the codec's large-M row from +`kq_largem_row` in `src/kquant_matmul.cpp`: the tile, the row-tile fold, the shape floor (M at or +above min M at either ceil(M/64) parity, N at or above min N) and the K ceiling, above which the +BM=64 tile runs at either parity. A codec whose tile and floor are not measured takes those of its +loader class (the 256-weight k-quant loaders, the 32-weight flat block loaders, or the IQ grid and +ternary loaders) and is marked inherited. The measured rows come from +`benchmarks/bench_dit_shapes.py` on an M5 Max at the diffusion transformer rows (M 4096 to 33792 +at the Z-Image and Qwen-Image projection shapes). There the BM=128 tile takes 0.85-0.99x the +BM=64 time at rows with an odd ceil(M/64) up to the K ceiling, and 1.01-1.09x above it (`q4_k` +and `q8_0` at K 12288), and per projection the median fold of 2 row tiles takes 1.00-1.26x the +unfolded time, of 4 row tiles 1.03-1.58x. + +The K ceiling is measured for every codec with NAX kernels at N 4096 and M 4128, 8256 and 16384: +it is the largest K before the first K (of 10240, 12288 and 16384) where BM=128 takes more than +1.01x the BM=64 time over the three M, measured on the codec's default K tile. Every codec wins at +K 3840 and 4096. On the K tile of 128, `q3_k`, `q5_k`, `q6_k`, `iq3_xxs`, `iq3_s`, `iq2_xs` and +`iq2_s` win at every K measured (`iq3_s` 0.91-0.94x, the other IQ codecs 0.93-1.00x), while the +flat block codecs lose at K 12288 and 16384: `q4_0` and `q4_1` 1.02-1.06x on the K tile of 128, +`q5_1` and `q8_0` 1.04-1.11x on 64. A ceiling moves only where the new tile takes at most the old +time at every M measured: on the K tile of 128 at K 12288, `q4_k` and `q5_0` take 1.00-1.02x over +the three M and 1.02-1.03x at M 4128, so they keep K 10240. + +| Codecs | Tile (bm, bn, wm, wn) | Fold log | Min M | Min N | Max K | Row | +|---|---|---|---|---|---|---| +| `q4_k`, `q8_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | 10240 | measured | +| `q5_k`, `q6_k` | 128, 64, 2, 2 | 0 | 4096 | 0 | none | measured | +| `q2_k` | 128, 64, 2, 2 | 0 | 4096 | 0 | 4096 | inherited, k-quant class | +| `q3_k` | 128, 64, 2, 2 | 0 | 4096 | 0 | none | inherited, k-quant class | +| `q4_0`, `q4_1`, `q5_0`, `q5_1` | 128, 64, 2, 2 | 0 | 4096 | 0 | 10240 | inherited, flat class | +| `iq4_nl` | 128, 64, 2, 2 | 0 | 4096 | 0 | 12288 | inherited, flat class | +| `pq2_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | 4096 | inherited, grid class | +| `q2_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | 4096 | inherited, grid class; the K ceiling is that of `pq2_0`, not measured | +| `stq1_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | 12288 | inherited, grid class | +| `iq2_s`, `iq2_xs`, `iq3_s`, `iq3_xxs` | 128, 64, 2, 2 | 0 | 4096 | 0 | none | inherited, grid class | +| `iq1_s`, `iq1_m`, `iq2_xxs`, `iq4_xs`, `ptq1_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | none | inherited, grid class | + `stq1_0` (structured-sparse ternary, llama.cpp PR #22836) ships the full ALU and NAX kernel set. Its pre-NAX floors are measured (M3 Max: plain split-K entry M 5), and so are its NAX small-M routes in `kq_nax_small_m`. The rest of its NAX policy is inherited from `iq1_s` and needs M5-silicon @@ -235,7 +445,15 @@ calibration: `bm128_min_m` (`benchmarks/bench_qmm_bm128_ab.py`), `kq_splitk_min_ candidacy (no `_db` instantiation yet). The PrismML codecs `pq2_0` and `ptq1_0` (128-wide blocks) ship the same kernel set. `q2_0` (64-wide blocks, the `pq2_0` code layout) runs the `pq2_0` kernels at its block width and takes the `pq2_0` route entries, which are not yet measured for -it. Their NAX split-K entry is M 9, where the verify kernels hand off. The +it. On the large-M and image ops `q2_0` inherits what `pq2_0` has and no kernel was built for it: +the grid class large-M row with the `pq2_0` K ceiling, the BM=32, BM=64 and BM=128 tiles on the K +tile of 64 (neither codec has a `_bk128` kernel, and a `KQ_NAX_Q2_0_BK128` variable is not read), +the split-K kernels with their float partial twins, and `verify_mma` at one row tile per simdgroup +with the float kernels by default and the half kernels under `KQ_VERIFY_MMA_HALF`. `KQ_NAX_TILE`, +`KQ_QMM_TILE` and the tile levers under `KQ_TUNE_LIVE` apply to it as to every NAX codec. Left out: +`quantized_matmul_glu` and `quantized_matmul_lora` have no `q2_0` or `pq2_0` kernel and build the +product from `quantized_matmul` and stock ops, and like every codec `q2_0` has no kernel on float +activations (float32 runs the bfloat16 kernels). Their NAX split-K entry is M 9, where the verify kernels hand off. The `bm128_min_m` and db64 floors are still inherited from `iq1_s`. The `ptq1_0` M=1 kernel gives each 28-byte block four lanes. Each lane reads its six trit bytes in two loads plus the word that holds the high-trit bytes and the scale, and decodes the trits with the exact-float base-3 coefficient @@ -244,26 +462,49 @@ instead of extracting each trit. A simdgroup computes four output rows, so the f eight rows per threadgroup. Measured on M5 Max at the Ternary Bonsai 2 27B projection shapes with the weights streamed from DRAM, it reads 360-440 GB/s, 1.2-1.3x the previous two-rows-per-simdgroup layout. The `pq2_0` M=1 kernel -masks each pair of 2-bit codes into the mantissas of a half2 and runs the dot as half2 fmas, two -weights per instruction, which brings it close to the rate of a kernel that only loads the bytes. +masks each pair of 2-bit codes into the mantissas of a half2. On float16 activations it runs the dot +as half2 fmas, two weights per instruction, which brings it close to the rate of a kernel that only +loads the bytes. On bfloat16 activations it widens each pair to float and runs the dot as float2 +fmas, since a bfloat16 value past 65504 does not fit in half. Both codecs have a `verify_qmv` sibling that decodes each row's block once and dots it against every activation row. It serves M 2 below the `verify_mma` entry on GPUs without NAX, where the `mv_ext` re-decode per row costs `ptq1_0` 3x. On NAX GPUs per-row qmv serves M 2. -The register-resident MMA verify kernels (`verify_mma`) serve `pq2_0`, `ptq1_0`, `q4_0` and `q8_0` -through M 8. Each simdgroup decodes a block of its weight rows straight into 8x8 simdgroup-matrix +The register-resident MMA verify kernels (`verify_mma`) serve `pq2_0`, `ptq1_0`, `q2_0`, `q4_0` and +`q8_0` through M 8. Each simdgroup decodes a block of its weight rows straight into 8x8 simdgroup-matrix fragments, with the k order inside the block permuted so every lane extracts the code pairs it holds cheapest, and multiplies them against the activations staged once per K chunk, so the weight bytes are read once for all M rows and the tile's threadgroup staging of the weights disappears. Split-K over the wire blocks feeds the same partial fold as `qmm_splitk` on the projection shapes; a head-sized N already fills the GPU with one split and writes the output directly. -The kernels stage the activations as half and sum each block in a half accumulator, so an activation -beyond the half range (magnitude 65504) overflows. The `pq2_0` and `ptq1_0` codes are -1, 0 and 1, -so a block's sum overflows only when its 128 activations add up past that range. The `q4_0` and -`q8_0` kernels decode their codes scaled by 1/16 and 1/256 and apply the factor in the float block -scale, so one channel near the half limit cannot overflow a block of 32 products. A float -accumulator removes the limit but runs `pq2_0` at half speed. +A threadgroup is 8 simdgroups, and each owns two 8-row tiles on `q4_0` and `q8_0` and one on +`pq2_0`, `q2_0` and `ptq1_0` (`q2_0` runs the `pq2_0` chain at 8 fragment MMAs per 64-weight block +and takes its one-tile layout, which is not measured for it). A `pq2_0` block runs its 16 fragment +MMAs into one accumulator per row tile, so the MMAs form a dependent chain whose latency, longer in +float than in half, bounds the kernel. At two row tiles per simdgroup the float kernel took about +twice the time of the half kernel on the M5 Max. One row tile doubles the threadgroups that hide the +latency, and the float kernel then runs in 0.45-0.52x of its two-tile time over the Qwen3.8-27B +projection and head shapes at M 2 to 8. Splitting each chain over two or four accumulators recovered +most of that too, and staging the activations as half with float accumulators did not. `q4_0` and +`ptq1_0` gained nothing from either change. + +The kernels decode the codes exactly in half and keep the activations in their own dtype in +threadgroup memory. Both widen to float in the fragment, so each block sums in float and the block +scale applies in float. Unsplit, or split with float partials (the float16 default, and +`KQ_SPLITK_F32=1`), the result rounds once, as the mat-vec kernels do with the exact weights, at any +activation magnitude. `KQ_VERIFY_MMA_HALF=1` runs the half kernels of `pq2_0`, `q2_0`, `q4_0` and +`q8_0`, which stage the activations as half and sum each block in half. There an activation beyond the half +range (magnitude 65504), such as a bfloat16 outlier, turns to inf, and so does a block sum past it. +On ordinary activations with float partials, the error of the half kernels runs 0.4 to 0.9 output +roundings beyond the rounding of the float64 result on bfloat16, and 4 to 7 on float16, and the +float kernels run within 0.01. Partials in the activation dtype (the bfloat16 default, and +`KQ_SPLITK_F32=0`) add about one output rounding to a split call of either kernel, which puts the +half kernels at 1.1 to 1.4 on bfloat16 and 4 to 7.5 on float16, and the float kernels at about 1 on +both. Against the half kernels over the same shapes and M on the M5 Max, the float kernels take +0.90-1.00x the time for `pq2_0` and 0.94-1.02x for `q4_0` on bfloat16, and 0.99-1.01x for `pq2_0` +and 1.00-1.05x for `q4_0` on float16. `ptq1_0` has no half kernels, since in float it runs in +0.90-1.03x the time of a half build. The entry M depends on the device. Without NAX, `pq2_0` and `ptq1_0` enter at M 3 (`pq2_0` at M 4 on float16 activations, where `verify_qmv` holds M 3) and `q4_0` at M 2. Measured on M3 Max at the @@ -285,7 +526,7 @@ wider. Each simdgroup owns 32 weight rows and decodes every block of them straig right operand of a 16x32x16 NAX matmul, with the block scale folded into the half weights. The left operand holds the activation rows padded to 16, read from device memory in the k order of the decoded weights, so nothing is staged in threadgroup memory and the loop has no barriers. The -accumulator is float32, which removes the half-range limit of `verify_mma`. +accumulator is float32. A simdgroup's device loads and NAX ops do not overlap, so the kernel hides its loads through the simdgroups resident beside it, and the host sizes the split count for that. The kernels differ in @@ -401,6 +642,320 @@ entries the kernel ran up to 1.10x slower at K 4096 and 1.21x slower at N 12288, entry under M 5 sits one row higher. At M 8 the kernel runs 1.05-1.39x faster than NAX split-K on the shapes of N 2048 and wider. +## GLU product of two quantized matmuls + +- `quantized_matmul_glu` - `y = act(x @ Wg^T) * (x @ Wu^T)` for two packed weights of one codec and + one shape, the gated feed-forward of a diffusion transformer block. `act` is `silu` or + `gelu_tanh`. The weights are two arrays, or the two row bands of one fused gate and up weight, + which the kernel reads in place. + + On a NAX GPU (M5), from 1024 activation rows, one dispatch computes the product for `q4_k`, + `q5_k` and `q6_k` weights whose inner dim is a multiple of 128, and for `q8_0` weights on + bfloat16 activations whose inner dim is a multiple of 64. Each threadgroup owns a 128 by 64 tile + of y. It walks K over the gate weight, applies the activation in f32 and stores the tile, then + walks K over the up weight and multiplies the stored tile into the sum before the final store. + y is the only buffer the op writes, so a call peaks at one output where the unfused ops peak at + two with `silu` and at three with `gelu_tanh`. + + Measured on an M5 Max against the two matmuls and the stock activation and product, as the median + of rotated rounds in one process: + + | K, N, activations | Rows | `q4_k` | `q5_k` | `q6_k` | `q8_0` | + |---|---|---|---|---|---| + | 3840, 10240, bfloat16 | 1056 | 0.91 | 0.99 | 0.96 | 0.99 | + | 3840, 10240, bfloat16 | 2080 | 0.89 | 0.88 | 0.87 | 0.93 | + | 3840, 10240, bfloat16 | 4128 | 0.96 | 0.94 | 0.96 | 0.96 | + | 3840, 10240, bfloat16 | 4608 | 0.97 | 0.96 | 0.97 | 0.96 | + | 3840, 10240, bfloat16 | 8208 | 0.96 | 0.96 | 0.97 | 0.97 | + | 3840, 10240, bfloat16 | 16416 | 0.97 | 0.97 | 0.99 | 0.98 | + | 3840, 10240, float16 | 1056 | 0.95 | 1.00 | 0.98 | - | + | 3840, 10240, float16 | 4128 | 0.94 | 0.96 | 0.97 | - | + | 3840, 10240, float16 | 16416 | 1.00 | 0.98 | 0.99 | - | + | 4096, 12288, bfloat16 | 4128 | 0.94 | 0.95 | 0.95 | 0.97 | + | 4096, 12288, bfloat16 | 16416 | 0.97 | 0.97 | 0.97 | 0.99 | + + From 4128 rows the gain is the activation and product pass that the kernel removes, 1 to 6 + percent. `q4_k` on float16 at 16416 rows is a tie. `q8_0` on float16 builds the unfused ops, so + it has no cell. `benchmarks/bench_glu.py` repeats the table for one codec. + + x is float16 or bfloat16, and float32 x runs in bfloat16 as in `quantized_matmul`. Leading batch + dims fold into rows. The activation and the product run in f32, with one rounding of the + activation and one of the product. A finite float16 result past the half range is stored as the + largest finite value with its sign, the activation before the product and the product at the + store, where the unfused ops return inf. A NaN or an infinity in x gives NaN and inf where the + unfused ops give them. + + Every other case builds the same product from `quantized_matmul` and stock ops: other codecs, + `q8_0` on float16, fewer rows, other inner dims, CPU streams and GPUs without NAX. + + Levers, each read once per process unless `KQ_TUNE_LIVE=1`: + + - `KQ_GLU` - `0` builds the unfused ops on every call. Unset (on) by default. + - `KQ_GLU_MIN_ROWS` - the row count from which the fused kernel runs, an integer of at least 1. + A value below 1 keeps the default, with a warning on stderr. Default 1024. + + A build with the CMake option `KQ_PROBE_KERNELS=ON`, which `KQ_PROBE_KERNELS=1` in the + environment of the build selects, also carries the tiles that the shipped ones were measured + against. It reads `KQ_GLU_BM`, `KQ_GLU_BK`, `KQ_GLU_WALKS` and `KQ_GLU_DUAL128` to select one, + and a tile it does not carry for the codec raises a `ValueError` that names the levers and the + kernel. A release build reads none of the four. + + `mx.grad` gives the gradient with respect to x at every row count: the gradient of the fused + kernel runs the unfused ops. The weights are frozen, as in `quantized_matmul`. For float32 x the + fused kernel sums the gate term and the up term of the gradient in bfloat16, where the unfused + ops sum them in float32, so the two gradients differ by one bfloat16 rounding, about 2e-3 + relative. `mx.jvp` and `mx.vmap` raise at every row count, since `quantized_matmul` has neither. + +## Quantized matmul with a low-rank term + +- `quantized_matmul_lora` - `y = x @ W^T + (strength * (x @ A)) @ B` for a packed weight and a + low-rank adapter, `lora_a` [K, R] and `lora_b` [R, N], at the row counts of a diffusion + transformer step. `strength` is a 0-D array or [1] for one scale, or [R] for one per rank, and + any other shape raises a `ValueError`. It is an input array, so a new value needs no reload and + no new trace of a compiled function. Several adapters are one call: join their A columns and + their B rows, and give one strength per rank. + + `bands` holds rows of (col_start, col_end, z_start, rank), in column order and disjoint. The + columns of a band take that block of ranks and that block of `lora_b`, and a column outside + every band takes no low-rank term. A fused q, k and v weight with one adapter per projection is + three bands, so no projection multiplies the ranks of another. `bands=None`, the default, is no + table: every column takes every rank. `bands=[]` is a table with no band: no column takes a + rank, and y is `quantized_matmul(x, w)`, bit for bit. + + On a NAX GPU (M5), for `q4_k`, `q5_k`, `q6_k` and `q8_0` weights on bfloat16 activations with K + up to 4096 and N of at least K, from 4096 activation rows where no column takes more than 128 + ranks, and from 1024 rows at every rank where N is at least 8192, stock ops compute + `z = strength * (x @ A)` and one dispatch computes the product and adds `z @ B`. Each + threadgroup owns a 128 by 64 tile of y. It walks the rank of its column band 32 ranks a step, + with z and B read from their buffers, then walks K over the weight, and stores the tile. Both + walks add into one f32 sum, which is rounded once at the store. y is the only buffer the + dispatch writes, where the unfused ops write the product and then read and write it again for + the sum. The kernel takes a rank of at most 256 without bands, and with bands at most 16 of + them, each starting on a multiple of 64 columns, ending on a multiple of 32 columns or at N, and + holding at most 256 ranks. The `q8_0` kernel walks K in tiles of 64, so a `q8_0` weight whose K + is no multiple of 64 runs the unfused ops. + + The bounds on K, N, rows and rank are where the kernel measured faster. It removes the second + pass over y, a cost in proportion to N. It waits for `x @ A`, a cost in proportion to K and to + the rank, where the unfused product runs beside `x @ A`. So the kernel gains where N is at least + K, and its gain is a share of the call that falls as K grows. Where N is near K the gain is + about one percent, and it holds from 4096 rows and up to a rank of 128. With bands, the rank + that counts is the rank of the largest band. + + Measured on an M5 Max against `quantized_matmul`, `x @ A` and one `mx.addmm`, one call an eval, + at 4608 rows on `q4_k`, as the median of four processes of rotated rounds, kernel time over + unfused time: + + | K, N | Rank 16 | Rank 64 | + |---|---|---| + | 2048, 2048 | 0.99 | 0.96 | + | 3072, 3072 | 0.98 | 0.97 | + | 3840, 3840 | 0.99 | 0.98 | + | 3840, 10240 | 0.97 | 0.97 | + | 3840, 11520, three bands | 0.96 | 0.94 | + | 4096, 4096 | 0.99 | 0.99 | + | 4096, 24576, two bands | 0.96 | 0.96 | + + Outside the bounds, with the kernel run under `KQ_QMM_LORA=2`: + + | K, N | Rank 16 | Rank 64 | + |---|---|---| + | 4096, 1024 | 1.03 | 1.03 | + | 4096, 2048 | 1.01 | 0.98 | + | 6144, 1536 | 1.04 | 1.02 | + | 10240, 3840 | 1.02 | 1.02 | + | 12288, 4096 | 1.01 | 1.02 | + | 5120, 5120 | 0.99 | 0.99 | + | 6144, 6144 | 0.99 | 0.99 | + | 8192, 8192 | 0.99 | 1.00 | + + The other codecs inside the bounds, at the same rows, rank 16 and rank 64: + + | K, N | `q5_k` | `q6_k` | `q8_0` | + |---|---|---|---| + | 3840, 3840 | 0.99, 0.98 | 0.99, 0.98 | 0.98, 0.98 | + | 3840, 11520, three bands | 0.96, 0.94 | 0.96, 0.94 | 0.96, 0.94 | + | 4096, 24576, two bands | 0.96, 0.96 | 0.97, 0.96 | 0.96, 0.96 | + + Below 4096 rows on `q4_k`, rank 16 and rank 64. The first three shapes are outside the bounds + there and ran under `KQ_QMM_LORA=2`: + + | K, N | 1056 rows | 2112 rows | + |---|---|---| + | 2048, 4096 | 1.08, 1.00 | 1.02, 0.99 | + | 3840, 3840 | 1.02, 0.99 | 1.03, 1.02 | + | 4096, 4096 | 1.02, 1.01 | 1.04, 1.03 | + | 4096, 8192 | 0.99, 0.98 | 0.98, 0.98 | + | 3840, 10240 | 0.99, 0.98 | 0.96, 0.96 | + | 3840, 11520, three bands | 0.97, 0.94 | 0.96, 0.93 | + | 4096, 24576, two bands | 0.96, 0.94 | 0.93, 0.93 | + + The tables from here on are the median of four processes that run the two forms in alternating + slots, one call an eval. Ranks 128 and 256 at 4608 rows on `q4_k`. At rank 256 the two shapes + with N of K are outside the bounds: + + | K, N | Rank 128 | Rank 256 | + |---|---|---| + | 3840, 3840 | 0.99 | 1.00 | + | 4096, 4096 | 0.98 | 1.01 | + | 3840, 10240 | 0.97 | 0.99 | + | 3840, 11520, three bands | 0.93 | 0.90 | + | 4096, 24576, two bands | 0.97 | 0.95 | + + From 4096 rows up on `q4_k`, rank 16 and rank 64: + + | K, N | 4096 rows | 8192 rows | 16384 rows | + |---|---|---|---| + | 3840, 3840 | 1.00, 1.00 | 0.98, 0.98 | 0.97, 0.97 | + | 4096, 4096 | 1.00, 0.99 | 0.97, 0.98 | 0.96, 0.97 | + | 3840, 11520, three bands | 0.97, 0.94 | 0.96, 0.94 | 0.97, 0.93 | + + The other codecs below 4096 rows, rank 16 and rank 64: + + | K, N | Rows | `q5_k` | `q6_k` | `q8_0` | + |---|---|---|---|---| + | 4096, 8192 | 1056 | 0.98, 0.97 | 0.95, 0.94 | 0.98, 0.98 | + | 4096, 8192 | 2112 | 0.97, 0.96 | 0.94, 0.94 | 0.98, 0.98 | + | 3840, 10240 | 1056 | 0.97, 0.96 | 0.94, 0.93 | 0.99, 1.00 | + | 3840, 10240 | 2112 | 0.96, 0.96 | 0.92, 0.92 | 0.97, 0.98 | + | 3840, 11520, three bands | 1056 | 0.96, 0.92 | 0.93, 0.92 | 0.99, 0.96 | + | 3840, 11520, three bands | 2112 | 0.94, 0.92 | 0.91, 0.90 | 0.96, 0.94 | + | 4096, 24576, two bands | 1056 | 0.95, 0.93 | 0.92, 0.91 | 0.95, 0.94 | + | 4096, 24576, two bands | 2112 | 0.93, 0.92 | 0.90, 0.89 | 0.93, 0.92 | + + K below 2048 on `q4_k`. A call here takes 0.2 to 2.4 ms, and two runs of the unfused ops differ + by up to 4 percent at the smallest: + + | K, N | Rows | Rank 16 | Rank 64 | + |---|---|---|---| + | 1536, 1536 | 4608 | 0.99 | 0.97 | + | 1024, 1024 | 4608 | 0.99 | 0.95 | + | 512, 512 | 4608 | 0.97 | 0.97 | + | 256, 256 | 4608 | 0.99 | 0.99 | + | 1024, 8192 | 1056 | 1.01 | 0.94 | + | 1024, 8192 | 4608 | 0.87 | 0.87 | + | 512, 8192 | 1056 | 0.95 | 0.95 | + | 512, 8192 | 4608 | 0.79 | 0.80 | + | 256, 8192 | 1056 | 0.94 | 0.91 | + | 256, 8192 | 4608 | 0.70 | 0.72 | + + x is float16 or bfloat16, and float32 x runs in bfloat16 as in `quantized_matmul`. Leading batch + dims fold into rows. `lora_a`, `lora_b` and `strength` take any float dtype and run in the dtype + of y. K, N and R are each below 2^25 and x holds fewer than 2^31 rows: the kernel computes + addresses inside a row tile in 32 bits, and a call past either limit raises a `ValueError`. + + Every other case builds the same sum from `quantized_matmul` and stock ops, which round the + product to the activation dtype before the add: other codecs, float16 activations, fewer rows, + K above 4096, N below K, N below 8192 under 4096 rows or with more than 128 ranks on a column, + a rank above 256, more bands, a band that starts off a multiple of 64 columns or ends before N + off a multiple of 32, a `q8_0` K off a multiple of 64, CPU streams and GPUs without NAX. A rank + of zero, an x of no rows and a weight of no rows return the product alone. + + Levers, each read once per process unless `KQ_TUNE_LIVE=1`. The op reads them when it is + called and builds its graph, so a compiled function keeps the kernel choice of its trace: + + - `KQ_QMM_LORA` - `0` builds the unfused ops on every call. `2` runs the kernel at every K and + N from 1024 rows, inside its other bounds. Unset, `1` or any other value keeps the bounds on + K, N, rows and rank. + - `KQ_QMM_LORA_MIN_ROWS` - the lowest row count at which the kernel runs, an integer of at + least 1. A value below 1 keeps the default, with a warning on stderr. Default 1024. + + A function compiled with `shapeless=True` keeps the path of its trace at every later shape. A + trace that ran the kernel runs it at every row count, below 1024 rows and outside the bounds on + rows too, and the result is right there. A trace that built the unfused ops keeps them. A + shapeless function cannot hold a call with `bands` that builds the unfused ops, since the stock + slices that cut the bands have no shape rule: that call raises a `ValueError` that names + `quantized_matmul_lora`. Compile such a function with fixed shapes. + + `mx.grad` gives the gradient with respect to x, `lora_a`, `lora_b` and `strength`, alone or in + any combination, at every row count: the gradient of the kernel runs the unfused ops. The weight + is frozen, as in `quantized_matmul`. `mx.jvp` and `mx.vmap` run with respect to `lora_a`, + `lora_b` and `strength` at every row count, on the unfused ops. With respect to x both raise + `Not implemented for KQuantMatmul` at every row count, since `quantized_matmul` has neither. + + The unfused ops add each band with `mx.addmm`. On the GPU that stock op misreads a first + argument of one column that is a view of a wider array, so the op hands it a dense copy of a + one-column band ([upstream-defects.md](upstream-defects.md)). Wider bands are read in place. + +## Quantized matmul with a residual epilogue + +- `quantized_matmul_residual` - `y = residual + gate * (x @ W^T)` for a packed weight: the + residual add after the output projection of a transformer block, with the gate of an adaptive + norm where the block has one. `residual` has the output shape. `gate` is [N], or x's batch dims, + then 1, then N (one gate row per batch item), and `gate=None` gives `residual + x @ W^T`. The + residual and the gate run in the output dtype, the dtype `quantized_matmul` gives for x. + + On a NAX GPU (M5), for `q4_k`, `q5_k`, `q6_k` and `q8_0` weights on bfloat16 activations, from + 1024 activation rows, one dispatch walks K into one f32 sum per output, then reads its tile of the + residual and its columns of the gate and stores `residual + gate * sum`, rounded once. The product + is never written, where the unfused ops write it, then read it and the residual and write the + sum, or two passes on stock elementwise ops. Each threadgroup owns a 128 by 64 tile of y, or 64 by + 64 for `q4_k` and `q8_0` above K 10240, the K ceiling of the large-row tile of + `quantized_matmul`. A gate per batch item runs the batch in the grid. The `q8_0` kernel walks K + in tiles of 64, so a `q8_0` weight whose K is no multiple of 64 runs the unfused ops. + + At Qwen-Image 2.1 step shapes (4096 rows a batch item, N 4096, bfloat16) on an M5 Max, the call + costs what the product alone costs: 0.928 of the unfused ops at K 4096 on `q5_k` and 0.973 at K + 12288 on `q4_k`, at batch 1, against 0.916 and 0.973 for the product alone. A compiled + `x + gate * y` after the product gives 0.982 and 0.984. + + Every other case builds the same sum from `quantized_matmul` and stock ops, which round the + product and the gated product to the output dtype before the add: other codecs, float16 + activations, fewer rows, a `q8_0` K off a multiple of 64, CPU streams and GPUs without NAX. + The op rounds once from 1024 rows counted across the batch, so x of shape [2, 512, K] runs the + kernel while each [512, K] half on its own runs the unfused ops, and the two can differ in the + last bit. + + Lever, read once per process unless `KQ_TUNE_LIVE=1`. The op reads it when it is called and + builds its graph, so a compiled function keeps the kernel choice of its trace: + + - `KQ_QMM_RESID` - `0` builds the unfused ops on every call. `2` runs the kernel at every row + count. Unset, `1` or any other value runs it from 1024 rows. + + `mx.grad` gives the gradient with respect to x, the residual and the gate at every row count: + the gradient of the kernel runs the unfused ops. The weight is frozen, as in `quantized_matmul`. + `mx.jvp` and `mx.vmap` run with respect to the residual and the gate. With respect to x both + raise `Not implemented for KQuantMatmul`, since `quantized_matmul` has neither. + +## 3x3 convolution + +- `conv2d_3x3` - `y = conv(x, w) + bias` for a 3x3 kernel, stride 1, padding 1, channels last: + the convolutions that carry the decoder of an image VAE. `x` is [B, H, W, C], `w` is + [O, 3, 3, C] in the layout of `nn.Conv2d`, and `bias` is [O] or None. `w` and `bias` run in the + dtype of x. The result is the cross-correlation `nn.Conv2d` computes, with zeros outside the + image. + + On a NAX GPU (M5), for bfloat16 with C a multiple of 16, from 1024 output pixels, one dispatch + computes y as an implicit GEMM: rows are output pixels, columns are the O filters, and K walks the + 9 taps and C channels, 32 channels a step, into one f32 sum per output. Each row reads its tap + from x in place, zeros where the tap falls off the image, so no patch buffer is written. The bias + is added in f32 and y is rounded once. Each threadgroup owns a 128 by 128 tile of y, or 128 by 64 + where a 64 column tile pads O less (O 144 and 288, for example). + + On an M5 Max, at the decoder shapes of Z-Image and Qwen-Image 2.1 (H and W 256 to 1024, C 128 to + 1152), the kernel runs 31 to 54 TFLOPS, 0.25 to 0.34 of the time of MLX's convolution without the + Winograd transform. MLX's Winograd path is faster where O is 256 to 1152, with an error to the + float64 result 35 times larger on a 256 channel layer. At O 128 the kernel takes about half the + time of the Winograd path, and at O 144 MLX takes no Winograd transform. A whole Z-Image decode + at 1024 pixels takes 0.36 of its time on the convolution without Winograd. + + Every other case builds MLX's `conv_general` on the flipped kernel, which takes its implicit GEMM + and not the Winograd transform, then a stock add of the bias, which rounds the convolution before + the add: float16 and float32, C off a multiple of 16, fewer pixels, CPU streams and GPUs without + NAX. + + Levers. The op reads `KQ_CONV_NAX` when it is called and builds its graph, so a compiled function + keeps the kernel choice of its trace. `KQ_CONV_TILE` is read at dispatch. Each is read once per + process unless `KQ_TUNE_LIVE=1`: + + - `KQ_CONV_NAX` - `0` builds the composition on every call. `2` runs the kernel at every pixel + count. Unset, `1` or any other value runs it from 1024 pixels. + - `KQ_CONV_TILE` - `bm,bn,wm,wn` pins the tile. The metallib carries `128,128,4,4`, + `128,64,4,2`, `64,128,2,4` and `64,64,2,2`; any other tile raises a `ValueError` that names the + missing kernel. + + `mx.grad`, `mx.jvp` and `mx.vmap` run the composition, so they give what MLX's convolution gives. + ## MoE GLU Fused gate/up expert matvecs with the GLU epilogue applied in the same dispatch, so each activation @@ -440,13 +995,14 @@ into a packed layout first. (scales and biases, bits 8, group 64) dequantize on the tile stage. `return_lse=True` adds per-row log-sum-exp. On NAX GPUs one query at head dim 512 with up to 8 query heads per KV head runs the chunk pass on the matrix units once KV heads x batch x keys reaches 3072. That needs K and V - strides in multiples of 8 elements and no `starts`, `ends` or q8 operands. Below 3072 the + strides in multiples of 8 elements and no `ends` or q8 operands. With `starts` it skips the + padded keys of each row as the simdgroup kernel does. Below 3072 the simdgroup kernel's finer splits keep more of the GPU busy. Eight simdgroups share each 256-key tile, each scoring 32 keys over the full head dim and then accumulating a 64-column slice of `P @ V`. The automatic split count is the largest power of two that keeps KV heads x batch x splits at 32 or under, and 1 once KV heads x batch passes 32. It ignores the depth, so it falls as the batch grows. The two kernels round differently, so the same rows are not bit-identical - between this pass, the simdgroup kernel (`starts`, `ends`, q8) and another batch size. + between this pass, the simdgroup kernel (`ends`, q8) and another batch size. `KQ_GQA_NAX=0` forces the simdgroup kernel. - **`sdpa_decode_gqa_cascade`** - shared-prefix batched decode: every row attends one common prefix plus its own private suffix. The prefix is walked once for all rows on the matrix-unit tile, private @@ -543,6 +1099,36 @@ Tuning levers (defaults are right for normal use): after attention or the next layer's input norm. Rows wider than 16384 and CPU streams run the composition. +## Attention q and k norm with RoPE + +- `qk_norm_rope` - the per-head RMSNorm and rotary embedding that a diffusion transformer applies + to q and k before attention, both tensors in one dispatch. Each head row `x` of q (weight `q_w`) + and of k (weight `k_w`) becomes `y = x * rsqrt(mean(x^2) + eps) * w`, and then each interleaved + pair `(y[2i], y[2i + 1])` of sequence row `l` rotates by the angle whose cosine and sine are + `cos[l, i]` and `sin[l, i]`. The tables hold one angle per row and pair, so any position scheme + works, such as the caption index and image row and column of a three-axis image RoPE. The angles + of `mx.fast.rope` follow the sequence index and cannot express these positions. + + The op reads q and k `[B, L, H, D]` in place through their strides, so the q and k slices of one + fused qkv output need no copy, and lazy transposed views work the same way. The batch shares the + `[L, D / 2]` tables. It returns `(q_out, k_out)`, each row-contiguous `[B, H, L, D]`, the layout + attention reads, so no transpose copy follows. + + q and k are float16 or bfloat16 and the outputs keep that dtype. The weights and tables may be any + float dtype and are used in f32. The norm and the rotation run in f32 with one rounding at the + write, so nearly every output element is the exact result rounded to the output dtype. The unfused + composition also rounds the normalized value before the rotation. + + One threadgroup per `(b, l)` row covers every head of q and k, with 16 lanes per head row. Each + lane loads its cosines, sines and weights once per row and reuses them across heads. Head dims 64, + 96, 128 and 256 have kernels in both dtypes. Views whose head-dim stride is 1 and whose rows start + on the vector width take vector loads, and any other view takes element loads through its strides. + + CPU streams, and GPU calls at other even head dims, run the same math as stock MLX ops, an f32 norm + and rotation followed by one cast to the input dtype and a transposed copy. The op is for + inference. On the GPU the fused kernel has no gradient or vmap, and `mx.grad` or `mx.vmap` through + it raises. + ## Hadamard rotation - **`hadamard_rotate`** - the signed block Walsh-Hadamard rotation a Hadamard-folded weight expects @@ -563,10 +1149,14 @@ Tuning levers (defaults are right for normal use): - **`codecs`** - the list of supported codec names. - **`metallib_loads`** / **`metallib_dir`** - whether the bundled metallib opened on the device, and where it lives. +- `metallib_kernels` - the sorted names of the kernels in the bundled metallib (empty without + Metal), the names `KQ_NAX_TILE` and `KQ_QMM_TILE` select from. - **`nax_available`** / **`nax_gather_enabled`** - whether the GPU exposes NAX tensor units, and whether the sorted-gather NAX GEMM kernels (`gather_qmm_rhs_nax`, the `gather_qmm_seg` NAX walk) are reachable for a codec. - **`cpu_neon_available`** - whether the arm64 NEON int8 GEMV path is compiled in. +- `gpu_core_count` / `max_threadgroup_memory` - the GPU core count from IOKit (`KQ_GPU_CORES` + overrides it) and the threadgroup memory limit in bytes, both 0 when unknown. ## Feeder-loop primitives diff --git a/docs/upstream-defects.md b/docs/upstream-defects.md new file mode 100644 index 00000000..26a8b5e7 --- /dev/null +++ b/docs/upstream-defects.md @@ -0,0 +1,42 @@ +# Upstream defects + +Defects of stock dependencies that mlx-kquant guards against in its own code. Each entry gives +the dependency and the versions seen, what goes wrong, the smallest repro, the guard and the test +that fails without the guard. An entry stays until no supported version of the dependency has the +defect. + +## `mx.addmm` misreads a first argument of one column on the GPU + +- Dependency: mlx 0.32.2 and 0.32.3, GPU streams. The CPU stream is right. +- Affected: `mx.addmm(c, a, b)` whose result has one column and more than one row, so c is + [..., M, 1] with M above 1, when c is a view whose M values are not consecutive in memory: a + column of a wider array, every second row of a taller one, reversed rows, a transposed row. + float32, float16 and bfloat16, lazy or evaluated c, with or without batch dims. +- What goes wrong: the op reads c at a fixed step from its first element that is not the row + stride of c. For a column of a wider array it returns the elements that follow c[0, 0] in its + row, one per output row. For a transposed row it reads past the buffer. No error is raised, and + `a @ b` is added to the wrong values. +- Not affected: a c of two or more columns, or of one row, in any view; a dense c, a band of rows + of a one-column array among them; a bias of one element broadcast to one column, as in + `nn.Linear` with one output; a c of one column added to a wider product; a and b in any view; + `a @ b + c`. + +Smallest repro: + +```python +import mlx.core as mx + +wide = mx.arange(21, dtype=mx.float32).reshape(3, 7) +c = wide[:, 3:4] # [3, 1], the values 3, 10, 17 +a, b = mx.zeros((3, 2)), mx.zeros((2, 1)) +print(mx.addmm(c, a, b).reshape(-1)) # GPU: [3, 4, 5] +print(mx.addmm(c, a, b, stream=mx.cpu).reshape(-1)) # [3, 10, 17] +print(mx.addmm(mx.contiguous(c), a, b).reshape(-1)) # GPU: [3, 10, 17] +``` + +- Guard: `lora_addmm` in `src/kquant_matmul_lora.cpp`. The unfused ops of + `quantized_matmul_lora` add each band of columns with `mx.addmm` on a slice of the base + product. A slice of one column goes in as `mx.contiguous(c)`, which copies it to dense memory + and passes a dense c through with no copy. A slice of two or more columns is read in place. +- Test: `test_lora_band_of_one_column` in `tests/test_matmul_lora.py`. It holds every column of + a call with one-column bands against a float64 sum. diff --git a/metal/kq_conv_nax.metal b/metal/kq_conv_nax.metal new file mode 100644 index 00000000..c39a3df2 --- /dev/null +++ b/metal/kq_conv_nax.metal @@ -0,0 +1,21 @@ +// clang-format off +// NAX 3x3 convolution as an implicit GEMM; see kq_conv_nax.h. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" +#include "mlx/backend/metal/kernels/kq_conv_nax.h" + +#define instantiate_kq_conv3x3_nax(type, bm, bn, wm, wn) \ + instantiate_kernel( \ + "kq_conv3x3_nax_" #type "_bm" #bm "_bn" #bn "_wm" #wm "_wn" #wn, \ + kq_conv3x3_nax, \ + type, \ + bm, \ + bn, \ + wm, \ + wn) + +instantiate_kq_conv3x3_nax(bfloat16_t, 128, 128, 4, 4) +instantiate_kq_conv3x3_nax(bfloat16_t, 128, 64, 4, 2) +instantiate_kq_conv3x3_nax(bfloat16_t, 64, 128, 2, 4) +instantiate_kq_conv3x3_nax(bfloat16_t, 64, 64, 2, 2) + // clang-format on diff --git a/metal/kq_lora_epilogue.metal b/metal/kq_lora_epilogue.metal index b18f4aae..b8946417 100644 --- a/metal/kq_lora_epilogue.metal +++ b/metal/kq_lora_epilogue.metal @@ -10,6 +10,7 @@ instantiate_kernel("kq_lora_mix_z_bfloat16_t", kq_lora_mix_z, bfloat16_t) instantiate_kernel("kq_lora_mix_z_float16_t", kq_lora_mix_z, float16_t) instantiate_kernel("kq_lora_mix_apply_bfloat16_t", kq_lora_mix_apply, bfloat16_t) instantiate_kernel("kq_lora_mix_apply_float16_t", kq_lora_mix_apply, float16_t) +instantiate_kernel("kq_lora_densify_1", kq_lora_densify, uint8_t) instantiate_kernel("kq_lora_densify_2", kq_lora_densify, uint16_t) instantiate_kernel("kq_lora_densify_4", kq_lora_densify, uint32_t) // clang-format on diff --git a/metal/kq_qk_norm_rope.metal b/metal/kq_qk_norm_rope.metal new file mode 100644 index 00000000..8a330d8e --- /dev/null +++ b/metal/kq_qk_norm_rope.metal @@ -0,0 +1,17 @@ +// clang-format off +// Fused per-head RMSNorm and RoPE instantiations; see kq_qk_norm_rope.h. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/kq_qk_norm_rope.h" + +#define instantiate_kq_qk_norm_rope(type, d) \ + instantiate_kernel("kq_qk_norm_rope_" #type "_" #d, kq_qk_norm_rope, type, d) + +instantiate_kq_qk_norm_rope(bfloat16_t, 64) +instantiate_kq_qk_norm_rope(bfloat16_t, 96) +instantiate_kq_qk_norm_rope(bfloat16_t, 128) +instantiate_kq_qk_norm_rope(bfloat16_t, 256) +instantiate_kq_qk_norm_rope(float16_t, 64) +instantiate_kq_qk_norm_rope(float16_t, 96) +instantiate_kq_qk_norm_rope(float16_t, 128) +instantiate_kq_qk_norm_rope(float16_t, 256) + // clang-format on diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index 35903df3..f15d3905 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -6,6 +6,10 @@ #include "mlx/backend/metal/kernels/quantized_utils.h" #include "mlx/backend/metal/kernels/kq_quantized.h" +// The ops cast float32 activations to bfloat16 before dispatch, so the +// kernels that take activations exist for bfloat16 and float16 only, and +// float keeps its dequantize kernels. + #define instantiate_kquant_batched(func, type, gs, bits, batched, codec) \ instantiate_kernel( \ "kquant_" #codec "_" #func "_" #type "_gs_" #gs "_b_" #bits \ @@ -27,6 +31,10 @@ aligned_N, \ batched) +// Split-K kernels come in pairs: the plain name stores T partials, and +// the _pf32 twin stores float partials for kquant_qmm_splitk_accum_*_pf32 +// (float16 by default, bfloat16 under KQ_SPLITK_F32=1). + #define instantiate_kquant_qmm_t_splitk(type, gs, bits, aligned_N, codec) \ instantiate_kernel( \ "kquant_" #codec "_qmm_t_splitk_" #type "_gs_" #gs "_b_" #bits \ @@ -35,7 +43,17 @@ type, \ gs, \ bits, \ - aligned_N) + aligned_N) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_splitk_" #type "_gs_" #gs "_b_" #bits \ + "_alN_" #aligned_N "_pf32", \ + kq_ ## codec ## _qmm_t_splitk, \ + type, \ + gs, \ + bits, \ + aligned_N, \ + 0, \ + float) // BM=16 and BM=8 split-K tiles for M <= 16 and M <= 8 (codecs with the // small_bm template arm). @@ -48,7 +66,17 @@ gs, \ bits, \ aligned_N, \ - 16) + 16) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_splitk_bm16_" #type "_gs_" #gs "_b_" #bits \ + "_alN_" #aligned_N "_pf32", \ + kq_ ## codec ## _qmm_t_splitk, \ + type, \ + gs, \ + bits, \ + aligned_N, \ + 16, \ + float) #define instantiate_kquant_qmm_t_splitk_bm8(type, gs, bits, aligned_N, codec) \ instantiate_kernel( \ "kquant_" #codec "_qmm_t_splitk_bm8_" #type "_gs_" #gs "_b_" #bits \ @@ -58,7 +86,17 @@ gs, \ bits, \ aligned_N, \ - 8) + 8) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_splitk_bm8_" #type "_gs_" #gs "_b_" #bits \ + "_alN_" #aligned_N "_pf32", \ + kq_ ## codec ## _qmm_t_splitk, \ + type, \ + gs, \ + bits, \ + aligned_N, \ + 8, \ + float) #define instantiate_kquant_qmm_n(type, gs, bits, batched, codec) \ instantiate_kernel( \ @@ -148,7 +186,7 @@ instantiate_kquant_gather_qmm_n(type, 32, 8, q8_0) \ instantiate_kquant_dequantize(type, 32, 8, q8_0) -instantiate_kquant_q8_0_for_type(float) +instantiate_kquant_dequantize(float, 32, 8, q8_0) instantiate_kquant_q8_0_for_type(bfloat16_t) instantiate_kquant_q8_0_for_type(float16_t) @@ -181,7 +219,7 @@ instantiate_kquant_q8_0_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 32, 5, q5_1) \ instantiate_kquant_dequantize(type, 32, 5, q5_1) -instantiate_kquant_q5_1_for_type(float) +instantiate_kquant_dequantize(float, 32, 5, q5_1) instantiate_kquant_q5_1_for_type(bfloat16_t) instantiate_kquant_q5_1_for_type(float16_t) @@ -214,7 +252,7 @@ instantiate_kquant_q5_1_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 32, 4, q4_0) \ instantiate_kquant_dequantize(type, 32, 4, q4_0) -instantiate_kquant_q4_0_for_type(float) +instantiate_kquant_dequantize(float, 32, 4, q4_0) instantiate_kquant_q4_0_for_type(bfloat16_t) instantiate_kquant_q4_0_for_type(float16_t) @@ -247,7 +285,7 @@ instantiate_kquant_q4_0_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 32, 4, q4_1) \ instantiate_kquant_dequantize(type, 32, 4, q4_1) -instantiate_kquant_q4_1_for_type(float) +instantiate_kquant_dequantize(float, 32, 4, q4_1) instantiate_kquant_q4_1_for_type(bfloat16_t) instantiate_kquant_q4_1_for_type(float16_t) @@ -280,7 +318,7 @@ instantiate_kquant_q4_1_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 32, 5, q5_0) \ instantiate_kquant_dequantize(type, 32, 5, q5_0) -instantiate_kquant_q5_0_for_type(float) +instantiate_kquant_dequantize(float, 32, 5, q5_0) instantiate_kquant_q5_0_for_type(bfloat16_t) instantiate_kquant_q5_0_for_type(float16_t) @@ -313,7 +351,7 @@ instantiate_kquant_q5_0_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 4, q4_k) \ instantiate_kquant_dequantize(type, 256, 4, q4_k) -instantiate_kquant_q4_k_for_type(float) +instantiate_kquant_dequantize(float, 256, 4, q4_k) instantiate_kquant_q4_k_for_type(bfloat16_t) instantiate_kquant_q4_k_for_type(float16_t) @@ -346,7 +384,7 @@ instantiate_kquant_q4_k_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 5, q5_k) \ instantiate_kquant_dequantize(type, 256, 5, q5_k) -instantiate_kquant_q5_k_for_type(float) +instantiate_kquant_dequantize(float, 256, 5, q5_k) instantiate_kquant_q5_k_for_type(bfloat16_t) instantiate_kquant_q5_k_for_type(float16_t) @@ -381,7 +419,7 @@ instantiate_kquant_q5_k_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 6, q6_k) \ instantiate_kquant_dequantize(type, 256, 6, q6_k) -instantiate_kquant_q6_k_for_type(float) +instantiate_kquant_dequantize(float, 256, 6, q6_k) instantiate_kquant_q6_k_for_type(bfloat16_t) instantiate_kquant_q6_k_for_type(float16_t) @@ -405,7 +443,6 @@ instantiate_kquant_q6_k_for_type(float16_t) instantiate_mv_ext(codec, type, gs, bits, 11) \ instantiate_mv_ext(codec, type, gs, bits, 12) #define instantiate_mv_ext_all(codec, gs, bits) \ - instantiate_mv_ext_for_type(codec, gs, bits, float) \ instantiate_mv_ext_for_type(codec, gs, bits, bfloat16_t) \ instantiate_mv_ext_for_type(codec, gs, bits, float16_t) instantiate_mv_ext_all(q8_0, 32, 8) @@ -452,7 +489,6 @@ instantiate_mv_ext_all(ptq1_0, 128, 1) instantiate_mv_ext_nr2(codec, type, gs, bits, 11) \ instantiate_mv_ext_nr2(codec, type, gs, bits, 12) #define instantiate_mv_ext_nr2_all(codec, gs, bits) \ - instantiate_mv_ext_nr2_for_type(codec, gs, bits, float) \ instantiate_mv_ext_nr2_for_type(codec, gs, bits, bfloat16_t) \ instantiate_mv_ext_nr2_for_type(codec, gs, bits, float16_t) instantiate_mv_ext_nr2_all(q6_k, 256, 6) @@ -478,7 +514,6 @@ instantiate_mv_ext_nr2_all(q6_k, 256, 6) instantiate_mv_ext_sb(codec, type, gs, bits, 11) \ instantiate_mv_ext_sb(codec, type, gs, bits, 12) #define instantiate_mv_ext_sb_all(codec, gs, bits) \ - instantiate_mv_ext_sb_for_type(codec, gs, bits, float) \ instantiate_mv_ext_sb_for_type(codec, gs, bits, bfloat16_t) \ instantiate_mv_ext_sb_for_type(codec, gs, bits, float16_t) instantiate_mv_ext_sb_all(q6_k, 256, 6) @@ -504,7 +539,6 @@ instantiate_mv_ext_sb_all(q6_k, 256, 6) instantiate_mv_ext_nx(codec, type, gs, bits, 11, nx) \ instantiate_mv_ext_nx(codec, type, gs, bits, 12, nx) #define instantiate_mv_ext_nx_all(codec, gs, bits, nx) \ - instantiate_mv_ext_nx_for_type(codec, gs, bits, float, nx) \ instantiate_mv_ext_nx_for_type(codec, gs, bits, bfloat16_t, nx) \ instantiate_mv_ext_nx_for_type(codec, gs, bits, float16_t, nx) instantiate_mv_ext_nx_all(q6_k, 256, 6, 16) @@ -565,7 +599,7 @@ instantiate_mv_ext_hd_all(q6_k, 256, 6) instantiate_kquant_gather_qmm_n(type, 256, 3, q3_k) \ instantiate_kquant_dequantize(type, 256, 3, q3_k) -instantiate_kquant_q3_k_for_type(float) +instantiate_kquant_dequantize(float, 256, 3, q3_k) instantiate_kquant_q3_k_for_type(bfloat16_t) instantiate_kquant_q3_k_for_type(float16_t) @@ -598,7 +632,7 @@ instantiate_kquant_q3_k_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 2, q2_k) \ instantiate_kquant_dequantize(type, 256, 2, q2_k) -instantiate_kquant_q2_k_for_type(float) +instantiate_kquant_dequantize(float, 256, 2, q2_k) instantiate_kquant_q2_k_for_type(bfloat16_t) instantiate_kquant_q2_k_for_type(float16_t) @@ -632,7 +666,7 @@ instantiate_kquant_q2_k_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 32, 4, iq4_nl) \ instantiate_kquant_dequantize(type, 32, 4, iq4_nl) -instantiate_kquant_iq4_nl_for_type(float) +instantiate_kquant_dequantize(float, 32, 4, iq4_nl) instantiate_kquant_iq4_nl_for_type(bfloat16_t) instantiate_kquant_iq4_nl_for_type(float16_t) @@ -647,8 +681,6 @@ instantiate_kquant_iq4_nl_for_type(float16_t) instantiate_kquant_qmm_t(type, 32, 4, true, 1, mxfp4) \ instantiate_kquant_qmm_t(type, 32, 4, false, 0, mxfp4) \ instantiate_kquant_qmm_t(type, 32, 4, false, 1, mxfp4) \ - instantiate_kquant_qmm_t_splitk(type, 32, 4, true, mxfp4) \ - instantiate_kquant_qmm_t_splitk(type, 32, 4, false, mxfp4) \ instantiate_kquant_qmm_n(type, 32, 4, 0, mxfp4) \ instantiate_kquant_qmm_n(type, 32, 4, 1, mxfp4) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 32, 4, mxfp4) \ @@ -660,7 +692,7 @@ instantiate_kquant_iq4_nl_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 32, 4, mxfp4) \ instantiate_kquant_dequantize(type, 32, 4, mxfp4) -instantiate_kquant_mxfp4_for_type(float) +instantiate_kquant_dequantize(float, 32, 4, mxfp4) instantiate_kquant_mxfp4_for_type(bfloat16_t) instantiate_kquant_mxfp4_for_type(float16_t) @@ -675,8 +707,6 @@ instantiate_kquant_mxfp4_for_type(float16_t) instantiate_kquant_qmm_t(type, 64, 4, true, 1, nvfp4) \ instantiate_kquant_qmm_t(type, 64, 4, false, 0, nvfp4) \ instantiate_kquant_qmm_t(type, 64, 4, false, 1, nvfp4) \ - instantiate_kquant_qmm_t_splitk(type, 64, 4, true, nvfp4) \ - instantiate_kquant_qmm_t_splitk(type, 64, 4, false, nvfp4) \ instantiate_kquant_qmm_n(type, 64, 4, 0, nvfp4) \ instantiate_kquant_qmm_n(type, 64, 4, 1, nvfp4) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 64, 4, nvfp4) \ @@ -688,7 +718,7 @@ instantiate_kquant_mxfp4_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 64, 4, nvfp4) \ instantiate_kquant_dequantize(type, 64, 4, nvfp4) -instantiate_kquant_nvfp4_for_type(float) +instantiate_kquant_dequantize(float, 64, 4, nvfp4) instantiate_kquant_nvfp4_for_type(bfloat16_t) instantiate_kquant_nvfp4_for_type(float16_t) @@ -720,7 +750,7 @@ instantiate_kquant_nvfp4_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 4, iq4_xs) \ instantiate_kquant_dequantize(type, 256, 4, iq4_xs) -instantiate_kquant_iq4_xs_for_type(float) +instantiate_kquant_dequantize(float, 256, 4, iq4_xs) instantiate_kquant_iq4_xs_for_type(bfloat16_t) instantiate_kquant_iq4_xs_for_type(float16_t) @@ -752,7 +782,7 @@ instantiate_kquant_iq4_xs_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 3, iq3_xxs) \ instantiate_kquant_dequantize(type, 256, 3, iq3_xxs) -instantiate_kquant_iq3_xxs_for_type(float) +instantiate_kquant_dequantize(float, 256, 3, iq3_xxs) instantiate_kquant_iq3_xxs_for_type(bfloat16_t) instantiate_kquant_iq3_xxs_for_type(float16_t) @@ -784,7 +814,7 @@ instantiate_kquant_iq3_xxs_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 3, iq3_s) \ instantiate_kquant_dequantize(type, 256, 3, iq3_s) -instantiate_kquant_iq3_s_for_type(float) +instantiate_kquant_dequantize(float, 256, 3, iq3_s) instantiate_kquant_iq3_s_for_type(bfloat16_t) instantiate_kquant_iq3_s_for_type(float16_t) @@ -816,7 +846,7 @@ instantiate_kquant_iq3_s_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 2, iq2_xxs) \ instantiate_kquant_dequantize(type, 256, 2, iq2_xxs) -instantiate_kquant_iq2_xxs_for_type(float) +instantiate_kquant_dequantize(float, 256, 2, iq2_xxs) instantiate_kquant_iq2_xxs_for_type(bfloat16_t) instantiate_kquant_iq2_xxs_for_type(float16_t) @@ -848,7 +878,7 @@ instantiate_kquant_iq2_xxs_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 2, iq2_xs) \ instantiate_kquant_dequantize(type, 256, 2, iq2_xs) -instantiate_kquant_iq2_xs_for_type(float) +instantiate_kquant_dequantize(float, 256, 2, iq2_xs) instantiate_kquant_iq2_xs_for_type(bfloat16_t) instantiate_kquant_iq2_xs_for_type(float16_t) @@ -880,7 +910,7 @@ instantiate_kquant_iq2_xs_for_type(float16_t) instantiate_kquant_gather_qmm_n(type, 256, 2, iq2_s) \ instantiate_kquant_dequantize(type, 256, 2, iq2_s) -instantiate_kquant_iq2_s_for_type(float) +instantiate_kquant_dequantize(float, 256, 2, iq2_s) instantiate_kquant_iq2_s_for_type(bfloat16_t) instantiate_kquant_iq2_s_for_type(float16_t) @@ -911,7 +941,7 @@ instantiate_kquant_iq2_s_for_type(float16_t) instantiate_kquant_gather_qmm_seg_t(type, 256, 1, false, iq1_s) \ instantiate_kquant_gather_qmm_n(type, 256, 1, iq1_s) \ instantiate_kquant_dequantize(type, 256, 1, iq1_s) -instantiate_kquant_iq1_s_for_type(float) +instantiate_kquant_dequantize(float, 256, 1, iq1_s) instantiate_kquant_iq1_s_for_type(bfloat16_t) instantiate_kquant_iq1_s_for_type(float16_t) @@ -942,7 +972,7 @@ instantiate_kquant_iq1_s_for_type(float16_t) instantiate_kquant_gather_qmm_seg_t(type, 256, 1, false, iq1_m) \ instantiate_kquant_gather_qmm_n(type, 256, 1, iq1_m) \ instantiate_kquant_dequantize(type, 256, 1, iq1_m) -instantiate_kquant_iq1_m_for_type(float) +instantiate_kquant_dequantize(float, 256, 1, iq1_m) instantiate_kquant_iq1_m_for_type(bfloat16_t) instantiate_kquant_iq1_m_for_type(float16_t) @@ -973,7 +1003,7 @@ instantiate_kquant_iq1_m_for_type(float16_t) instantiate_kquant_gather_qmm_seg_t(type, 256, 1, false, stq1_0) \ instantiate_kquant_gather_qmm_n(type, 256, 1, stq1_0) \ instantiate_kquant_dequantize(type, 256, 1, stq1_0) -instantiate_kquant_stq1_0_for_type(float) +instantiate_kquant_dequantize(float, 256, 1, stq1_0) instantiate_kquant_stq1_0_for_type(bfloat16_t) instantiate_kquant_stq1_0_for_type(float16_t) @@ -1006,25 +1036,54 @@ instantiate_kquant_stq1_0_for_type(float16_t) instantiate_kquant_gather_qmm_seg_t(type, gs, bits, false, codec) \ instantiate_kquant_gather_qmm_n(type, gs, bits, codec) \ instantiate_kquant_dequantize(type, gs, bits, codec) -instantiate_kquant_prism_for_type(pq2_0, 128, 2, float) -instantiate_kquant_prism_for_type(q2_0, 64, 2, float) +instantiate_kquant_dequantize(float, 128, 2, pq2_0) +instantiate_kquant_dequantize(float, 64, 2, q2_0) instantiate_kquant_prism_for_type(pq2_0, 128, 2, bfloat16_t) instantiate_kquant_prism_for_type(q2_0, 64, 2, bfloat16_t) instantiate_kquant_prism_for_type(pq2_0, 128, 2, float16_t) instantiate_kquant_prism_for_type(q2_0, 64, 2, float16_t) -instantiate_kquant_prism_for_type(ptq1_0, 128, 1, float) +instantiate_kquant_dequantize(float, 128, 1, ptq1_0) instantiate_kquant_prism_for_type(ptq1_0, 128, 1, bfloat16_t) instantiate_kquant_prism_for_type(ptq1_0, 128, 1, float16_t) -// Register-resident MMA verify kernels (kq_verify_mma.h), M <= 8, half -// activations staged from bf16 or f16. +// Register-resident MMA verify kernels (kq_verify_mma.h), M <= 8. The +// activations stay in bf16 or f16 and each block sums in float. The _h +// twins stage half activations and sum in half (KQ_VERIFY_MMA_HALF); +// ptq1_0 has none. #define instantiate_kquant_verify_mma(type, gs, bits, codec) \ instantiate_kernel( \ "kquant_" #codec "_verify_mma_" #type "_gs_" #gs "_b_" #bits, \ kq_ ## codec ## _verify_mma, \ type, \ gs, \ - bits) + bits) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_mma_" #type "_gs_" #gs "_b_" #bits \ + "_pf32", \ + kq_ ## codec ## _verify_mma, \ + type, \ + gs, \ + bits, \ + float) +#define instantiate_kquant_verify_mma_h(type, gs, bits, codec) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_mma_" #type "_gs_" #gs "_b_" #bits \ + "_h", \ + kq_ ## codec ## _verify_mma, \ + type, \ + gs, \ + bits, \ + type, \ + true) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_mma_" #type "_gs_" #gs "_b_" #bits \ + "_h_pf32", \ + kq_ ## codec ## _verify_mma, \ + type, \ + gs, \ + bits, \ + float, \ + true) instantiate_kquant_verify_mma(bfloat16_t, 128, 2, pq2_0) instantiate_kquant_verify_mma(bfloat16_t, 64, 2, q2_0) instantiate_kquant_verify_mma(float16_t, 128, 2, pq2_0) @@ -1035,6 +1094,14 @@ instantiate_kquant_verify_mma(bfloat16_t, 32, 4, q4_0) instantiate_kquant_verify_mma(float16_t, 32, 4, q4_0) instantiate_kquant_verify_mma(bfloat16_t, 32, 8, q8_0) instantiate_kquant_verify_mma(float16_t, 32, 8, q8_0) +instantiate_kquant_verify_mma_h(bfloat16_t, 128, 2, pq2_0) +instantiate_kquant_verify_mma_h(float16_t, 128, 2, pq2_0) +instantiate_kquant_verify_mma_h(bfloat16_t, 64, 2, q2_0) +instantiate_kquant_verify_mma_h(float16_t, 64, 2, q2_0) +instantiate_kquant_verify_mma_h(bfloat16_t, 32, 4, q4_0) +instantiate_kquant_verify_mma_h(float16_t, 32, 4, q4_0) +instantiate_kquant_verify_mma_h(bfloat16_t, 32, 8, q8_0) +instantiate_kquant_verify_mma_h(float16_t, 32, 8, q8_0) // clang-format on // clang-format off @@ -1058,7 +1125,6 @@ instantiate_kquant_verify_mma(float16_t, 32, 8, q8_0) instantiate_kquant_gather_qmm_rhs(type, gs, bits, false, 16, codec) #define instantiate_kquant_gather_qmm_rhs_codec(gs, bits, codec) \ - instantiate_kquant_gather_qmm_rhs_for_type(float, gs, bits, codec) \ instantiate_kquant_gather_qmm_rhs_for_type(float16_t, gs, bits, codec) \ instantiate_kquant_gather_qmm_rhs_for_type(bfloat16_t, gs, bits, codec) @@ -1143,12 +1209,14 @@ instantiate_kquant_gather_qmm_rhs_codec(128, 1, ptq1_0) } // Split-K partial fold for qmm_t_splitk: y[i] = (T)(f32 sum over the -// splits axis of partials[z * stride + i]). Partials carry one T rounding -// per slice; the fold accumulates in f32 and rounds once more. One thread -// per output element; cost is noise next to the GEMM pass. -template +// splits axis of partials[z * stride + i]). The fold accumulates in f32 +// and rounds once. Float partials (P = float, the _pf32 kernels) keep the +// result at that one output rounding; T partials add one T rounding per +// slice. One thread per output element; cost is noise next to the GEMM +// pass. +template [[kernel]] void kq_qmm_splitk_accum( - const device T* partials [[buffer(0)]], + const device P* partials [[buffer(0)]], device T* y [[buffer(1)]], const constant int& n_elems [[buffer(2)]], const constant int& splits [[buffer(3)]], @@ -1171,7 +1239,15 @@ template kq_qmm_splitk_accum, \ type) -instantiate_kq_qmm_splitk_accum(float) +#define instantiate_kq_qmm_splitk_accum_pf32(type) \ + instantiate_kernel( \ + "kquant_qmm_splitk_accum_" #type "_pf32", \ + kq_qmm_splitk_accum, \ + type, \ + float) + instantiate_kq_qmm_splitk_accum(bfloat16_t) instantiate_kq_qmm_splitk_accum(float16_t) +instantiate_kq_qmm_splitk_accum_pf32(bfloat16_t) +instantiate_kq_qmm_splitk_accum_pf32(float16_t) // clang-format on diff --git a/metal/kq_quantized_nax.metal b/metal/kq_quantized_nax.metal index 7cf5c2ba..1378f89d 100644 --- a/metal/kq_quantized_nax.metal +++ b/metal/kq_quantized_nax.metal @@ -90,7 +90,6 @@ instantiate_kquant_nax_gather_qmm_rhs(type, gs, bits, false, nn, 64, 64, 2, 2, codec) #define instantiate_kquant_nax_codec(codec, gs, bits) \ - instantiate_kquant_nax_codec_for_type(codec, float, gs, bits) \ instantiate_kquant_nax_codec_for_type(codec, float16_t, gs, bits) \ instantiate_kquant_nax_codec_for_type(codec, bfloat16_t, gs, bits) @@ -129,7 +128,12 @@ "kquant_" #codec "_qmm_t_nax_splitk_" #type "_gs_" #gs "_b_" #bits \ "_bm" #bm "_bn64_bk64_wm2_wn2_alN_" #aligned_N, \ kq_ ## codec ## _qmm_t_nax_splitk, \ - type, gs, bits, aligned_N, bm, 64, 2, 2) + type, gs, bits, aligned_N, bm, 64, 2, 2) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_splitk_" #type "_gs_" #gs "_b_" #bits \ + "_bm" #bm "_bn64_bk64_wm2_wn2_alN_" #aligned_N "_pf32", \ + kq_ ## codec ## _qmm_t_nax_splitk, \ + type, gs, bits, aligned_N, bm, 64, 2, 2, float) #define instantiate_kquant_nax_splitk_bm(codec, gs, bits, bm) \ instantiate_kquant_nax_qmm_t_splitk(float16_t, gs, bits, true, bm, codec) \ instantiate_kquant_nax_qmm_t_splitk(float16_t, gs, bits, false, bm, codec) \ @@ -200,10 +204,14 @@ instantiate_kquant_nax_db64(q5_0, 32, 5) // per-threadgroup work (loader dequant + MMA issue, real DRAM ~half // peak); BM=128 halves the row-tile count and with it total loader // dequant per unique weight. A loads straight from device, so the cost -// is Dtile registers: 8 accumulator fragments vs 4. Dispatched only when -// ceil(M/64) is even and M >= 193 (padding parity with BM=64; q6_k ABBA -// at three shapes: +9-25% M224-512, prefill M512-2048 +8-10% hot, -// +31-40% cold). KQ_NAX_BM128=0 kills. +// is Dtile registers: 8 accumulator fragments vs 4. Dispatched when +// ceil(M/64) is even and M is at or above the per-codec floor (padding +// parity with BM=64; q6_k ABBA at three shapes: +9-25% M224-512, prefill +// M512-2048 +8-10% hot, +31-40% cold), and at every M from +// KQ_NAX_BM128_ODD_M, where the row-live guard in kq_qmm_t_nax_tgp_impl +// runs a last row tile of 64 rows or fewer with the loads and MMA of a +// BM=64 tile. +// KQ_NAX_BM128=0 kills. #define instantiate_kquant_nax_qmm_t_bm128(codec, type, gs, bits) \ instantiate_kquant_nax_qmm_t(type, gs, bits, true, 1, 128, 64, 2, 2, codec) \ instantiate_kquant_nax_qmm_t(type, gs, bits, true, 0, 128, 64, 2, 2, codec) \ @@ -236,6 +244,45 @@ instantiate_kquant_nax_bm128(stq1_0, 256, 1) instantiate_kquant_nax_bm128(pq2_0, 128, 2) instantiate_kquant_nax_bm128(q2_0, 64, 2) instantiate_kquant_nax_bm128(ptq1_0, 128, 1) +// The BM=64 and BM=128 tiles on a K tile of 128 (half super-block or pair +// loaders), per codec. KQ_NAX__BK128=0 keeps the BK=64 kernels. +#define instantiate_kquant_nax_bk128(codec, gs, b, type, al, batched, bm) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_" #type "_gs_" #gs "_b_" #b \ + "_bm" #bm "_bn64_bk128_wm2_wn2" \ + "_alN_" #al "_batch_" #batched, \ + kq_##codec##_qmm_t_nax_bk128, \ + type, \ + gs, \ + b, \ + al, \ + batched, \ + bm, \ + 64, \ + 2, \ + 2) +#define instantiate_kquant_nax_bk128_bm(codec, gs, b, type, bm) \ + instantiate_kquant_nax_bk128(codec, gs, b, type, true, 1, bm) \ + instantiate_kquant_nax_bk128(codec, gs, b, type, true, 0, bm) \ + instantiate_kquant_nax_bk128(codec, gs, b, type, false, 1, bm) \ + instantiate_kquant_nax_bk128(codec, gs, b, type, false, 0, bm) +#define instantiate_kquant_nax_bk128_codec(codec, gs, b) \ + instantiate_kquant_nax_bk128_bm(codec, gs, b, float16_t, 64) \ + instantiate_kquant_nax_bk128_bm(codec, gs, b, bfloat16_t, 64) \ + instantiate_kquant_nax_bk128_bm(codec, gs, b, float16_t, 128) \ + instantiate_kquant_nax_bk128_bm(codec, gs, b, bfloat16_t, 128) +instantiate_kquant_nax_bk128_codec(q6_k, 256, 6) +instantiate_kquant_nax_bk128_codec(q5_k, 256, 5) +instantiate_kquant_nax_bk128_codec(q3_k, 256, 3) +instantiate_kquant_nax_bk128_codec(q4_k, 256, 4) +instantiate_kquant_nax_bk128_codec(q4_0, 32, 4) +instantiate_kquant_nax_bk128_codec(q4_1, 32, 4) +instantiate_kquant_nax_bk128_codec(q5_0, 32, 5) +instantiate_kquant_nax_bk128_codec(iq4_nl, 32, 4) +instantiate_kquant_nax_bk128_codec(iq3_xxs, 256, 3) +instantiate_kquant_nax_bk128_codec(iq3_s, 256, 3) +instantiate_kquant_nax_bk128_codec(iq2_xs, 256, 2) +instantiate_kquant_nax_bk128_codec(iq2_s, 256, 2) // Small-BM gather_qmm_rhs tile for the few-rows-per-expert prefill regime // (MoE top-k at chat chunk sizes: rows/E ~ 4-32). Every expert segment in @@ -366,7 +413,12 @@ instantiate_kquant_nax_gather_seg_codec(ptq1_0, 128, 1) instantiate_kernel( \ "kquant_" #codec "_verify_nax_" #type "_gs_" #gs "_b_" #bits, \ kq_ ## codec ## _verify_nax, \ - type, gs, bits) + type, gs, bits) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_nax_" #type "_gs_" #gs "_b_" #bits \ + "_pf32", \ + kq_ ## codec ## _verify_nax, \ + type, gs, bits, float) instantiate_kquant_verify_nax(bfloat16_t, 128, 2, pq2_0) instantiate_kquant_verify_nax(bfloat16_t, 64, 2, q2_0) instantiate_kquant_verify_nax(float16_t, 128, 2, pq2_0) diff --git a/metal/kq_quantized_nax_epi.metal b/metal/kq_quantized_nax_epi.metal new file mode 100644 index 00000000..828f6b1f --- /dev/null +++ b/metal/kq_quantized_nax_epi.metal @@ -0,0 +1,34 @@ +// clang-format off +// NAX qmm_t kernels with an epilogue; see kq_quantized_nax_epi.h. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/gemm.h" +#include "mlx/backend/metal/kernels/quantized_utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/loader.h" +#include "mlx/backend/metal/kernels/kq_quantized_nax_epi.h" + +#define instantiate_kquant_nax_glu_walks(codec, type, al, bm, bk) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_glu_" #type \ + "_bm" #bm "_bn64_bk" #bk "w_wm2_wn2_alN_" #al, \ + kq_##codec##_qmm_t_nax_glu_walks, \ + type, \ + al, \ + bm, \ + 64, \ + bk, \ + 2, \ + 2) +#define instantiate_kquant_nax_glu_walks_tile(codec, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, float16_t, true, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, float16_t, false, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, bfloat16_t, true, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, bfloat16_t, false, bm, bk) + +instantiate_kquant_nax_glu_walks_tile(q4_k, 128, 128) +instantiate_kquant_nax_glu_walks_tile(q5_k, 128, 128) +instantiate_kquant_nax_glu_walks_tile(q6_k, 128, 128) +// q8_0 on bfloat16 only: on float16 it measured slower than the unfused ops. +instantiate_kquant_nax_glu_walks(q8_0, bfloat16_t, true, 128, 64) +instantiate_kquant_nax_glu_walks(q8_0, bfloat16_t, false, 128, 64) + // clang-format on diff --git a/metal/kq_quantized_nax_epi_probe.metal b/metal/kq_quantized_nax_epi_probe.metal new file mode 100644 index 00000000..023621dd --- /dev/null +++ b/metal/kq_quantized_nax_epi_probe.metal @@ -0,0 +1,75 @@ +// clang-format off +// Probe tiles of the epilogue kernels, built with KQ_PROBE_KERNELS=ON only: +// the one-walk GLU form and the two-walk tiles the op does not run. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/gemm.h" +#include "mlx/backend/metal/kernels/quantized_utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/loader.h" +#include "mlx/backend/metal/kernels/kq_quantized_nax_epi.h" + +#define instantiate_kquant_nax_glu(codec, type, al, bm, bk) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_glu_" #type \ + "_bm" #bm "_bn64_bk" #bk "_wm2_wn2_alN_" #al, \ + kq_##codec##_qmm_t_nax_glu, \ + type, \ + al, \ + bm, \ + 64, \ + bk, \ + 2, \ + 2) +#define instantiate_kquant_nax_glu_d128(codec, type, al, bm) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_glu_" #type \ + "_bm" #bm "_bn64_bk128d_wm2_wn2_alN_" #al, \ + kq_##codec##_qmm_t_nax_glu_d128, \ + type, \ + al, \ + bm, \ + 64, \ + 2, \ + 2) +#define instantiate_kquant_nax_glu_d128_bm(codec, bm) \ + instantiate_kquant_nax_glu_d128(codec, float16_t, true, bm) \ + instantiate_kquant_nax_glu_d128(codec, float16_t, false, bm) \ + instantiate_kquant_nax_glu_d128(codec, bfloat16_t, true, bm) \ + instantiate_kquant_nax_glu_d128(codec, bfloat16_t, false, bm) +#define instantiate_kquant_nax_glu_walks(codec, type, al, bm, bk) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_glu_" #type \ + "_bm" #bm "_bn64_bk" #bk "w_wm2_wn2_alN_" #al, \ + kq_##codec##_qmm_t_nax_glu_walks, \ + type, \ + al, \ + bm, \ + 64, \ + bk, \ + 2, \ + 2) +#define instantiate_kquant_nax_glu_walks_tile(codec, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, float16_t, true, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, float16_t, false, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, bfloat16_t, true, bm, bk) \ + instantiate_kquant_nax_glu_walks(codec, bfloat16_t, false, bm, bk) +#define instantiate_kquant_nax_glu_tile(codec, bm, bk) \ + instantiate_kquant_nax_glu(codec, float16_t, true, bm, bk) \ + instantiate_kquant_nax_glu(codec, float16_t, false, bm, bk) \ + instantiate_kquant_nax_glu(codec, bfloat16_t, true, bm, bk) \ + instantiate_kquant_nax_glu(codec, bfloat16_t, false, bm, bk) +#define instantiate_kquant_nax_glu_probe_codec(codec) \ + instantiate_kquant_nax_glu_tile(codec, 64, 64) \ + instantiate_kquant_nax_glu_tile(codec, 64, 128) \ + instantiate_kquant_nax_glu_tile(codec, 128, 64) \ + instantiate_kquant_nax_glu_tile(codec, 128, 128) \ + instantiate_kquant_nax_glu_d128_bm(codec, 64) \ + instantiate_kquant_nax_glu_d128_bm(codec, 128) \ + instantiate_kquant_nax_glu_walks_tile(codec, 64, 64) \ + instantiate_kquant_nax_glu_walks_tile(codec, 64, 128) \ + instantiate_kquant_nax_glu_walks_tile(codec, 128, 64) + +instantiate_kquant_nax_glu_probe_codec(q4_k) +instantiate_kquant_nax_glu_walks_tile(q5_k, 128, 64) +instantiate_kquant_nax_glu_walks_tile(q6_k, 128, 64) + // clang-format on diff --git a/metal/kq_quantized_nax_lora.metal b/metal/kq_quantized_nax_lora.metal new file mode 100644 index 00000000..7f356ad9 --- /dev/null +++ b/metal/kq_quantized_nax_lora.metal @@ -0,0 +1,30 @@ +// clang-format off +// NAX qmm_t kernels with a low-rank term; see kq_quantized_nax_lora.h. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/gemm.h" +#include "mlx/backend/metal/kernels/quantized_utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/loader.h" +#include "mlx/backend/metal/kernels/kq_quantized_nax_lora.h" + +#define instantiate_kquant_nax_lora(codec, type, al, bm, bk) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_lora_" #type \ + "_bm" #bm "_bn64_bk" #bk "_wm2_wn2_alN_" #al, \ + kq_##codec##_qmm_t_nax_lora, \ + type, \ + al, \ + bm, \ + 64, \ + bk, \ + 2, \ + 2) + +#define instantiate_kquant_nax_lora_codec(codec, bk) \ + instantiate_kquant_nax_lora(codec, bfloat16_t, true, 128, bk) \ + instantiate_kquant_nax_lora(codec, bfloat16_t, false, 128, bk) +instantiate_kquant_nax_lora_codec(q4_k, 128) +instantiate_kquant_nax_lora_codec(q5_k, 128) +instantiate_kquant_nax_lora_codec(q6_k, 128) +instantiate_kquant_nax_lora_codec(q8_0, 64) + // clang-format on diff --git a/metal/kq_quantized_nax_resid.metal b/metal/kq_quantized_nax_resid.metal new file mode 100644 index 00000000..6f232c29 --- /dev/null +++ b/metal/kq_quantized_nax_resid.metal @@ -0,0 +1,38 @@ +// clang-format off +// NAX qmm_t kernels with a residual epilogue; see kq_quantized_nax_resid.h. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/gemm.h" +#include "mlx/backend/metal/kernels/quantized_utils.h" +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/loader.h" +#include "mlx/backend/metal/kernels/kq_quantized_nax_resid.h" + +#define instantiate_kquant_nax_resid(codec, type, al, gt, bm, bk) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_nax_resid_" #type \ + "_bm" #bm "_bn64_bk" #bk "_wm2_wn2_alN_" #al "_gate_" #gt, \ + kq_##codec##_qmm_t_nax_resid, \ + type, \ + al, \ + gt, \ + bm, \ + 64, \ + bk, \ + 2, \ + 2) + +#define instantiate_kquant_nax_resid_bm(codec, bm, bk) \ + instantiate_kquant_nax_resid(codec, bfloat16_t, true, true, bm, bk) \ + instantiate_kquant_nax_resid(codec, bfloat16_t, false, true, bm, bk) \ + instantiate_kquant_nax_resid(codec, bfloat16_t, true, false, bm, bk) \ + instantiate_kquant_nax_resid(codec, bfloat16_t, false, false, bm, bk) + +// BM 64 for the codecs whose large-row tile has a K ceiling (q4_k and +// q8_0 above K 10240), BM 128 for every codec. +instantiate_kquant_nax_resid_bm(q4_k, 128, 128) +instantiate_kquant_nax_resid_bm(q4_k, 64, 128) +instantiate_kquant_nax_resid_bm(q5_k, 128, 128) +instantiate_kquant_nax_resid_bm(q6_k, 128, 128) +instantiate_kquant_nax_resid_bm(q8_0, 128, 64) +instantiate_kquant_nax_resid_bm(q8_0, 64, 64) + // clang-format on diff --git a/metal/mlx/backend/metal/kernels/kq_conv_nax.h b/metal/mlx/backend/metal/kernels/kq_conv_nax.h new file mode 100644 index 00000000..04db9544 --- /dev/null +++ b/metal/mlx/backend/metal/kernels/kq_conv_nax.h @@ -0,0 +1,204 @@ +// NAX 3x3 convolution, stride 1, padding 1, as an implicit GEMM. +// +// y[p, o] = bias[o] + sum_{dy, dx, c} x[p + (dy - 1, dx - 1), c] * w[o, dy, +// dx, c] +// +// x [B, H, W, C] and y [B, H, W, O] channels last, w [O, 3, 3, C] (rows of +// 9 C), bias [O], all in one float dtype. Rows of the GEMM are the M = B H W +// output pixels, columns the O output channels, and K the 9 C taps and +// channels, walked tap by tap, 32 channels a step. A threadgroup owns one +// [BM, BN] tile of y. Each simdgroup loads its A tile straight from x: a row +// is one output pixel, and its 32 columns are 32 channels of the input +// pixel under the tap, or zeros where the tap falls outside the image. The B +// tile loads from w as in a GEMM with B transposed. Sums are f32 and y is +// rounded once at the store, with the bias added in f32. +// +// C is a multiple of 16, so a 16-channel fragment column is all inside C or +// all past it. M times C and M times O stay below 2^31. +#pragma once + +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" + +using namespace mlx::steel; + +template +METAL_FUNC void kq_conv3x3_nax_impl( + const device T* x, + const device T* w, + const device T* bias, + device T* y, + const constant int& H, + const constant int& W, + const constant int& C, + const constant int& O, + const constant int& M, + const constant int& has_bias, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]]) { + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + using Frag = typename NAXTile::NAXFrag_t; + constexpr short ER = Frag::kElemRows; + constexpr short EC = Frag::kElemCols; + constexpr short RJ = Frag::kElemRowsJump; + constexpr short NR = TM * ER; + + const int y_row = int(tid.y) * BM; + const int y_col = int(tid.x) * BN; + const short tm = SM * (simd_gid / WN); + const short tn = SN * (simd_gid % WN); + if (y_row + tm >= M) { + return; + } + const short sgp_sm = short(min(int(SM), M - (y_row + tm))); + const short sgp_sn = short(clamp(O - (y_col + tn), 0, int(SN))); + if (sgp_sn <= 0) { + return; + } + + // The rows of this thread: pixel, its row and column in the image. + const short2 sc = Frag::get_coord(); + const int hw = H * W; + int pix[NR]; + int py[NR]; + int px[NR]; + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; mm++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < ER; i++) { + const short r = mm * ER + i; + const int p = y_row + tm + mm * 16 + sc.y + i * RJ; + pix[r] = p; + if (p < M) { + const int q = p % hw; + py[r] = q / W; + px[r] = q % W; + } else { + // Every tap of a pixel past M falls outside the image. + py[r] = -4; + px[r] = -4; + } + } + } + + const int ldw = 9 * C; + const device T* wt = w + int64_t(y_col + tn) * ldw; + + NAXTile Dtile; + Dtile.clear(); + + const bool n_full = sgp_sn == SN; + + STEEL_PRAGMA_NO_UNROLL + for (short tap = 0; tap < 9; tap++) { + const short dy = tap / 3 - 1; + const short dx = tap % 3 - 1; + // Element offset of the input pixel under the tap, -1 outside. + int off[NR]; + STEEL_PRAGMA_UNROLL + for (short r = 0; r < NR; r++) { + const int yy = py[r] + dy; + const int xx = px[r] + dx; + const bool in = yy >= 0 && yy < H && xx >= 0 && xx < W; + off[r] = in ? (pix[r] + dy * W + dx) * C + sc.x : -1; + } + const device T* wk = wt + tap * C; + + STEEL_PRAGMA_NO_UNROLL + for (int c0 = 0; c0 < C; c0 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short ck = short(min(int(SK), C - c0)); + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; mm++) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; kk++) { + thread auto& f = Atile.frag_at(mm, kk); + const bool col_in = kk * 16 < ck; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < ER; i++) { + const int o = off[mm * ER + i]; + if (col_in && o >= 0) { + const device T* src = x + o + c0 + kk * 16; + STEEL_PRAGMA_UNROLL + for (short j = 0; j < EC; j++) { + f[i * EC + j] = src[j]; + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < EC; j++) { + f[i * EC + j] = T(0); + } + } + } + } + } + + if (n_full && ck == SK) { + Btile.load(wk + c0, ldw); + } else { + Btile.load_safe(wk + c0, ldw, short2(ck, sgp_sn)); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; + } + } + + device T* ys = y + int64_t(y_row + tm) * O + y_col + tn; + const bool full = n_full && sgp_sm == SM; + if (has_bias) { + NAXTile Ctile; + const device T* bs = bias + y_col + tn; + if (n_full) { + Ctile.load(bs, 0); + } else { + Ctile.load_safe(bs, 0, short2(sgp_sn, SM)); + } + constexpr short NE = decltype(Dtile)::kElemsPerFrag; + constexpr short NF = decltype(Dtile)::kNumFrags; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < NF; ++f) { + STEEL_PRAGMA_UNROLL + for (short e = 0; e < NE; ++e) { + Dtile.val_frags[f][e] += Ctile.val_frags[f][e]; + } + } + } + if (full) { + Dtile.store(ys, O); + } else { + Dtile.store_safe(ys, O, short2(sgp_sn, sgp_sm)); + } +} + +template +[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void kq_conv3x3_nax( + const device T* x [[buffer(0)]], + const device T* w [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* y [[buffer(3)]], + const constant int& H [[buffer(4)]], + const constant int& W [[buffer(5)]], + const constant int& C [[buffer(6)]], + const constant int& O [[buffer(7)]], + const constant int& M [[buffer(8)]], + const constant int& has_bias [[buffer(9)]], + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]]) { + kq_conv3x3_nax_impl( + x, w, bias, y, H, W, C, O, M, has_bias, tid, simd_gid); +} diff --git a/metal/mlx/backend/metal/kernels/kq_lora_epilogue.h b/metal/mlx/backend/metal/kernels/kq_lora_epilogue.h index 7f7a0840..abe0bf65 100644 --- a/metal/mlx/backend/metal/kernels/kq_lora_epilogue.h +++ b/metal/mlx/backend/metal/kernels/kq_lora_epilogue.h @@ -317,7 +317,9 @@ kq_lora_mix_apply( // Densify a strided LoRA operand into a row-major temporary (host side // kq_lora_view_gpu). Element-size templated: 2-byte (half/bfloat) and 4-byte -// (float/uint32/int32) operands share one kernel each. +// (float/uint32/int32) operands share one kernel each. The 1-byte form +// copies a packed weight view whose rows are not dense (host side +// kquant_dense.cpp). struct KqLoraDensifyArgs { int shape[8]; int64_t strides[8]; diff --git a/metal/mlx/backend/metal/kernels/kq_moe_glu.h b/metal/mlx/backend/metal/kernels/kq_moe_glu.h index 3f9743a2..53192888 100644 --- a/metal/mlx/backend/metal/kernels/kq_moe_glu.h +++ b/metal/mlx/backend/metal/kernels/kq_moe_glu.h @@ -158,8 +158,12 @@ template if (simd_lid == 0) { g += gb[row0 + r]; u += ub[row0 + r]; - g = min(g, (U)limit); - u = metal::clamp(u, -(U)limit, (U)limit); + // As mx.minimum and mx.clip: a NaN passes. min() and clamp() return + // a bound for a NaN, so each side is one comparison. + const U lim = (U)limit; + g = g > lim ? lim : g; + u = u > lim ? lim : u; + u = u < -lim ? -lim : u; const U sig = 1.0f / (1.0f + metal::exp(-alpha * g)); out[r] = static_cast((g * sig) * (u + 1.0f)); } diff --git a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h index d02fa0c7..daf5bd9e 100644 --- a/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h +++ b/metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h @@ -31,6 +31,20 @@ #define KQ_GLU_ACT_SILU_LIMIT 2 #define KQ_GLU_ACT_SWIGLU_CLAMP 3 +// The clamps of the limited activations, as mx.minimum and mx.clip give +// them: a NaN passes. min() and clamp() return a bound for a NaN, which +// would turn a NaN gate or up value into a finite output, so each side is +// one comparison. +inline float kq_glu_cap(float g, float limit) { + return g > limit ? limit : g; +} + +inline float kq_glu_clip(float u, float limit) { + u = u > limit ? limit : u; + u = u < -limit ? -limit : u; + return u; +} + template inline float kq_glu_epilogue(float g, float u, float limit, float alpha = 1.0f) { @@ -44,8 +58,8 @@ kq_glu_epilogue(float g, float u, float limit, float alpha = 1.0f) { // gpt-oss clamped SwiGLU: gate clamped from above only, up clamped both // sides, sigmoid slope alpha, and a (u + 1) linear term -- matches the // packed kq_moe_glu_gather epilogue (kq_moe_glu.h). - g = metal::min(g, limit); - u = metal::clamp(u, -limit, limit); + g = kq_glu_cap(g, limit); + u = kq_glu_clip(u, limit); const float sig = 1.0f / (1.0f + metal::exp(-alpha * g)); return (g * sig) * (u + 1.0f); } @@ -56,8 +70,8 @@ kq_glu_epilogue(float g, float u, float limit, float alpha = 1.0f) { // The silu instantiations take the same clamp from a positive run-time // limit; the shexp family dispatches silu_limit that way (no _silu_limit // shexp kernels). Hosts pass 0 for plain silu. - g = metal::min(g, limit); - u = metal::clamp(u, -limit, limit); + g = kq_glu_cap(g, limit); + u = kq_glu_clip(u, limit); } const float sig = 1.0f / (1.0f + metal::exp(-g)); return (g * sig) * u; @@ -1977,12 +1991,14 @@ template // changes per chunk: the grid entry comes back as half4 pairs (no u64 byte // unpack), the activation row is read from a half copy staged once per // threadgroup in dynamic threadgroup memory (buffer threadgroup(0), K halves, -// host-sized), and the four 4-wide dots run in half with a float sum per -// chunk before the chunk scale. Expert rows only: the shared-expert slot in -// the shexp / mix kernels keeps the float path over the device activation -// row, so the shexp numerics are unchanged. Outputs differ from the float -// kernels at the half rounding level (about 1e-3 relative on the fused GLU -// output at K=4096, against a bf16 floor near 9e-4). +// host-sized), and the four 4-wide dots widen both halves to float, so +// each product and the chunk sum run in float before the chunk scale: a +// grid code (up to 62) times any half activation stays in range, and only +// an activation past 65504 turns to inf in the staged row. Expert rows +// only: the shared-expert slot in the shexp / mix kernels keeps the float +// path over the device activation row, so the shexp numerics are +// unchanged. Outputs differ from the float kernels at the half rounding +// level of the staged activations. // --------------------------------------------------------------------------- #define KQ_EXT_STAGE_LUTS_H(CodecT, name) \ @@ -2028,8 +2044,8 @@ METAL_FUNC float kq_ext_row_partial_h( KqTgLutsH::deq_chunk16h(block, short(ich % chpb), lw, luts, sc); const threadgroup half4* xp = xs + ich * 4; acc += sc * - (float(dot(lw[0], xp[0])) + float(dot(lw[1], xp[1])) + - float(dot(lw[2], xp[2])) + float(dot(lw[3], xp[3]))); + (dot(float4(lw[0]), float4(xp[0])) + dot(float4(lw[1]), float4(xp[1])) + + dot(float4(lw[2]), float4(xp[2])) + dot(float4(lw[3]), float4(xp[3]))); } return acc; } @@ -2053,20 +2069,20 @@ METAL_FUNC float2 kq_ext_glu_row_partial_h( const int64_t boff = (int64_t)(ich / chpb) * Codec::block_bytes; const short cch = short(ich % chpb); const threadgroup half4* xp = xs + ich * 4; - const half4 a0 = xp[0]; - const half4 a1 = xp[1]; - const half4 a2 = xp[2]; - const half4 a3 = xp[3]; + const float4 a0 = float4(xp[0]); + const float4 a1 = float4(xp[1]); + const float4 a2 = float4(xp[2]); + const float4 a3 = float4(xp[3]); half4x4 lw; float sc; KqTgLutsH::deq_chunk16h(g_row + boff, cch, lw, luts, sc); acc.x += sc * - (float(dot(lw[0], a0)) + float(dot(lw[1], a1)) + float(dot(lw[2], a2)) + - float(dot(lw[3], a3))); + (dot(float4(lw[0]), a0) + dot(float4(lw[1]), a1) + + dot(float4(lw[2]), a2) + dot(float4(lw[3]), a3)); KqTgLutsH::deq_chunk16h(u_row + boff, cch, lw, luts, sc); acc.y += sc * - (float(dot(lw[0], a0)) + float(dot(lw[1], a1)) + float(dot(lw[2], a2)) + - float(dot(lw[3], a3))); + (dot(float4(lw[0]), a0) + dot(float4(lw[1]), a1) + + dot(float4(lw[2]), a2) + dot(float4(lw[3]), a3)); } return acc; } @@ -2205,8 +2221,8 @@ template < } // Score-mixed down gather with the shared-expert slot folded in: the routed -// slots restage their h row per slot (two barriers each) and dot in half; -// the shared slot stays on the float path. +// slots restage their h row as half per slot (two barriers each) and dot +// each chunk in float; the shared slot stays on the float path. template [[kernel]] void kq_ext_gather_qmv_mix_h( const device uint8_t* w [[buffer(0)]], diff --git a/metal/mlx/backend/metal/kernels/kq_qk_norm_rope.h b/metal/mlx/backend/metal/kernels/kq_qk_norm_rope.h new file mode 100644 index 00000000..55a75132 --- /dev/null +++ b/metal/mlx/backend/metal/kernels/kq_qk_norm_rope.h @@ -0,0 +1,163 @@ +// Fused per-head RMSNorm and RoPE of the DiT attention q and k, both in one +// dispatch: +// +// y = x * rsqrt(mean(x^2) + eps) * w +// out[2i] = y[2i] * cos[l, i] - y[2i + 1] * sin[l, i] +// out[2i + 1] = y[2i] * sin[l, i] + y[2i + 1] * cos[l, i] +// +// q and k are [B, L, H, D], read in place through their b, l, h and d +// strides (a slice of a fused qkv output needs no copy). q_out and k_out +// are [B, H, L, D] row-contiguous, the layout attention reads. cos and sin +// are float [L, D / 2] tables with one angle per row and pair, so any +// position scheme works (multi-axis image positions included); the norm +// weights are float [D]. The batch shares the tables. +// +// One threadgroup per (b, l) row. A group of KQ_QKNR_LPR lanes covers one +// head row of D, EPL = D / KQ_QKNR_LPR elements per lane, so a simdgroup +// covers 32 / KQ_QKNR_LPR head rows per pass and the simdgroups walk the +// 2 * H head rows of q, then k. A lane loads its cos, sin and weight values +// once per row and reuses them across heads. The mean square is a +// shuffle-xor sum inside the lane group. All math in f32, one round at the +// write. +// +// kq_qknr_vec (function constant 350) selects vector loads, which the host +// sets when the d stride is 1 and every row start is aligned to the vector +// width; otherwise the kernel reads element by element through the strides. + +#pragma once + +#define KQ_QKNR_LPR 16 +#define KQ_QKNR_UNROLL 2 + +constant bool kq_qknr_vec [[function_constant(350)]]; + +// Mirror of the host-side KQQkNormRopeParams (kquant_qk_norm_rope.cpp); the +// layouts must match field for field. Strides are in elements. +struct KQQkNormRopeParams { + int64_t q_b; + int64_t q_l; + int64_t q_h; + int64_t q_d; + int64_t k_b; + int64_t k_l; + int64_t k_h; + int64_t k_d; + int L; + int H; + float eps; +}; + +template +METAL_FUNC void kq_qknr_load(const device T* src, int64_t sd, thread float* x) { + if (kq_qknr_vec) { + using V = metal::vec; + for (int i = 0; i < EPL; i += VW) { + const V v = *reinterpret_cast(src + i); + for (int t = 0; t < VW; ++t) { + x[i + t] = float(v[t]); + } + } + } else { + for (int i = 0; i < EPL; ++i) { + x[i] = float(src[i * sd]); + } + } +} + +template +[[kernel]] void kq_qk_norm_rope( + const device T* q [[buffer(0)]], + const device T* k [[buffer(1)]], + const device float* qw [[buffer(2)]], + const device float* kw [[buffer(3)]], + const device float* cosv [[buffer(4)]], + const device float* sinv [[buffer(5)]], + device T* q_out [[buffer(6)]], + device T* k_out [[buffer(7)]], + const constant KQQkNormRopeParams& p [[buffer(8)]], + uint row [[threadgroup_position_in_grid]], + uint sg [[simdgroup_index_in_threadgroup]], + uint nsg [[simdgroups_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + constexpr int LPR = KQ_QKNR_LPR; + constexpr int EPL = D / LPR; + constexpr int VW = (EPL % 4 == 0) ? 4 : 2; + constexpr int RPS = 32 / LPR; + constexpr int U = KQ_QKNR_UNROLL; + static_assert(D % (2 * LPR) == 0 && EPL % VW == 0, "unsupported head dim"); + + const int L = p.L; + const int H = p.H; + const int b = int(row) / L; + const int l = int(row) - b * L; + const int c0 = (int(lane) % LPR) * EPL; + + float cs[EPL / 2], sn[EPL / 2], wq[EPL], wk[EPL]; + const device float* cr = cosv + int64_t(l) * (D / 2) + c0 / 2; + const device float* sr = sinv + int64_t(l) * (D / 2) + c0 / 2; + for (int i = 0; i < EPL / 2; ++i) { + cs[i] = cr[i]; + sn[i] = sr[i]; + } + for (int i = 0; i < EPL; ++i) { + wq[i] = qw[c0 + i]; + wk[i] = kw[c0 + i]; + } + + const device T* qr = q + b * p.q_b + l * p.q_l + c0 * p.q_d; + const device T* kr = k + b * p.k_b + l * p.k_l + c0 * p.k_d; + // Output row (b, 0, l) of [B, H, L, D]; head h adds h * L * D. + const int64_t obase = (int64_t(b) * H * L + l) * D + c0; + const int64_t ohead = int64_t(L) * D; + const int step = int(nsg) * RPS; + + for (int j0 = int(sg) * RPS + int(lane) / LPR; j0 < 2 * H; j0 += step * U) { + float x[U][EPL]; + bool act[U], isk[U]; + int hh[U]; + for (int u = 0; u < U; ++u) { + const int j = j0 + u * step; + act[u] = j < 2 * H; + isk[u] = j >= H; + hh[u] = isk[u] ? j - H : j; + if (act[u]) { + const device T* src = isk[u] ? kr + hh[u] * p.k_h : qr + hh[u] * p.q_h; + kq_qknr_load(src, isk[u] ? p.k_d : p.q_d, x[u]); + } else { + for (int i = 0; i < EPL; ++i) { + x[u][i] = 0.0f; + } + } + } + for (int u = 0; u < U; ++u) { + float ss = 0.0f; + for (int i = 0; i < EPL; ++i) { + ss = metal::fma(x[u][i], x[u][i], ss); + } + for (int o = LPR / 2; o >= 1; o >>= 1) { + ss += simd_shuffle_xor(ss, ushort(o)); + } + if (!act[u]) { + continue; + } + const float inv = metal::precise::rsqrt(ss / float(D) + p.eps); + float y[EPL]; + for (int i = 0; i < EPL; ++i) { + y[i] = x[u][i] * inv * (isk[u] ? wk[i] : wq[i]); + } + device T* dst = (isk[u] ? k_out : q_out) + obase + hh[u] * ohead; + using V = metal::vec; + for (int i = 0; i < EPL; i += VW) { + V v; + for (int t = 0; t < VW; t += 2) { + const int e = i + t; + const float a = y[e]; + const float c = y[e + 1]; + v[t] = T(a * cs[e / 2] - c * sn[e / 2]); + v[t + 1] = T(a * sn[e / 2] + c * cs[e / 2]); + } + *reinterpret_cast(dst + i) = v; + } + } + } +} diff --git a/metal/mlx/backend/metal/kernels/kq_quantized.h b/metal/mlx/backend/metal/kernels/kq_quantized.h index ae2b7841..2528f03f 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized.h @@ -107,6 +107,9 @@ MLX_MTL_CONST int SIMD_SIZE = 32; struct kq_empty {}; +// U is the type of y: T for a whole matmul and for the split-K partials +// of the plain kernels, float for the _pf32 kernels (float16 by default, +// bfloat16 under KQ_SPLITK_F32=1). template < typename T, typename LoaderW, @@ -115,11 +118,12 @@ template < const int BK = 32, const int BN = 32, const int WM = 2, - const int WN = 2> + const int WN = 2, + typename U = T> METAL_FUNC void kq_qmm_t_impl( const device uint8_t* w, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -135,12 +139,11 @@ METAL_FUNC void kq_qmm_t_impl( (void)lid; - static_assert(WM * WN == 4, "the tile runs on four simdgroups"); constexpr int BK_padded = (BK + 16 / sizeof(T)); using mma_t = mlx::steel::BlockMMA< T, - T, + U, BM, BN, BK, @@ -1470,7 +1473,16 @@ struct KqQ8_0BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q8_0_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1507,7 +1519,7 @@ template static_assert( group_size == KQ_Q8_0_GROUP, "Q8_0 kernel requires group_size=32"); static_assert(bits == 8, "Q8_0 kernel requires bits=8"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -1517,8 +1529,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -1527,12 +1539,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q8_0_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -1565,7 +1578,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -1987,10 +2000,10 @@ template // Register-resident MMA verify (kq_verify_mma.h). Lane L owns qs bytes // 8L..8L+7: fragments (b0, b2), (b1, b3), (b4, b6), (b5, b7). Each pair // has its sign bits flipped and is masked into the mantissas of a half2 -// (1024 + 128 + q), which one fma maps to q / 256, exact in half. The -// 1/256 keeps every product within half its activation, so one channel -// near the top of the half range stays finite in the block sum, and -// d_scale restores the factor. +// (1024 + 128 + q), which one fma maps to q / 256, exact in half, and +// d_scale restores the factor. In the half kernels the 1/256 keeps every +// product within half its activation, so one channel near the top of the +// half range stays finite in the block sum. struct KqQ8_0Mma { static constant constexpr int block_k = KQ_Q8_0_GROUP; static constant constexpr int block_bytes = KQ_Q8_0_BLOCK_BYTES; @@ -1999,14 +2012,14 @@ struct KqQ8_0Mma { static METAL_FUNC int perm(int f, int col) { return 8 * (col / 2) + 4 * (f >> 1) + (f & 1) + 2 * (col & 1); } - template + template static METAL_FUNC void block( thread const KqVmmaRows& rows, int boff, short L, short fm, - const threadgroup half* xb, - thread simdgroup_half8x8 (&acc)[NT]) { + const threadgroup X* xb, + thread simdgroup_matrix (&acc)[NT]) { uint wv[NT][2]; for (short t = 0; t < NT; ++t) { const packed_ushort4 v = @@ -2469,7 +2482,16 @@ struct KqQ5_1BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q5_1_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2506,7 +2528,7 @@ template static_assert( group_size == KQ_Q5_1_GROUP, "Q5_1 kernel requires group_size=32"); static_assert(bits == 5, "Q5_1 kernel requires bits=5"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -2516,8 +2538,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -2526,12 +2548,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q5_1_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -2564,7 +2587,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_fp.h b/metal/mlx/backend/metal/kernels/kq_quantized_fp.h index e0cf92be..1c1e3cd9 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_fp.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_fp.h @@ -300,7 +300,16 @@ struct KqMxfp4BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_mxfp4_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -336,60 +345,16 @@ template } static_assert(group_size == KQ_MXFP4_GROUP, "MXFP4 requires gs=32"); static_assert(bits == 4, "MXFP4 requires bits=4"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqMxfp4BlockLoader; - kq_qmm_t_impl( + KqMxfp4BlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template -[[kernel]] void kq_mxfp4_qmm_t_splitk( - const device uint8_t* w, - const device uint8_t* /* scales */, - const device T* x, - device T* y, - const constant int& K, - const constant int& N, - const constant int& M, - const constant int& k_partition_size, - const constant int& split_k_partition_stride, - uint3 tid [[threadgroup_position_in_grid]], - uint lid [[thread_index_in_threadgroup]], - uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]]) { - static_assert(group_size == KQ_MXFP4_GROUP, "MXFP4 requires gs=32"); - static_assert(bits == 4, "MXFP4 requires bits=4"); - constexpr int BM = 32, BK = 32, BN = 32; - constexpr int BK_padded = (BK + 16 / sizeof(T)); - threadgroup T Xs[BM * BK_padded]; - threadgroup T Ws[BN * BK_padded]; - using LoaderW = - KqMxfp4BlockLoader; - const int k_start = tid.z * k_partition_size; - x += k_start; - auto wl = w; - wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; - y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( - wl, - x, - y, - Xs, - Ws, - K, - N, - M, - k_partition_size, - tid, - lid, - simd_gid, - simd_lid); -} - template [[kernel]] void kq_mxfp4_qmm_n( const device uint8_t* w, @@ -810,7 +775,16 @@ struct KqNvfp4BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_nvfp4_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -846,60 +820,16 @@ template } static_assert(group_size == KQ_NVFP4_SUPERBLOCK, "NVFP4 requires gs=64"); static_assert(bits == 4, "NVFP4 requires bits=4"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqNvfp4BlockLoader; - kq_qmm_t_impl( + KqNvfp4BlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template -[[kernel]] void kq_nvfp4_qmm_t_splitk( - const device uint8_t* w, - const device uint8_t* /* scales */, - const device T* x, - device T* y, - const constant int& K, - const constant int& N, - const constant int& M, - const constant int& k_partition_size, - const constant int& split_k_partition_stride, - uint3 tid [[threadgroup_position_in_grid]], - uint lid [[thread_index_in_threadgroup]], - uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]]) { - static_assert(group_size == KQ_NVFP4_SUPERBLOCK, "NVFP4 requires gs=64"); - static_assert(bits == 4, "NVFP4 requires bits=4"); - constexpr int BM = 32, BK = 32, BN = 32; - constexpr int BK_padded = (BK + 16 / sizeof(T)); - threadgroup T Xs[BM * BK_padded]; - threadgroup T Ws[BN * BK_padded]; - using LoaderW = - KqNvfp4BlockLoader; - const int k_start = tid.z * k_partition_size; - x += k_start; - auto wl = w; - wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; - y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( - wl, - x, - y, - Xs, - Ws, - K, - N, - M, - k_partition_size, - tid, - lid, - simd_gid, - simd_lid); -} - template [[kernel]] void kq_nvfp4_qmm_n( const device uint8_t* w, diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index fd49a3e2..8a7ba2bc 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h @@ -1274,7 +1274,16 @@ struct KqIq4_xsBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq4_xs_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1310,13 +1319,13 @@ template } static_assert(group_size == KQ_IQ4_XS_SUPERBLOCK, "IQ4_XS requires gs=256"); static_assert(bits == 4, "IQ4_XS requires bits=4"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq4_xsBlockLoader; - kq_qmm_t_impl( + KqIq4_xsBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -1325,12 +1334,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq4_xs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -1355,7 +1365,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -1745,7 +1755,16 @@ struct KqIq3_xxsBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq3_xxs_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1781,13 +1800,13 @@ template } static_assert(group_size == KQ_IQ3_XXS_SUPERBLOCK, "IQ3_XXS requires gs=256"); static_assert(bits == 3, "IQ3_XXS requires bits=3"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq3_xxsBlockLoader; - kq_qmm_t_impl( + KqIq3_xxsBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -1796,12 +1815,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq3_xxs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -1826,7 +1846,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -2223,7 +2243,16 @@ struct KqIq3_sBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq3_s_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2259,13 +2288,13 @@ template } static_assert(group_size == KQ_IQ3_S_SUPERBLOCK, "IQ3_S requires gs=256"); static_assert(bits == 3, "IQ3_S requires bits=3"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq3_sBlockLoader; - kq_qmm_t_impl( + KqIq3_sBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -2274,12 +2303,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq3_s_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -2304,7 +2334,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -2697,7 +2727,16 @@ struct KqIq2_xxsBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq2_xxs_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2733,13 +2772,13 @@ template } static_assert(group_size == KQ_IQ2_XXS_SUPERBLOCK, "IQ2_XXS requires gs=256"); static_assert(bits == 2, "IQ2_XXS requires bits=2"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq2_xxsBlockLoader; - kq_qmm_t_impl( + KqIq2_xxsBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -2748,12 +2787,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq2_xxs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -2778,7 +2818,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -3174,7 +3214,16 @@ struct KqIq2_xsBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq2_xs_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -3210,13 +3259,13 @@ template } static_assert(group_size == KQ_IQ2_XS_SUPERBLOCK, "IQ2_XS requires gs=256"); static_assert(bits == 2, "IQ2_XS requires bits=2"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq2_xsBlockLoader; - kq_qmm_t_impl( + KqIq2_xsBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -3225,12 +3274,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq2_xs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -3255,7 +3305,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -3653,7 +3703,16 @@ struct KqIq2_sBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq2_s_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -3689,13 +3748,13 @@ template } static_assert(group_size == KQ_IQ2_S_SUPERBLOCK, "IQ2_S requires gs=256"); static_assert(bits == 2, "IQ2_S requires bits=2"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq2_sBlockLoader; - kq_qmm_t_impl( + KqIq2_sBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -3704,12 +3763,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq2_s_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -3734,7 +3794,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -4120,7 +4180,16 @@ struct KqIq1_sBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq1_s_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -4156,13 +4225,13 @@ template } static_assert(group_size == KQ_IQ1_S_SUPERBLOCK, "IQ1_S requires gs=256"); static_assert(bits == 1, "IQ1_S requires bits=1"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq1_sBlockLoader; - kq_qmm_t_impl( + KqIq1_sBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -4171,12 +4240,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq1_s_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -4201,7 +4271,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -4610,7 +4680,16 @@ struct KqIq1_mBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq1_m_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -4646,13 +4725,13 @@ template } static_assert(group_size == KQ_IQ1_M_SUPERBLOCK, "IQ1_M requires gs=256"); static_assert(bits == 1, "IQ1_M requires bits=1"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq1_mBlockLoader; - kq_qmm_t_impl( + KqIq1_mBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -4661,12 +4740,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq1_m_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -4691,7 +4771,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -5180,13 +5260,13 @@ struct KqTgLuts { // --------------------------------------------------------------------------- // Half-dot decode tables (KQ_MOE_HALF): the grid codecs' 8-bit magnitude // entries pre-expanded to half4 pairs in threadgroup memory, so the chunk -// dequant is a table read plus a sign select in half and the dot runs in -// half against a half-staged activation row. Specialized for the codecs -// where the float path is ALU-bound at decode (iq2_xs, iq2_xxs, iq3_xxs); -// every other codec reports has_half = false and never reaches these -// kernels (host gate). deq_chunk16h mirrors deq_chunk16s: 16 unscaled -// weights in natural order plus the chunk scale, applied by the caller in -// float once per chunk. +// dequant is a table read plus a sign select in half. The caller widens +// the weights and a half-staged activation row to float for each chunk +// dot. Specialized for the codecs where the float path is ALU-bound at +// decode (iq2_xs, iq2_xxs, iq3_xxs); every other codec reports has_half = +// false and never reaches these kernels (host gate). deq_chunk16h mirrors +// deq_chunk16s: 16 unscaled weights in natural order plus the chunk scale, +// applied by the caller in float once per chunk. // --------------------------------------------------------------------------- // Negates lane j of a half4 where bit j of s (low 4 bits) is set. diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h index ea44e5df..36cbe46c 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h @@ -563,7 +563,16 @@ struct KqQ4_KBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q4_k_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -600,7 +609,7 @@ template static_assert( group_size == KQ_Q4_K_SUPERBLOCK, "Q4_K kernel requires group_size=256"); static_assert(bits == 4, "Q4_K kernel requires bits=4"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -610,8 +619,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -623,12 +632,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q4_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -661,7 +671,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -1637,7 +1647,16 @@ struct KqQ5_KBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q5_k_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1674,7 +1693,7 @@ template static_assert( group_size == KQ_Q5_K_SUPERBLOCK, "Q5_K kernel requires group_size=256"); static_assert(bits == 5, "Q5_K kernel requires bits=5"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -1684,8 +1703,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -1694,12 +1713,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q5_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -1732,7 +1752,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -2673,7 +2693,16 @@ struct KqQ6_KBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q6_k_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2710,7 +2739,7 @@ template static_assert( group_size == KQ_Q6_K_SUPERBLOCK, "Q6_K kernel requires group_size=256"); static_assert(bits == 6, "Q6_K kernel requires bits=6"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -2720,8 +2749,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -2730,12 +2759,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q6_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -2768,7 +2798,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -3808,7 +3838,16 @@ struct KqQ3_KBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q3_k_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -3845,7 +3884,7 @@ template static_assert( group_size == KQ_Q3_K_SUPERBLOCK, "Q3_K kernel requires group_size=256"); static_assert(bits == 3, "Q3_K kernel requires bits=3"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -3855,8 +3894,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -3865,12 +3904,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q3_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -3903,7 +3943,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -4706,7 +4746,16 @@ struct KqQ2_KBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q2_k_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -4743,7 +4792,7 @@ template static_assert( group_size == KQ_Q2_K_SUPERBLOCK, "Q2_K kernel requires group_size=256"); static_assert(bits == 2, "Q2_K kernel requires bits=2"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -4753,8 +4802,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -4763,12 +4812,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q2_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -4801,7 +4851,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h b/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h index 71e577ab..69b0880e 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h @@ -393,7 +393,16 @@ struct KqQ4_0BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q4_0_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -430,7 +439,7 @@ template static_assert( group_size == KQ_Q4_0_GROUP, "Q4_0 kernel requires group_size=32"); static_assert(bits == 4, "Q4_0 kernel requires bits=4"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -440,8 +449,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -450,12 +459,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q4_0_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -488,7 +498,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -758,8 +768,9 @@ template // 4L..4L+3 of qs: fragments (n0, n2), (n1, n3), (h0, h2), (h1, h3) of // its four low and four high nibbles, each pair masked into the // mantissas of a half2 (1024 + q) and mapped by one fma to (q - 8) / 16, -// exact in half. The 1/16 keeps a block of 32 products inside half range -// for activations up to the half limit; d_scale restores it. +// exact in half, and d_scale restores the 1/16. In the half kernels it +// keeps one channel near the half limit, times a code of up to 8, inside +// half range. struct KqQ4_0Mma { static constant constexpr int block_k = KQ_Q4_0_GROUP; static constant constexpr int block_bytes = KQ_Q4_0_BLOCK_BYTES; @@ -768,14 +779,14 @@ struct KqQ4_0Mma { static METAL_FUNC int perm(int f, int col) { return 16 * ((f % 4) >> 1) + 4 * (col / 2) + (f & 1) + 2 * (col & 1); } - template + template static METAL_FUNC void block( thread const KqVmmaRows& rows, int boff, short L, short fm, - const threadgroup half* xb, - thread simdgroup_half8x8 (&acc)[NT]) { + const threadgroup X* xb, + thread simdgroup_matrix (&acc)[NT]) { uint lo[NT], hi[NT]; for (short t = 0; t < NT; ++t) { const packed_ushort2 v2 = @@ -1210,7 +1221,16 @@ struct KqQ4_1BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q4_1_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1247,7 +1267,7 @@ template static_assert( group_size == KQ_Q4_1_GROUP, "Q4_1 kernel requires group_size=32"); static_assert(bits == 4, "Q4_1 kernel requires bits=4"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -1257,8 +1277,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -1267,12 +1287,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q4_1_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -1305,7 +1326,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -1994,7 +2015,16 @@ struct KqQ5_0BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_q5_0_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2031,7 +2061,7 @@ template static_assert( group_size == KQ_Q5_0_GROUP, "Q5_0 kernel requires group_size=32"); static_assert(bits == 5, "Q5_0 kernel requires bits=5"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -2041,8 +2071,8 @@ template BK, BK_padded, /*reduction_dim=*/1, - /*tgp_size=*/2 * 2 * SIMD_SIZE>; - kq_qmm_t_impl( + /*tgp_size=*/WM * WN * SIMD_SIZE>; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -2051,12 +2081,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_q5_0_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -2089,7 +2120,7 @@ template < wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, @@ -2641,7 +2672,16 @@ struct KqIq4_nlBlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_iq4_nl_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2677,13 +2717,13 @@ template } static_assert(group_size == KQ_IQ4_NL_GROUP, "IQ4_NL requires gs=32"); static_assert(bits == 4, "IQ4_NL requires bits=4"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqIq4_nlBlockLoader; - kq_qmm_t_impl( + KqIq4_nlBlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -2692,12 +2732,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_iq4_nl_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -2722,7 +2763,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h index c66a4cf0..b9c15b03 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h @@ -22,11 +22,12 @@ template < const int WM = 2, const int WN = 2, const bool use_db = false, - const bool tail_ksplit = false> + const bool tail_ksplit = false, + typename U = T> METAL_FUNC void kq_qmm_t_nax_tgp_impl( const device uint8_t* w, const device T* x, - device T* y, + device U* y, threadgroup T* Ws, const constant int& K, const constant int& N, @@ -101,6 +102,12 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( // than either uniform path. The safe path costs ~3%, so branch on the // tile, clamp per SG. const bool is_unaligned_sm = (M - y_row) < BM; + // Row-live guard: a simdgroup row whose rows all lie at or beyond M keeps + // its share of the weight decode and every barrier, and skips the A load, + // the MMA and the store. Uniform per simdgroup, since a simdgroup row + // shares one tm. Lets a BM=128 tile carry a last row tile of 64 rows or + // fewer with the loads and MMA of one BM=64 tile. + const bool sg_live = tm < M - y_row; const short sgp_sn = aligned_N ? SN : short(min(int(SN), N - (y_col + tn))); @@ -166,7 +173,7 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( const threadgroup T* Wcur = Ws + (step & 1) * WS_STRIDE; STEEL_PRAGMA_NO_UNROLL - for (int kk1 = kk_first; kk1 < BK; kk1 += kk_stride) { + for (int kk1 = kk_first; sg_live && kk1 < BK; kk1 += kk_stride) { NAXTile Atile; NAXTile Btile; @@ -206,7 +213,7 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( threadgroup_barrier(mem_flags::mem_threadgroup); STEEL_PRAGMA_NO_UNROLL - for (int kk1 = kk_first; kk1 < BK; kk1 += kk_stride) { + for (int kk1 = kk_first; sg_live && kk1 < BK; kk1 += kk_stride) { NAXTile Atile; NAXTile Btile; @@ -272,7 +279,7 @@ METAL_FUNC void kq_qmm_t_nax_tgp_impl( Dtile.store(y + tm * N + tn, N); } else if (kAlignedM.value && sgp_sn == SN) { Dtile.store(y + tm * N + tn, N); - } else { + } else if (sg_live) { Dtile.store_safe(y + tm * N + tn, N, short2(sgp_sn, sgp_sm)); } }); @@ -1613,6 +1620,100 @@ struct KqNaxQ4_KBlockLoader { // Q5_K: 176 bytes per 256-weight super-block. Q4_K layout + qh[32] high bits. +// Q4_K on a K tile of 128: one tile is the sub-blocks 4h to 4h + 3 of half +// h, the qs pairs 2h and 2h + 1. Two threads share a row, 16 positions l +// each. A thread reads qs[64h + l] and qs[64h + 32 + l] and the four scale +// and min pairs of its half once per tile, where the BK=64 loader reads two +// pairs on each of its two tiles. +template < + typename T, + short BROWS, + short BCOLS, + short dst_ld, + short reduction_dim, + short tgp_size> +struct KqNaxQ4_KHalfLoader { + MLX_MTL_CONST int weights_per_block = KQ_Q4_K_SUPERBLOCK; + MLX_MTL_CONST int bytes_per_block = KQ_Q4_K_BLOCK_BYTES; + static_assert(BCOLS == 128, "Q4_K half loader requires BCOLS == 128."); + static_assert(reduction_dim == 1, "Q4_K half loader serves qmm_t only."); + static_assert(BROWS * 2 == tgp_size, "Q4_K half loader: 2 threads per row."); + + MLX_MTL_CONST short TCOLS = 2; + MLX_MTL_CONST short n_pos = 16; + // load_unsafe depends on half_idx only. + MLX_MTL_CONST bool db_safe = true; + + const short thread_idx; + const short bi; + const short lj; + + threadgroup T* dst; + const device uint8_t* src; + short half_idx; + + KqNaxQ4_KHalfLoader( + const device uint8_t* src_, + const int src_ld_, + threadgroup T* dst_, + ushort simd_group_id [[simdgroup_index_in_threadgroup]], + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + : thread_idx(simd_group_id* SIMD_SIZE + simd_lane_id), + bi(thread_idx / TCOLS), + lj((thread_idx % TCOLS) * n_pos), + dst(dst_ + bi * dst_ld + lj), + src(src_ + bi * (src_ld_ * bytes_per_block / weights_per_block)), + half_idx(0) {} + + void load_unsafe() thread { + const float d = float(*(const device half*)(src + KQ_Q4_K_D_OFFSET)); + const float dmin = float(*(const device half*)(src + KQ_Q4_K_DMIN_OFFSET)); + const device uint8_t* scales12 = src + KQ_Q4_K_SCALES_OFFSET; + float es[4], em[4]; +#pragma unroll + for (short q = 0; q < 4; q++) { + uint8_t sc6, mn6; + kq_get_scale_min_k4(half_idx * 4 + q, scales12, sc6, mn6); + es[q] = d * float(sc6); + em[q] = dmin * float(mn6); + } + const device uint8_t* qs = src + KQ_Q4_K_QS_OFFSET + half_idx * 64 + lj; + const uint4 a_v = *reinterpret_cast(qs); + const uint4 b_v = *reinterpret_cast(qs + 32); + const thread uint8_t* a = reinterpret_cast(&a_v); + const thread uint8_t* b = reinterpret_cast(&b_v); +#pragma unroll + for (short j = 0; j < n_pos; j++) { + dst[j] = T(es[0] * float(a[j] & 0x0F) - em[0]); + dst[32 + j] = T(es[1] * float(a[j] >> 4) - em[1]); + dst[64 + j] = T(es[2] * float(b[j] & 0x0F) - em[2]); + dst[96 + j] = T(es[3] * float(b[j] >> 4) - em[3]); + } + } + + void load_safe(short2 src_tile_dim) thread { + if (bi >= src_tile_dim.y) { +#pragma unroll + for (short j = 0; j < n_pos; j++) { + dst[j] = T(0); + dst[32 + j] = T(0); + dst[64 + j] = T(0); + dst[96 + j] = T(0); + } + return; + } + load_unsafe(); + } + + void next() thread { + half_idx += 1; + if (half_idx == 2) { + half_idx = 0; + src += bytes_per_block; + } + } +}; + template < typename T, short BROWS, @@ -1744,6 +1845,109 @@ struct KqNaxQ5_KBlockLoader { } }; +// Q5_K on a K tile of 128: one tile is the sub-blocks 4h to 4h + 3 of half +// h, two qs pairs, plus the high bits. Two threads share a row, 16 positions +// l each. A thread reads qs[64h + l] and qs[64h + 32 + l] for the low bits +// and qh[l] once per tile for bits 4h to 4h + 3, where the BK=64 loader +// reads qh[l] on all four tiles of the super-block. +template < + typename T, + short BROWS, + short BCOLS, + short dst_ld, + short reduction_dim, + short tgp_size> +struct KqNaxQ5_KHalfLoader { + MLX_MTL_CONST int weights_per_block = KQ_Q5_K_SUPERBLOCK; + MLX_MTL_CONST int bytes_per_block = KQ_Q5_K_BLOCK_BYTES; + static_assert(BCOLS == 128, "Q5_K half loader requires BCOLS == 128."); + static_assert(reduction_dim == 1, "Q5_K half loader serves qmm_t only."); + static_assert(BROWS * 2 == tgp_size, "Q5_K half loader: 2 threads per row."); + + MLX_MTL_CONST short TCOLS = 2; + MLX_MTL_CONST short n_pos = 16; + // load_unsafe depends on half_idx only. + MLX_MTL_CONST bool db_safe = true; + + const short thread_idx; + const short bi; + const short lj; + + threadgroup T* dst; + const device uint8_t* src; + short half_idx; + + KqNaxQ5_KHalfLoader( + const device uint8_t* src_, + const int src_ld_, + threadgroup T* dst_, + ushort simd_group_id [[simdgroup_index_in_threadgroup]], + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + : thread_idx(simd_group_id* SIMD_SIZE + simd_lane_id), + bi(thread_idx / TCOLS), + lj((thread_idx % TCOLS) * n_pos), + dst(dst_ + bi * dst_ld + lj), + src(src_ + bi * (src_ld_ * bytes_per_block / weights_per_block)), + half_idx(0) {} + + void load_unsafe() thread { + const float d = float(*(const device half*)(src + KQ_Q5_K_D_OFFSET)); + const float dmin = float(*(const device half*)(src + KQ_Q5_K_DMIN_OFFSET)); + const device uint8_t* scales12 = src + KQ_Q5_K_SCALES_OFFSET; + float es[4], em[4]; +#pragma unroll + for (short q = 0; q < 4; q++) { + uint8_t sc6, mn6; + kq_get_scale_min_k4(half_idx * 4 + q, scales12, sc6, mn6); + es[q] = d * float(sc6); + em[q] = dmin * float(mn6); + } + const device uint8_t* qs = src + KQ_Q5_K_QS_OFFSET + half_idx * 64 + lj; + const device uint8_t* qh = src + KQ_Q5_K_QH_OFFSET + lj; + const uint4 a_v = *reinterpret_cast(qs); + const uint4 b_v = *reinterpret_cast(qs + 32); + const uint4 h_v = *reinterpret_cast(qh); + const thread uint8_t* a = reinterpret_cast(&a_v); + const thread uint8_t* b = reinterpret_cast(&b_v); + const thread uint8_t* hb = reinterpret_cast(&h_v); + const short hs = half_idx * 4; +#pragma unroll + for (short j = 0; j < n_pos; j++) { + const uint8_t h = hb[j] >> hs; + const uint8_t q0 = (a[j] & 0x0F) | ((h & 1u) << 4); + const uint8_t q1 = (a[j] >> 4) | (((h >> 1) & 1u) << 4); + const uint8_t q2 = (b[j] & 0x0F) | (((h >> 2) & 1u) << 4); + const uint8_t q3 = (b[j] >> 4) | (((h >> 3) & 1u) << 4); + dst[j] = T(es[0] * float(q0) - em[0]); + dst[32 + j] = T(es[1] * float(q1) - em[1]); + dst[64 + j] = T(es[2] * float(q2) - em[2]); + dst[96 + j] = T(es[3] * float(q3) - em[3]); + } + } + + void load_safe(short2 src_tile_dim) thread { + if (bi >= src_tile_dim.y) { +#pragma unroll + for (short j = 0; j < n_pos; j++) { + dst[j] = T(0); + dst[32 + j] = T(0); + dst[64 + j] = T(0); + dst[96 + j] = T(0); + } + return; + } + load_unsafe(); + } + + void next() thread { + half_idx += 1; + if (half_idx == 2) { + half_idx = 0; + src += bytes_per_block; + } + } +}; + // Q6_K: 210 bytes per 256-weight super-block. Reversed field order: // [ql[128]][qh[64]][int8 scales[16]][fp16 d] @@ -1863,6 +2067,109 @@ struct KqNaxQ6_KBlockLoader { } }; +// Q6_K on a K tile of 128: one tile is one half of a super-block (64 bytes +// of ql, 32 of qh, 8 scales). Two threads share a row, 16 positions l each. +// A thread reads ql[l], ql[32 + l] and qh[l] once and writes the weights l, +// l + 32, l + 64 and l + 96 of the half. The BK=64 loader reads each of +// these bytes on two tiles. The dequant arithmetic is the same, so the +// weights in Ws are bit-identical. +template < + typename T, + short BROWS, + short BCOLS, + short dst_ld, + short reduction_dim, + short tgp_size> +struct KqNaxQ6_KHalfLoader { + MLX_MTL_CONST int weights_per_block = KQ_Q6_K_SUPERBLOCK; + MLX_MTL_CONST int bytes_per_block = KQ_Q6_K_BLOCK_BYTES; + static_assert(BCOLS == 128, "Q6_K half loader requires BCOLS == 128."); + static_assert(reduction_dim == 1, "Q6_K half loader serves qmm_t only."); + static_assert(BROWS * 2 == tgp_size, "Q6_K half loader: 2 threads per row."); + + MLX_MTL_CONST short TCOLS = 2; + MLX_MTL_CONST short n_pos = 16; + // load_unsafe depends on half_idx only. + MLX_MTL_CONST bool db_safe = true; + + const short thread_idx; + const short bi; + const short lj; + + threadgroup T* dst; + const device uint8_t* src; + short half_idx; + + KqNaxQ6_KHalfLoader( + const device uint8_t* src_, + const int src_ld_, + threadgroup T* dst_, + ushort simd_group_id [[simdgroup_index_in_threadgroup]], + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + : thread_idx(simd_group_id* SIMD_SIZE + simd_lane_id), + bi(thread_idx / TCOLS), + lj((thread_idx % TCOLS) * n_pos), + dst(dst_ + bi * dst_ld + lj), + src(src_ + bi * (src_ld_ * bytes_per_block / weights_per_block)), + half_idx(0) {} + + void load_unsafe() thread { + const float d = float(*(const device half*)(src + KQ_Q6_K_D_OFFSET)); + // Scale 2q + (l >= 16) of the half serves weight 32q + l. + const device int8_t* scales = + (const device int8_t*)(src + KQ_Q6_K_SCALES_OFFSET) + half_idx * 8 + + lj / 16; + float es[4]; +#pragma unroll + for (short q = 0; q < 4; q++) { + es[q] = d * float(scales[q * 2]); + } + const device uint8_t* ql = src + KQ_Q6_K_QL_OFFSET + half_idx * 64 + lj; + const device uint8_t* qh = src + KQ_Q6_K_QH_OFFSET + half_idx * 32 + lj; + uint8_t la[n_pos], lb[n_pos], hb[n_pos]; +#pragma unroll + for (short j = 0; j < n_pos; j++) { + la[j] = ql[j]; + lb[j] = ql[32 + j]; + hb[j] = qh[j]; + } +#pragma unroll + for (short j = 0; j < n_pos; j++) { + const uint h = hb[j]; + const int q0 = int((la[j] & 0x0F) | ((h & 0x03) << 4)) - 32; + const int q1 = int((lb[j] & 0x0F) | (((h >> 2) & 0x03) << 4)) - 32; + const int q2 = int((la[j] >> 4) | (((h >> 4) & 0x03) << 4)) - 32; + const int q3 = int((lb[j] >> 4) | (((h >> 6) & 0x03) << 4)) - 32; + dst[j] = T(es[0] * float(q0)); + dst[32 + j] = T(es[1] * float(q1)); + dst[64 + j] = T(es[2] * float(q2)); + dst[96 + j] = T(es[3] * float(q3)); + } + } + + void load_safe(short2 src_tile_dim) thread { + if (bi >= src_tile_dim.y) { +#pragma unroll + for (short j = 0; j < n_pos; j++) { + dst[j] = T(0); + dst[32 + j] = T(0); + dst[64 + j] = T(0); + dst[96 + j] = T(0); + } + return; + } + load_unsafe(); + } + + void next() thread { + half_idx += 1; + if (half_idx == 2) { + half_idx = 0; + src += bytes_per_block; + } + } +}; + // Q3_K: 110 bytes per 256-weight super-block. // [hmask[32]][qs[64]][scales[12]][fp16 d] @@ -1991,6 +2298,109 @@ struct KqNaxQ3_KBlockLoader { } }; +// Q3_K on a K tile of 128: one tile is one half of a super-block (32 bytes +// of qs, the 32 hmask bytes, 8 scales). Two threads share a row, 16 +// positions l each. A thread reads qs[32h + l] and hmask[l] once and writes +// the weights l, l + 32, l + 64 and l + 96 of the half. The BK=64 loader +// reads each qs byte on two tiles and each hmask byte on four. +template < + typename T, + short BROWS, + short BCOLS, + short dst_ld, + short reduction_dim, + short tgp_size> +struct KqNaxQ3_KHalfLoader { + MLX_MTL_CONST int weights_per_block = KQ_Q3_K_SUPERBLOCK; + MLX_MTL_CONST int bytes_per_block = KQ_Q3_K_BLOCK_BYTES; + static_assert(BCOLS == 128, "Q3_K half loader requires BCOLS == 128."); + static_assert(reduction_dim == 1, "Q3_K half loader serves qmm_t only."); + static_assert(BROWS * 2 == tgp_size, "Q3_K half loader: 2 threads per row."); + + MLX_MTL_CONST short TCOLS = 2; + MLX_MTL_CONST short n_pos = 16; + // load_unsafe depends on half_idx only. + MLX_MTL_CONST bool db_safe = true; + + const short thread_idx; + const short bi; + const short lj; + + threadgroup T* dst; + const device uint8_t* src; + short half_idx; + + KqNaxQ3_KHalfLoader( + const device uint8_t* src_, + const int src_ld_, + threadgroup T* dst_, + ushort simd_group_id [[simdgroup_index_in_threadgroup]], + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + : thread_idx(simd_group_id* SIMD_SIZE + simd_lane_id), + bi(thread_idx / TCOLS), + lj((thread_idx % TCOLS) * n_pos), + dst(dst_ + bi * dst_ld + lj), + src(src_ + bi * (src_ld_ * bytes_per_block / weights_per_block)), + half_idx(0) {} + + void load_unsafe() thread { + const float d = float(*(const device half*)(src + KQ_Q3_K_D_OFFSET)); + // Scale 2kt + (l >= 16) serves weight 32kt + l. + const short scale_off = lj / 16; + float es[4]; +#pragma unroll + for (short q = 0; q < 4; q++) { + const short kt = half_idx * 4 + q; + es[q] = d * + float((int)kq_q3_k_unpack_scale( + kt * 2 + scale_off, src + KQ_Q3_K_SCALES_OFFSET) - + 32); + } + const device uint8_t* qs = src + KQ_Q3_K_QS_OFFSET + half_idx * 32 + lj; + const device uint8_t* hm = src + KQ_Q3_K_HMASK_OFFSET + lj; + uint8_t qb[n_pos], hb[n_pos]; +#pragma unroll + for (short j = 0; j < n_pos; j++) { + qb[j] = qs[j]; + hb[j] = hm[j]; + } + const short hs = half_idx * 4; +#pragma unroll + for (short j = 0; j < n_pos; j++) { + const uint8_t q = qb[j]; + const uint8_t h = hb[j] >> hs; +#pragma unroll + for (short p = 0; p < 4; p++) { + const uint8_t q2 = (q >> (p * 2)) & 0x03; + const int q3 = (int)q2 - (((h >> p) & 1) ? 0 : 4); + dst[p * 32 + j] = T(es[p] * float(q3)); + } + } + } + + void load_safe(short2 src_tile_dim) thread { + if (bi >= src_tile_dim.y) { +#pragma unroll + for (short j = 0; j < n_pos; j++) { + dst[j] = T(0); + dst[32 + j] = T(0); + dst[64 + j] = T(0); + dst[96 + j] = T(0); + } + return; + } + load_unsafe(); + } + + void next() thread { + half_idx += 1; + if (half_idx == 2) { + half_idx = 0; + src += bytes_per_block; + } + } +}; + // Q2_K: 84 bytes per 256-weight super-block. [scales[16]][qs[64]][fp16 d][fp16 // dmin] @@ -3651,18 +4061,173 @@ KQ_NAX_DEFINE_KERNELS(q5_0, 32, 5, KqNaxQ5_0BlockLoader) KQ_NAX_DEFINE_KERNELS(q4_k, 256, 4, KqNaxQ4_KBlockLoader) KQ_NAX_DEFINE_KERNELS(q5_k, 256, 5, KqNaxQ5_KBlockLoader) KQ_NAX_DEFINE_KERNELS(q6_k, 256, 6, KqNaxQ6_KBlockLoader) + KQ_NAX_DEFINE_KERNELS(q3_k, 256, 3, KqNaxQ3_KBlockLoader) KQ_NAX_DEFINE_KERNELS(q2_k, 256, 2, KqNaxQ2_KBlockLoader) +// BK=128 tile from two BK=64 loaders of a db_safe codec: lo fills K columns +// 0-63 of the tile and hi columns 64-127, one BK=64 position apart, and +// next() moves both two positions. The tile holds the same weights as two +// BK=64 tiles, so the output is bit-identical. Half the barriers and loop +// steps of BK=64 per K. +template +struct KqNaxPairLoader { + static_assert(BCOLS == 128, "pair loader requires BCOLS == 128."); + static_assert(reduction_dim == 1, "pair loader serves qmm_t only."); + static_assert(Half::db_safe, "pair loader requires a db_safe loader."); + MLX_MTL_CONST int weights_per_block = Half::weights_per_block; + MLX_MTL_CONST int bytes_per_block = Half::bytes_per_block; + MLX_MTL_CONST bool db_safe = true; + + Half lo; + Half hi; + + KqNaxPairLoader( + const device uint8_t* src, + const int src_ld, + threadgroup T* dst, + ushort simd_group_id, + ushort simd_lane_id) thread + : lo(src, src_ld, dst, simd_group_id, simd_lane_id), + hi(src, src_ld, dst + 64, simd_group_id, simd_lane_id) { + hi.next(); + } + + void load_unsafe() const thread { + lo.load_unsafe(); + hi.load_unsafe(); + } + + void load_safe(short2 src_tile_dim) const thread { + lo.load_safe(src_tile_dim); + hi.load_safe(src_tile_dim); + } + + void next() thread { + lo.next(); + lo.next(); + hi.next(); + hi.next(); + } +}; + +#define KQ_NAX_DEFINE_PAIR_LOADER(NAME, HALF) \ + template < \ + typename T, \ + short BROWS, \ + short BCOLS, \ + short dst_ld, \ + short reduction_dim, \ + short tgp_size> \ + using NAME = KqNaxPairLoader< \ + T, \ + BCOLS, \ + reduction_dim, \ + HALF>; + +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxQ4_0PairLoader, KqNaxQ4_0BlockLoader) +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxQ4_1PairLoader, KqNaxQ4_1BlockLoader) +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxQ5_0PairLoader, KqNaxQ5_0BlockLoader) +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxIq4_nlPairLoader, KqNaxIq4_nlBlockLoader) +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxIq3_xxsPairLoader, KqNaxIq3_xxsBlockLoader) +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxIq3_sPairLoader, KqNaxIq3_sBlockLoader) +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxIq2_xsPairLoader, KqNaxIq2_xsBlockLoader) +KQ_NAX_DEFINE_PAIR_LOADER(KqNaxIq2_sPairLoader, KqNaxIq2_sBlockLoader) +#undef KQ_NAX_DEFINE_PAIR_LOADER + +// qmm_t on a K tile of 128 with a half-super-block or pair loader (one Ws +// buffer). +// The MMA walks K in the same 32-wide steps as on BK=64, so the output is +// bit-identical to the BK=64 kernel of the same tile. +#define KQ_NAX_DEFINE_BK128_KERNEL(codec, GROUP_CONST, bits_val, LOADER) \ + template < \ + typename T, \ + int group_size, \ + int bits, \ + bool aligned_N, \ + bool batched, \ + int BM, \ + int BN, \ + int WM, \ + int WN> \ + [[kernel]] void kq_##codec##_qmm_t_nax_bk128( \ + const device uint8_t* w, \ + const device uint8_t* /* scales */, \ + const device T* x, \ + device T* y, \ + const constant int& K, \ + const constant int& N, \ + const constant int& M, \ + const constant int& x_batch_ndims, \ + const constant int* x_shape, \ + const constant int64_t* x_strides, \ + const constant int& w_batch_ndims, \ + const constant int* w_shape, \ + const constant int64_t* w_strides, \ + const constant int64_t* /* s_strides */, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 tgp_grid [[threadgroups_per_grid]], \ + uint lid [[thread_index_in_threadgroup]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + static_assert(group_size == GROUP_CONST, #codec " group_size"); \ + static_assert(bits == bits_val, #codec " bits"); \ + tid = kq_swizzle_tid(tid, tgp_grid, N, BN); \ + if (int(tid.y) * BM >= M) { \ + return; \ + } \ + constexpr int BK = 128; \ + constexpr int BK_padded = (BK + 16 / sizeof(T)); \ + using LoaderW = LOADER< \ + T, \ + BN, \ + BK, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + threadgroup T Ws[BN * BK_padded]; \ + if constexpr (batched) { \ + kq_adjust_matrix_offsets( \ + x, \ + w, \ + y, \ + M * N, \ + x_batch_ndims, \ + x_shape, \ + x_strides, \ + w_batch_ndims, \ + w_shape, \ + w_strides, \ + tid); \ + } \ + kq_qmm_t_nax_tgp_impl( \ + w, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); \ + } + +KQ_NAX_DEFINE_BK128_KERNEL(q6_k, 256, 6, KqNaxQ6_KHalfLoader) +KQ_NAX_DEFINE_BK128_KERNEL(q5_k, 256, 5, KqNaxQ5_KHalfLoader) +KQ_NAX_DEFINE_BK128_KERNEL(q3_k, 256, 3, KqNaxQ3_KHalfLoader) +KQ_NAX_DEFINE_BK128_KERNEL(q4_k, 256, 4, KqNaxQ4_KHalfLoader) +KQ_NAX_DEFINE_BK128_KERNEL(q4_0, 32, 4, KqNaxQ4_0PairLoader) +KQ_NAX_DEFINE_BK128_KERNEL(q4_1, 32, 4, KqNaxQ4_1PairLoader) +KQ_NAX_DEFINE_BK128_KERNEL(q5_0, 32, 5, KqNaxQ5_0PairLoader) +KQ_NAX_DEFINE_BK128_KERNEL(iq4_nl, 32, 4, KqNaxIq4_nlPairLoader) +KQ_NAX_DEFINE_BK128_KERNEL(iq3_xxs, 256, 3, KqNaxIq3_xxsPairLoader) +KQ_NAX_DEFINE_BK128_KERNEL(iq3_s, 256, 3, KqNaxIq3_sPairLoader) +KQ_NAX_DEFINE_BK128_KERNEL(iq2_xs, 256, 2, KqNaxIq2_xsPairLoader) +KQ_NAX_DEFINE_BK128_KERNEL(iq2_s, 256, 2, KqNaxIq2_sPairLoader) +#undef KQ_NAX_DEFINE_BK128_KERNEL + // Split-K qmm_t on the NAX tile (KQ_QMM_SPLITK_NAX): // grid.z indexes K-slices; each slice walks k_partition_size weights from a -// superblock-aligned start and stores a T partial tile at -// tid.z * split_k_partition_stride. The shared kquant_qmm_splitk_accum pass -// folds slices in f32. The host guarantees k_partition_size is a multiple -// of both the codec superblock and BK, so every slice starts the loader at -// kt_base 0. The last slice stops at K, so it may be shorter than the -// others. Non-batched transpose shapes only; no swizzle (grid.y is a -// single row tile in the target band). +// superblock-aligned start and stores a U partial tile (float on float16 +// activations, T on bfloat16; KQ_SPLITK_F32 overrides both) at +// tid.z * split_k_partition_stride. The shared +// kquant_qmm_splitk_accum pass folds slices in f32. The host guarantees +// k_partition_size is a multiple of both the codec superblock and BK, so every +// slice starts the loader at kt_base 0. The last slice stops at K, so it may be +// shorter than the others. Non-batched transpose shapes only; no swizzle +// (grid.y is a single row tile in the target band). #define KQ_NAX_DEFINE_SPLITK_KERNEL(codec, GROUP_CONST, bits_val, LOADER) \ template < \ typename T, \ @@ -3672,12 +4237,13 @@ KQ_NAX_DEFINE_KERNELS(q2_k, 256, 2, KqNaxQ2_KBlockLoader) int BM, \ int BN, \ int WM, \ - int WN> \ + int WN, \ + typename U = T> \ [[kernel]] void kq_##codec##_qmm_t_nax_splitk( \ const device uint8_t* w, \ const device uint8_t* /* scales */, \ const device T* x, \ - device T* y, \ + device U* y, \ const constant int& K, \ const constant int& N, \ const constant int& M, \ @@ -3722,7 +4288,8 @@ KQ_NAX_DEFINE_KERNELS(q2_k, 256, 2, KqNaxQ2_KBlockLoader) WM, \ WN, \ kWsBufs == 2, \ - true>(wl, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid, k_len); \ + true, \ + U>(wl, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid, k_len); \ } KQ_NAX_DEFINE_SPLITK_KERNEL(q6_k, 256, 6, KqNaxQ6_KBlockLoader) diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax_epi.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax_epi.h new file mode 100644 index 00000000..1838eb63 --- /dev/null +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax_epi.h @@ -0,0 +1,566 @@ +// NAX qmm_t kernels with an epilogue, built on the tile and loaders of +// kq_quantized_nax.h. Separate entry points, so the plain qmm_t kernels +// keep their code. +// +// GLU product: y = act(x @ Wg^T) * (x @ Wu^T) for two quantized weights of +// one codec and one shape. A threadgroup owns one [BM, BN] output tile of +// y, the only output buffer. The activation and the product run in f32. +// +// Two forms. The two-walk form (kq_qmm_t_nax_glu_walks_impl) is the one +// the op runs. The one-walk form (kq_qmm_t_nax_glu_impl) keeps both sums +// live and measured slower; it builds as a probe kernel only. +#include "mlx/backend/metal/kernels/kq_quantized_nax.h" + +// 0 = silu, 1 = gelu (tanh form). +constant int kq_epi_act [[function_constant(360)]]; +// Saturate at the finite half range before the store (float16 outputs). +constant bool kq_epi_saturate [[function_constant(361)]]; + +// A finite value past the half range stores the largest finite half with +// its sign. A NaN and an infinity pass: clamp() returns its lower bound for +// a NaN, so the two sides are compared one at a time. +METAL_FUNC float kq_epi_saturate16(float v) { + v = (v > 65504.0f && v <= FLT_MAX) ? 65504.0f : v; + v = (v < -65504.0f && v >= -FLT_MAX) ? -65504.0f : v; + return v; +} + +METAL_FUNC float kq_epi_activation(float g) { + if (kq_epi_act == 1) { + // 0.5 g (1 + tanh(sqrt(2 / pi) (g + 0.044715 g^3))) + const float u = 0.7978845608028654f * (g + 0.044715f * g * g * g); + return 0.5f * g * (1.0f + precise::tanh(u)); + } + return g / (1.0f + precise::exp(-g)); +} + +// GLU product in one K walk with two f32 tiles, rounded once at the store. +// dual = true keeps the two decoded weight tiles in two Ws buffers: one +// barrier pair per K step serves both weights, and each A tile is loaded +// once for both products. dual = false decodes the two weights one after +// the other into one Ws buffer. WS_LD is the row stride of a tile in Ws: +// BK plus the pad of the plain kernels, or BK alone where two padded +// buffers would pass the 32 KB of threadgroup memory. +template < + typename T, + typename LoaderW, + const bool aligned_N, + const int BM, + const int BK, + const int BN, + const int WM, + const int WN, + const bool dual, + const int WS_LD = (BK + 16 / sizeof(T))> +METAL_FUNC void kq_qmm_t_nax_glu_impl( + const device uint8_t* wg, + const device uint8_t* wu, + const device T* x, + device T* y, + threadgroup T* Ws, + const constant int& K, + const constant int& N, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BK % 32 == 0, "BK must be a multiple of the MMA K step"); + + // Row stride of a decoded weight tile in Ws. + constexpr int BK_padded = WS_LD; + constexpr int WS_STRIDE = BN * BK_padded; + + const int K_w = (K / LoaderW::weights_per_block) * LoaderW::bytes_per_block; + const int y_row = tid.y * BM; + const int y_col = tid.x * BN; + + x += y_row * static_cast(K); + wg += static_cast(y_col) * K_w; + wu += static_cast(y_col) * K_w; + y += y_row * static_cast(N) + y_col; + + threadgroup T* Wg = Ws; + threadgroup T* Wu = dual ? Ws + WS_STRIDE : Ws; + LoaderW loader_g(wg, K, Wg, simd_gid, simd_lid); + LoaderW loader_u(wu, K, Wu, simd_gid, simd_lid); + + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + const short tm = SM * (simd_gid / WN); + const short tn = SN * (simd_gid % WN); + + constexpr bool transpose_a = false; + constexpr bool transpose_b = true; + + // Spans are clamped in int before narrowing: M and N pass 32767. + const short sgp_sm = short(min(int(SM), M - (y_row + tm))); + const bool is_unaligned_sm = (M - y_row) < BM; + // A simdgroup row at or beyond M decodes its share of the weights and + // keeps every barrier, and skips the A loads, the MMAs and the store. + const bool sg_live = tm < M - y_row; + + const short sgp_sn = aligned_N ? SN : short(min(int(SN), N - (y_col + tn))); + const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); + const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); + + NAXTile Dg; + NAXTile Du; + Dg.clear(); + Du.clear(); + + x += tm * K; + + dispatch_bool(!is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(aligned_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_g.load_unsafe(); + } else { + loader_g.load_safe(short2(BK, tgp_bn)); + } + if constexpr (dual) { + if constexpr (kAlignedN.value) { + loader_u.load_unsafe(); + } else { + loader_u.load_safe(short2(BK, tgp_bn)); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; sg_live && kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } + + Btile.template load(Wg + tn * BK_padded + kk1); + tile_matmad_nax( + Dg, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + if constexpr (dual) { + Btile.template load(Wu + tn * BK_padded + kk1); + tile_matmad_nax( + Du, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + } + + (void)compiler_barrier; + } + + if constexpr (!dual) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_u.load_unsafe(); + } else { + loader_u.load_safe(short2(BK, tgp_bn)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; sg_live && kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } + + Btile.template load(Wu + tn * BK_padded + kk1); + tile_matmad_nax( + Du, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; + } + } + + x += BK; + loader_g.next(); + loader_u.next(); + } + + if (!sg_live) { + return; + } + + constexpr short NE = decltype(Dg)::kElemsPerFrag; + constexpr short NF = decltype(Dg)::kNumFrags; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < NF; ++f) { + STEEL_PRAGMA_UNROLL + for (short e = 0; e < NE; ++e) { + float v = kq_epi_activation(Dg.val_frags[f][e]) * Du.val_frags[f][e]; + if (kq_epi_saturate) { + v = kq_epi_saturate16(v); + } + Dg.val_frags[f][e] = v; + } + } + + if constexpr (kAlignedM.value && kAlignedN.value) { + Dg.store(y + tm * N + tn, N); + } else if (kAlignedM.value && sgp_sn == SN) { + Dg.store(y + tm * N + tn, N); + } else { + Dg.store_safe(y + tm * N + tn, N, short2(sgp_sn, sgp_sm)); + } + }); + }); +} + +// GLU product in two K walks with one f32 tile: the gate walk stores +// act(x @ Wg^T) to y in the output dtype, the up walk reads that tile back +// and stores the product. One accumulator tile is live at a time, as in +// the plain kernels. With the saturating store a finite activation past +// the half range is stored as the largest finite half before the product, +// and a finite product after it. A NaN or an infinity passes. +template < + typename T, + typename LoaderW, + const bool aligned_N, + const int BM, + const int BK, + const int BN, + const int WM, + const int WN> +METAL_FUNC void kq_qmm_t_nax_glu_walks_impl( + const device uint8_t* wg, + const device uint8_t* wu, + const device T* x, + device T* y, + threadgroup T* Ws, + const constant int& K, + const constant int& N, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BK % 32 == 0, "BK must be a multiple of the MMA K step"); + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + + const int K_w = (K / LoaderW::weights_per_block) * LoaderW::bytes_per_block; + const int y_row = tid.y * BM; + const int y_col = tid.x * BN; + + x += y_row * static_cast(K); + wg += static_cast(y_col) * K_w; + wu += static_cast(y_col) * K_w; + y += y_row * static_cast(N) + y_col; + + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + const short tm = SM * (simd_gid / WN); + const short tn = SN * (simd_gid % WN); + + constexpr bool transpose_a = false; + constexpr bool transpose_b = true; + + // Spans are clamped in int before narrowing: M and N pass 32767. + const short sgp_sm = short(min(int(SM), M - (y_row + tm))); + const bool is_unaligned_sm = (M - y_row) < BM; + const bool sg_live = tm < M - y_row; + + const short sgp_sn = aligned_N ? SN : short(min(int(SN), N - (y_col + tn))); + const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); + const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); + + x += tm * K; + y += tm * N + tn; + + dispatch_bool(!is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(aligned_N || !is_unaligned_bn, [&](auto kAlignedN) { + STEEL_PRAGMA_NO_UNROLL + for (int pass = 0; pass < 2; pass++) { + LoaderW loader_w(pass == 0 ? wg : wu, K, Ws, simd_gid, simd_lid); + const device T* xk = x; + NAXTile Dtile; + Dtile.clear(); + + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe(short2(BK, tgp_bn)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; sg_live && kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xk + kk1, K); + } else { + Atile.load_safe(xk + kk1, K, short2(SK, sgp_sm)); + } + + Btile.template load(Ws + tn * BK_padded + kk1); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; + } + + xk += BK; + loader_w.next(); + } + + if (!sg_live) { + continue; + } + + constexpr short NE = decltype(Dtile)::kElemsPerFrag; + constexpr short NF = decltype(Dtile)::kNumFrags; + if (pass == 0) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < NF; ++f) { + STEEL_PRAGMA_UNROLL + for (short e = 0; e < NE; ++e) { + float v = kq_epi_activation(Dtile.val_frags[f][e]); + if (kq_epi_saturate) { + v = kq_epi_saturate16(v); + } + Dtile.val_frags[f][e] = v; + } + } + } else { + NAXTile Ctile; + if constexpr (kAlignedM.value && kAlignedN.value) { + Ctile.load(y, N); + } else if (kAlignedM.value && sgp_sn == SN) { + Ctile.load(y, N); + } else { + Ctile.load_safe(y, N, short2(sgp_sn, sgp_sm)); + } + STEEL_PRAGMA_UNROLL + for (short f = 0; f < NF; ++f) { + STEEL_PRAGMA_UNROLL + for (short e = 0; e < NE; ++e) { + float v = Ctile.val_frags[f][e] * Dtile.val_frags[f][e]; + if (kq_epi_saturate) { + v = kq_epi_saturate16(v); + } + Dtile.val_frags[f][e] = v; + } + } + } + + if constexpr (kAlignedM.value && kAlignedN.value) { + Dtile.store(y, N); + } else if (kAlignedM.value && sgp_sn == SN) { + Dtile.store(y, N); + } else { + Dtile.store_safe(y, N, short2(sgp_sn, sgp_sm)); + } + } + }); + }); +} + +#define KQ_NAX_DEFINE_GLU_WALKS_KERNEL(codec, LOADER64, LOADER128) \ + template < \ + typename T, \ + bool aligned_N, \ + int BM, \ + int BN, \ + int BK, \ + int WM, \ + int WN> \ + [[kernel]] void kq_##codec##_qmm_t_nax_glu_walks( \ + const device uint8_t* wg, \ + const device uint8_t* wu, \ + const device T* x, \ + device T* y, \ + const constant int& K, \ + const constant int& N, \ + const constant int& M, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 tgp_grid [[threadgroups_per_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + tid = kq_swizzle_tid(tid, tgp_grid, N, BN); \ + if (int(tid.y) * BM >= M) { \ + return; \ + } \ + constexpr int BK_padded = (BK + 16 / sizeof(T)); \ + threadgroup T Ws[BN * BK_padded]; \ + if constexpr (BK == 64) { \ + using LoaderW = LOADER64< \ + T, \ + BN, \ + 64, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_glu_walks_impl( \ + wg, wu, x, y, Ws, K, N, M, tid, simd_gid, simd_lid); \ + } else { \ + using LoaderW = LOADER128< \ + T, \ + BN, \ + 128, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_glu_walks_impl( \ + wg, wu, x, y, Ws, K, N, M, tid, simd_gid, simd_lid); \ + } \ + } + +// GLU kernel of one codec. LOADER64 decodes a K tile of 64 and LOADER128 a +// K tile of 128. BK 64 keeps two Ws buffers; BK 128 keeps one, since two +// would pass 32 KB. +#define KQ_NAX_DEFINE_GLU_KERNEL(codec, LOADER64, LOADER128) \ + template < \ + typename T, \ + bool aligned_N, \ + int BM, \ + int BN, \ + int BK, \ + int WM, \ + int WN> \ + [[kernel]] void kq_##codec##_qmm_t_nax_glu( \ + const device uint8_t* wg, \ + const device uint8_t* wu, \ + const device T* x, \ + device T* y, \ + const constant int& K, \ + const constant int& N, \ + const constant int& M, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 tgp_grid [[threadgroups_per_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + tid = kq_swizzle_tid(tid, tgp_grid, N, BN); \ + if (int(tid.y) * BM >= M) { \ + return; \ + } \ + constexpr int BK_padded = (BK + 16 / sizeof(T)); \ + constexpr bool dual = BK == 64; \ + threadgroup T Ws[(dual ? 2 : 1) * BN * BK_padded]; \ + if constexpr (BK == 64) { \ + using LoaderW = LOADER64< \ + T, \ + BN, \ + 64, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_glu_impl( \ + wg, wu, x, y, Ws, K, N, M, tid, simd_gid, simd_lid); \ + } else { \ + using LoaderW = LOADER128< \ + T, \ + BN, \ + 128, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_glu_impl< \ + T, \ + LoaderW, \ + aligned_N, \ + BM, \ + 128, \ + BN, \ + WM, \ + WN, \ + false>(wg, wu, x, y, Ws, K, N, M, tid, simd_gid, simd_lid); \ + } \ + } + +// GLU kernel on a K tile of 128 with two unpadded Ws buffers, 32 KB at +// BN 64 and a 2-byte T. +#define KQ_NAX_DEFINE_GLU_D128_KERNEL(codec, LOADER128) \ + template \ + [[kernel]] void kq_##codec##_qmm_t_nax_glu_d128( \ + const device uint8_t* wg, \ + const device uint8_t* wu, \ + const device T* x, \ + device T* y, \ + const constant int& K, \ + const constant int& N, \ + const constant int& M, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 tgp_grid [[threadgroups_per_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + tid = kq_swizzle_tid(tid, tgp_grid, N, BN); \ + if (int(tid.y) * BM >= M) { \ + return; \ + } \ + static_assert(2 * BN * 128 * sizeof(T) <= 32768, "Ws over 32 KB"); \ + threadgroup T Ws[2 * BN * 128]; \ + using LoaderW = LOADER128< \ + T, \ + BN, \ + 128, \ + 128, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_glu_impl< \ + T, \ + LoaderW, \ + aligned_N, \ + BM, \ + 128, \ + BN, \ + WM, \ + WN, \ + true, \ + 128>(wg, wu, x, y, Ws, K, N, M, tid, simd_gid, simd_lid); \ + } + +KQ_NAX_DEFINE_GLU_KERNEL(q4_k, KqNaxQ4_KBlockLoader, KqNaxQ4_KHalfLoader) +KQ_NAX_DEFINE_GLU_D128_KERNEL(q4_k, KqNaxQ4_KHalfLoader) +KQ_NAX_DEFINE_GLU_WALKS_KERNEL(q4_k, KqNaxQ4_KBlockLoader, KqNaxQ4_KHalfLoader) +KQ_NAX_DEFINE_GLU_WALKS_KERNEL(q5_k, KqNaxQ5_KBlockLoader, KqNaxQ5_KHalfLoader) +KQ_NAX_DEFINE_GLU_WALKS_KERNEL(q6_k, KqNaxQ6_KBlockLoader, KqNaxQ6_KHalfLoader) +// q8_0 has a K tile of 64 only. +KQ_NAX_DEFINE_GLU_WALKS_KERNEL(q8_0, KqNaxQ8_0BlockLoader, KqNaxQ8_0BlockLoader) +#undef KQ_NAX_DEFINE_GLU_WALKS_KERNEL +#undef KQ_NAX_DEFINE_GLU_KERNEL +#undef KQ_NAX_DEFINE_GLU_D128_KERNEL diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax_lora.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax_lora.h new file mode 100644 index 00000000..52ba7ce1 --- /dev/null +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax_lora.h @@ -0,0 +1,321 @@ +// NAX qmm_t kernel with a low-rank term in the same reduction, built on the +// tile and loaders of kq_quantized_nax.h. Separate entry points, so the +// plain qmm_t kernels keep their code. +// +// y = x @ W^T + z @ B +// +// x [M, K], W a quantized weight [N, K], z [M, R] and B [R, N] in the +// dtype of x. A threadgroup owns one [BM, BN] tile of y. It walks the rank +// of its column band in steps of 32, with the A tile loading z and the B +// tile loading B straight from their buffers, and then walks K as the +// plain kernel does. Both walks add into one f32 accumulator, which is +// rounded once at the store, so y is written once and never read. +#include "mlx/backend/metal/kernels/kq_quantized_nax.h" + +// The walk of a band's rank, 32 ranks a step, into the accumulator of a +// simdgroup. sg_cols is how many of the simdgroup's columns lie in the +// band, and sg_rows how many of its rows lie in M. A simdgroup with no +// column in the band runs no step, so no value of z reaches its columns. +// Past the band's rank the A tile and the B tile load zero, so no z column +// and no B row of another band enters the sum, whatever it holds. Past the +// band's end the B tile loads zero. The host runs this kernel only for +// bands that end on a simdgroup's column tile or at N, so a simdgroup holds +// no stored column outside its band. +template < + typename T, + const bool aligned_M, + const short TM, + const short TN, + const short TK> +METAL_FUNC void kq_qmm_t_nax_lora_tail( + thread NAXTile& Dtile, + const device T* z, + const device T* b, + const int R, + const int N, + const int band_rank, + const short sg_cols, + const short sg_rows) { + constexpr short SK = TK * 16; + constexpr short SN = TN * 16; + STEEL_PRAGMA_NO_UNROLL + for (int r0 = 0; sg_cols > 0 && r0 < band_rank; r0 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short zk = short(min(int(SK), band_rank - r0)); + if constexpr (aligned_M) { + if (zk == SK) { + Atile.load(z + r0, R); + } else { + Atile.load_safe(z + r0, R, short2(zk, sg_rows)); + } + } else { + Atile.load_safe(z + r0, R, short2(zk, sg_rows)); + } + + const device T* br = b + r0 * static_cast(N); + if (zk == SK && sg_cols == SN) { + Btile.load(br, N); + } else { + Btile.load_safe(br, N, short2(sg_cols, zk)); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; + } +} + +// band_z, band_rank and band_cols describe the column band of the +// threadgroup's tile: the first z column, the rank, and how many of the +// tile's columns lie in the band. A tile with no column in a band has +// band_cols at or below 0 and runs the plain walk alone. +template < + typename T, + typename LoaderW, + const bool aligned_N, + const int BM, + const int BK, + const int BN, + const int WM, + const int WN> +METAL_FUNC void kq_qmm_t_nax_lora_impl( + const device uint8_t* w, + const device T* x, + const device T* z, + const device T* b, + device T* y, + threadgroup T* Ws, + const constant int& K, + const constant int& N, + const constant int& M, + const constant int& R, + const int band_z, + const int band_rank, + const int band_cols, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BK % 32 == 0, "BK must be a multiple of the MMA K step"); + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + + const int K_w = (K / LoaderW::weights_per_block) * LoaderW::bytes_per_block; + const int y_row = tid.y * BM; + const int y_col = tid.x * BN; + + x += y_row * static_cast(K); + w += static_cast(y_col) * K_w; + y += y_row * static_cast(N) + y_col; + z += y_row * static_cast(R) + band_z; + b += band_z * static_cast(N) + y_col; + + LoaderW loader_w(w, K, Ws, simd_gid, simd_lid); + + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + const short tm = SM * (simd_gid / WN); + const short tn = SN * (simd_gid % WN); + + constexpr bool transpose_a = false; + constexpr bool transpose_b = true; + + // Spans are clamped in int before narrowing: M and N pass 32767. + const short sgp_sm = short(min(int(SM), M - (y_row + tm))); + const bool is_unaligned_sm = (M - y_row) < BM; + // A simdgroup row at or beyond M keeps its share of the weight decode + // and every barrier, and skips the A loads, the MMAs and the store. + const bool sg_live = tm < M - y_row; + + const short sgp_sn = aligned_N ? SN : short(min(int(SN), N - (y_col + tn))); + const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); + const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); + + NAXTile Dtile; + Dtile.clear(); + + x += tm * K; + z += tm * R; + // Columns of the simdgroup inside the band. + const short sg_cols = short(clamp(band_cols - int(tn), 0, int(SN))); + + dispatch_bool(!is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(aligned_N || !is_unaligned_bn, [&](auto kAlignedN) { + // The low-rank term first, then the K walk and the store, in the + // order of the plain kernel. + if (sg_live) { + kq_qmm_t_nax_lora_tail( + Dtile, z, b + tn, R, N, band_rank, sg_cols, sgp_sm); + } + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe(short2(BK, tgp_bn)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; sg_live && kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } + + Btile.template load(Ws + tn * BK_padded + kk1); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; + } + + x += BK; + loader_w.next(); + } + + if (!sg_live) { + return; + } + + if constexpr (kAlignedM.value && kAlignedN.value) { + Dtile.store(y + tm * N + tn, N); + } else if (kAlignedM.value && sgp_sn == SN) { + Dtile.store(y + tm * N + tn, N); + } else { + Dtile.store_safe(y + tm * N + tn, N, short2(sgp_sn, sgp_sm)); + } + }); + }); +} + +// Kernel of one codec. LOADER64 decodes a K tile of 64 and LOADER128 a K +// tile of 128; a codec with one K tile names its loader twice. bands holds +// n_bands rows of (col_start, col_end, z_start, rank), disjoint and in +// column order, with every col_start on a column tile. The tile takes the +// last band that starts at or before it: past that band's end band_cols is +// at or below 0. +#define KQ_NAX_DEFINE_LORA_KERNEL(codec, LOADER64, LOADER128) \ + template < \ + typename T, \ + bool aligned_N, \ + int BM, \ + int BN, \ + int BK, \ + int WM, \ + int WN> \ + [[kernel]] void kq_##codec##_qmm_t_nax_lora( \ + const device uint8_t* w, \ + const device T* x, \ + const device T* z, \ + const device T* b, \ + device T* y, \ + const constant int& K, \ + const constant int& N, \ + const constant int& M, \ + const constant int& R, \ + const constant int& n_bands, \ + const constant int* bands, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 tgp_grid [[threadgroups_per_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + tid = kq_swizzle_tid(tid, tgp_grid, N, BN); \ + if (int(tid.y) * BM >= M) { \ + return; \ + } \ + const int y_col = int(tid.x) * BN; \ + int band_z = 0; \ + int band_rank = 0; \ + int band_cols = 0; \ + for (int i = 0; i < n_bands; i++) { \ + if (y_col >= bands[4 * i]) { \ + band_z = bands[4 * i + 2]; \ + band_rank = bands[4 * i + 3]; \ + band_cols = min(BN, bands[4 * i + 1] - y_col); \ + } \ + } \ + constexpr int BK_padded = (BK + 16 / sizeof(T)); \ + threadgroup T Ws[BN * BK_padded]; \ + if constexpr (BK == 64) { \ + using LoaderW = LOADER64< \ + T, \ + BN, \ + 64, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_lora_impl( \ + w, \ + x, \ + z, \ + b, \ + y, \ + Ws, \ + K, \ + N, \ + M, \ + R, \ + band_z, \ + band_rank, \ + band_cols, \ + tid, \ + simd_gid, \ + simd_lid); \ + } else { \ + using LoaderW = LOADER128< \ + T, \ + BN, \ + 128, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_lora_impl( \ + w, \ + x, \ + z, \ + b, \ + y, \ + Ws, \ + K, \ + N, \ + M, \ + R, \ + band_z, \ + band_rank, \ + band_cols, \ + tid, \ + simd_gid, \ + simd_lid); \ + } \ + } + +KQ_NAX_DEFINE_LORA_KERNEL(q4_k, KqNaxQ4_KBlockLoader, KqNaxQ4_KHalfLoader) +KQ_NAX_DEFINE_LORA_KERNEL(q5_k, KqNaxQ5_KBlockLoader, KqNaxQ5_KHalfLoader) +KQ_NAX_DEFINE_LORA_KERNEL(q6_k, KqNaxQ6_KBlockLoader, KqNaxQ6_KHalfLoader) +KQ_NAX_DEFINE_LORA_KERNEL(q8_0, KqNaxQ8_0BlockLoader, KqNaxQ8_0BlockLoader) +#undef KQ_NAX_DEFINE_LORA_KERNEL diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax_resid.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax_resid.h new file mode 100644 index 00000000..dcf69e24 --- /dev/null +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax_resid.h @@ -0,0 +1,255 @@ +// NAX qmm_t kernel with a residual epilogue, built on the tile and loaders +// of kq_quantized_nax.h. Separate entry points, so the plain qmm_t kernels +// keep their code. +// +// y = r + g * (x @ W^T) +// +// x [B, M, K], W a quantized weight [N, K], r and y [B, M, N], g [B, N] or +// [N] (g_stride N or 0), all in the dtype of x. With gated false there is +// no g and y = r + x @ W^T. A threadgroup owns one [BM, BN] tile of y in +// batch tid.z. It walks K as the plain kernel does into one f32 +// accumulator, then reads its tile of r and its columns of g and stores +// r + g * acc, rounded once. +#include "mlx/backend/metal/kernels/kq_quantized_nax.h" + +template < + typename T, + typename LoaderW, + const bool aligned_N, + const bool gated, + const int BM, + const int BK, + const int BN, + const int WM, + const int WN> +METAL_FUNC void kq_qmm_t_nax_resid_impl( + const device uint8_t* w, + const device T* x, + const device T* r, + const device T* g, + device T* y, + threadgroup T* Ws, + const constant int& K, + const constant int& N, + const constant int& M, + const constant int& g_stride, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BK % 32 == 0, "BK must be a multiple of the MMA K step"); + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + + const int K_w = (K / LoaderW::weights_per_block) * LoaderW::bytes_per_block; + const int y_row = tid.y * BM; + const int y_col = tid.x * BN; + const int64_t batch = tid.z; + + x += (batch * M + y_row) * static_cast(K); + w += static_cast(y_col) * K_w; + y += (batch * M + y_row) * static_cast(N) + y_col; + r += (batch * M + y_row) * static_cast(N) + y_col; + if constexpr (gated) { + g += batch * g_stride + y_col; + } + + LoaderW loader_w(w, K, Ws, simd_gid, simd_lid); + + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + const short tm = SM * (simd_gid / WN); + const short tn = SN * (simd_gid % WN); + + constexpr bool transpose_a = false; + constexpr bool transpose_b = true; + + // Spans are clamped in int before narrowing: M and N pass 32767. + const short sgp_sm = short(min(int(SM), M - (y_row + tm))); + const bool is_unaligned_sm = (M - y_row) < BM; + // A simdgroup row at or beyond M keeps its share of the weight decode + // and every barrier, and skips the A loads, the MMAs and the store. + const bool sg_live = tm < M - y_row; + + const short sgp_sn = aligned_N ? SN : short(min(int(SN), N - (y_col + tn))); + const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); + const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); + + NAXTile Dtile; + Dtile.clear(); + + x += tm * K; + + dispatch_bool(!is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(aligned_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe(short2(BK, tgp_bn)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; sg_live && kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } + + Btile.template load(Ws + tn * BK_padded + kk1); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; + } + + x += BK; + loader_w.next(); + } + + if (!sg_live) { + return; + } + + const device T* rs = r + tm * N + tn; + device T* ys = y + tm * N + tn; + const bool full = (kAlignedM.value && kAlignedN.value) || + (kAlignedM.value && sgp_sn == SN); + + // The residual tile, and the gate as a tile whose rows all read + // the same columns (a row stride of 0). + NAXTile Rtile; + if (full) { + Rtile.load(rs, N); + } else { + Rtile.load_safe(rs, N, short2(sgp_sn, sgp_sm)); + } + constexpr short NE = decltype(Dtile)::kElemsPerFrag; + constexpr short NF = decltype(Dtile)::kNumFrags; + if constexpr (gated) { + NAXTile Gtile; + if (full) { + Gtile.load(g + tn, 0); + } else { + Gtile.load_safe(g + tn, 0, short2(sgp_sn, sgp_sm)); + } + STEEL_PRAGMA_UNROLL + for (short f = 0; f < NF; ++f) { + STEEL_PRAGMA_UNROLL + for (short e = 0; e < NE; ++e) { + Dtile.val_frags[f][e] = Rtile.val_frags[f][e] + + Gtile.val_frags[f][e] * Dtile.val_frags[f][e]; + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < NF; ++f) { + STEEL_PRAGMA_UNROLL + for (short e = 0; e < NE; ++e) { + Dtile.val_frags[f][e] += Rtile.val_frags[f][e]; + } + } + } + + if (full) { + Dtile.store(ys, N); + } else { + Dtile.store_safe(ys, N, short2(sgp_sn, sgp_sm)); + } + }); + }); +} + +// Kernel of one codec. LOADER64 decodes a K tile of 64 and LOADER128 a K +// tile of 128; a codec with one K tile names its loader twice. +#define KQ_NAX_DEFINE_RESID_KERNEL(codec, LOADER64, LOADER128) \ + template < \ + typename T, \ + bool aligned_N, \ + bool gated, \ + int BM, \ + int BN, \ + int BK, \ + int WM, \ + int WN> \ + [[kernel]] void kq_##codec##_qmm_t_nax_resid( \ + const device uint8_t* w, \ + const device T* x, \ + const device T* r, \ + const device T* g, \ + device T* y, \ + const constant int& K, \ + const constant int& N, \ + const constant int& M, \ + const constant int& g_stride, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 tgp_grid [[threadgroups_per_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + tid = kq_swizzle_tid(tid, tgp_grid, N, BN); \ + if (int(tid.y) * BM >= M) { \ + return; \ + } \ + constexpr int BK_padded = (BK + 16 / sizeof(T)); \ + threadgroup T Ws[BN * BK_padded]; \ + if constexpr (BK == 64) { \ + using LoaderW = LOADER64< \ + T, \ + BN, \ + 64, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_resid_impl< \ + T, \ + LoaderW, \ + aligned_N, \ + gated, \ + BM, \ + 64, \ + BN, \ + WM, \ + WN>(w, x, r, g, y, Ws, K, N, M, g_stride, tid, simd_gid, simd_lid); \ + } else { \ + using LoaderW = LOADER128< \ + T, \ + BN, \ + 128, \ + BK_padded, \ + /*reduction_dim=*/1, \ + /*tgp_size=*/WM * WN * SIMD_SIZE>; \ + kq_qmm_t_nax_resid_impl< \ + T, \ + LoaderW, \ + aligned_N, \ + gated, \ + BM, \ + 128, \ + BN, \ + WM, \ + WN>(w, x, r, g, y, Ws, K, N, M, g_stride, tid, simd_gid, simd_lid); \ + } \ + } + +KQ_NAX_DEFINE_RESID_KERNEL(q4_k, KqNaxQ4_KBlockLoader, KqNaxQ4_KHalfLoader) +KQ_NAX_DEFINE_RESID_KERNEL(q5_k, KqNaxQ5_KBlockLoader, KqNaxQ5_KHalfLoader) +KQ_NAX_DEFINE_RESID_KERNEL(q6_k, KqNaxQ6_KBlockLoader, KqNaxQ6_KHalfLoader) +KQ_NAX_DEFINE_RESID_KERNEL(q8_0, KqNaxQ8_0BlockLoader, KqNaxQ8_0BlockLoader) +#undef KQ_NAX_DEFINE_RESID_KERNEL diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_prism.h b/metal/mlx/backend/metal/kernels/kq_quantized_prism.h index 7572c6b2..dd7be87c 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_prism.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_prism.h @@ -136,8 +136,9 @@ METAL_FUNC void kq_pq2_0_decode_pairs(uint wv, thread half2* h) { // M=1 mat-vec: SB / 16 lanes per block (eight for PQ2_0, four for Q2_0), 32 // lanes of blocks per simdgroup pass, 16 contiguous weights per lane decoded -// as half2 pairs. Two half MACs per instruction against the activation pair -// (y_j, y_j+8) keep the decode near the load-only rate. +// as half2 pairs. For float16 activations, two half MACs per instruction +// against the activation pair (y_j, y_j+8) keep the decode near the +// load-only rate. template METAL_FUNC void kq_pq2_0_qmv_sb_impl( const device uint8_t* w, @@ -151,9 +152,9 @@ METAL_FUNC void kq_pq2_0_qmv_sb_impl( constexpr int num_simdgroups = 2; constexpr int lanes = SB / 16; constexpr int blocks_per_pass = 32 / lanes; - // fp32 activations keep an fp32 pair accumulator; half inputs stay in - // the half pipe. - typedef metal::conditional_t, float2, half2> A2; + // float16 activations stay in the half pipe. Other activations keep a + // float pair accumulator: bfloat16 values past 65504 do not fit in half. + typedef metal::conditional_t, half2, float2> A2; typedef float U; const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; @@ -226,7 +227,7 @@ METAL_FUNC void kq_pq2_0_verify_qmv_sb_impl( constexpr int num_simdgroups = 2; constexpr int lanes = SB / 16; constexpr int blocks_per_pass = 32 / lanes; - typedef metal::conditional_t, float2, half2> A2; + typedef metal::conditional_t, half2, float2> A2; typedef float U; const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; @@ -934,7 +935,11 @@ using KqPtq1_0BlockLoader = KqPrismBlockLoader< int group_size, \ int bits, \ bool aligned_N, \ - bool batched> \ + bool batched, \ + int BM = 64, \ + int BN = 64, \ + int WM = 2, \ + int WN = 2> \ [[kernel]] void kq_##CODEC##_qmm_t( \ const device uint8_t* w, \ const device uint8_t* /* scales */, \ @@ -970,12 +975,12 @@ using KqPtq1_0BlockLoader = KqPrismBlockLoader< } \ static_assert(group_size == SB, #CODEC " requires gs=" #SB); \ static_assert(bits == BITS, #CODEC " requires bits=" #BITS); \ - constexpr int BM = 64, BK = 32, BN = 64; \ + constexpr int BK = 32; \ constexpr int BK_padded = (BK + 16 / sizeof(T)); \ threadgroup T Xs[BM * BK_padded]; \ threadgroup T Ws[BN * BK_padded]; \ - using LoaderW = LOADER; \ - kq_qmm_t_impl( \ + using LoaderW = LOADER; \ + kq_qmm_t_impl( \ w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); \ } \ \ @@ -984,12 +989,13 @@ using KqPtq1_0BlockLoader = KqPrismBlockLoader< int group_size, \ int bits, \ bool aligned_N, \ - int small_bm = 0> \ + int small_bm = 0, \ + typename U = T> \ [[kernel]] void kq_##CODEC##_qmm_t_splitk( \ const device uint8_t* w, \ const device uint8_t* /* scales */, \ const device T* x, \ - device T* y, \ + device U* y, \ const constant int& K, \ const constant int& N, \ const constant int& M, \ @@ -1013,7 +1019,7 @@ using KqPtq1_0BlockLoader = KqPrismBlockLoader< auto wl = w; \ wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; \ y += tid.z * static_cast(split_k_partition_stride); \ - kq_qmm_t_impl( \ + kq_qmm_t_impl( \ wl, \ x, \ y, \ @@ -1226,14 +1232,14 @@ struct KqPq2_0Mma { static METAL_FUNC int perm(int f, int col) { return 32 * (col / 2) + 16 * (f / 8) + (f % 8) + 8 * (col & 1); } - template + template static METAL_FUNC void block( thread const KqVmmaRows& rows, int boff, short L, short fm, - const threadgroup half* xb, - thread simdgroup_half8x8 (&acc)[NT]) { + const threadgroup X* xb, + thread simdgroup_matrix (&acc)[NT]) { uint wv[NT][2]; for (short t = 0; t < NT; ++t) { const packed_ushort4 v = @@ -1272,14 +1278,14 @@ struct KqQ2_0Mma { static METAL_FUNC int perm(int f, int col) { return 16 * (col / 2) + f + 8 * (col & 1); } - template + template static METAL_FUNC void block( thread const KqVmmaRows& rows, int boff, short L, short fm, - const threadgroup half* xb, - thread simdgroup_half8x8 (&acc)[NT]) { + const threadgroup X* xb, + thread simdgroup_matrix (&acc)[NT]) { uint wv[NT]; for (short t = 0; t < NT; ++t) { const packed_ushort2 v = @@ -1322,14 +1328,14 @@ struct KqPtq1_0Mma { const int level = f % 5; return byte < 16 ? 16 * level + byte : 80 + 8 * level + (byte - 16); } - template + template static METAL_FUNC void block( thread const KqVmmaRows& rows, int boff, short L, short fm, - const threadgroup half* xb, - thread simdgroup_half8x8 (&acc)[NT]) { + const threadgroup X* xb, + thread simdgroup_matrix (&acc)[NT]) { for (short p = 0; p < 3; ++p) { half2 u[NT]; for (short t = 0; t < NT; ++t) { @@ -1363,6 +1369,13 @@ struct KqPtq1_0Mma { } }; -KQ_DEFINE_VERIFY_MMA_KERNEL(pq2_0, KqPq2_0Mma, 2) +// One row tile per simdgroup (NT 1). A pq2_0 block runs its 16 fragment +// MMAs into one accumulator per row tile, a chain of dependent MMAs whose +// latency, longer in float than in half, bounds the kernel; one row tile +// per simdgroup doubles the threadgroups that hide it. On an M5 Max over +// the Qwen3.8-27B shapes at M 2 to 8 the float kernel ran in 0.45-0.52x +// its time at two row tiles. q2_0 runs the same chain, 8 fragment MMAs per +// 64-weight block, and takes the pq2_0 row tile count; it is not measured. +KQ_DEFINE_VERIFY_MMA_KERNEL(pq2_0, KqPq2_0Mma, 1) KQ_DEFINE_VERIFY_MMA_KERNEL(ptq1_0, KqPtq1_0Mma, 1) -KQ_DEFINE_VERIFY_MMA_KERNEL(q2_0, KqQ2_0Mma, 2) +KQ_DEFINE_VERIFY_MMA_KERNEL(q2_0, KqQ2_0Mma, 1) diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_stq.h b/metal/mlx/backend/metal/kernels/kq_quantized_stq.h index eb61c855..96495d9e 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_stq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_stq.h @@ -359,7 +359,16 @@ struct KqStq1_0BlockLoader { } }; -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool batched, + int BM = 64, + int BN = 64, + int WM = 2, + int WN = 2> [[kernel]] void kq_stq1_0_qmm_t( const device uint8_t* w, const device uint8_t* /* scales */, @@ -395,13 +404,13 @@ template } static_assert(group_size == KQ_STQ1_0_SUPERBLOCK, "STQ1_0 requires gs=256"); static_assert(bits == 1, "STQ1_0 requires bits=1"); - constexpr int BM = 64, BK = 32, BN = 64; + constexpr int BK = 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; using LoaderW = - KqStq1_0BlockLoader; - kq_qmm_t_impl( + KqStq1_0BlockLoader; + kq_qmm_t_impl( w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } @@ -410,12 +419,13 @@ template < int group_size, int bits, bool aligned_N, - int small_bm = 0> + int small_bm = 0, + typename U = T> [[kernel]] void kq_stq1_0_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -440,7 +450,7 @@ template < auto wl = w; wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; y += tid.z * static_cast(split_k_partition_stride); - kq_qmm_t_impl( + kq_qmm_t_impl( wl, x, y, diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h b/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h index 5c402e94..b7a5dc64 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h @@ -30,6 +30,7 @@ using namespace metal; using namespace mlx::steel; constant int gqa_splits [[function_constant(2)]]; +constant bool gqa_has_starts [[function_constant(4)]]; namespace kq_sdpa_nax { @@ -338,6 +339,11 @@ kq_sdpa_fa_indexed_nax_2pass_1( // threadgroup. Inside a fragment the 16-wide contraction index is a // permutation of the head dim, and Q uses the same one: lane quad g reads // dims 16g..16g+15 of each 64-dim block, 32 contiguous bytes per row. +// +// With starts (a left-padded batched KV cache), batch row b attends keys +// [starts[b], N). Whole tiles below the start are skipped, keys below it +// in the first tile are masked, and K and V rows are read from the start +// up, so the pad rows are never read. template [[kernel, max_total_threads_per_threadgroup(256)]] void kq_sdpa_gqa_nax_2pass_1( const device T* queries [[buffer(0)]], @@ -355,6 +361,7 @@ template const constant int& gqa [[buffer(12)]], const constant size_t& k_batch_stride [[buffer(13)]], const constant size_t& v_batch_stride [[buffer(14)]], + const device int* starts [[buffer(15)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -391,12 +398,15 @@ template const int chunk = ((N + gqa_splits * BK - 1) / (gqa_splits * BK)) * BK; const int k0 = split_idx * chunk; const int k1 = min(k0 + chunk, N); + // First key of this batch row, and the first tile that holds one. + const int row_start = gqa_has_starts ? max(0, starts[tid.y]) : 0; + const int kt0 = max(k0, (row_start / BK) * BK); // Partials row for head fm of this group ([B, Hq, 1, splits, D]). const size_t po = (hb * gqa + fm) * gqa_splits + split_idx; device float* dst = out + po * D + sg * DC + goff; - if (k0 >= k1) { + if (max(k0, row_start) >= k1) { // Empty split: pass 2 folds it at zero weight. if (fm < gqa) { STEEL_PRAGMA_UNROLL @@ -433,17 +443,21 @@ template float m_run = Limits::finite_min; // row fm, log2 domain float l_run = 0; // this lane's share of row fm's sum - for (int kt = k0; kt < k1; kt += BK) { + // The first tile holds a key at or above row_start, so the running max is + // finite from the first tile on. + for (int kt = kt0; kt < k1; kt += BK) { // S (heads x 32 keys) for this simdgroup's keys. Lane rows are keys - // kb + fm, + 8, + 16, + 24; rows past N read row N - 1 and mask. + // kb + fm, + 8, + 16, + 24; rows past N read row N - 1, rows below + // row_start read row row_start, and both mask. const int kb = kt + sg * KS; ffrag_t S0 = ffrag_t(0); ffrag_t S1 = ffrag_t(0); - if (kb < k1) { + if (kb < k1 && kb + KS > row_start) { const device T* kr[4]; STEEL_PRAGMA_UNROLL for (short r = 0; r < 4; r++) { - kr[r] = kbase + (size_t)min(kb + fm + 8 * r, N - 1) * k_seq_stride; + kr[r] = kbase + + (size_t)min(max(kb + fm + 8 * r, row_start), N - 1) * k_seq_stride; } STEEL_PRAGMA_UNROLL for (short blk = 0; blk < D / 64; blk++) { @@ -481,9 +495,12 @@ template float tmax = Limits::finite_min; STEEL_PRAGMA_UNROLL for (short j = 0; j < 4; j++) { - sv[j] = kb + fn + j < k1 ? S0[j] * scale2 : Limits::finite_min; - sv[4 + j] = - kb + 16 + fn + j < k1 ? S1[j] * scale2 : Limits::finite_min; + const int ka = kb + fn + j; + const int kc = ka + 16; + sv[j] = ka < k1 && ka >= row_start ? S0[j] * scale2 + : Limits::finite_min; + sv[4 + j] = kc < k1 && kc >= row_start ? S1[j] * scale2 + : Limits::finite_min; tmax = max(tmax, max(sv[j], sv[4 + j])); } tmax = max(tmax, simd_shuffle_xor(tmax, ushort(1))); @@ -522,10 +539,12 @@ template // rows are keys kt + 16 kc + fm and + 8; masked keys carry P = 0. STEEL_PRAGMA_UNROLL for (short kc = 0; kc < BK / 16; kc++) { - if (kt + kc * 16 < k1) { + if (kt + kc * 16 < k1 && kt + kc * 16 + 16 > row_start) { const int r0 = kt + kc * 16 + fm; - const device T* vr0 = vbase + (size_t)min(r0, N - 1) * v_seq_stride; - const device T* vr1 = vbase + (size_t)min(r0 + 8, N - 1) * v_seq_stride; + const device T* vr0 = + vbase + (size_t)min(max(r0, row_start), N - 1) * v_seq_stride; + const device T* vr1 = + vbase + (size_t)min(max(r0 + 8, row_start), N - 1) * v_seq_stride; T4 v0[4], v1[4]; STEEL_PRAGMA_UNROLL for (short s = 0; s < 4; s++) { diff --git a/metal/mlx/backend/metal/kernels/kq_verify_mma.h b/metal/mlx/backend/metal/kernels/kq_verify_mma.h index 40734064..f10c21a1 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_mma.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_mma.h @@ -7,10 +7,14 @@ // lane-owned bytes, the activations are staged once per K chunk as B^T // (rows k, cols m) in threadgroup memory in natural k order, and one // simdgroup_multiply_accumulate per fragment accumulates D (rows n, cols -// m). The per-block accumulator is half; the block scale is applied to it -// in float. A codec whose code magnitudes could carry a block's sum out of -// half range decodes scaled-down codes and restores the factor through -// d_scale. Split-K over grid z writes T partials that +// m). The codes decode exactly in half. The activations stay in their own +// dtype in threadgroup memory, both operands widen to float in the +// fragment, and each block sums in float; the block scale is applied to it +// in float. The KQ_VERIFY_MMA_HALF kernels (HALF, name suffix _h) stage +// the activations as half and sum each block in half, whose range stops +// at 65504; for them, a codec whose code magnitudes could carry a block's +// sum out of half range decodes scaled-down codes and restores the factor +// through d_scale. Split-K over grid z writes U partials that // kquant_qmm_splitk_accum folds. // // The k order inside a block is a fixed per-codec permutation @@ -24,9 +28,9 @@ // block_k, block_bytes, d_offset, d_scale (factor on the block scale) // perm(f, col): element index of fragment f (0..block_k/8-1), column col // block(rows, boff, L, fm, xb, acc): decode block boff of each row -// and accumulate its block_k/8 fragments; L = fn / 2 is the lane's -// column pair, xb the staged B base of this block (fragment f at -// xb + 8 * perm(f, fm)). +// and accumulate its block_k/8 fragments through kq_vmma_step; L = +// fn / 2 is the lane's column pair, xb the staged B base of this block +// (fragment f at xb + 8 * perm(f, fm)). MLX_MTL_CONST int KQ_VMMA_THREADS = 256; MLX_MTL_CONST int KQ_VMMA_NSG = KQ_VMMA_THREADS / 32; @@ -37,35 +41,49 @@ struct KqVmmaRows { const device uint8_t* p[NT]; }; -template +// Staged activation type X and accumulator type A: T and float, or half +// and half for the KQ_VERIFY_MMA_HALF kernels. +template +struct KqVmmaTypes { + using X = metal::conditional_t; + using A = metal::conditional_t; +}; + +// One fragment step: the lane's two staged activations at bsrc and its +// decoded code pair a[t] of each row tile, widened to A. +template METAL_FUNC void kq_vmma_step( - const threadgroup half* bsrc, + const threadgroup X* bsrc, thread const half2 (&a)[NT], - thread simdgroup_half8x8 (&acc)[NT]) { - simdgroup_half8x8 B; - reinterpret_cast(B.thread_elements()) = - *(const threadgroup half2*)bsrc; + thread simdgroup_matrix (&acc)[NT]) { + simdgroup_matrix B; + const vec b = *(const threadgroup vec*)bsrc; + reinterpret_cast&>(B.thread_elements()) = + vec(static_cast(b.x), static_cast(b.y)); for (short t = 0; t < NT; ++t) { - simdgroup_half8x8 A; - reinterpret_cast(A.thread_elements()) = a[t]; - simdgroup_multiply_accumulate(acc[t], A, B, acc[t]); + simdgroup_matrix W; + reinterpret_cast&>(W.thread_elements()) = + vec(static_cast(a[t].x), static_cast(a[t].y)); + simdgroup_multiply_accumulate(acc[t], W, B, acc[t]); } } -template +template METAL_FUNC void kq_verify_mma_impl( const device uint8_t* w, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, const constant int& k_partition_size, - threadgroup half* XsT, + threadgroup typename KqVmmaTypes::X* XsT, uint3 tid, uint lid, uint simd_gid, uint simd_lid) { + using X = typename KqVmmaTypes::X; + using A = typename KqVmmaTypes::A; constexpr int KC = KQ_VMMA_KC; constexpr int BK = Codec::block_k; const int row_bytes = (K / BK) * Codec::block_bytes; @@ -89,26 +107,31 @@ METAL_FUNC void kq_verify_mma_impl( threadgroup_barrier(mem_flags::mem_threadgroup); for (int kk = lid; kk < kn; kk += KQ_VMMA_THREADS) { const device T* xe = x + kc + kk; - vec v; + vec v; for (short m = 0; m < 8; ++m) { - v[m] = m < M ? half(float(xe[m * K])) : half(0); + if constexpr (HALF) { + v[m] = m < M ? half(float(xe[m * K])) : half(0); + } else { + v[m] = m < M ? xe[m * K] : T(0); + } } - *(threadgroup vec*)(XsT + kk * 8) = v; + *(threadgroup vec*)(XsT + kk * 8) = v; } threadgroup_barrier(mem_flags::mem_threadgroup); for (int kb = 0; kb < kn; kb += BK) { const int boff = ((kc + kb) / BK) * Codec::block_bytes; - const threadgroup half* xb = XsT + kb * 8 + fn; - simdgroup_half8x8 acc[NT]; + const threadgroup X* xb = XsT + kb * 8 + fn; + simdgroup_matrix acc[NT]; for (short t = 0; t < NT; ++t) { - acc[t] = make_filled_simdgroup_matrix(half(0)); + acc[t] = make_filled_simdgroup_matrix(A(0)); } Codec::template block(rows, boff, L, fm, xb, acc); for (short t = 0; t < NT; ++t) { const float d = Codec::d_scale * float(*(const device half*)(rows.p[t] + boff + Codec::d_offset)); c[t] += d * - float2(reinterpret_cast(acc[t].thread_elements())); + float2(reinterpret_cast&>( + acc[t].thread_elements())); } } } @@ -116,10 +139,10 @@ METAL_FUNC void kq_verify_mma_impl( const int nt = n + 8 * t; if (nt < N) { if (fn < M) { - y[fn * N + nt] = static_cast(c[t].x); + y[fn * N + nt] = static_cast(c[t].x); } if (fn + 1 < M) { - y[(fn + 1) * N + nt] = static_cast(c[t].y); + y[(fn + 1) * N + nt] = static_cast(c[t].y); } } } @@ -128,12 +151,17 @@ METAL_FUNC void kq_verify_mma_impl( // Entry point: qmm_t_splitk's buffer layout, grid (ceil(N / (64 NT)), 1, // splits), 256 threads. #define KQ_DEFINE_VERIFY_MMA_KERNEL(CODEC, TRAITS, NT) \ - template \ + template < \ + typename T, \ + int group_size, \ + int bits, \ + typename U = T, \ + bool HALF = false> \ [[kernel]] void kq_##CODEC##_verify_mma( \ const device uint8_t* w, \ const device uint8_t* /* scales */, \ const device T* x, \ - device T* y, \ + device U* y, \ const constant int& K, \ const constant int& N, \ const constant int& M, \ @@ -144,8 +172,8 @@ METAL_FUNC void kq_verify_mma_impl( uint simd_gid [[simdgroup_index_in_threadgroup]], \ uint simd_lid [[thread_index_in_simdgroup]]) { \ static_assert(group_size == TRAITS::block_k, #CODEC " block width"); \ - threadgroup half XsT[KQ_VMMA_KC * 8]; \ - kq_verify_mma_impl( \ + threadgroup typename KqVmmaTypes::X XsT[KQ_VMMA_KC * 8]; \ + kq_verify_mma_impl( \ w, \ x, \ y + tid.z * static_cast(split_k_partition_stride), \ diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index 663a7dcd..dc32757d 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -14,7 +14,8 @@ // 8..15 as padding, loaded from device memory per unit in the codec's k // order. Nothing is staged in threadgroup memory and there are no // barriers. Each simdgroup keeps one f32 accumulator. Split-K over grid z -// writes T partials that kquant_qmm_splitk_accum folds. +// writes U partials, T or float as in kq_qmm_t_impl, that +// kquant_qmm_splitk_accum folds. // // A unit is one or more wire blocks, or half of one, and step s of a unit // is its s-th 16-wide NAX op. For pq2_0 and q4_0, step s consumes the @@ -794,11 +795,11 @@ struct KqIq4NlNax { } }; -template +template METAL_FUNC void kq_verify_nax_impl( const device uint8_t* w, const device T* x, - device T* y, + device U* y, const constant int& K, const constant int& N, const constant int& M, @@ -906,7 +907,7 @@ METAL_FUNC void kq_verify_nax_impl( for (short j = 0; j < 4; ++j) { const int n = n0 + 16 * f + fn + j; if (n < N) { - y[static_cast(fm) * N + n] = static_cast(cc[f * 8 + j]); + y[static_cast(fm) * N + n] = static_cast(cc[f * 8 + j]); } } } @@ -915,12 +916,12 @@ METAL_FUNC void kq_verify_nax_impl( // Entry point: qmm_t_splitk's buffer layout, grid (ceil(N / 128), 1, // splits), 128 threads. #define KQ_DEFINE_VERIFY_NAX_KERNEL(CODEC, TRAITS) \ - template \ + template \ [[kernel]] void kq_##CODEC##_verify_nax( \ const device uint8_t* w, \ const device uint8_t* /* scales */, \ const device T* x, \ - device T* y, \ + device U* y, \ const constant int& K, \ const constant int& N, \ const constant int& M, \ @@ -930,7 +931,7 @@ METAL_FUNC void kq_verify_nax_impl( uint simd_gid [[simdgroup_index_in_threadgroup]], \ uint simd_lid [[thread_index_in_simdgroup]]) { \ static_assert(group_size == TRAITS::group, #CODEC " block width"); \ - kq_verify_nax_impl( \ + kq_verify_nax_impl( \ w, \ x, \ y + tid.z * static_cast(split_k_partition_stride), \ diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index 6dd5dcd4..c757ce51 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -5,7 +5,7 @@ Current API: * ``codecs`` / ``metallib_dir`` / ``metallib_loads`` / - ``cpu_neon_available`` - toolchain self-checks. + ``metallib_kernels`` / ``cpu_neon_available`` - toolchain self-checks. * ``dequantize`` - GGUF K-quant wire bytes -> float array. * ``quantized_matmul`` - x @ dequant(w) for K-quant weights. * ``quantized_matmul_qmv_bias`` - decode-only (M=1) bias-fused variant of @@ -28,6 +28,7 @@ codec_has_matmul, codec_has_moe_glu, codecs, + conv2d_3x3, cpu_neon_available, dequantize, dsa_indexer_qat, @@ -70,7 +71,9 @@ latent_fp4_pack, latent_fp4_unpack, load_gguf, + max_threadgroup_memory, metallib_dir, + metallib_kernels, metallib_loads, moe_glu_gather, moe_glu_gather_kq, @@ -78,10 +81,14 @@ moe_router_topk, nax_available, nax_gather_enabled, + qk_norm_rope, qmv_fine_max_n, quantize, quantized_matmul, + quantized_matmul_glu, + quantized_matmul_lora, quantized_matmul_qmv_bias, + quantized_matmul_residual, residency_commit, residency_erase, residency_insert, @@ -132,6 +139,7 @@ "codec_has_matmul", "codec_has_moe_glu", "codecs", + "conv2d_3x3", "cpu_neon_available", "dequantize", "dsa_indexer_qat", @@ -172,18 +180,24 @@ "latent_fp4_unpack", "load_gguf", "metallib_dir", + "metallib_kernels", "metallib_loads", "moe_glu_gather", "moe_glu_gather_kq", "moe_glu_gather_shexp_kq", "moe_router_topk", "nax_available", + "qk_norm_rope", "qmv_fine_max_n", "gpu_core_count", + "max_threadgroup_memory", "nax_gather_enabled", "quantize", "quantized_matmul", + "quantized_matmul_glu", + "quantized_matmul_lora", "quantized_matmul_qmv_bias", + "quantized_matmul_residual", "rmsnorm2_add", "rmsnorm_gate", "rmsnorm_multi3", diff --git a/mlx_kquant/_deps.py b/mlx_kquant/_deps.py index 6c2c26f6..32b1621e 100644 --- a/mlx_kquant/_deps.py +++ b/mlx_kquant/_deps.py @@ -20,7 +20,7 @@ def require_tools() -> None: or broken. A present-but-broken mlx-lm raises whatever its import died on (e.g. - transformers 5.13 rejects mlx-lm <= 0.31.3's tokenizer registration with + transformers 5.13.0 rejects mlx-lm <= 0.31.3's tokenizer registration with an ``AttributeError``); surface that as a version-conflict hint instead of a raw traceback from a dependency's internals. """ @@ -32,6 +32,6 @@ def require_tools() -> None: raise ImportError( "mlx-kquant: mlx-lm is installed but failed to import, which " "usually means an incompatible dependency version (for example " - "transformers 5.13 breaks mlx-lm <= 0.31.3; fix with: pip install " - f"'transformers<5.13'). Underlying error: {e!r}" + "transformers 5.13.0 breaks mlx-lm <= 0.31.3; fix with: pip install " + f"'transformers!=5.13.0'). Underlying error: {e!r}" ) from e diff --git a/mlx_kquant/_version.py b/mlx_kquant/_version.py index 2dd55360..3d187266 100644 --- a/mlx_kquant/_version.py +++ b/mlx_kquant/_version.py @@ -1 +1 @@ -__version__ = "0.4.18" +__version__ = "0.5.0" diff --git a/mlx_kquant/loader.py b/mlx_kquant/loader.py index 5083485d..8aea08a5 100644 --- a/mlx_kquant/loader.py +++ b/mlx_kquant/loader.py @@ -23,8 +23,9 @@ import importlib import importlib.util import json +from collections.abc import Callable from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import mlx.core as mx @@ -116,6 +117,7 @@ def load( trust_remote_code: bool = False, model_config: dict | None = None, revision: str | None = None, + get_model_classes: Callable[..., tuple[type, Any]] | None = None, ) -> tuple[nn.Module, dict]: """Load a kquant checkpoint and return ``(model, config)``. @@ -128,6 +130,11 @@ def load( imported and executed (arbitrary code). Off by default. model_config: optional overrides merged into the loaded config. revision: optional HF revision (branch / tag / commit). + get_model_classes: optional ``get_model_classes(config=config)`` that + returns ``(Model, ModelArgs)``, as mlx-lm's ``load_model`` takes; + by default the mlx-lm class for ``config["model_type"]``. A + ``model_file`` the checkpoint declares takes precedence, as in + mlx-lm. Returns: ``(model, config)`` with ``KQuant*`` modules installed and weights loaded. @@ -182,6 +189,8 @@ def load( arch = importlib.util.module_from_spec(spec) spec.loader.exec_module(arch) model_class, args_class = arch.Model, arch.ModelArgs + elif get_model_classes is not None: + model_class, args_class = get_model_classes(config=config) else: model_class, args_class = _get_classes(config) diff --git a/mlx_kquant/mlx_lm_patch.py b/mlx_kquant/mlx_lm_patch.py index c2348a7a..596def69 100644 --- a/mlx_kquant/mlx_lm_patch.py +++ b/mlx_kquant/mlx_lm_patch.py @@ -109,6 +109,7 @@ def _patched_load_model( strict=strict, model_config=model_config, trust_remote_code=_trust_remote_code, + get_model_classes=get_model_classes, ) kwargs = {} diff --git a/pyproject.toml b/pyproject.toml index bd351535..93492c79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,10 +2,10 @@ requires = [ "setuptools>=77", "cmake>=3.27", - "mlx==0.32.1", + "mlx==0.32.3", # Must match mlx's nanobind ABI exactly; CMakeLists.txt (KQ_NANOBIND_PIN) # enforces it at configure time for non-isolated builds - keep in sync. - "nanobind==2.13.0", + "nanobind==3.0.1", ] build-backend = "setuptools.build_meta" @@ -41,15 +41,16 @@ classifiers = [ "Environment :: GPU", "Topic :: Scientific/Engineering :: Artificial Intelligence", ] -dependencies = ["mlx==0.32.1"] +dependencies = ["mlx==0.32.3"] # Model-level helpers (the create-weights toolchain, loader, and CLI) need the # mlx-lm model zoo and friends; the raw ``kq.*`` ops do not. Keep that weight # behind an extra so ``import mlx_kquant`` stays a thin op layer. [project.optional-dependencies] -# transformers<5.13: 5.13.0 rejects mlx-lm 0.31.3's string-key -# AutoTokenizer.register at import; lift when a fixed mlx-lm ships. -tools = ["mlx-lm>=0.27", "transformers<5.13", "gguf", "numpy"] +# transformers!=5.13.0: that release rejects the string key mlx-lm 0.31.3 +# passes to AutoTokenizer.register, so importing mlx-lm fails. 5.12.1 and +# 5.13.1 through 5.17.0 import it cleanly. +tools = ["mlx-lm>=0.27", "transformers!=5.13.0", "gguf", "numpy"] dev = ["pytest", "ruff", "numpy", "gguf"] [project.scripts] @@ -96,11 +97,11 @@ ignore = ["E741"] # ambiguous var names (l/I) appear in math-y kernel glue # --------------------------------------------------------------------------- # Wheel building (cibuildwheel). macOS arm64 only - the extension is a Metal # module. Build isolation pulls [build-system].requires (incl. the ABI-pinned -# mlx==0.32.1 + nanobind + the Metal toolchain from the runner's Xcode). +# mlx==0.32.3 + nanobind + the Metal toolchain from the runner's Xcode). # # NOTE: this config is a skeleton - it has NOT yet produced a wheel validated on # a second machine. Before the first PyPI publish, build locally, then -# `pip install` the wheel into a CLEAN venv with only mlx==0.32.1 and confirm +# `pip install` the wheel into a CLEAN venv with only mlx==0.32.3 and confirm # `python -c "import mlx_kquant as kq; kq.metallib_loads()"` on a GPU box. # --------------------------------------------------------------------------- [tool.cibuildwheel] @@ -120,7 +121,7 @@ archs = ["arm64"] # https://github.com/ml-explore/mlx/issues/3586 environment = { MACOSX_DEPLOYMENT_TARGET = "26.2" } # Do NOT vendor libmlx into the wheel: the extension is ABI-pinned to the user's -# installed mlx==0.32.1 and resolves it via the @loader_path/../mlx/lib rpath +# installed mlx==0.32.3 and resolves it via the @loader_path/../mlx/lib rpath # (see CMakeLists.txt). Bundling a second libmlx copy would split MLX global # state and resurrect the gguflib<->libmlx symbol collision. --exclude keeps it # out of the wheel; --ignore-missing-dependencies stops delocate from erroring diff --git a/setup.py b/setup.py index bd02599f..7648e4d1 100644 --- a/setup.py +++ b/setup.py @@ -16,6 +16,14 @@ f"{os.environ.get('CMAKE_ARGS', '')} -DPython_EXECUTABLE={sys.executable}".strip() ) +# KQ_PROBE_KERNELS is a cached CMake option, so every build names its value: +# without it a build after a probe build keeps the probe kernels. The +# environment variable KQ_PROBE_KERNELS=1 asks for a probe build; a +# -DKQ_PROBE_KERNELS in CMAKE_ARGS is left alone. +if "-DKQ_PROBE_KERNELS" not in os.environ["CMAKE_ARGS"]: + probe = os.environ.get("KQ_PROBE_KERNELS", "").lower() in ("1", "on", "true", "yes") + os.environ["CMAKE_ARGS"] += f" -DKQ_PROBE_KERNELS={'ON' if probe else 'OFF'}" + # Force the Metal deployment target to 26.2 on macOS. # https://github.com/ml-explore/mlx/issues/3586 if sys.platform == "darwin": diff --git a/src/kquant.h b/src/kquant.h index 95bf3d81..27d53f04 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -27,6 +27,14 @@ std::string metallib_dir(); // Returns false only on a non-Metal build. bool metallib_loads(); +// Sorted names of the kernels in the bundled metallib (empty on non-Metal +// builds). +std::vector metallib_kernels(); + +// True in a build with the CMake option KQ_PROBE_KERNELS=ON: the metallib +// carries the probe tiles and the levers that select them are read. +bool probe_kernels(); + // True when the GPU supports the NAX (tensor-core) matmul kernels (false on // non-Metal builds and pre-NAX GPUs). bool nax_available(); @@ -35,11 +43,52 @@ bool nax_available(); // it); 0 when unknown. int gpu_core_count(); +// The GPU's threadgroup memory limit in bytes; 0 on non-Metal builds. +int max_threadgroup_memory(); + // Largest N at which the M=1 mat-vec takes the fine (2 rows per threadgroup) // tiling for this codec; 0 means coarse. The calibrated ceiling unless // KQ_GPU_CORES is set, which scales it by cores/40. On Metal only. int qmv_fine_max_n(const std::string& kquant_type); +// The NAX qmm_t large-M row of a codec: bm, bn, wm, wn, swizzle_log, min_m, +// min_n, max_k and inherited (1 when the tile and floor are its loader +// class's); empty on non-Metal builds. +std::vector nax_large_row(const std::string& kquant_type); + +// The row that runs a plain transposed quantized_matmul as dequantize then +// the stock matmul by default, for an architecture string (as +// get_architecture gives it), NAX, a codec, the rows of x, N and K: min_m, +// min_k, max_k, min_n and max_n. Empty where the qmm_t kernel runs: on NAX +// GPUs, outside applegpu_g13 to g16 with a suffix other than p, for a codec +// without a row, below its floor in M and outside the K and N range. +std::vector alu_dequantize_row( + const std::string& arch, + bool nax, + const std::string& kquant_type, + int M, + int N, + int K); + +// Whether a plain transposed quantized_matmul on a 2-D weight on a GPU +// stream takes the dequantize route, for an architecture string, NAX, the +// KQ_DEQUANT_ROUTE mode ("auto", "off" or "force"), whether KQ_QMM_TILE +// pins the tile, a codec, the rows of x, N and K. The op decides with this. +bool dequantize_route_takes( + const std::string& arch, + bool nax, + const std::string& mode, + bool pinned, + const std::string& kquant_type, + int M, + int N, + int K); + +// Every kernel name the split-K dispatchers can build (qmm_splitk, +// qmm_nax_splitk, verify_mma, verify_nax and their fold), for the kernel +// name tests; empty on non-Metal builds. +std::vector split_kernel_names(); + // True when gather_qmm's sorted-rhs NAX GEMM leaf can serve `kquant_type` // here: NAX hardware present, the codec ships NAX kernels, and KQ_DISABLE_NAX // is unset (read live). Callers with their own sorted-prefill arms (e.g. the @@ -1071,6 +1120,43 @@ class KQuantDequantize : public mx::Primitive { int bits_; }; +// Lays out the last two dims of a weight operand as dense rows, at eval. +// An op builder cannot do this from strides: an array without memory +// reports a dense layout whatever view it becomes. An input whose last two +// dims are dense rows passes through with its buffer, strides and offset. +// Any other input is copied row-major. With `whole`, an input passes +// through only when all of its dims are row-major. Inference-only. +class KQuantMatrixDense : public mx::Primitive { + public: + explicit KQuantMatrixDense(mx::Stream stream, bool whole = false) + : mx::Primitive(stream), whole_(whole) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override { + return {inputs[0].shape()}; + } + + const char* name() const override { + return "KQuantMatrixDense"; + } + bool is_equivalent(const mx::Primitive& other) const override { + return whole_ == static_cast(other).whole_; + } + + private: + // True when the input passes through as it is. + bool passes(const mx::array& in) const; + + bool whole_; +}; + // Quantized matmul. vjp implements only the gradient wrt x (re-dispatch with // the transpose flipped) so LoRA can train on a frozen kquant base; jvp/vmap // and the gradient wrt the quantized weights/scales inherit the base-class @@ -1129,6 +1215,56 @@ class KQuantMatmul : public mx::Primitive { std::vector& outputs); }; +// The output of a matmul that runs as dequantize then the stock matmul, +// with the vjp of KQuantMatmul. inputs = {y, x, w, scales}, where y is the +// product under stop_gradient; the output shares y's buffer. The gradient +// wrt x runs the quantized matmul the other way, so no backward pass keeps +// the dense copy, and the gradient wrt w or scales throws, as on the kernel. +// jvp and vmap throw the base-class text with the name KQuantMatmul. +class KQuantDequantRoute : public mx::Primitive { + public: + explicit KQuantDequantRoute(mx::Stream stream, std::string kquant_type) + : mx::Primitive(stream), kquant_type_(std::move(kquant_type)) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override { + return {inputs[0].shape()}; + } + + const char* name() const override { + return "KQuantDequantRoute"; + } + bool is_equivalent(const mx::Primitive& other) const override { + return kquant_type_ == + static_cast(other).kquant_type_; + } + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + // jvp and vmap refuse by the name KQuantMatmul, as the kernel path does, + // so the error text does not depend on which path a call takes. + std::vector jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) override; + std::pair, std::vector> vmap( + const std::vector& inputs, + const std::vector& axes) override; + + private: + std::string kquant_type_; +}; + // Bias-fused decode-only (M=1) qmv/qmv_fast dispatch -- see // quantized_matmul_qmv_bias() above for the shape contract eval_gpu enforces. // Inference-only, like the SDPA primitive below: no vjp override, so a grad @@ -1462,6 +1598,310 @@ class KQuantRMSNormGate : public mx::Primitive { float eps_; }; +// Fused per-head RMSNorm and RoPE of the DiT attention q and k: +// y = rms_norm(x, w, eps), then each interleaved pair (y[2i], y[2i + 1]) +// rotates by the angle whose cos and sin are cos[l, i] and sin[l, i]. +// q and k [B, L, H, D] (float16/bfloat16, any strides: a slice of a fused +// qkv output is read in place), q_w and k_w [D], cos and sin [L, D / 2] +// (any float dtype, used in f32; one angle per row, so any position scheme +// works). Returns (q_out, k_out) as row-contiguous [B, H, L, D] in the +// input dtype. f32 math, one round at the write. On the GPU head dims 64, +// 96, 128 and 256 run one fused dispatch; other even head dims and CPU +// streams run the same math as stock ops. +std::pair qk_norm_rope( + mx::array q, + mx::array k, + mx::array q_w, + mx::array k_w, + mx::array cos, + mx::array sin, + float eps, + mx::StreamOrDevice s = {}); + +// Fused qk norm and RoPE primitive (see qk_norm_rope). Inference-only. +class KQuantQkNormRope : public mx::Primitive { + public: + KQuantQkNormRope(mx::Stream stream, float eps) + : mx::Primitive(stream), eps_(eps) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override; + + const char* name() const override { + return "KQuantQkNormRope"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + private: + float eps_; +}; + +// GLU product of two quantized matmuls that share x: +// y = act(x @ Wg^T) * (x @ Wu^T) +// x [..., K] (float16/bfloat16; float32 runs in bfloat16), w_gate and w_up +// 2-D uint8 wire bytes of one codec and one shape [N, K bytes], which may +// be row slices of one wider weight. activation is "silu" or "gelu_tanh". +// Returns [..., N] in the dtype quantized_matmul gives for x. On a NAX GPU, +// for codecs with a GLU kernel (q8_0 on bfloat16 only), K a multiple of the +// codec's K tile (64 for q8_0, 128 for the others) and at least 1024 rows, +// one dispatch walks K for the gate weight and stores the activation of +// that sum in y, then walks K for the up weight and multiplies the stored +// activation into that sum at the final store. The sums, the activation and +// the product run in f32, with one rounding of the activation to the dtype +// of y and one of the product. A float16 store keeps a finite value in the +// finite half range, and a NaN or an infinity passes. Every other case +// builds the product from quantized_matmul and stock ops. +mx::array quantized_matmul_glu( + mx::array x, + mx::array w_gate, + mx::array w_up, + const std::string& kquant_type, + const std::string& activation = "silu", + mx::StreamOrDevice s = {}); + +// Quantized matmul with an epilogue fused into the store (see +// quantized_matmul_glu). vjp and jvp are those of the unfused ops, wrt x +// only: the quantized weights are frozen, as in KQuantMatmul. +class KQuantMatmulFused : public mx::Primitive { + public: + KQuantMatmulFused(mx::Stream stream, std::string kquant_type, int act) + : mx::Primitive(stream), + kquant_type_(std::move(kquant_type)), + act_(act) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override; + + const char* name() const override { + return "KQuantMatmulFused"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + std::vector jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) override; + + private: + std::string kquant_type_; + int act_; // 0 = silu, 1 = gelu (tanh form) +}; + +// Quantized matmul with a low-rank term: +// y = x @ W^T + (strength * (x @ A)) @ B +// x [..., K] (float16/bfloat16; float32 runs in bfloat16), w 2-D uint8 wire +// bytes [N, K bytes], lora_a [K, R] and lora_b [R, N] in any float dtype +// (they run in the dtype of y). strength scales the R columns of x @ A: a +// 0-D array, [1] or [R]. bands holds rows of (col_start, col_end, z_start, +// rank): the columns [col_start, col_end) of y take the ranks [z_start, +// z_start + rank) and that block of lora_b, and a column outside every band +// takes no low-rank term. Bands are in column order and disjoint. Without +// bands (nullopt) every column takes every rank. With a table of no row no +// column takes one, and y is the product alone. Several adapters are one +// call: their A columns and B rows are joined, with one strength per +// column. K, N and R are below 2^25 and x holds fewer than 2^31 rows. +// Returns [..., N] in the dtype quantized_matmul gives for x. On a NAX GPU, +// for q4_k, q5_k, q6_k and q8_0 weights on bfloat16 with K up to 4096 and N +// of at least K, at least 4096 rows and at most 128 ranks per column (from +// 1024 rows and at every rank where N is at least 8192), every band +// starting on a multiple of 64 columns and ending on a multiple of 32 +// columns or at N, at most 16 bands and a rank of at most 256 per band +// (without bands, a rank of at most 256), stock ops compute +// z = strength * (x @ A) and one dispatch walks the rank for z and B and +// then K for the weight in the same f32 sums, with y as its only output +// buffer and one rounding at the store. Every other case builds the sum +// from quantized_matmul and stock ops. +mx::array quantized_matmul_lora( + mx::array x, + mx::array w, + const std::string& kquant_type, + mx::array lora_a, + mx::array lora_b, + const std::optional& strength = std::nullopt, + const std::optional>& bands = std::nullopt, + mx::StreamOrDevice s = {}); + +// Quantized matmul with a low-rank term in the reduction (see +// quantized_matmul_lora). Inputs are x, w, z and B. vjp, jvp and vmap are +// those of the unfused ops. The vjp is wrt x, z and B: the quantized weight +// is frozen, as in KQuantMatmul. jvp and vmap run for z and B and refuse x +// and the weight by the name of KQuantMatmul, which has neither. +class KQuantMatmulLora : public mx::Primitive { + public: + KQuantMatmulLora( + mx::Stream stream, + std::string kquant_type, + std::vector bands) + : mx::Primitive(stream), + kquant_type_(std::move(kquant_type)), + bands_(std::move(bands)) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override; + + const char* name() const override { + return "KQuantMatmulLora"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + std::vector jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) override; + std::pair, std::vector> vmap( + const std::vector& inputs, + const std::vector& axes) override; + + private: + std::string kquant_type_; + // Rows of (col_start, col_end, z_start, rank); empty for one band over + // every column and every rank. + std::vector bands_; +}; + +// 3x3 convolution, stride 1, padding 1, channels last: +// y = conv(x, w) + bias +// x [B, H, W, C] in a float dtype, w [O, 3, 3, C] and bias [O] in any float +// dtype (they run in the dtype of x). Returns [B, H, W, O] in the dtype of +// x. On a NAX GPU, for bfloat16 with C a multiple of 16, from 1024 output +// pixels, one dispatch computes y as an implicit GEMM with f32 sums and +// one rounding. Every other case builds MLX's convolution on the flipped +// kernel (its implicit GEMM, no Winograd transform) and a stock add. +mx::array conv2d_3x3( + mx::array x, + mx::array w, + const std::optional& bias = std::nullopt, + mx::StreamOrDevice s = {}); + +// 3x3 convolution as an implicit GEMM (see conv2d_3x3). Inputs are x, w +// and, with a bias, the bias. vjp, jvp and vmap are those of the composed +// convolution. +class KQuantConv3x3 : public mx::Primitive { + public: + explicit KQuantConv3x3(mx::Stream stream) : mx::Primitive(stream) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override; + + const char* name() const override { + return "KQuantConv3x3"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + std::vector jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) override; + std::pair, std::vector> vmap( + const std::vector& inputs, + const std::vector& axes) override; +}; + +// Quantized matmul with a residual epilogue: +// y = residual + gate * (x @ W^T) +// x [..., K] (float16/bfloat16; float32 runs in bfloat16), w 2-D uint8 wire +// bytes [N, K bytes], residual of the output shape [..., N], gate optional: +// [N], or x's batch dims, then 1, then N (one gate row per batch, as an +// adaptive norm gives). With no gate y = residual + x @ W^T. The residual +// and the gate run in the output dtype. K and N are below 2^25 and x holds +// fewer than 2^31 rows. Returns [..., N] in the dtype quantized_matmul +// gives for x. On a NAX GPU, for q4_k, q5_k, q6_k and q8_0 weights on +// bfloat16, from 1024 rows, one dispatch walks K into f32 sums and stores +// residual + gate * sum, rounded once. Every other case builds the sum from +// quantized_matmul and stock ops, which round the product and the gated +// product to the output dtype before the add. +mx::array quantized_matmul_residual( + mx::array x, + mx::array w, + const std::string& kquant_type, + mx::array residual, + const std::optional& gate = std::nullopt, + mx::StreamOrDevice s = {}); + +// Quantized matmul with a residual epilogue (see quantized_matmul_residual). +// Inputs are x, w, the residual and, when gated, the gate. vjp, jvp and +// vmap are those of the unfused ops. The vjp is wrt x, the residual and the +// gate: the quantized weight is frozen, as in KQuantMatmul. +class KQuantMatmulResid : public mx::Primitive { + public: + KQuantMatmulResid(mx::Stream stream, std::string kquant_type) + : mx::Primitive(stream), kquant_type_(std::move(kquant_type)) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override; + + const char* name() const override { + return "KQuantMatmulResid"; + } + bool is_equivalent(const mx::Primitive& other) const override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + std::vector jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) override; + std::pair, std::vector> vmap( + const std::vector& inputs, + const std::vector& axes) override; + + private: + std::string kquant_type_; +}; + // Fused causal short convolution for gated-delta prefill (KDA, GDN): // y[t] = silu(sum_j w[:, j] * in[t - K + 1 + j]), in = [state; x] // over x [B, T, C] (float16/bfloat16), state [B, K - 1, C] (the carried diff --git a/src/kquant_conv_nax.cpp b/src/kquant_conv_nax.cpp new file mode 100644 index 00000000..3d42e733 --- /dev/null +++ b/src/kquant_conv_nax.cpp @@ -0,0 +1,306 @@ +// 3x3 convolution, stride 1, padding 1, channels last (see kq_conv_nax.h +// for the kernel). +// +// y = conv(x, w) + bias +// +// On a NAX GPU, for bfloat16 with C a multiple of 16, from kConvMinRows +// output pixels, one dispatch computes y as an implicit GEMM with f32 sums, +// rounded once. Every other case builds the same convolution from MLX's +// conv_general on the flipped kernel (its implicit GEMM, f32 sums, no +// Winograd transform) and a stock add of the bias. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kquant.h" +#include "kquant_internal.h" + +#include "mlx/allocator.h" +#include "mlx/ops.h" +#include "mlx/transforms.h" +#include "mlx/utils.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +const char* kConvOp = "[mlx_kquant.conv2d_3x3]"; + +// Output pixels from which the kernel runs under the default mode. +constexpr int64_t kConvMinRows = 1024; +// The kernel holds element offsets of x and y in an int. +constexpr int64_t kConvMaxElems = int64_t{1} << 31; + +bool conv_tune_live() { + static const bool live = kq_env_lever("KQ_TUNE_LIVE", 0) != 0; + return live; +} + +// KQ_CONV_NAX: 0 = the composition on every call, 2 = the kernel at every +// pixel count, unset, 1 or any other value = from kConvMinRows pixels. +int conv_mode() { + static const int cached = kq_env_lever("KQ_CONV_NAX", 1); + return conv_tune_live() ? kq_env_lever("KQ_CONV_NAX", 1) : cached; +} + +bool conv_log() { + static const bool cached = std::getenv("KQ_QMM_LOG") != nullptr; + return conv_tune_live() ? std::getenv("KQ_QMM_LOG") != nullptr : cached; +} + +// The composition: MLX's convolution on the flipped kernel with flip set, +// which on the GPU takes the implicit GEMM and not the Winograd transform, +// then the bias. +mx::array conv_composed( + const mx::array& x, + const mx::array& w, + const std::optional& bias, + mx::Stream s) { + auto rev = mx::array({2, 1, 0}); + auto wf = mx::take(mx::take(w, rev, 1, s), rev, 2, s); + auto y = mx::conv_general( + x, wf, {1, 1}, {1, 1}, {1, 1}, {1, 1}, {1, 1}, 1, true, s); + if (bias.has_value()) { + y = mx::add(y, *bias, s); + } + return y; +} + +} // namespace + +#ifdef _METAL_ + +void KQuantConv3x3::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + const mx::array x = kq_row_major_gpu(inputs[0], s); + const mx::array w = kq_row_major_gpu(inputs[1], s); + const bool has_bias = inputs.size() > 2; + const mx::array b = has_bias ? kq_row_major_gpu(inputs[2], s) : w; + auto& out = outputs[0]; + out.set_data(mx::allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + const int H = x.shape(1); + const int W = x.shape(2); + const int C = x.shape(3); + const int O = w.shape(0); + const int M = x.shape(0) * H * W; + const int hb = has_bias ? 1 : 0; + + // Column tile: 128, or 64 where 64 pads O less. KQ_CONV_TILE + // (bm,bn,wm,wn) pins a tile the metallib carries, read once per process + // unless KQ_TUNE_LIVE=1. + int bm = 128, bn = 128, wm = 4, wn = 4; + const int pad128 = (O + 127) / 128 * 128; + const int pad64 = (O + 63) / 64 * 64; + if (pad64 < pad128) { + bn = 64; + wn = 2; + } + static const char* tile_cached = std::getenv("KQ_CONV_TILE"); + const char* tile_env = + conv_tune_live() ? std::getenv("KQ_CONV_TILE") : tile_cached; + if (tile_env != nullptr && tile_env[0] != '\0') { + if (std::sscanf(tile_env, "%d,%d,%d,%d", &bm, &bn, &wm, &wn) != 4) { + throw std::invalid_argument( + std::string(kConvOp) + " KQ_CONV_TILE=" + tile_env + + " is not bm,bn,wm,wn."); + } + } + const std::string kname = "kq_conv3x3_nax_" + kq_type_string(x.dtype()) + + "_bm" + std::to_string(bm) + "_bn" + std::to_string(bn) + "_wm" + + std::to_string(wm) + "_wn" + std::to_string(wn); + if (tile_env != nullptr && tile_env[0] != '\0') { + kq_require_kernel(d, std::string("KQ_CONV_TILE=") + tile_env, kname); + } + + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims( + (O + bn - 1) / bn, (static_cast(M) + bm - 1) / bm, 1); + if (conv_log()) { + std::fprintf( + stderr, + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + kname.c_str(), + M, + O, + 9 * C, + static_cast(grid_dims.width), + static_cast(grid_dims.height), + static_cast(grid_dims.depth)); + } + + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + int c = 0; + ce.set_input_array(x, c++); + ce.set_input_array(w, c++); + ce.set_input_array(b, c++); + ce.set_output_array(out, c++); + ce.set_bytes(H, c++); + ce.set_bytes(W, c++); + ce.set_bytes(C, c++); + ce.set_bytes(O, c++); + ce.set_bytes(M, c++); + ce.set_bytes(hb, c++); + ce.dispatch_threadgroups(grid_dims, group_dims); +} + +#else // !_METAL_ + +void KQuantConv3x3::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error(std::string(kConvOp) + " requires Metal."); +} + +#endif + +void KQuantConv3x3::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + std::string(kConvOp) + " CPU streams build the composed ops."); +} + +namespace { + +std::optional conv_bias_of(const std::vector& in) { + return in.size() > 2 ? std::optional(in[2]) : std::nullopt; +} + +} // namespace + +// The kernel computes the function of the composition, so its transforms +// are the composition's. +std::vector KQuantConv3x3::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector&) { + auto fun = [&](const std::vector& in) { + return std::vector{ + conv_composed(in[0], in[1], conv_bias_of(in), stream())}; + }; + auto vjps = mx::vjp(fun, primals, cotangents).second; + std::vector out; + for (auto arg : argnums) { + out.push_back(vjps[arg]); + } + return out; +} + +std::vector KQuantConv3x3::jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) { + auto fun = [&](const std::vector& in) { + std::vector full = primals; + for (size_t i = 0; i < argnums.size(); i++) { + full[argnums[i]] = in[i]; + } + return std::vector{ + conv_composed(full[0], full[1], conv_bias_of(full), stream())}; + }; + std::vector at; + for (auto arg : argnums) { + at.push_back(primals[arg]); + } + return mx::jvp(fun, at, tangents).second; +} + +std::pair, std::vector> KQuantConv3x3::vmap( + const std::vector& inputs, + const std::vector& axes) { + auto fun = [&](const std::vector& in) { + return std::vector{ + conv_composed(in[0], in[1], conv_bias_of(in), stream())}; + }; + return {mx::vmap(fun, axes, {0})(inputs), {0}}; +} + +bool KQuantConv3x3::is_equivalent(const mx::Primitive&) const { + return true; +} + +std::vector KQuantConv3x3::output_shapes( + const std::vector& inputs) { + auto shape = inputs[0].shape(); + shape.back() = inputs[1].shape(0); + return {shape}; +} + +mx::array conv2d_3x3( + mx::array x, + mx::array w, + const std::optional& bias, + mx::StreamOrDevice s_) { + const std::string op = kConvOp; + if (x.ndim() != 4 || !mx::issubdtype(x.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " x must be a float array [B, H, W, C] (channels last)."); + } + const int C = x.shape(3); + if (w.ndim() != 4 || w.shape(1) != 3 || w.shape(2) != 3 || w.shape(3) != C || + !mx::issubdtype(w.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " w must be a float array [O, 3, 3, C] with C " + + std::to_string(C) + "."); + } + const int O = w.shape(0); + if (bias.has_value() && + (bias->ndim() != 1 || bias->shape(0) != O || + !mx::issubdtype(bias->dtype(), mx::floating))) { + throw std::invalid_argument( + op + " bias must be a float array [O] with O " + std::to_string(O) + + "."); + } + auto s = mx::to_stream(s_); + const mx::Dtype dt = x.dtype(); + auto wc = mx::astype(w, dt, s); + std::optional bc; + if (bias.has_value()) { + bc = mx::astype(*bias, dt, s); + } + const int64_t pixels = static_cast(x.size()) / std::max(C, 1); + const int mode = conv_mode(); + const bool fused = s.device == mx::Device::gpu && nax_available() && + mode != 0 && dt == mx::bfloat16 && C > 0 && C % 16 == 0 && O > 0 && + static_cast(x.size()) < kConvMaxElems && + pixels * O < kConvMaxElems && (mode == 2 || pixels >= kConvMinRows); + if (!fused) { + return conv_composed(x, wc, bc, s); + } + auto out_shape = x.shape(); + out_shape.back() = O; + std::vector in = {std::move(x), std::move(wc)}; + if (bc.has_value()) { + in.push_back(*bc); + } + return mx::array( + std::move(out_shape), + dt, + std::make_shared(s), + std::move(in)); +} + +} // namespace mlx_kquant diff --git a/src/kquant_cpu_decode.cpp b/src/kquant_cpu_decode.cpp index 9805d218..f509658a 100644 --- a/src/kquant_cpu_decode.cpp +++ b/src/kquant_cpu_decode.cpp @@ -1343,6 +1343,15 @@ void kquant_qmm_cpu( throw std::runtime_error( "[mlx_kquant] quantized_matmul: unsupported codec: " + kquant_type); } + // No rows or no columns leave nothing to write, and an inner dim of 0 + // gives zeros. The GEMM below takes no leading dim of 0. + if (M <= 0 || N <= 0) { + return; + } + if (K <= 0) { + std::fill_n(result, static_cast(M) * N, T(0.0f)); + return; + } const int w_rows = transpose_w ? N : K; const int w_cols = transpose_w ? K : N; const std::size_t row_bytes = diff --git a/src/kquant_dense.cpp b/src/kquant_dense.cpp new file mode 100644 index 00000000..fdeaac68 --- /dev/null +++ b/src/kquant_dense.cpp @@ -0,0 +1,227 @@ +// Dense matrix rows for weight operands, decided at eval. +// +// The matmul kernels walk the last two dims of a packed weight as dense +// rows. An op builder can check that only on an array that has memory: an +// unevaluated view reports the default dense strides at build time, so a +// lazy w[0::2] looks dense and the kernel then reads the rows in between. +// kq_matrix_dense keeps the build-time check for an array with memory and +// gives any other array a KQuantMatrixDense node. The node aliases an input +// whose matrix rows are dense, whatever the stride of its leading dims and +// the size of its buffer, and copies any other input row-major. +// kq_row_major is the same for an operand that a kernel reads as one +// row-major run, such as the bytes dequantize decodes: there the node +// aliases only an input that is row-major in all of its dims. +// +// The _gpu forms decide inside a primitive's own eval, where every input +// has memory: they return the input when its layout is the one the kernel +// reads and a row-major copy otherwise. An op uses them for an operand that +// is unevaluated on most calls, such as an activation, which then costs no +// node when it is dense. +#include +#include +#include +#include + +#include "kquant.h" +#include "kquant_internal.h" + +#include "mlx/allocator.h" +#include "mlx/backend/cpu/encoder.h" +#include "mlx/ops.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" // kq_get_kernel +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +// Dtypes and ranks the copy kernels take (kq_lora_densify_). +bool dense_copy_ok(const mx::array& a) { + const size_t isz = a.itemsize(); + return (isz == 1 || isz == 2 || isz == 4) && a.ndim() <= 8; +} + +} // namespace + +mx::array kq_matrix_dense(const mx::array& x, const mx::Stream& s) { + if (x.status() != mx::array::Status::unscheduled) { + return kq_matrix_rows_dense(x) ? x : mx::contiguous(x, false, s); + } + // One element or none has no layout to get wrong. + if (x.size() <= 1) { + return x; + } + if (!dense_copy_ok(x)) { + return mx::contiguous(x, false, s); + } + return mx::array( + x.shape(), x.dtype(), std::make_shared(s), {x}); +} + +mx::array kq_row_major(const mx::array& x, const mx::Stream& s) { + if (x.status() != mx::array::Status::unscheduled) { + return x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + } + if (x.size() <= 1) { + return x; + } + if (!dense_copy_ok(x)) { + return mx::contiguous(x, false, s); + } + return mx::array( + x.shape(), x.dtype(), std::make_shared(s, true), {x}); +} + +bool KQuantMatrixDense::passes(const mx::array& in) const { + if (in.size() == 0) { + return true; + } + return whole_ ? in.flags().row_contiguous : kq_matrix_rows_dense(in); +} + +#ifdef _METAL_ + +namespace { + +// Copies a (any strides) into out (row-major, allocated) with one thread +// per element. +void dense_copy_gpu(const mx::array& a, mx::array& out, const mx::Stream& s) { + if (!dense_copy_ok(a)) { + throw std::invalid_argument( + "[mlx_kquant] a strided view must have a 1-, 2- or 4-byte dtype " + "and at most 8 dims."); + } + // The copy kernel indexes its threads with 32 bits. + if (a.size() > UINT32_MAX) { + throw std::invalid_argument( + "[mlx_kquant] a strided view of more than 4294967295 elements " + "cannot be copied at eval. Evaluate a contiguous copy of it first, " + "for example with mx.contiguous(w)."); + } + auto& d = mx::metal::device(s.device); + auto& ce = mx::metal::get_command_encoder(s); + // KqLoraDensifyArgs in kq_lora_epilogue.h. + struct { + int shape[8]; + int64_t strides[8]; + int ndim; + int64_t size; + } args{}; + for (int i = 0; i < static_cast(a.ndim()); i++) { + args.shape[i] = a.shape(i); + args.strides[i] = a.strides()[i]; + } + args.ndim = static_cast(a.ndim()); + args.size = static_cast(a.size()); + auto kernel = + kq_get_kernel(d, "kq_lora_densify_" + std::to_string(a.itemsize())); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(a, 0); + ce.set_output_array(out, 1); + ce.set_bytes(args, 2); + const size_t tg = a.size() < 1024 ? a.size() : 1024; + ce.dispatch_threads(MTL::Size(a.size(), 1, 1), MTL::Size(tg, 1, 1)); +} + +} // namespace + +mx::array kq_dense_copy_gpu(const mx::array& a, const mx::Stream& s) { + mx::array t(a.shape(), a.dtype(), nullptr, {}); + t.set_data(mx::allocator::malloc(t.nbytes())); + dense_copy_gpu(a, t, s); + mx::metal::get_command_encoder(s).add_temporary(t); + return t; +} + +mx::array kq_matrix_dense_gpu(const mx::array& a, const mx::Stream& s) { + if (a.size() == 0 || kq_matrix_rows_dense(a)) { + return a; + } + return kq_dense_copy_gpu(a, s); +} + +mx::array kq_row_major_gpu(const mx::array& a, const mx::Stream& s) { + if (a.size() == 0 || a.flags().row_contiguous) { + return a; + } + return kq_dense_copy_gpu(a, s); +} + +mx::array kq_unit_stride_gpu(const mx::array& a, const mx::Stream& s) { + if (a.size() == 0 || a.ndim() == 0 || a.strides().back() == 1) { + return a; + } + return kq_dense_copy_gpu(a, s); +} + +void KQuantMatrixDense::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + const auto& in = inputs[0]; + auto& out = outputs[0]; + if (passes(in)) { + out.copy_shared_buffer(in); + return; + } + out.set_data(mx::allocator::malloc(out.nbytes())); + dense_copy_gpu(in, out, stream()); +} + +#else // !_METAL_ + +mx::array kq_matrix_dense_gpu(const mx::array&, const mx::Stream&) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + +mx::array kq_dense_copy_gpu(const mx::array&, const mx::Stream&) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + +mx::array kq_row_major_gpu(const mx::array&, const mx::Stream&) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + +mx::array kq_unit_stride_gpu(const mx::array&, const mx::Stream&) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + +void KQuantMatrixDense::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + +#endif // _METAL_ + +void KQuantMatrixDense::eval_cpu( + const std::vector& inputs, + std::vector& outputs) { + const auto& in = inputs[0]; + auto& out = outputs[0]; + if (passes(in)) { + out.copy_shared_buffer(in); + return; + } + out.set_data(mx::allocator::malloc(out.nbytes())); + auto& encoder = mx::cpu::get_command_encoder(stream()); + encoder.set_input_array(in); + encoder.set_output_array(out); + encoder.dispatch([src = mx::array::unsafe_weak_copy(in), + dst = mx::array::unsafe_weak_copy(out)]() mutable { + const size_t isz = src.itemsize(); + const char* ip = src.data(); + char* op = dst.data(); + const int64_t n = static_cast(src.size()); + for (int64_t i = 0; i < n; i++) { + const int64_t loc = elem_to_loc64(i, src.shape(), src.strides()); + std::memcpy(op + i * isz, ip + loc * static_cast(isz), isz); + } + }); +} + +} // namespace mlx_kquant diff --git a/src/kquant_device.cpp b/src/kquant_device.cpp index fa7928e4..3c3b3183 100644 --- a/src/kquant_device.cpp +++ b/src/kquant_device.cpp @@ -1,9 +1,18 @@ // GPU core count for the occupancy policies that scale with it. Probed once // from the IOKit accelerator entry; KQ_GPU_CORES overrides it and is read per -// call so one process can A/B a policy. +// call so one process can A/B a policy. Also the environment lever parser +// (kq_env_lever, kq_env_flag). +#include +#include +#include +#include #include +#include +#include +#include #include "kquant.h" +#include "kquant_internal.h" #ifdef __APPLE__ #include @@ -48,8 +57,94 @@ int probe_gpu_core_count() { #endif } +// Warns once per lever and value that `value` is not one `name` takes. +void warn_env_lever(const char* name, const char* value, const char* takes) { + static std::mutex mu; + static std::set seen; + std::lock_guard lock(mu); + if (seen.insert(std::string(name) + "=" + value).second) { + std::fprintf( + stderr, + "[mlx_kquant] %s=%s is not %s; %s keeps its default.\n", + name, + value, + takes, + name); + } +} + +int parse_env_lever(const char* name, int unset, bool flag) { + const char* e = std::getenv(name); + if (e == nullptr || e[0] == '\0') { + return unset; + } + std::string v; + for (const char* c = e; *c != '\0'; ++c) { + v += static_cast(std::tolower(static_cast(*c))); + } + if (v == "false" || v == "off" || v == "no") { + return 0; + } + if (v == "true" || v == "on" || v == "yes") { + return 1; + } + const char* takes = flag ? "0, 1, true, false, on, off, yes or no" + : "an integer, true, false, on, off, yes or no"; + char* end = nullptr; + errno = 0; + const long n = std::strtol(e, &end, 10); + if (end == e || *end != '\0' || errno != 0 || n < INT_MIN || n > INT_MAX || + (flag && n != 0 && n != 1)) { + warn_env_lever(name, e, takes); + return unset; + } + return static_cast(n); +} + } // namespace +int kq_env_lever(const char* name, int unset) { + return parse_env_lever(name, unset, false); +} + +int kq_env_flag(const char* name, int unset) { + return parse_env_lever(name, unset, true); +} + +int kq_env_lever_from(const char* name, int unset, int lowest) { + const int v = parse_env_lever(name, unset, false); + if (v >= lowest) { + return v; + } + const char* e = std::getenv(name); + const std::string takes = "an integer of at least " + std::to_string(lowest); + warn_env_lever(name, e != nullptr ? e : "", takes.c_str()); + return unset; +} + +int kq_env_choice( + const char* name, + const std::vector& choices, + int unset) { + const char* e = std::getenv(name); + if (e == nullptr || e[0] == '\0') { + return unset; + } + std::string v; + for (const char* c = e; *c != '\0'; ++c) { + v += static_cast(std::tolower(static_cast(*c))); + } + std::string takes = "one of"; + for (size_t i = 0; i < choices.size(); ++i) { + if (v == choices[i]) { + return static_cast(i); + } + takes += (i == 0 ? " " : ", ") + choices[i]; + } + warn_env_lever(name, e, takes.c_str()); + return unset; +} + int gpu_core_count() { const char* e = std::getenv("KQ_GPU_CORES"); if (e != nullptr) { diff --git a/src/kquant_dsa_indexer.cpp b/src/kquant_dsa_indexer.cpp index 2ba8f920..875fa911 100644 --- a/src/kquant_dsa_indexer.cpp +++ b/src/kquant_dsa_indexer.cpp @@ -231,7 +231,10 @@ void KQDsaIndexerScoreDecode::eval_gpu( auto& ce = mx::metal::get_command_encoder(s); const auto& q = inputs[0]; - const auto& k = inputs[1]; // [B, P * D], the row-major key block + // [B, P * D], the row-major key block. A batch is row-major from the + // builder. One sequence comes through a reshape that can keep a view + // with an element stride other than 1, which is copied here. + const mx::array k = kq_unit_stride_gpu(inputs[1], s); const auto& weights = inputs[2]; auto& out = outputs[0]; @@ -594,9 +597,10 @@ mx::array dsa_indexer_score_decode( // a no-op at eval when the input already is. At B == 1 the keys go // through a reshape to [1, P * D] instead: a cache's prefix slice is a // row-contiguous view that Reshape keeps as a view at eval, while - // Contiguous would copy it (its buffer is larger than the view); any - // other layout Reshape copies. The kernel steps batches by P * D, so a - // batch of slices still needs the contiguous copy. + // Contiguous would copy it (its buffer is larger than the view). + // Reshape keeps other views too, and the eval copies one whose rows + // are not dense. The kernel steps batches by P * D, so a batch of + // slices still needs the contiguous copy. auto q = mx::contiguous(mx::astype(queries, final_type, s), false, s); auto k = mx::reshape( mx::astype(keys, final_type, s), diff --git a/src/kquant_dsa_sparse_attn.cpp b/src/kquant_dsa_sparse_attn.cpp index d927bc46..415ac71e 100644 --- a/src/kquant_dsa_sparse_attn.cpp +++ b/src/kquant_dsa_sparse_attn.cpp @@ -102,7 +102,10 @@ void KQDsaSparseAttention::eval_gpu( // memory (two threadgroups per core) plus finer live-arm buckets win there; // dense 512-entry lists keep the 256-wide tile (fewer per-tile softmax // rescales / selected[] builds). KQ_DSA_BK=128/256 overrides for tuning. - static const char* bk_env = std::getenv("KQ_DSA_BK"); + static const char* bk_env = [] { + const char* e = std::getenv("KQ_DSA_BK"); + return e != nullptr && e[0] != '\0' ? e : nullptr; + }(); const std::string bk_tag = bk_env ? ((std::string(bk_env) == "128") ? "bk128" : "bk256") : ((topkN <= 128) ? "bk128" : "bk256"); diff --git a/src/kquant_internal.h b/src/kquant_internal.h index 7d130399..2c5e85e9 100644 --- a/src/kquant_internal.h +++ b/src/kquant_internal.h @@ -5,6 +5,7 @@ #include #include #include +#include #include "mlx/array.h" #include "mlx/stream.h" @@ -18,6 +19,48 @@ namespace mlx_kquant { // Throws on other dtypes. std::string kq_type_string(mx::Dtype d); +// Integer value of the environment lever `name`. Absent or empty reads as +// `unset`; 0, false, off and no read as 0; 1, true, on and yes read as 1; +// any other integer reads as itself. Any other value reads as `unset`, with +// a warning on stderr the first time the lever holds it. +int kq_env_lever(const char* name, int unset); + +// kq_env_lever for a lever whose set values are 0 and 1: any other integer +// also reads as `unset`, with the same warning. +int kq_env_flag(const char* name, int unset); + +// kq_env_lever for a lever whose set values start at `lowest`: a smaller +// integer also reads as `unset`, with the same warning. +int kq_env_lever_from(const char* name, int unset, int lowest); + +// Index in `choices` (lower case) of the environment lever `name`, read in +// any case. Absent or empty reads as `unset`; any other value reads as +// `unset`, with the same warning. +int kq_env_choice( + const char* name, + const std::vector& choices, + int unset); + +// K ceiling of the codec's large-M NAX tile: above it the plain qmm_t +// runs the BM=64 tile in place of BM=128. 0 = no ceiling. +int kq_nax_large_max_k(const std::string& kquant_type); + +// True when KQ_QMM_TILE pins the plain qmm_t tile, which also keeps a call +// off the dequantize route. Metal builds only. +bool kq_qmm_tile_pinned(); + +// KQ_DEQUANT_ROUTE: auto (the chip classes of the dequantize route), off +// (never) or force (on any GPU, for the codecs and shapes of the route). +enum KqDequantRoute { kDequantAuto = 0, kDequantOff = 1, kDequantForce = 2 }; + +// The KQ_DEQUANT_ROUTE mode, read once per process unless KQ_TUNE_LIVE=1. +// Metal builds only. +int kq_dequant_route_mode(); + +// True when KQ_QMM_LOG is set, read as the qmm tiles read it. Metal builds +// only. +bool kq_qmm_log_on(); + // KVarN key/value widths the record codec supports. inline bool kq_kvarn_bits_ok(int bits) { return bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || @@ -93,6 +136,69 @@ inline void kq_check_weight_base_at_build( } } +// True when the last two dims of a are laid out as dense rows: the last +// dim has stride 1 and the dim before it has the row length as its stride. +// The matmul kernels walk a weight that way and take the strides of any +// leading dims. The answer means something only for an array with memory. +inline bool kq_matrix_rows_dense(const mx::array& a) { + const int nd = static_cast(a.ndim()); + if (nd == 0) { + return true; + } + if (nd < 2) { + return a.strides()[0] == 1; + } + return a.strides()[nd - 1] == 1 && + a.strides()[nd - 2] == static_cast(a.shape(nd - 1)); +} + +// Op-builder form. An x with memory comes back as it is when its matrix +// rows are dense, and as a row-major copy otherwise. An x without memory +// has no strides to read yet, so it gets a KQuantMatrixDense node, which +// decides at eval. +mx::array kq_matrix_dense(const mx::array& x, const mx::Stream& s); + +// Op-builder form for an operand a kernel reads as one row-major run: an x +// with memory comes back as it is when it is row-major, and as a row-major +// copy otherwise. An x without memory gets a KQuantMatrixDense node that +// decides at eval, and passes a row-major x through with its buffer. +mx::array kq_row_major(const mx::array& x, const mx::Stream& s); + +// Eval form for a GPU primitive: a itself when its matrix rows are dense, +// else a row-major copy made by one dispatch on s, which the command buffer +// keeps alive. a has a 1-, 2- or 4-byte dtype and at most 8 dims. +mx::array kq_matrix_dense_gpu(const mx::array& a, const mx::Stream& s); + +// A row-major copy of a in a buffer the command buffer keeps alive, +// whatever the layout of a. The dtype and dim limits are those above. +mx::array kq_dense_copy_gpu(const mx::array& a, const mx::Stream& s); + +// The same for an operand a kernel reads as one row-major run: a itself +// when it is row-major, else the row-major copy. An op whose operand is +// unevaluated on most calls, such as an activation, binds the result of +// this in its eval and adds no node for the operand. +mx::array kq_row_major_gpu(const mx::array& a, const mx::Stream& s); + +// The same for an operand a kernel reads with a unit stride in its last +// dim and the operand's own strides in the dims before it: a itself when +// the last stride is 1, else the row-major copy. +mx::array kq_unit_stride_gpu(const mx::array& a, const mx::Stream& s); + +// Eval guard for an operand an op does not copy, such as an expert stack: +// throws std::invalid_argument(what + " must be row-contiguous.") when a +// is not row-major. An op builder can run it only on an array that has +// memory, since any other array reports the default dense strides. +inline void kq_check_row_major(const mx::array& a, const std::string& what) { + if (!a.flags().row_contiguous) { + throw std::invalid_argument(what + " must be row-contiguous."); + } +} + +// True when a has memory, so that its strides and flags are its own. +inline bool kq_has_layout(const mx::array& a) { + return a.status() != mx::array::Status::unscheduled; +} + // 64-bit variant of mx::elem_to_loc (mlx/backend/common/utils.h), whose linear // index parameter is a 32-bit int. That truncates offsets into tensors whose // flattened element count exceeds 2^31 - e.g. a stacked MoE expert weight whose @@ -112,4 +218,21 @@ inline int64_t elem_to_loc64( return loc; } +// Eval form for a CPU primitive, called inside its dispatched task: the +// elements of a in row-major order, a's own memory when a is row-major and +// else a copy held by store. +template +const T* kq_row_major_cpu(const mx::array& a, std::vector& store) { + const T* p = a.data(); + if (a.flags().row_contiguous) { + return p; + } + const int64_t n = static_cast(a.size()); + store.resize(static_cast(n)); + for (int64_t i = 0; i < n; i++) { + store[static_cast(i)] = p[elem_to_loc64(i, a.shape(), a.strides())]; + } + return store.data(); +} + } // namespace mlx_kquant diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index fba751ff..f6ab404e 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -10,12 +10,17 @@ // decode-latency-critical path); the decode-only bias-fused fast path lives in // the KQuantQmvBias primitive below (qmv_bias), which reuses this file's qmv // dispatch helpers. +#include #include #include +#include #include +#include +#include #include #include #include +#include #include "kquant.h" #include "kquant_codec.h" @@ -76,13 +81,38 @@ using mx::Stream; // +13% (q5_k, q3_k) -- the k-quants, q8_0, and the cheap-loader IQs. // - 961: the flat quartet only wins the deepest bands (+1-7% at M1024) // and loses or washes below. +// Large-M row: the tile, the row-tile fold (swizzle_log, 0 = none), the +// shape floor (M >= min_m at either ceil(M/64) parity, N >= min_n; min_m 0 = +// never) and the K ceiling (max_k: above it the row runs BM=64; 0 = none) +// of the image-generation band. A codec whose tile and floor are not +// measured takes those of its loader class and is marked inherited; max_k +// is measured for every codec with NAX kernels. +struct KqLargeM { + int bm; + int bn; + int wm; + int wn; + int swizzle_log; + int min_m; + int min_n; + int max_k; + bool inherited; +}; struct KqSmallBmPolicy { bool bm32; int route_min; int db64_min_n; int bm128_min_m; + KqLargeM large; }; +static KqLargeM kq_largem_row(const std::string& t); +static KqSmallBmPolicy kq_smallbm_base(const std::string& t); static KqSmallBmPolicy kq_smallbm_policy(const std::string& t) { + KqSmallBmPolicy p = kq_smallbm_base(t); + p.large = kq_largem_row(t); + return p; +} +static KqSmallBmPolicy kq_smallbm_base(const std::string& t) { if (t == "q4_k" || t == "q3_k") { return {true, 7, 0, 449}; } @@ -140,6 +170,64 @@ static KqSmallBmPolicy kq_smallbm_policy(const std::string& t) { return {false, 0, 0, 193}; // iq2_xs, iq2_s, iq1_m } +// Loader classes of the large-M rows: the 256-weight k-quant loaders, the +// 32-weight flat block loaders, and the rest (the IQ grid and ternary +// loaders). +enum class KqLoaderClass { kquant, flat, grid }; +static KqLoaderClass kq_loader_class(const std::string& t) { + if (t == "q2_k" || t == "q3_k" || t == "q4_k" || t == "q5_k" || t == "q6_k") { + return KqLoaderClass::kquant; + } + if (t == "q4_0" || t == "q4_1" || t == "q5_0" || t == "q5_1" || t == "q8_0" || + t == "iq4_nl") { + return KqLoaderClass::flat; + } + return KqLoaderClass::grid; +} + +// K above which BM=128 runs slower than BM=64 from M 4096, the largest K +// before the first measured loss of more than 1% (geometric mean over M +// 4128, 8256 and 16384 at N 4096, K 10240, 12288 and 16384; K 3840 and +// 4096 win for every codec). 0: no loss up to K 16384. On the BK=128 tiles +// q3_k, q5_k, q6_k, iq2_s, iq2_xs, iq3_s and iq3_xxs win at every K, and so +// do iq1_s, iq1_m, iq2_xxs, iq4_xs and ptq1_0 on BK=64. A ceiling holds for +// the codec's default K tile. q2_0 takes the pq2_0 ceiling; it is not +// measured. +static int kq_largem_max_k(const std::string& t) { + if (t == "q2_k" || t == "pq2_0" || t == "q2_0") { + return 4096; + } + if (t == "q4_k" || t == "q8_0" || t == "q4_0" || t == "q4_1" || t == "q5_0" || + t == "q5_1") { + return 10240; + } + if (t == "iq4_nl" || t == "stq1_0") { + return 12288; + } + return 0; +} + +// Measured on M5 Max with benchmarks/bench_dit_shapes.py at the DiT rows +// (4096 to 33792). A row with inherited = true takes its loader class's +// tile and floor. Measured rows: BM=128 is faster than BM=64 at every +// odd-tile row from 4096 up to the K ceiling, and every row-tile fold is +// slower than none. +static KqLargeM kq_largem_row(const std::string& t) { + const int max_k = kq_largem_max_k(t); + if (t == "q4_k" || t == "q5_k" || t == "q6_k" || t == "q8_0") { + return {128, 64, 2, 2, 0, 4096, 0, max_k, false}; + } + switch (kq_loader_class(t)) { + case KqLoaderClass::kquant: + return {128, 64, 2, 2, 0, 4096, 0, max_k, true}; + case KqLoaderClass::flat: + return {128, 64, 2, 2, 0, 4096, 0, max_k, true}; + case KqLoaderClass::grid: + return {128, 64, 2, 2, 0, 4096, 0, max_k, true}; + } + return {128, 64, 2, 2, 0, 4096, 0, max_k, true}; +} + // Shape-adjusted sub-13 route threshold. The table crossovers were tuned // at [17408x5120]; they hold for N >= 8192 but sit too low for strong-mv // codecs at decode projection shapes, where the qmm grid is only @@ -522,23 +610,241 @@ static int kq_splitk_nax_min_m(const std::string& t, int N) { return N <= 1024 ? p.splitk_small : p.splitk; } +// Integer value of an environment lever (kq_env_lever): `unset` when the +// variable is absent, empty or not a value the lever takes. +static int kq_env_int(const char* name, int unset) { + return kq_env_lever(name, unset); +} + +// KQ_TUNE_LIVE=1, read once per process, makes the tile levers below read +// their variables on every dispatch, so one process can switch tiles +// between calls on one resident copy of the weights. Without it each tile +// lever keeps its first read for the life of the process. +static bool kq_tune_live() { + static const bool live = kq_env_int("KQ_TUNE_LIVE", 0) != 0; + return live; +} + +// KQ_NAX_SMALL_BM: 0 = off, 2 = force bm32 for policy-excluded codecs, +// unset or 1 = per-codec policy. Read by the BM=32 tile pick in qmm_nax +// and by the small-M qmm route in eval_gpu_base. +static int kq_nax_small_bm_mode() { + static const int cached = kq_env_int("KQ_NAX_SMALL_BM", 1); + return kq_tune_live() ? kq_env_int("KQ_NAX_SMALL_BM", 1) : cached; +} + +// KQ_NAX_BM128: 0 = off, 1 = floor 193 for all codecs, 2 = no floor, +// unset (-1) = per-codec floor. +static int kq_nax_bm128_mode() { + static const int cached = kq_env_int("KQ_NAX_BM128", -1); + return kq_tune_live() ? kq_env_int("KQ_NAX_BM128", -1) : cached; +} + +// KQ_NAX_BM128_ODD_M: lowest M at which the BM=128 tile also runs when +// ceil(M/64) is odd, 0 = never (the parity rule at every M), unset (-1) = +// min_m of the codec's large-M row. The row-live guard in the tile skips +// the A loads and MMA of the dead simdgroup row of the last row tile, so +// there the half-live tile does the loads and MMA of one BM=64 tile. +static int kq_nax_bm128_odd_m() { + static const int cached = kq_env_int("KQ_NAX_BM128_ODD_M", -1); + return kq_tune_live() ? kq_env_int("KQ_NAX_BM128_ODD_M", -1) : cached; +} + +// KQ_NAX__BK128 (KQ_NAX_Q6K_BK128, KQ_NAX_Q4_0_BK128, +// KQ_NAX_IQ3_S_BK128, ...): K tile of the codec on the BM=64 and BM=128 +// tiles, 1 = 128, 0 = 64, unset = the codec's default. Codecs not listed +// keep BK=64: q8_0, q5_1, iq4_xs, iq2_xxs, iq1_s and iq1_m measured slower +// on BK=128 at diffusion model rows. +struct KqBk128Lever { + const char* type; + const char* env; + int dflt; +}; +static constexpr KqBk128Lever kq_nax_bk128_levers[] = { + {"q6_k", "KQ_NAX_Q6K_BK128", 1}, + {"q5_k", "KQ_NAX_Q5K_BK128", 1}, + {"q3_k", "KQ_NAX_Q3K_BK128", 1}, + {"q4_k", "KQ_NAX_Q4K_BK128", 1}, + {"q4_0", "KQ_NAX_Q4_0_BK128", 1}, + {"q4_1", "KQ_NAX_Q4_1_BK128", 1}, + {"q5_0", "KQ_NAX_Q5_0_BK128", 1}, + {"iq4_nl", "KQ_NAX_IQ4_NL_BK128", 1}, + {"iq3_xxs", "KQ_NAX_IQ3_XXS_BK128", 1}, + {"iq3_s", "KQ_NAX_IQ3_S_BK128", 1}, + {"iq2_xs", "KQ_NAX_IQ2_XS_BK128", 1}, + {"iq2_s", "KQ_NAX_IQ2_S_BK128", 1}, +}; +static constexpr int kq_nax_bk128_codecs = + sizeof(kq_nax_bk128_levers) / sizeof(kq_nax_bk128_levers[0]); +static int kq_nax_bk128_slot(const std::string& t) { + for (int j = 0; j < kq_nax_bk128_codecs; j++) { + if (t == kq_nax_bk128_levers[j].type) { + return j; + } + } + return -1; +} +static bool kq_nax_bk128_on(const std::string& t) { + const int i = kq_nax_bk128_slot(t); + if (i < 0) { + return false; + } + static const std::array cached = [] { + std::array a{}; + for (int j = 0; j < kq_nax_bk128_codecs; j++) { + a[j] = kq_env_flag( + kq_nax_bk128_levers[j].env, kq_nax_bk128_levers[j].dflt) != 0; + } + return a; + }(); + const KqBk128Lever& l = kq_nax_bk128_levers[i]; + return kq_tune_live() ? kq_env_flag(l.env, l.dflt) != 0 : cached[i]; +} + +// KQ_NAX_DB64: 0 = off, 1 = no N floor, unset (-1) = per-codec N floor. +static int kq_nax_db64_mode() { + static const int cached = kq_env_int("KQ_NAX_DB64", -1); + return kq_tune_live() ? kq_env_int("KQ_NAX_DB64", -1) : cached; +} + +// KQ_NAX_SWIZZLE: row-tile fold log on the BM=64 and BM=128 tiles, 0 = +// off, unset (-1) = swizzle_log of the codec's large-M row where that row +// applies and off elsewhere. +static int kq_nax_swizzle_log() { + static const int cached = kq_env_int("KQ_NAX_SWIZZLE", -1); + return kq_tune_live() ? kq_env_int("KQ_NAX_SWIZZLE", -1) : cached; +} + +// KQ_SPLITK_F32 picks the dtype of the K-slice partials of the split-K +// routes (qmm_splitk, verify_mma, verify_nax, qmm_nax_splitk). Unset +// (default): float16 activations store float partials, since a float16 +// partial can pass 65504 where the sum does not, and bfloat16 activations +// store bfloat16 partials, which round every slice once more before the +// fold. 1 = float partials for both, so the fold rounds the result once; +// 0 = partials in the activation dtype for both. Read once per process +// unless KQ_TUNE_LIVE=1. +static int kq_splitk_f32_mode() { + static const int cached = kq_env_flag("KQ_SPLITK_F32", -1); + return kq_tune_live() ? kq_env_flag("KQ_SPLITK_F32", -1) : cached; +} + +// Partial dtype and kernel name suffix of a split-K call. The ops cast +// float32 activations to bfloat16, so dt is bfloat16 or float16, the +// only types with split kernels. +struct KqSplitkPartials { + mx::Dtype dtype; + std::string suffix; +}; + +static KqSplitkPartials kq_splitk_partials(mx::Dtype dt, bool f32) { + if (f32) { + return {mx::float32, "_pf32"}; + } + return {dt, ""}; +} + +static KqSplitkPartials kq_splitk_partials(mx::Dtype dt) { + const int mode = kq_splitk_f32_mode(); + return kq_splitk_partials(dt, mode < 0 ? dt == mx::float16 : mode != 0); +} + +// KQ_QMM_LOG set (any value): the qmm tiles log the kernel name and shape +// of each dispatch to stderr. +static bool kq_qmm_log() { + static const bool cached = std::getenv("KQ_QMM_LOG") != nullptr; + return kq_tune_live() ? std::getenv("KQ_QMM_LOG") != nullptr : cached; +} + +// A qmm tile: rows and columns per threadgroup, and the grid of +// simdgroups over them. +struct KqTile { + int bm; + int bn; + int wm; + int wn; +}; + +static std::string kq_tile_string(const KqTile& t) { + return std::to_string(t.bm) + "," + std::to_string(t.bn) + "," + + std::to_string(t.wm) + "," + std::to_string(t.wn); +} + +// Reads a tile override "bm,bn,wm,wn". Unset or empty gives none, and a +// malformed value raises with the lever name. +static std::optional kq_read_tile(const char* lever) { + const char* e = std::getenv(lever); + if (e == nullptr || e[0] == '\0') { + return std::nullopt; + } + KqTile t{}; + char tail = 0; + const int n = + std::sscanf(e, "%d,%d,%d,%d%c", &t.bm, &t.bn, &t.wm, &t.wn, &tail); + if (n != 4 || t.bm <= 0 || t.bn <= 0 || t.wm <= 0 || t.wn <= 0) { + throw std::invalid_argument( + std::string("[mlx_kquant] ") + lever + "=" + e + + " is not bm,bn,wm,wn (four positive integers)."); + } + return t; +} + +// KQ_NAX_TILE pins the NAX qmm_t tile of transpose calls in place of the +// BM=32, BM=64 and BM=128 policy. Read once per process unless +// KQ_TUNE_LIVE=1. +static std::optional kq_nax_tile() { + static const std::optional cached = kq_read_tile("KQ_NAX_TILE"); + return kq_tune_live() ? kq_read_tile("KQ_NAX_TILE") : cached; +} + +// KQ_QMM_TILE pins the qmm_t tile of transpose calls on the path without +// NAX. Read once per process unless KQ_TUNE_LIVE=1. +static std::optional kq_qmm_tile() { + static const std::optional cached = kq_read_tile("KQ_QMM_TILE"); + return kq_tune_live() ? kq_read_tile("KQ_QMM_TILE") : cached; +} + +// Raises when a tile override selects a kernel the metallib does not +// carry, naming the kernel. +static void kq_require_tile_kernel( + Device& d, + const char* lever, + const KqTile& t, + const std::string& kname) { + kq_require_kernel(d, std::string(lever) + "=" + kq_tile_string(t), kname); +} + +// One KQ_QMM_LOG line: kernel name, shape and grid of a qmm tile dispatch. +static void kq_log_qmm( + const std::string& kname, + int M, + int N, + int K, + const MTL::Size& grid) { + if (kq_qmm_log()) { + std::fprintf( + stderr, + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + kname.c_str(), + M, + N, + K, + static_cast(grid.width), + static_cast(grid.height), + static_cast(grid.depth)); + } +} + // KQ_VERIFY_EXT, read once per process. 0 forces verify_qmv and 1 forces // mv_ext on the mat-vec band (M 2-12), -1 leaves the per-codec default. static int kq_verify_ext_env() { - static const int v = []() { - const char* e = std::getenv("KQ_VERIFY_EXT"); - return e != nullptr ? std::atoi(e) : -1; - }(); + static const int v = kq_env_int("KQ_VERIFY_EXT", -1); return v; } // KQ_QMM_SPLITK, read once per process. A target above 1 forces the ALU // split-K route for M <= 32, 0 disables it, -1 leaves the default. static int kq_qmm_splitk_env() { - static const int v = []() { - const char* e = std::getenv("KQ_QMM_SPLITK"); - return e != nullptr ? std::atoi(e) : -1; - }(); + static const int v = kq_env_int("KQ_QMM_SPLITK", -1); return v; } @@ -572,10 +878,7 @@ void qmm_nax( // are db_safe and have smallbm instantiations only. KQ_NAX_SMALL_BM: // 0 = off, 2 = force bm32 for policy-excluded codecs (crossover // finding), unset or 1 = per-codec policy. - static const int small_bm_mode = []() { - const char* e = std::getenv("KQ_NAX_SMALL_BM"); - return e == nullptr ? 1 : std::atoi(e); - }(); + const int small_bm_mode = kq_nax_small_bm_mode(); // Window ends at 32: above it grid.y goes to 2 and every weight column // tile streams once per M-tile, which halves per-weight bandwidth at // DRAM-bound (measured: q6_k M48 146 GB/s on bm32 vs 184 on bm64). @@ -595,18 +898,40 @@ void qmm_nax( // policy table). KQ_NAX_BM128: 0 = off, 1 = force floor 193 for all // codecs (crossover finding), 2 = drop the floor entirely (probes the // M65-128 wash band, symmetric with KQ_NAX_SMALL_BM=2), unset = - // per-codec policy. - static const int bm128_mode = []() { - const char* e = std::getenv("KQ_NAX_BM128"); - return e == nullptr ? -1 : std::atoi(e); - }(); - int bm128_min_m = kq_smallbm_policy(kquant_type).bm128_min_m; - if (bm128_mode != 0 && transpose && ((M + 63) / 64) % 2 == 0 && + // per-codec policy. The loss cells predate the row-live guard. From the + // codec's large-M row (min_m, min_n) the row's tile runs at either + // parity, and above the row's K ceiling the BM=64 tile runs instead; + // KQ_NAX_BM128_ODD_M replaces min_m, and 0 keeps the parity rule at + // every M. + const int bm128_mode = kq_nax_bm128_mode(); + const KqSmallBmPolicy policy = kq_smallbm_policy(kquant_type); + const KqLargeM& large = policy.large; + const int odd_m_env = kq_nax_bm128_odd_m(); + const int bm128_odd_m = odd_m_env >= 0 ? odd_m_env : large.min_m; + const bool large_rows = transpose && bm128_mode != 0 && bm128_odd_m > 0 && + M >= bm128_odd_m && N >= large.min_n; + const bool bm128_rows = ((M + 63) / 64) % 2 == 0 || large_rows; + if (bm128_mode != 0 && transpose && bm128_rows && (bm128_mode == 1 ? M >= 193 : bm128_mode == 2 ? true - : M >= bm128_min_m)) { + : M >= policy.bm128_min_m)) { bm = 128; } + if (large_rows) { + bm = large.max_k > 0 && K > large.max_k ? 64 : large.bm; + bn = large.bn; + wm = large.wm; + wn = large.wn; + } + // KQ_NAX_TILE replaces the tile the policy above picked. + const std::optional tile = + transpose ? kq_nax_tile() : std::optional(); + if (tile) { + bm = tile->bm; + bn = tile->bn; + wm = tile->wm; + wn = tile->wn; + } // M33-64 band: the name-suffixed _db BM=64 variant double-buffers Ws // (bm32 here would stream weights twice via grid.y=2; blanket DB@64 // regressed M96+ -3-15% and prefill -3-7%, so it is band-gated). @@ -615,14 +940,20 @@ void qmm_nax( // (db64_min_n > 0) carry _db instantiations, so 1 is bounded by // availability; probing another codec needs its instantiation // restored and a metallib rebuild. - static const int db64_mode = []() { - const char* e = std::getenv("KQ_NAX_DB64"); - return e == nullptr ? -1 : std::atoi(e); - }(); + const int db64_mode = kq_nax_db64_mode(); int db64_min_n = kq_smallbm_policy(kquant_type).db64_min_n; - bool use_db64 = transpose && bm == 64 && M >= 33 && M <= 64 && - db64_min_n > 0 && + bool use_db64 = transpose && bm == 64 && bn == 64 && wm == 2 && wn == 2 && + M >= 33 && M <= 64 && db64_min_n > 0 && (db64_mode == 1 || (db64_mode == -1 && N >= db64_min_n)); + // Codecs with a BK=128 lever on the BM=64 and BM=128 tiles walk K 128 + // weights per step: the k-quant half loaders take one half super-block per + // tile, where a code byte that spans 128 weights is read once instead of on + // two or four tiles, and the pair loaders run two BK=64 loaders per tile. + // Bit-identical output. A 32-weight block codec needs K % 128 == 0. + if (transpose && !use_db64 && bn == 64 && wm == 2 && wn == 2 && + (bm == 64 || bm == 128) && K % 128 == 0 && kq_nax_bk128_on(kquant_type)) { + bk = 128; + } MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); // Row-tile traversal swizzle (M > 64): fold row-tiles into grid.x so @@ -630,19 +961,23 @@ void qmm_nax( // row-tiles 2..2^log read it from SLC instead of re-streaming DRAM // (weight-normalized bandwidth otherwise divides by the row-tile // count). kq_swizzle_tid in the kernel derives the fold factor from - // threadgroups_per_grid; padded row-tiles early-return. - // KQ_NAX_SWIZZLE: 1 = on (experiment lever; default off). - static const bool swizzle_on = []() { - const char* e = std::getenv("KQ_NAX_SWIZZLE"); - return e != nullptr && std::atoi(e) != 0; - }(); - if (swizzle_on && transpose && bm == 64) { - int y_tiles = (M + bm - 1) / bm; - if (y_tiles >= 2) { - int log = y_tiles >= 4 ? 2 : 1; + // threadgroups_per_grid; padded row-tiles early-return. The fold is + // capped at the row-tile count. KQ_NAX_SWIZZLE sets the log; unset, the + // large-M row's swizzle_log applies on the row's own tile (not under + // KQ_NAX_TILE). + const int swizzle_env = kq_nax_swizzle_log(); + int swizzle_log = swizzle_env >= 0 ? swizzle_env + : (large_rows && !tile) ? large.swizzle_log + : 0; + if (swizzle_log > 0 && transpose && (bm == 64 || bm == 128)) { + const int y_tiles = (M + bm - 1) / bm; + while (swizzle_log > 0 && (1 << swizzle_log) > y_tiles) { + swizzle_log--; + } + if (swizzle_log > 0) { grid_dims = MTL::Size( - static_cast((N + bn - 1) / bn) << log, - (y_tiles + (1 << log) - 1) >> log, + static_cast((N + bn - 1) / bn) << swizzle_log, + (y_tiles + (1 << swizzle_log) - 1) >> swizzle_log, B); } } @@ -674,6 +1009,11 @@ void qmm_nax( batched ? "_batch_1" : "_batch_0", use_db64 ? "_db" : ""); + if (tile) { + kq_require_tile_kernel(d, "KQ_NAX_TILE", *tile, kname); + } + kq_log_qmm(kname, M, N, K, grid_dims); + auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -731,6 +1071,22 @@ void qmm( int wm = 2, wn = 2; int bm = 64; int bn = transpose ? 64 : 32; + // KQ_QMM_TILE replaces the tile. The shipped bm64 bn64 wm2 wn2 tile keeps + // its plain kernel name, and any other tile carries its geometry in the + // name. + const std::optional tile = + transpose ? kq_qmm_tile() : std::optional(); + std::string tile_tag; + if (tile) { + bm = tile->bm; + bn = tile->bn; + wm = tile->wm; + wn = tile->wn; + if (bm != 64 || bn != 64 || wm != 2 || wn != 2) { + tile_tag = "_bm" + std::to_string(bm) + "_bn" + std::to_string(bn) + + "_wm" + std::to_string(wm) + "_wn" + std::to_string(wn); + } + } MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); @@ -747,14 +1103,14 @@ void qmm( group_size, "_b_", bits, + tile_tag, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", batched ? "_batch_1" : "_batch_0"); - static const bool qmm_log = std::getenv("KQ_QMM_LOG") != nullptr; - if (qmm_log) { - std::fprintf( - stderr, "[kq qmm] %s M=%d N=%d K=%d\n", kname.c_str(), M, N, K); + if (tile) { + kq_require_tile_kernel(d, "KQ_QMM_TILE", *tile, kname); } + kq_log_qmm(kname, M, N, K, grid_dims); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -772,16 +1128,162 @@ void qmm( ce.dispatch_threadgroups(grid_dims, group_dims); } +// qmm_splitk tile. The bm16 tile halves the MMA row-padding waste at +// small M and widens BN to 64 so each K-step does more MMA work per +// barrier; the bm8 tile halves it again for M <= 8 (one simdgroup along M, +// four along N). On M3 Max the bm16 tile is bound per weight, not per byte +// (q8_0 ~240 GB/s and q6_k ~180 GB/s both ~220 G weights/s), and the BM32 +// tile costs 1.45x per weight for 2x the MMA, so MMA issue is ~45% of the +// tile. KQ_QMM_SPLITK_BM8=0 keeps M <= 8 on the bm16 tile. +struct KqSplitkTile { + int bm; + int bn; +}; +static KqSplitkTile kq_splitk_tile(int M, const std::string& t) { + static const bool bm8_enabled = kq_env_int("KQ_QMM_SPLITK_BM8", 1) != 0; + const bool small = M <= 16 && kq_splitk_codec(t); + const bool bm8 = small && bm8_enabled && M <= 8; + return {bm8 ? 8 : (small ? 16 : 32), small ? 64 : 32}; +} + +// NAX split-K tile rows: BM=32 through M 48, where the last row tile of +// the mid-M band holds at most 16 rows, and BM=64 above. +static int kq_nax_splitk_bm(int M) { + return M <= 48 ? 32 : 64; +} + +// Kernel names of the split-K routes and their fold. The dispatchers and +// split_kernel_names build the names here, so a test of the names that +// split_kernel_names lists covers every name a dispatch can ask for. +static std::string kq_qmm_splitk_kname( + const std::string& t, + const std::string& ts, + int gs, + int bits, + KqSplitkTile tile, + bool aligned, + const std::string& suffix) { + std::string kname; + kname.reserve(64); + mx::concatenate( + kname, + kq_kname_prefix(t) + + (tile.bm == 8 + ? "qmm_t_splitk_bm8_" + : (tile.bn == 64 ? "qmm_t_splitk_bm16_" : "qmm_t_splitk_")), + ts, + "_gs_", + gs, + "_b_", + bits, + aligned ? "_alN_true" : "_alN_false", + suffix); + return kname; +} + +static std::string kq_qmm_nax_splitk_kname( + const std::string& t, + const std::string& ts, + int gs, + int bits, + int bm, + bool aligned, + const std::string& suffix) { + std::string kname; + kname.reserve(80); + mx::concatenate( + kname, + kq_kname_prefix(t) + "qmm_t_nax_splitk_", + ts, + "_gs_", + gs, + "_b_", + bits, + "_bm", + bm, + "_bn64_bk64_wm2_wn2", + aligned ? "_alN_true" : "_alN_false", + suffix); + return kname; +} + +// verify_mma codecs with half kernels; ptq1_0 runs as fast in float and +// has none. q2_0 has them as pq2_0 does. +static bool kq_verify_mma_has_half(const std::string& t) { + return t == "pq2_0" || t == "q2_0" || t == "q4_0" || t == "q8_0"; +} + +// KQ_VERIFY_MMA_HALF=1 (probe lever, read live) runs the half kernels on +// the codecs that have them, for timing against the float kernels. They +// stage the activations as half and sum each block in half, so their range +// stops at 65504, which a pq2_0 block sum of unscaled codes reaches at +// activations of 256. +static bool kq_verify_mma_half(const std::string& t) { + return kq_verify_mma_has_half(t) && kq_env_int("KQ_VERIFY_MMA_HALF", 0) == 1; +} + +// half selects the KQ_VERIFY_MMA_HALF kernels. +static std::string kq_verify_mma_kname( + const std::string& t, + const std::string& ts, + int gs, + int bits, + bool half, + const std::string& suffix) { + std::string kname; + kname.reserve(64); + mx::concatenate( + kname, + kq_kname_prefix(t) + "verify_mma_", + ts, + "_gs_", + gs, + "_b_", + bits, + half ? "_h" : "", + suffix); + return kname; +} + +// kstep 256 on q4_0 selects its eight-block kernel. +static std::string kq_verify_nax_kname( + const std::string& t, + const std::string& ts, + int gs, + int bits, + int kstep, + const std::string& suffix) { + std::string kname; + kname.reserve(64); + mx::concatenate( + kname, + (t == "q4_0" && kstep == 256 ? "kquant_q4_0_sb_" : kq_kname_prefix(t)) + + "verify_nax_", + ts, + "_gs_", + gs, + "_b_", + bits, + suffix); + return kname; +} + +static std::string kq_splitk_accum_kname( + const std::string& ts, + const std::string& suffix) { + return "kquant_qmm_splitk_accum_" + ts + suffix; +} + // Split-K qmm_t for the small-M decode band and the mid-M band on non-NAX // GPUs (see kq_midm_splitk). The plain tile grid is only // ceil(N/64) x 1 threadgroups at decode shapes with M <= 32 (84 at // [5376 x 21504]) and the in-tile K walk serializes, capping qmm/NAX at // 160-257 GB/s while mv_ext decays past M~4 on L2 activation re-reads. // Partitioning K across grid.z multiplies threadgroup count by `splits`; -// each slice writes a T partial tile and a second tiny pass folds them in -// f32. Slice starts must land on wire-block boundaries, so the partition -// is a multiple of group_size (the caller guarantees splits divides -// K / group_size). Non-batched transpose shapes only. +// each slice writes a partial tile, float or T by kq_splitk_partials, and +// a second tiny pass folds them in f32. Slice starts must land on wire-block +// boundaries, so the partition is a multiple of group_size (the caller +// guarantees splits divides K / group_size). Non-batched transpose shapes only. void qmm_splitk( const array& x, const array& w, @@ -797,45 +1299,22 @@ void qmm_splitk( const Stream& s, const std::string& kquant_type) { constexpr int wm = 2, wn = 2; - // The bm16 tile halves the MMA row-padding waste at small M and - // widens BN to 64 so each K-step does more MMA work per barrier; the - // bm8 tile halves it again for M <= 8 (one simdgroup along M, four - // along N). On M3 Max the bm16 tile is bound per weight, not per - // byte (q8_0 ~240 GB/s and q6_k ~180 GB/s both ~220 G weights/s), and - // the BM32 tile costs 1.45x per weight for 2x the MMA, so MMA issue is - // ~45% of the tile. KQ_QMM_SPLITK_BM8=0 keeps M <= 8 on the bm16 tile. - static const bool bm8_enabled = []() { - const char* e = std::getenv("KQ_QMM_SPLITK_BM8"); - return e == nullptr || std::atoi(e) != 0; - }(); - const bool small = M <= 16 && kq_splitk_codec(kquant_type); - const bool bm8 = small && bm8_enabled && M <= 8; - const int bm = bm8 ? 8 : (small ? 16 : 32); - const int bn = small ? 64 : 32; + const KqSplitkTile tile = kq_splitk_tile(M, kquant_type); + const int bm = tile.bm; + const int bn = tile.bn; const int k_partition = (K / group_size / splits) * group_size; const int part_stride = M * N; + const KqSplitkPartials pp = kq_splitk_partials(x.dtype()); - array partials({splits, M, N}, x.dtype(), nullptr, {}); + array partials({splits, M, N}, pp.dtype, nullptr, {}); partials.set_data(mx::allocator::malloc(partials.nbytes())); auto& ce = mx::metal::get_command_encoder(s); ce.add_temporary(partials); std::string type_string = kq_type_string(x.dtype()); - bool aligned = N % bn == 0; - std::string kname; - kname.reserve(64); - mx::concatenate( - kname, - kq_kname_prefix(kquant_type) + - (bm8 ? "qmm_t_splitk_bm8_" - : (small ? "qmm_t_splitk_bm16_" : "qmm_t_splitk_")), - type_string, - "_gs_", - group_size, - "_b_", - bits, - aligned ? "_alN_true" : "_alN_false"); + const std::string kname = kq_qmm_splitk_kname( + kquant_type, type_string, group_size, bits, tile, N % bn == 0, pp.suffix); auto kernel = kq_get_kernel(d, kname); ce.set_compute_pipeline_state(kernel); @@ -852,9 +1331,10 @@ void qmm_splitk( ce.set_bytes(part_stride, c++); MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, splits); + kq_log_qmm(kname, M, N, K, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); - std::string aname = "kquant_qmm_splitk_accum_" + type_string; + const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); auto accum = kq_get_kernel(d, aname); ce.set_compute_pipeline_state(accum); const int n_elems = M * N; @@ -875,7 +1355,9 @@ void qmm_splitk( // once per K chunk, so the weight read is amortized over all M rows // without the tile's threadgroup staging of the weights. Split-K over // grid z with qmm_splitk's partial/fold shape; slices are whole wire -// blocks. Non-batched transpose shapes, bf16 or f16 activations. +// blocks. Non-batched transpose shapes, bf16 or f16 activations, which +// stay in their dtype while each block sums in float. kq_verify_mma_half +// selects the half kernels. void verify_mma( const array& x, const array& w, @@ -897,23 +1379,22 @@ void verify_mma( auto& ce = mx::metal::get_command_encoder(s); // One split writes out directly; the partials and the fold exist only // for splits > 1. - array partials({splits, M, N}, x.dtype(), nullptr, {}); + const KqSplitkPartials pp = splits > 1 ? kq_splitk_partials(x.dtype()) + : KqSplitkPartials{x.dtype(), ""}; + array partials({splits, M, N}, pp.dtype, nullptr, {}); if (splits > 1) { partials.set_data(mx::allocator::malloc(partials.nbytes())); ce.add_temporary(partials); } std::string type_string = kq_type_string(x.dtype()); - std::string kname; - kname.reserve(64); - mx::concatenate( - kname, - kq_kname_prefix(kquant_type) + "verify_mma_", + const std::string kname = kq_verify_mma_kname( + kquant_type, type_string, - "_gs_", group_size, - "_b_", - bits); + bits, + kq_verify_mma_half(kquant_type), + pp.suffix); auto kernel = kq_get_kernel(d, kname); ce.set_compute_pipeline_state(kernel); @@ -930,12 +1411,13 @@ void verify_mma( ce.set_bytes(part_stride, c++); MTL::Size group_dims(256, 1, 1); MTL::Size grid_dims((N + rows - 1) / rows, 1, splits); + kq_log_qmm(kname, M, N, K, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; } - std::string aname = "kquant_qmm_splitk_accum_" + type_string; + const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); auto accum = kq_get_kernel(d, aname); ce.set_compute_pipeline_state(accum); const int n_elems = M * N; @@ -956,8 +1438,7 @@ void verify_mma( // the two-block kernel (probe lever, read live per call). static int kq_verify_nax_call_kstep(int K, const std::string& t) { if (t == "q4_0" && K % 256 == 0) { - const char* e = std::getenv("KQ_VERIFY_NAX_Q4_0_SB"); - if (e == nullptr || std::atoi(e) != 0) { + if (kq_env_int("KQ_VERIFY_NAX_Q4_0_SB", 1) != 0) { return 256; } } @@ -1010,25 +1491,17 @@ void verify_nax( const int part_stride = M * N; auto& ce = mx::metal::get_command_encoder(s); - array partials({splits, M, N}, x.dtype(), nullptr, {}); + const KqSplitkPartials pp = splits > 1 ? kq_splitk_partials(x.dtype()) + : KqSplitkPartials{x.dtype(), ""}; + array partials({splits, M, N}, pp.dtype, nullptr, {}); if (splits > 1) { partials.set_data(mx::allocator::malloc(partials.nbytes())); ce.add_temporary(partials); } std::string type_string = kq_type_string(x.dtype()); - std::string kname; - kname.reserve(64); - mx::concatenate( - kname, - (kquant_type == "q4_0" && kstep == 256 ? "kquant_q4_0_sb_" - : kq_kname_prefix(kquant_type)) + - "verify_nax_", - type_string, - "_gs_", - group_size, - "_b_", - bits); + const std::string kname = kq_verify_nax_kname( + kquant_type, type_string, group_size, bits, kstep, pp.suffix); auto kernel = kq_get_kernel(d, kname); ce.set_compute_pipeline_state(kernel); @@ -1045,12 +1518,13 @@ void verify_nax( ce.set_bytes(part_stride, c++); MTL::Size group_dims(128, 1, 1); MTL::Size grid_dims((N + 127) / 128, 1, splits); + kq_log_qmm(kname, M, N, K, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; } - std::string aname = "kquant_qmm_splitk_accum_" + type_string; + const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); auto accum = kq_get_kernel(d, aname); ce.set_compute_pipeline_state(accum); const int n_elems = M * N; @@ -1087,34 +1561,22 @@ void qmm_nax_splitk( Device& d, const Stream& s, const std::string& kquant_type, - int k_partition, - int bm = 32) { + int k_partition) { + const int bm = kq_nax_splitk_bm(M); constexpr int bn = 64; constexpr int wm = 2, wn = 2; const int part_stride = M * N; + const KqSplitkPartials pp = kq_splitk_partials(x.dtype()); - array partials({splits, M, N}, x.dtype(), nullptr, {}); + array partials({splits, M, N}, pp.dtype, nullptr, {}); partials.set_data(mx::allocator::malloc(partials.nbytes())); auto& ce = mx::metal::get_command_encoder(s); ce.add_temporary(partials); std::string type_string = kq_type_string(x.dtype()); - bool aligned = N % bn == 0; - std::string kname; - kname.reserve(80); - mx::concatenate( - kname, - kq_kname_prefix(kquant_type) + "qmm_t_nax_splitk_", - type_string, - "_gs_", - group_size, - "_b_", - bits, - "_bm", - bm, - "_bn64_bk64_wm2_wn2", - aligned ? "_alN_true" : "_alN_false"); + const std::string kname = kq_qmm_nax_splitk_kname( + kquant_type, type_string, group_size, bits, bm, N % bn == 0, pp.suffix); auto kernel = kq_get_kernel(d, kname); ce.set_compute_pipeline_state(kernel); @@ -1131,9 +1593,10 @@ void qmm_nax_splitk( ce.set_bytes(part_stride, c++); MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, splits); + kq_log_qmm(kname, M, N, K, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); - std::string aname = "kquant_qmm_splitk_accum_" + type_string; + const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); auto accum = kq_get_kernel(d, aname); ce.set_compute_pipeline_state(accum); const int n_elems = M * N; @@ -1279,13 +1742,11 @@ void qmv( // forces fine, =0 forces coarse (A/B lever). Read live (KQ_DISABLE_NAX // precedent) so one process can interleave coarse/fine for // thermally-paired A/Bs; getenv cost is noise per dispatch. - const char* qmv_fine_e = std::getenv("KQ_QMV_FINE"); - const int qmv_fine_env = qmv_fine_e != nullptr ? std::atoi(qmv_fine_e) : -1; + const int qmv_fine_env = kq_env_int("KQ_QMV_FINE", -1); // The split-K kernel takes M == 1, B == 1 unless KQ_QMV_SPLITK=0 or // KQ_QMV_FINE names a per-row tiling. Read live like KQ_QMV_FINE. - const char* splitk_e = std::getenv("KQ_QMV_SPLITK"); - const bool splitk_on = splitk_e == nullptr || std::atoi(splitk_e) != 0; + const bool splitk_on = kq_env_int("KQ_QMV_SPLITK", 1) != 0; if (M == 1 && B == 1 && splitk_on && qmv_fine_env == -1 && qmv_splitk( x, w, scales, out, group_size, bits, N, K, d, s, kquant_type)) { @@ -1357,8 +1818,7 @@ void qmv_bias( // Same fine-tiling policy as qmv() above; M == 1 is the shape contract // here, so only the codec ceiling gates the default. - const char* qmv_fine_e = std::getenv("KQ_QMV_FINE"); - const int qmv_fine_env = qmv_fine_e != nullptr ? std::atoi(qmv_fine_e) : -1; + const int qmv_fine_env = kq_env_int("KQ_QMV_FINE", -1); const bool use_fine = codec_has_qmv_fine(kquant_type) && (qmv_fine_env == 1 || (qmv_fine_env != 0 && N <= kquant_qmv_fine_max_n(kquant_type))); @@ -1433,10 +1893,8 @@ void verify_qmv( // default. Only q8_0 and q6_k have the fine kernel instantiated. // KQ_VERIFY_FINE=1 forces fine for both, KQ_VERIFY_FINE=0 forces coarse for // both (A/B lever). - static const int verify_fine = []() { - const char* e = std::getenv("KQ_VERIFY_FINE"); - return e != nullptr ? std::atoi(e) : -1; // -1 = per-codec default - }(); + // -1 = per-codec default. + static const int verify_fine = kq_env_int("KQ_VERIFY_FINE", -1); bool codec_has_fine = (kquant_type == "q8_0" || kquant_type == "q6_k"); bool default_fine = (kquant_type == "q8_0"); // q6_k neutral -> coarse default bool use_fine = codec_has_fine && @@ -1628,16 +2086,106 @@ static bool kq_midm_splitk(int M, int N) { // the NAX tile weighted by calls per forward, and no cell loses more than // 2%. At four or five per core, N 10240 at M 33-64 (one full wave) loses // up to 1.6x, and M 97-128 wins or loses by codec. -static int kq_midm_nax_target(int M, int N) { +// q4_0 and q8_0 cut K in units of 64 weights, so at K 5120 these targets +// write 16 or 8 partials where the 256-weight codecs write 10 or 5. With +// float partials (f32) that traffic costs more than the added threadgroups +// gain, and the two codecs target 10 from M 33 to 64 and 5 from M 65 to 96. +// On an M5 Max over the Qwen3.8-27B projections on float16 activations, +// per cell against 16 and 8: 0.79-1.01x the time at M 33-64, and +// 0.67-1.11x at M 65-96, where every cell above 1.03x has K 6144 or 17408. +// There a target of 5 leaves 4 slices (96 and 272 units), so above K 5120 +// the target from M 65 to 96 stays 8: 0.88-0.97x the time of 4 slices, and +// 0.99-1.02x at K 6144 M 96, where the CI spans 1. Measured at K 5120, 6144 +// and 17408 only. bfloat16 with KQ_SPLITK_F32=1 takes the same targets, not +// measured there. +static int +kq_midm_nax_target(int M, int N, int K, const std::string& t, bool f32) { + const int cores = gpu_core_count(); + const int grid = ((N + 63) / 64) * ((M + 63) / 64); + int target = 0; if (kq_midm_splitk(M, N)) { + target = kq_splitk_nax_target; + } else if (M >= 65 && M <= 96 && grid < 16 * (cores > 0 ? cores : 40)) { + target = 8; + } + if (target > 0 && f32 && M <= 96 && (t == "q4_0" || t == "q8_0")) { + return M <= 64 ? 10 : (K > 5120 ? 8 : 5); + } + return target; +} + +// KQ_SPLITK_CAP= (probe lever, read live) caps the default slice target +// of the small-M split-K routes (NAX split-K and the ALU split route, M <= +// 32) so the split grid holds at most n threadgroups per GPU core +// (KQ_GPU_CORES sets the count). A grid that already holds n per core +// takes one slice, and the call runs what it runs with the split route +// off: from M 13 usually the qmm tile (BM=32 on NAX GPUs), or per-row qmv +// below get_qmv_batch_limit, and at M 9 to 12 the small-M qmm route, +// mv_ext or per-row qmv. verify_mma and verify_nax are not capped. 0 keeps +// the fixed target of 16 at every N on both routes. Unset keeps it on the +// ALU route and takes kq_splitk_nax_default on the NAX route. A target +// forced through KQ_QMM_SPLITK_NAX or KQ_QMM_SPLITK is not capped. +static int kq_splitk_cap_env() { + return kq_env_int("KQ_SPLITK_CAP", -1); +} + +static int kq_splitk_capped(int target, int grid, int cap) { + if (cap <= 0) { + return target; + } + const int cores = gpu_core_count(); + const int64_t room = + int64_t(cap) * (cores > 0 ? cores : 40) / std::max(grid, 1); + return static_cast( + std::max(1, std::min(target, room))); +} + +// True when a call of M <= 32 rows with the split route off runs the +// BM=32 NAX tile: the codec takes that tile, and below M 13 it also takes +// the small-M qmm route at this M. Otherwise the call runs an mv route or +// the BM=64 tile. +static bool kq_unsplit_takes_bm32(int M, int N, int K, const std::string& t) { + const int mode = kq_nax_small_bm_mode(); + if (mode == 0 || (mode != 2 && !kq_smallbm_policy(t).bm32)) { + return false; + } + if (M > 12) { + return true; + } + const int m = kq_smallm_route_min(t, N, K); + return m > 0 && M >= m; +} + +// Default slice target of the NAX split-K route at M <= 32. A grid that +// already holds 16 threadgroups per core, such as a vocab head, takes one +// slice where the call without the split runs the same BM=32 tile, except +// q6_k and iq2_xxs below M 32, where the split still wins. Where that call +// would leave the tile for an mv route or the BM=64 tile, the split stays. +// Below a full grid, q4_0 and q8_0 on float partials hold the split grid +// to 16 threadgroups per core, and the rest keep the target of 16. q8_0 +// below M 13 keeps the target on every grid, where both the single slice +// and the float-partial cap measured slower. +static int kq_splitk_nax_default( + int M, + int N, + int K, + int grid, + const std::string& t, + bool f32) { + if (t == "q8_0" && M <= 12) { return kq_splitk_nax_target; } const int cores = gpu_core_count(); - const int grid = ((N + 63) / 64) * ((M + 63) / 64); - if (M >= 65 && M <= 96 && grid < 16 * (cores > 0 ? cores : 40)) { - return 8; + const int64_t full = int64_t(16) * (cores > 0 ? cores : 40); + if (grid >= full) { + const bool keep = (M < 32 && (t == "q6_k" || t == "iq2_xxs")) || + !kq_unsplit_takes_bm32(M, N, K, t); + return keep ? kq_splitk_nax_target : 1; } - return 0; + if (f32 && (t == "q4_0" || t == "q8_0")) { + return kq_splitk_capped(kq_splitk_nax_target, grid, 16); + } + return kq_splitk_nax_target; } // Largest divisor of nblk at or under target (1 = no split). @@ -1666,8 +2214,7 @@ static int kq_nax_split(int target, int nblk, int& per) { } const int sp = kq_split_count(target, nblk); per = nblk / sp; - const char* e = std::getenv("KQ_SPLITK_RAGGED"); - const int mode = e != nullptr ? std::atoi(e) : 1; + const int mode = kq_env_int("KQ_SPLITK_RAGGED", 1); if (mode == 0) { return sp; } @@ -1683,8 +2230,13 @@ static int kq_nax_split(int target, int nblk, int& per) { // verify_mma split count. Split-K fills the grid for the projection // shapes; a head-sized N already has more threadgroups than the GPU holds, -// and there the partials only cost traffic. +// and there the partials only cost traffic. KQ_VERIFY_MMA_SPLITS= sets +// the target instead (probe lever, read live; 1 runs unsplit). static int kq_verify_mma_splits(int N, int K, int group_size, int rows) { + const int env = kq_env_int("KQ_VERIFY_MMA_SPLITS", 0); + if (env > 0) { + return kq_split_count(env, K / group_size); + } return (N + rows - 1) / rows >= 1024 ? 1 : kq_split_count(16, K / group_size); } @@ -1880,6 +2432,27 @@ bool kq_forced_route( } // namespace +int kq_nax_large_max_k(const std::string& kquant_type) { + return kq_largem_max_k(kquant_type); +} + +bool kq_qmm_tile_pinned() { + return kq_qmm_tile().has_value(); +} + +static int kq_read_dequant_route() { + return kq_env_choice("KQ_DEQUANT_ROUTE", {"auto", "off", "force"}, 0); +} + +int kq_dequant_route_mode() { + static const int cached = kq_read_dequant_route(); + return kq_tune_live() ? kq_read_dequant_route() : cached; +} + +bool kq_qmm_log_on() { + return kq_qmm_log(); +} + #endif // _METAL_ std::vector KQuantMatmul::output_shapes( @@ -1914,6 +2487,11 @@ void KQuantMatmul::eval_cpu( out.set_data(mx::allocator::malloc(out.nbytes())); auto& encoder = mx::cpu::get_command_encoder(stream()); + // An empty output has nothing to compute. The op returns zeros for an + // inner dim of 0, so the row size below is not 0. + if (out.size() == 0) { + return; + } encoder.set_input_array(x); encoder.set_input_array(w); encoder.set_output_array(out); @@ -2048,6 +2626,9 @@ void KQuantMatmul::eval_gpu_base( const auto& scales = inputs[2]; kq_check_weight_base(w, kquant_type_, "quantized_matmul"); out.set_data(mx::allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } bool non_batched = w.ndim() == 2 && x.flags().row_contiguous; int K = x.shape(-1); @@ -2152,15 +2733,11 @@ void KQuantMatmul::eval_gpu_base( if (transpose_ && non_batched && M >= 2 && M <= 32 && x.dtype() != mx::float32 && kq_is_nax_available() && codec_has_nax(kquant_type_)) { - const char* qmv_e = std::getenv("KQ_NAX_QMV"); - int qmv_max_m = qmv_e != nullptr ? std::atoi(qmv_e) : -1; + int qmv_max_m = kq_env_int("KQ_NAX_QMV", -1); if (qmv_max_m == -1) { - const char* vnax_f = std::getenv("KQ_VERIFY_NAX"); - const int vnax_force = vnax_f != nullptr ? std::atoi(vnax_f) : -1; - const char* vmma_f = std::getenv("KQ_VERIFY_MMA"); - const int vmma_force = vmma_f != nullptr ? std::atoi(vmma_f) : -1; - const char* sk_f = std::getenv("KQ_QMM_SPLITK_NAX"); - const int sk_force = sk_f != nullptr ? std::atoi(sk_f) : -1; + const int vnax_force = kq_env_int("KQ_VERIFY_NAX", -1); + const int vmma_force = kq_env_int("KQ_VERIFY_MMA", -1); + const int sk_force = kq_env_int("KQ_QMM_SPLITK_NAX", -1); const int ext = kq_verify_ext_env(); const bool forced = (vnax_force >= 1 && M >= vnax_force && codec_verify_nax_kstep(kquant_type_) > 0) || @@ -2186,12 +2763,9 @@ void KQuantMatmul::eval_gpu_base( if (transpose_ && non_batched && M <= 8 && x.dtype() != mx::float32 && kq_is_nax_available() && codec_has_nax(kquant_type_)) { const int kstep = kq_verify_nax_call_kstep(K, kquant_type_); - const char* vnax_e = std::getenv("KQ_VERIFY_NAX"); - const int vnax_env = vnax_e != nullptr ? std::atoi(vnax_e) : -1; - const char* vmma_f = std::getenv("KQ_VERIFY_MMA"); - const int vmma_force = vmma_f != nullptr ? std::atoi(vmma_f) : -1; - const char* sk_f = std::getenv("KQ_QMM_SPLITK_NAX"); - const int sk_force = sk_f != nullptr ? std::atoi(sk_f) : -1; + const int vnax_env = kq_env_int("KQ_VERIFY_NAX", -1); + const int vmma_force = kq_env_int("KQ_VERIFY_MMA", -1); + const int sk_force = kq_env_int("KQ_QMM_SPLITK_NAX", -1); const bool other_forced = (vmma_force >= 2 && M >= vmma_force && codec_verify_mma_rows(kquant_type_) > 0) || sk_force >= 1; @@ -2228,8 +2802,7 @@ void KQuantMatmul::eval_gpu_base( // KQ_VERIFY_MMA= forces the entry at M >= m, 0 or 1 disables; read // live so an A/B can flip arms on one generator. Decided ahead of NAX // split-K, whose default entry yields to it. - const char* vmma_e = std::getenv("KQ_VERIFY_MMA"); - const int vmma_env = vmma_e != nullptr ? std::atoi(vmma_e) : -1; + const int vmma_env = kq_env_int("KQ_VERIFY_MMA", -1); const int vmma_rows = codec_verify_mma_rows(kquant_type_); int vmma_min_m = 0; if (vmma_env >= 2) { @@ -2250,8 +2823,7 @@ void KQuantMatmul::eval_gpu_base( // on one generator. The default entry yields to verify_mma, to a // forced KQ_QMM_SPLITK, and to a set KQ_VERIFY_EXT through M 12, so // those arms time the route they name. - const char* sk_nax_e = std::getenv("KQ_QMM_SPLITK_NAX"); - const int sk_nax_env = sk_nax_e != nullptr ? std::atoi(sk_nax_e) : -1; + const int sk_nax_env = kq_env_int("KQ_QMM_SPLITK_NAX", -1); const int sk_nax_min_m = kq_splitk_nax_min_m(kquant_type_, N); const int ext_env = kq_verify_ext_env(); const bool sk_nax_yields = vmma_route || kq_qmm_splitk_env() > 1 || @@ -2262,8 +2834,18 @@ void KQuantMatmul::eval_gpu_base( if (sk_nax_route && transpose_ && non_batched && M <= 32 && codec_has_nax(kquant_type_) && kq_is_nax_available() && (K % 64 == 0) && x.dtype() != mx::float32) { - const int qmm_splitk_nax_env = - sk_nax_env > 1 ? sk_nax_env : kq_splitk_nax_target; + // The BM32 tile's grid, ceil(N/64) x ceil(M/32), sizes the cap. + const int grid = ((N + 63) / 64) * ((M + 31) / 32); + const int cap = kq_splitk_cap_env(); + const int qmm_splitk_nax_env = sk_nax_env > 1 ? sk_nax_env + : cap >= 0 ? kq_splitk_capped(kq_splitk_nax_target, grid, cap) + : kq_splitk_nax_default( + M, + N, + K, + grid, + kquant_type_, + kq_splitk_partials(x.dtype()).dtype == mx::float32); const int q = std::max(group_size_, 64); int per = 0; const int sp = kq_nax_split(qmm_splitk_nax_env, K / q, per); @@ -2311,7 +2893,12 @@ void KQuantMatmul::eval_gpu_base( !nax_path); if (splitk_route && transpose_ && non_batched && M <= 32 && kq_splitk_codec(kquant_type_)) { - const int splitk_target = qmm_splitk_env > 1 ? qmm_splitk_env : 16; + const KqSplitkTile tile = kq_splitk_tile(M, kquant_type_); + const int grid = + ((N + tile.bn - 1) / tile.bn) * ((M + tile.bm - 1) / tile.bm); + const int splitk_target = qmm_splitk_env > 1 + ? qmm_splitk_env + : kq_splitk_capped(16, grid, kq_splitk_cap_env()); const int sp = kq_split_count(splitk_target, K / group_size_); if (sp > 1) { qmm_splitk( @@ -2338,10 +2925,7 @@ void KQuantMatmul::eval_gpu_base( // 249-287 vs 149-256, q2_k M>=8 133-194 vs 88-135, iq3_xxs M>=8 // 163-175 vs 102-158). KQ_NAX_SMALL_BM=0 restores the old routing // (same switch that picks the BM=32 tile in qmm_nax). - static const bool nax_smallm_route = []() { - const char* e = std::getenv("KQ_NAX_SMALL_BM"); - return e == nullptr || std::atoi(e) != 0; - }(); + const bool nax_smallm_route = kq_nax_small_bm_mode() != 0; int smallm_min = kq_smallm_route_min(kquant_type_, N, K); if (nax_smallm_route && transpose_ && non_batched && smallm_min > 0 && M >= smallm_min && M <= 12 && kq_is_nax_available() && (K % 64 == 0) && @@ -2369,8 +2953,7 @@ void KQuantMatmul::eval_gpu_base( if (transpose_ && non_batched && M > 32 && M <= 128 && !nax_path && qmm_splitk_env != 0 && kq_splitk_codec(kquant_type_) && x.dtype() != mx::float32) { - const char* me = std::getenv("KQ_QMM_MIDM"); - const int midm = me != nullptr ? std::atoi(me) : -1; + const int midm = kq_env_int("KQ_QMM_MIDM", -1); if (midm == 1 || (midm == -1 && kq_midm_splitk(M, N))) { const int target = qmm_splitk_env > 1 ? qmm_splitk_env : 16; const int sp = kq_split_count(target, K / group_size_); @@ -2394,18 +2977,20 @@ void KQuantMatmul::eval_gpu_base( } } - // Mid-M split-K on NAX GPUs where kq_midm_nax_target gives a target. - // BM32 through M 48, where the last row tile holds at most 16 rows and - // runs K-split, and BM64 above. KQ_QMM_MIDM (read live) is the same probe + // Mid-M split-K on NAX GPUs where kq_midm_nax_target gives a target, on + // the tile kq_nax_splitk_bm picks. KQ_QMM_MIDM (read live) is the same probe // lever. KQ_QMM_SPLITK_NAX=0 turns it off with the small-M route, and a // larger value sets its target. if (transpose_ && non_batched && M > 32 && M <= 128 && nax_path && (K % 64 == 0) && x.dtype() != mx::float32) { - const char* me = std::getenv("KQ_QMM_MIDM"); - const int midm = me != nullptr ? std::atoi(me) : -1; - const char* ske = std::getenv("KQ_QMM_SPLITK_NAX"); - const int sk_env = ske != nullptr ? std::atoi(ske) : -1; - const int rule = kq_midm_nax_target(M, N); + const int midm = kq_env_int("KQ_QMM_MIDM", -1); + const int sk_env = kq_env_int("KQ_QMM_SPLITK_NAX", -1); + const int rule = kq_midm_nax_target( + M, + N, + K, + kquant_type_, + kq_splitk_partials(x.dtype()).dtype == mx::float32); if (sk_env != 0 && (midm == 1 || (midm == -1 && rule > 0))) { const int target = sk_env > 1 ? sk_env : rule > 0 ? rule @@ -2428,8 +3013,7 @@ void KQuantMatmul::eval_gpu_base( d, s, kquant_type_, - per * q, - M <= 48 ? 32 : 64); + per * q); return; } } @@ -2656,6 +3240,102 @@ int qmv_fine_max_n(const std::string& kquant_type) { : 0; } +// Every kernel name the split-K dispatchers can build, over the codecs +// their gates admit and the bfloat16 and float16 activations the ops pass, +// both N alignments, both partial dtypes and the verify_mma float and half +// kernels a codec has. The tiles, NAX tile rows and verify_nax steps come +// from the dispatch helpers over every M the split routes serve (1 to 128) +// and every K kind, under this process's levers. The gates read the codec +// table, not KQ_DISABLE_NAX. Names repeat across routes. +std::vector nax_large_row(const std::string& kquant_type) { + KqLargeM r = kq_largem_row(kquant_type); + return { + r.bm, + r.bn, + r.wm, + r.wn, + r.swizzle_log, + r.min_m, + r.min_n, + r.max_k, + r.inherited}; +} + +std::vector split_kernel_names() { + std::vector out; + auto each_suffix = [](mx::Dtype dt, auto&& f) { + f(kq_splitk_partials(dt, false).suffix); + f(kq_splitk_partials(dt, true).suffix); + }; + for (mx::Dtype dt : {mx::bfloat16, mx::float16}) { + const std::string ts = kq_type_string(dt); + each_suffix(dt, [&](const std::string& sfx) { + out.push_back(kq_splitk_accum_kname(ts, sfx)); + }); + } + for (const std::string& t : codec_names()) { + const KQuantCodec* codec = codec_by_name(t); + const int gs = codec->weights_per_block; + const int bits = codec->bits; + std::set> tiles; + std::set nax_bms; + for (int M = 1; M <= 128; ++M) { + const KqSplitkTile tile = kq_splitk_tile(M, t); + tiles.insert({tile.bm, tile.bn}); + nax_bms.insert(kq_nax_splitk_bm(M)); + } + // K of whole codec steps, and a multiple of 256 steps, which covers + // every step kq_verify_nax_call_kstep picks. + const int kstep = codec_verify_nax_kstep(t); + std::set ksteps; + for (int K : {kstep, 256 * kstep}) { + ksteps.insert(kq_verify_nax_call_kstep(K, t)); + } + for (mx::Dtype dt : {mx::bfloat16, mx::float16}) { + const std::string ts = kq_type_string(dt); + if (kq_splitk_codec(t)) { + for (const auto& [bm, bn] : tiles) { + const KqSplitkTile tile{bm, bn}; + for (bool aligned : {true, false}) { + each_suffix(dt, [&](const std::string& sfx) { + out.push_back( + kq_qmm_splitk_kname(t, ts, gs, bits, tile, aligned, sfx)); + }); + } + } + } + if (codec->has_nax_kernel) { + for (int bm : nax_bms) { + for (bool aligned : {true, false}) { + each_suffix(dt, [&](const std::string& sfx) { + out.push_back( + kq_qmm_nax_splitk_kname(t, ts, gs, bits, bm, aligned, sfx)); + }); + } + } + } + if (codec_verify_mma_rows(t) > 0) { + for (bool half : {false, true}) { + if (half && !kq_verify_mma_has_half(t)) { + continue; + } + each_suffix(dt, [&](const std::string& sfx) { + out.push_back(kq_verify_mma_kname(t, ts, gs, bits, half, sfx)); + }); + } + } + if (codec->has_nax_kernel && kstep > 0) { + for (int k : ksteps) { + each_suffix(dt, [&](const std::string& sfx) { + out.push_back(kq_verify_nax_kname(t, ts, gs, bits, k, sfx)); + }); + } + } + } + } + return out; +} + #else void KQuantQmvBias::eval_gpu( @@ -2669,6 +3349,14 @@ int qmv_fine_max_n(const std::string&) { return 0; } +std::vector nax_large_row(const std::string&) { + return {}; +} + +std::vector split_kernel_names() { + return {}; +} + #endif // _METAL_ } // namespace mlx_kquant diff --git a/src/kquant_matmul_fused.cpp b/src/kquant_matmul_fused.cpp new file mode 100644 index 00000000..565aac63 --- /dev/null +++ b/src/kquant_matmul_fused.cpp @@ -0,0 +1,464 @@ +// Quantized matmuls with an epilogue fused into the store (see +// kq_quantized_nax_epi.h for the kernels). +// +// GLU product: y = act(x @ Wg^T) * (x @ Wu^T) over two packed weights of +// one codec and one shape. On a NAX GPU, for the codecs with a GLU kernel, +// from kGluMinRows rows and with K a whole number of the codec's K tiles, +// one dispatch computes the product with y as its only output buffer. +// Every other case builds the same product from quantized_matmul and stock +// ops. +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kquant.h" +#include "kquant_codec.h" +#include "kquant_internal.h" + +#include "mlx/allocator.h" +#include "mlx/compile.h" +#include "mlx/ops.h" +#include "mlx/transforms.h" +#include "mlx/utils.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +const char* kGluOp = "[mlx_kquant.quantized_matmul_glu]"; + +// Activation kinds, the values of the kernel's function constant. +int glu_act_kind(const std::string& activation) { + if (activation == "silu") { + return 0; + } + if (activation == "gelu_tanh") { + return 1; + } + throw std::invalid_argument( + std::string(kGluOp) + " activation must be 'silu' or 'gelu_tanh', got '" + + activation + "'."); +} + +// The activation of the composition, built the way mlx.nn builds silu and +// gelu_approx: the same ops in the same order, with the constants a Python +// scalar gives, as one compiled elementwise function while mx's compile +// switch is on and as separate ops while it is off. The composition then +// has the bits and the peak memory of nn.silu(g) * u and +// nn.gelu_approx(g) * u in either state: the compiled function writes into +// g's buffer, where separate ops hold a buffer per step. +std::vector glu_silu(const std::vector& in) { + const auto& x = in[0]; + return {mx::multiply(x, mx::sigmoid(x))}; +} + +std::vector glu_gelu_tanh(const std::vector& in) { + // 0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x ** 3))) + const auto& x = in[0]; + const auto dt = x.dtype(); + auto x3 = mx::power(x, mx::array(3, dt)); + auto sum = mx::add(x, mx::multiply(mx::array(0.044715f, dt), x3)); + auto t = mx::tanh(mx::multiply(mx::array(0.7978845608028654f, dt), sum)); + return {mx::multiply( + mx::multiply(mx::array(0.5f, dt), x), mx::add(mx::array(1, dt), t))}; +} + +mx::array glu_activation(const mx::array& g, int act, mx::Stream s) { + // mx::compile reads the compile switch when it is called, so the call is + // made per activation, as a call of nn.silu does. With the switch off it + // returns the function as it is. With it on it returns a wrapper around + // the compiled function, which the compile cache keys by the function + // pointer: the function is compiled once per thread, and the wrapper + // owns no cache entry, so nothing is erased when it goes. + auto fn = mx::compile(act == 1 ? glu_gelu_tanh : glu_silu, true); + // The ops of the function run on the default stream. + mx::StreamContext on_stream(s); + return fn({g})[0]; +} + +// K tile of the GLU kernel for a codec and an activation dtype, 0 where +// the op builds the composition. The q8_0 decoder has a K tile of 64 only, +// and its kernel runs on bfloat16 only: on float16 it measured slower than +// the composition. +int glu_kernel_bk(const std::string& t, mx::Dtype dt) { + if (t == "q4_k" || t == "q5_k" || t == "q6_k") { + return 128; + } + if (t == "q8_0" && dt == mx::bfloat16) { + return 64; + } + return 0; +} + +bool glu_tune_live() { + static const bool live = kq_env_lever("KQ_TUNE_LIVE", 0) != 0; + return live; +} + +// KQ_GLU=0 builds the composition on every call. +bool glu_on() { + static const bool cached = kq_env_flag("KQ_GLU", 1) != 0; + return glu_tune_live() ? kq_env_flag("KQ_GLU", 1) != 0 : cached; +} + +// Rows from which the fused kernel runs. The kernel is measured against +// the composition from 1056 rows up; fewer rows build the composition. +// KQ_GLU_MIN_ROWS replaces the bound with a value of at least 1. A value +// below 1 keeps the bound, with a warning. +constexpr int kGluMinRows = 1024; +int glu_min_rows() { + static const int cached = + kq_env_lever_from("KQ_GLU_MIN_ROWS", kGluMinRows, 1); + return glu_tune_live() ? kq_env_lever_from("KQ_GLU_MIN_ROWS", kGluMinRows, 1) + : cached; +} + +// The op runs the two-walk kernel on a row tile of 128 and the codec's K +// tile. A build with KQ_PROBE_KERNELS=ON carries more tiles and reads the +// levers that select them: KQ_GLU_BM and KQ_GLU_BK (64 or 128), +// KQ_GLU_WALKS=0 (the one-walk kernel) and KQ_GLU_DUAL128=1. They are read +// once per process unless KQ_TUNE_LIVE=1. A release build has one tile per +// codec and reads none of them. +#ifdef KQ_PROBE_KERNELS +int glu_lever(const char* name) { + const int v = kq_env_lever(name, -1); + return v == 64 || v == 128 ? v : -1; +} +int glu_bm_lever() { + static const int cached = glu_lever("KQ_GLU_BM"); + return glu_tune_live() ? glu_lever("KQ_GLU_BM") : cached; +} +int glu_bk_lever() { + static const int cached = glu_lever("KQ_GLU_BK"); + return glu_tune_live() ? glu_lever("KQ_GLU_BK") : cached; +} +bool glu_walks() { + static const bool cached = kq_env_flag("KQ_GLU_WALKS", 1) != 0; + return glu_tune_live() ? kq_env_flag("KQ_GLU_WALKS", 1) != 0 : cached; +} +// The one-walk K tile of 128 on two unpadded Ws buffers. +bool glu_dual128() { + static const bool cached = kq_env_flag("KQ_GLU_DUAL128", 0) != 0; + return glu_tune_live() ? kq_env_flag("KQ_GLU_DUAL128", 0) != 0 : cached; +} +#else +int glu_bm_lever() { + return -1; +} +int glu_bk_lever() { + return -1; +} +bool glu_walks() { + return true; +} +bool glu_dual128() { + return false; +} +#endif + +bool glu_log() { + static const bool cached = std::getenv("KQ_QMM_LOG") != nullptr; + return glu_tune_live() ? std::getenv("KQ_QMM_LOG") != nullptr : cached; +} + +// The product from quantized_matmul and stock ops. The op builds this +// where it has no fused kernel, and the gradient of the fused kernel is +// the gradient of this. +mx::array glu_composed( + const mx::array& x, + const mx::array& w_gate, + const mx::array& w_up, + const std::string& kquant_type, + int act, + mx::Stream s) { + auto scales = mx::zeros({1}, mx::uint8, s); + auto g = quantized_matmul( + x, + w_gate, + scales, + kquant_type, + true, + std::nullopt, + std::nullopt, + std::nullopt, + s); + auto u = quantized_matmul( + x, + w_up, + scales, + kquant_type, + true, + std::nullopt, + std::nullopt, + std::nullopt, + s); + return mx::multiply(glu_activation(g, act, s), u, s); +} + +} // namespace + +#ifdef _METAL_ + +void KQuantMatmulFused::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + const auto& x = inputs[0]; + // A weight whose rows are dense binds as it is, at its own offset, so a + // row band of a wider weight is read in place. Any other view is copied + // to dense rows first. + const mx::array wg = kq_matrix_dense_gpu(inputs[1], s); + const mx::array wu = kq_matrix_dense_gpu(inputs[2], s); + auto& out = outputs[0]; + kq_check_weight_base(wg, kquant_type_, "quantized_matmul_glu"); + kq_check_weight_base(wu, kquant_type_, "quantized_matmul_glu"); + out.set_data(mx::allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + + const int K = x.shape(-1); + const int N = wg.shape(0); + // x is dense rows, so its batch dims fold into M. + const int M = static_cast(x.size() / K); + + int bm = 128; + int bk = glu_kernel_bk(kquant_type_, x.dtype()); + if (const int v = glu_bm_lever(); v > 0) { + bm = v; + } + if (const int v = glu_bk_lever(); v > 0 && v <= bk && K % v == 0) { + bk = v; + } + const int bn = 64; + + const bool aligned = N % bn == 0; + bool saturate = x.dtype() == mx::float16; + const bool walks = glu_walks(); + const bool dual = !walks && bk == 128 && glu_dual128(); + const std::string kname = kq_kname_prefix(kquant_type_) + "qmm_t_nax_glu_" + + kq_type_string(x.dtype()) + "_bm" + std::to_string(bm) + "_bn64_bk" + + std::to_string(bk) + + (walks ? "w" + : dual ? "d" + : "") + + "_wm2_wn2_alN_" + (aligned ? "true" : "false"); +#ifdef KQ_PROBE_KERNELS + // A probe lever can name a tile that this codec or dtype has no kernel + // for. Raise with the levers and the kernel name before the dispatch. + if (bm != 128 || bk != glu_kernel_bk(kquant_type_, x.dtype()) || !walks) { + kq_require_kernel( + d, + "KQ_GLU_BM=" + std::to_string(bm) + " KQ_GLU_BK=" + std::to_string(bk) + + " KQ_GLU_WALKS=" + (walks ? "1" : "0") + + " KQ_GLU_DUAL128=" + (dual ? "1" : "0"), + kname); + } +#endif + const std::string hash_name = + kname + "_act" + std::to_string(act_) + (saturate ? "_sat" : ""); + mx::metal::MTLFCList func_consts = { + {&act_, MTL::DataType::DataTypeInt, 360}, + {&saturate, MTL::DataType::DataTypeBool, 361}, + }; + MTL::Size group_dims(32, 2, 2); + MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1); + if (glu_log()) { + std::fprintf( + stderr, + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + kname.c_str(), + M, + N, + K, + static_cast(grid_dims.width), + static_cast(grid_dims.height), + static_cast(grid_dims.depth)); + } + + auto kernel = kq_get_kernel(d, kname, hash_name, func_consts); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + int c = 0; + ce.set_input_array(wg, c++); + ce.set_input_array(wu, c++); + ce.set_input_array(x, c++); + ce.set_output_array(out, c++); + ce.set_bytes(K, c++); + ce.set_bytes(N, c++); + ce.set_bytes(M, c++); + ce.dispatch_threadgroups(grid_dims, group_dims); +} + +#else // !_METAL_ + +void KQuantMatmulFused::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error(std::string(kGluOp) + " requires Metal."); +} + +#endif + +void KQuantMatmulFused::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + std::string(kGluOp) + " CPU streams build the composed ops."); +} + +// The fused kernel computes the function of the composition, so its +// derivatives are the composition's: the gradient wrt x, with the quantized +// weights frozen as in quantized_matmul. For float16 and bfloat16 x a +// derivative then does not depend on which side of the row bound a call +// falls. float32 x is cast to bfloat16 before this node, so the gate term +// and the up term of the gradient are summed in bfloat16 here and the sum +// is cast to float32. The composition casts x once per matmul, so it casts +// each term to float32 and sums there. The two gradients differ by the +// rounding of that one bfloat16 sum, and neither is closer to the exact +// gradient. +std::vector KQuantMatmulFused::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector&) { + for (auto arg : argnums) { + if (arg != 0) { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul_glu vjp: no gradient wrt the " + "quantized weights (the kquant base is frozen)."); + } + } + auto fun = [&](const std::vector& in) { + return std::vector{glu_composed( + in[0], primals[1], primals[2], kquant_type_, act_, stream())}; + }; + auto vjps = + mx::vjp(fun, std::vector{primals[0]}, cotangents).second; + return std::vector(argnums.size(), vjps[0]); +} + +std::vector KQuantMatmulFused::jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) { + for (auto arg : argnums) { + if (arg != 0) { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul_glu jvp: no derivative wrt the " + "quantized weights (the kquant base is frozen)."); + } + } + auto fun = [&](const std::vector& in) { + return std::vector{glu_composed( + in[0], primals[1], primals[2], kquant_type_, act_, stream())}; + }; + return mx::jvp( + fun, + std::vector{primals[0]}, + std::vector{tangents[0]}) + .second; +} + +bool KQuantMatmulFused::is_equivalent(const mx::Primitive& other) const { + const auto& o = static_cast(other); + return kquant_type_ == o.kquant_type_ && act_ == o.act_; +} + +std::vector KQuantMatmulFused::output_shapes( + const std::vector& inputs) { + auto shape = inputs[0].shape(); + shape.back() = inputs[1].shape(0); + return {shape}; +} + +mx::array quantized_matmul_glu( + mx::array x, + mx::array w_gate, + mx::array w_up, + const std::string& kquant_type, + const std::string& activation, + mx::StreamOrDevice s_) { + const std::string op = kGluOp; + const int act = glu_act_kind(activation); + const KQuantCodec* codec = codec_by_name(kquant_type); + if (codec == nullptr) { + throw std::invalid_argument( + op + " Unknown kquant_type: '" + kquant_type + "'."); + } + for (const auto* w : {&w_gate, &w_up}) { + if (w->dtype() != mx::uint8 || w->ndim() != 2) { + throw std::invalid_argument( + op + " w_gate and w_up must be 2-D uint8 (raw GGUF wire bytes)."); + } + } + if (w_gate.shape() != w_up.shape()) { + throw std::invalid_argument(op + " w_gate and w_up must have one shape."); + } + if (w_gate.shape(-1) % codec->bytes_per_block != 0) { + throw std::invalid_argument( + op + " weight last dim (" + std::to_string(w_gate.shape(-1)) + + " bytes) is not a whole number of " + + std::to_string(codec->bytes_per_block) + "-byte " + codec->name + + " blocks."); + } + const int K = + (w_gate.shape(-1) / codec->bytes_per_block) * codec->weights_per_block; + if (x.ndim() < 1 || x.shape(-1) != K) { + throw std::invalid_argument( + op + " x last dim does not match the weight inner dim (" + + std::to_string(K) + ")."); + } + if (!mx::issubdtype(x.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " Only real floating x dtypes are supported."); + } + auto s = mx::to_stream(s_); + + // K must be a whole number of the codec's K tiles (of the lever's tile + // when KQ_GLU_BK selects a smaller one). + // float32 x runs in bfloat16, as in quantized_matmul. + const mx::Dtype out_type = + x.dtype() == mx::float32 ? mx::bfloat16 : x.dtype(); + int bk = glu_kernel_bk(kquant_type, out_type); + if (const int v = glu_bk_lever(); v > 0 && v <= bk) { + bk = v; + } + const int64_t rows = K > 0 ? static_cast(x.size()) / K : 0; + bool fused = s.device == mx::Device::gpu && nax_available() && glu_on() && + bk > 0 && K % bk == 0 && x.ndim() >= 2 && rows >= glu_min_rows() && + rows > 0 && w_gate.shape(0) > 0; + if (!fused) { + return glu_composed(x, w_gate, w_up, kquant_type, act, s); + } + + kq_check_weight_base_at_build(w_gate, kquant_type, s, "quantized_matmul_glu"); + kq_check_weight_base_at_build(w_up, kquant_type, s, "quantized_matmul_glu"); + // x gets a Contiguous node: the strides of an unevaluated array mean + // nothing at build time. The weights go in as they are and the eval + // looks at their strides. A Contiguous node would copy a row band of a + // wider weight on every call, since its buffer is larger than the band. + auto x_c = mx::contiguous(mx::astype(x, out_type, s), false, s); + auto out_shape = x.shape(); + out_shape.back() = w_gate.shape(0); + return mx::array( + std::move(out_shape), + out_type, + std::make_shared(s, kquant_type, act), + {std::move(x_c), std::move(w_gate), std::move(w_up)}); +} + +} // namespace mlx_kquant diff --git a/src/kquant_matmul_lora.cpp b/src/kquant_matmul_lora.cpp new file mode 100644 index 00000000..eb74db4f --- /dev/null +++ b/src/kquant_matmul_lora.cpp @@ -0,0 +1,630 @@ +// Quantized matmul with a low-rank term in the same reduction (see +// kq_quantized_nax_lora.h for the kernel). +// +// y = x @ W^T + (strength * (x @ A)) @ B +// +// On a NAX GPU, for the codecs with a kernel (q4_k, q5_k, q6_k and q8_0 on +// bfloat16), from kLoraMinRows rows, on the shapes of lora_shape_rule, +// stock ops compute z = strength * (x @ A) and one dispatch computes the +// product and adds z @ B in its accumulator, with y as its only output +// buffer. Every other case builds the same sum from quantized_matmul and +// stock ops. +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kquant.h" +#include "kquant_codec.h" +#include "kquant_internal.h" + +#include "mlx/allocator.h" +#include "mlx/ops.h" +#include "mlx/transforms.h" +#include "mlx/utils.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +const char* kLoraOp = "[mlx_kquant.quantized_matmul_lora]"; + +// Column tile of the kernel. A band starts on a column tile. +constexpr int kLoraBn = 64; +// Column tile of a simdgroup. A band ends on one, or at N. +constexpr int kLoraSn = 32; +// Largest rank of one band the kernel runs. +constexpr int kLoraMaxRank = 256; +// Most bands the kernel looks through per threadgroup. +constexpr int kLoraMaxBands = 16; +// Sizes the op refuses. The kernel offsets a row tile by a 32-bit product +// of 64 rows and a row length (K, N or R), which wraps from 2^25, and +// holds the row count in an int. +constexpr int64_t kLoraMaxCols = int64_t{1} << 25; +constexpr int64_t kLoraMaxRows = int64_t{1} << 31; + +// K tile of the kernel for a codec and an activation dtype, 0 where the op +// builds the composition. +int lora_kernel_bk(const std::string& t, mx::Dtype dt) { + if (dt != mx::bfloat16) { + return 0; + } + if (t == "q4_k" || t == "q5_k" || t == "q6_k") { + return 128; + } + if (t == "q8_0") { + return 64; + } + return 0; +} + +bool lora_tune_live() { + static const bool live = kq_env_lever("KQ_TUNE_LIVE", 0) != 0; + return live; +} + +// KQ_QMM_LORA: 0 = the composition on every call, 2 = the kernel wherever +// it can run, unset, 1 or any other value = the shape rule +// (lora_shape_rule). +int lora_mode() { + static const int cached = kq_env_lever("KQ_QMM_LORA", 1); + return lora_tune_live() ? kq_env_lever("KQ_QMM_LORA", 1) : cached; +} + +// Rows from which the kernel runs. KQ_QMM_LORA_MIN_ROWS replaces the bound +// with a value of at least 1. A value below 1 keeps the bound, with a +// warning. +constexpr int kLoraMinRows = 1024; +int lora_min_rows() { + static const int cached = + kq_env_lever_from("KQ_QMM_LORA_MIN_ROWS", kLoraMinRows, 1); + return lora_tune_live() + ? kq_env_lever_from("KQ_QMM_LORA_MIN_ROWS", kLoraMinRows, 1) + : cached; +} + +// The shapes on which the kernel runs under the default mode: those where +// it measured faster than the unfused ops. The kernel removes the second +// pass over y, a cost in proportion to N. It waits for z, a cost in +// proportion to K that the unfused product does not pay, since x @ A runs +// beside it. So the kernel is faster where N is at least K, and its share +// of the call falls as K grows: above kLoraRuleK the two forms tie. From N +// of kLoraRuleWideN the kernel runs at every row count and rank it takes. +// Where N is nearer K the gain is about one percent: it holds from +// kLoraRuleRows rows and up to a rank of kLoraRuleRank, since the wait for +// z grows with the rank. rank is the most ranks a column takes. +constexpr int kLoraRuleK = 4096; +constexpr int kLoraRuleRows = 4096; +constexpr int kLoraRuleWideN = 8192; +constexpr int kLoraRuleRank = 128; +bool lora_shape_rule(int K, int N, int64_t rows, int rank) { + return K <= kLoraRuleK && N >= K && + (N >= kLoraRuleWideN || (rows >= kLoraRuleRows && rank <= kLoraRuleRank)); +} + +bool lora_log() { + static const bool cached = std::getenv("KQ_QMM_LOG") != nullptr; + return lora_tune_live() ? std::getenv("KQ_QMM_LOG") != nullptr : cached; +} + +mx::array last_axis_slice(const mx::array& a, int lo, int hi, mx::Stream s) { + mx::Shape start(a.ndim(), 0); + mx::Shape stop = a.shape(); + start.back() = lo; + stop.back() = hi; + return mx::slice(a, std::move(start), std::move(stop), s); +} + +// c + a @ b by the stock op, with c copied to dense memory first when it +// has one column. On the GPU the stock op reads a first argument of one +// column and several rows at a step that is not its row stride, so a +// column of a wider array comes out wrong (docs/upstream-defects.md). A +// dense c is passed through with no copy, and a c of two or more columns +// is read in place, as before. +mx::array +lora_addmm(mx::array c, const mx::array& a, const mx::array& b, mx::Stream s) { + if (c.shape(-1) == 1) { + c = mx::contiguous(c, false, s); + } + return mx::addmm(std::move(c), a, b, 1.0f, 1.0f, s); +} + +// The base product and z on their way into the banded sum: the identity, +// on the buffer of its input, under every transform. The stock Slice ops +// that cut the bands have no shape rule, so a shapeless compiled function +// cannot hold the banded sum. Each array that a Slice cuts first passes +// through this node, so the node comes before every Slice in the order of +// a trace and refuses by the name of the op. +class LoraBandedInput : public mx::UnaryPrimitive { + public: + explicit LoraBandedInput(mx::Stream stream) : mx::UnaryPrimitive(stream) {} + + void eval_cpu(const std::vector& inputs, mx::array& out) override { + out.copy_shared_buffer(inputs[0]); + } + void eval_gpu(const std::vector& inputs, mx::array& out) override { + out.copy_shared_buffer(inputs[0]); + } + + const char* name() const override { + return "KQuantLoraBandedInput"; + } + bool is_equivalent(const mx::Primitive&) const override { + return true; + } + + std::vector vjp( + const std::vector&, + const std::vector& cotangents, + const std::vector&, + const std::vector&) override { + return {cotangents[0]}; + } + std::vector jvp( + const std::vector&, + const std::vector& tangents, + const std::vector&) override { + return {tangents[0]}; + } + std::pair, std::vector> vmap( + const std::vector& inputs, + const std::vector& axes) override { + return {{inputs[0]}, axes}; + } + + std::vector output_shapes(const std::vector&) override { + throw std::invalid_argument( + std::string(kLoraOp) + + " A shapeless compiled function cannot hold a call with bands that " + "builds the unfused ops. Compile it with fixed shapes, or trace it " + "on a call that runs the kernel."); + } +}; + +// base + z @ B over the bands. With no band table every column takes +// every rank. With one, a band's columns take its ranks and its block of +// B, and a column outside every band takes none. +mx::array lora_add( + const mx::array& base_, + const mx::array& z_, + const mx::array& b, + const std::vector& bands, + mx::Stream s) { + if (bands.empty()) { + return lora_addmm(base_, z_, b, s); + } + auto banded_input = [&s](const mx::array& a) { + return mx::array( + a.shape(), a.dtype(), std::make_shared(s), {a}); + }; + const mx::array base = banded_input(base_); + const mx::array z = banded_input(z_); + const int N = base.shape(-1); + std::vector parts; + int col = 0; + for (size_t i = 0; i < bands.size(); i += 4) { + const int cs = bands[i]; + const int ce = bands[i + 1]; + const int z0 = bands[i + 2]; + const int rb = bands[i + 3]; + if (cs > col) { + parts.push_back(last_axis_slice(base, col, cs, s)); + } + auto y = last_axis_slice(base, cs, ce, s); + if (rb > 0) { + auto bb = mx::slice(b, {z0, cs}, {z0 + rb, ce}, s); + y = lora_addmm(std::move(y), last_axis_slice(z, z0, z0 + rb, s), bb, s); + } + parts.push_back(std::move(y)); + col = ce; + } + if (col < N) { + parts.push_back(last_axis_slice(base, col, N, s)); + } + return parts.size() == 1 ? parts[0] + : mx::concatenate(std::move(parts), -1, s); +} + +// The sum from quantized_matmul and stock ops, given z. The op builds this +// where it has no kernel, and the derivatives of the kernel are the +// derivatives of this. +mx::array lora_composed( + const mx::array& x, + const mx::array& w, + const mx::array& z, + const mx::array& b, + const std::string& kquant_type, + const std::vector& bands, + mx::Stream s) { + auto scales = mx::zeros({1}, mx::uint8, s); + auto base = quantized_matmul( + x, + w, + scales, + kquant_type, + true, + std::nullopt, + std::nullopt, + std::nullopt, + s); + return lora_add(base, z, b, bands, s); +} + +} // namespace + +#ifdef _METAL_ + +void KQuantMatmulLora::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + const auto& x = inputs[0]; + // A weight whose rows are dense binds as it is, at its own offset. Any + // other view is copied to dense rows first. z and B bind as they are + // when they are row-major. + const mx::array w = kq_matrix_dense_gpu(inputs[1], s); + const mx::array z = kq_row_major_gpu(inputs[2], s); + const mx::array b = kq_row_major_gpu(inputs[3], s); + auto& out = outputs[0]; + kq_check_weight_base(w, kquant_type_, "quantized_matmul_lora"); + out.set_data(mx::allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + + const int K = x.shape(-1); + const int N = w.shape(0); + const int R = b.shape(0); + // x is dense rows, so its batch dims fold into M. + const int M = static_cast(x.size() / K); + + const int bm = 128; + const int bk = lora_kernel_bk(kquant_type_, x.dtype()); + const bool aligned = N % kLoraBn == 0; + const std::string kname = kq_kname_prefix(kquant_type_) + "qmm_t_nax_lora_" + + kq_type_string(x.dtype()) + "_bm" + std::to_string(bm) + "_bn64_bk" + + std::to_string(bk) + "_wm2_wn2_alN_" + (aligned ? "true" : "false"); + + // With no band table, one band over every column and every rank. + std::vector bands = bands_; + if (bands.empty()) { + bands = {0, N, 0, R}; + } + const int n_bands = static_cast(bands.size() / 4); + + MTL::Size group_dims(32, 2, 2); + // M may be within a row tile of the largest int. + MTL::Size grid_dims( + (N + kLoraBn - 1) / kLoraBn, (static_cast(M) + bm - 1) / bm, 1); + if (lora_log()) { + std::fprintf( + stderr, + "[kq qmm] %s M=%d N=%d K=%d R=%d bands=%d grid=%lu,%lu,%lu\n", + kname.c_str(), + M, + N, + K, + R, + n_bands, + static_cast(grid_dims.width), + static_cast(grid_dims.height), + static_cast(grid_dims.depth)); + } + + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + int c = 0; + ce.set_input_array(w, c++); + ce.set_input_array(x, c++); + ce.set_input_array(z, c++); + ce.set_input_array(b, c++); + ce.set_output_array(out, c++); + ce.set_bytes(K, c++); + ce.set_bytes(N, c++); + ce.set_bytes(M, c++); + ce.set_bytes(R, c++); + ce.set_bytes(n_bands, c++); + ce.set_vector_bytes(bands, c++); + ce.dispatch_threadgroups(grid_dims, group_dims); +} + +#else // !_METAL_ + +void KQuantMatmulLora::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error(std::string(kLoraOp) + " requires Metal."); +} + +#endif + +void KQuantMatmulLora::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + std::string(kLoraOp) + " CPU streams build the composed ops."); +} + +// The kernel computes the function of the composition, so its vjp is the +// composition's: with respect to x, z and B, with the quantized weight +// frozen as in quantized_matmul. z is built from x, A and the strength by +// stock ops before this node, so their gradients follow from the one with +// respect to z. +namespace { + +constexpr int kLoraArgOfInput[] = {0, -1, 1, 2}; + +void lora_check_argnums(const std::vector& argnums) { + for (auto arg : argnums) { + if (arg == 1) { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul_lora vjp: no gradient wrt the " + "quantized weight (the kquant base is frozen)."); + } + } +} + +} // namespace + +std::vector KQuantMatmulLora::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector&) { + lora_check_argnums(argnums); + auto fun = [&](const std::vector& in) { + return std::vector{lora_composed( + in[0], primals[1], in[1], in[2], kquant_type_, bands_, stream())}; + }; + auto vjps = mx::vjp( + fun, + std::vector{primals[0], primals[2], primals[3]}, + cotangents) + .second; + std::vector out; + for (auto arg : argnums) { + out.push_back(vjps[kLoraArgOfInput[arg]]); + } + return out; +} + +// jvp and vmap are the transforms of the composition, over the inputs that +// carry a tangent or a batch axis. The others enter as constants. So a +// tangent or a batch of z or B (of lora_a, lora_b or the strength) runs on +// stock ops, and one of x or of the weight reaches quantized_matmul, which +// has neither transform and refuses by its own name, as it does where the +// op builds the composition. +std::vector KQuantMatmulLora::jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) { + auto fun = [&](const std::vector& in) { + std::vector full = primals; + for (size_t i = 0; i < argnums.size(); i++) { + full[argnums[i]] = in[i]; + } + return std::vector{lora_composed( + full[0], full[1], full[2], full[3], kquant_type_, bands_, stream())}; + }; + std::vector at; + for (auto arg : argnums) { + at.push_back(primals[arg]); + } + return mx::jvp(fun, at, tangents).second; +} + +std::pair, std::vector> KQuantMatmulLora::vmap( + const std::vector& inputs, + const std::vector& axes) { + auto fun = [&](const std::vector& in) { + return std::vector{lora_composed( + in[0], in[1], in[2], in[3], kquant_type_, bands_, stream())}; + }; + return {mx::vmap(fun, axes, {0})(inputs), {0}}; +} + +bool KQuantMatmulLora::is_equivalent(const mx::Primitive& other) const { + const auto& o = static_cast(other); + return kquant_type_ == o.kquant_type_ && bands_ == o.bands_; +} + +std::vector KQuantMatmulLora::output_shapes( + const std::vector& inputs) { + auto shape = inputs[0].shape(); + shape.back() = inputs[1].shape(0); + return {shape}; +} + +mx::array quantized_matmul_lora( + mx::array x, + mx::array w, + const std::string& kquant_type, + mx::array lora_a, + mx::array lora_b, + const std::optional& strength, + const std::optional>& bands_, + mx::StreamOrDevice s_) { + const std::string op = kLoraOp; + // No table: every column takes every rank, which the code below holds as + // an empty vector. A table with no row: no column takes a rank. + const bool no_band = bands_.has_value() && bands_->empty(); + const std::vector bands = bands_.value_or(std::vector{}); + const KQuantCodec* codec = codec_by_name(kquant_type); + if (codec == nullptr) { + throw std::invalid_argument( + op + " Unknown kquant_type: '" + kquant_type + "'."); + } + if (w.dtype() != mx::uint8 || w.ndim() != 2) { + throw std::invalid_argument( + op + " w must be 2-D uint8 (raw GGUF wire bytes)."); + } + if (w.shape(-1) % codec->bytes_per_block != 0) { + throw std::invalid_argument( + op + " weight last dim (" + std::to_string(w.shape(-1)) + + " bytes) is not a whole number of " + + std::to_string(codec->bytes_per_block) + "-byte " + codec->name + + " blocks."); + } + const int K = + (w.shape(-1) / codec->bytes_per_block) * codec->weights_per_block; + const int N = w.shape(0); + if (x.ndim() < 1 || x.shape(-1) != K) { + throw std::invalid_argument( + op + " x last dim does not match the weight inner dim (" + + std::to_string(K) + ")."); + } + if (!mx::issubdtype(x.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " Only real floating x dtypes are supported."); + } + if (lora_a.ndim() != 2 || lora_a.shape(0) != K || + !mx::issubdtype(lora_a.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " lora_a must be a float array [K, R] with K " + + std::to_string(K) + "."); + } + const int R = lora_a.shape(1); + if (lora_b.ndim() != 2 || lora_b.shape(0) != R || lora_b.shape(1) != N || + !mx::issubdtype(lora_b.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " lora_b must be a float array [R, N] with R " + + std::to_string(R) + " and N " + std::to_string(N) + "."); + } + if (K >= kLoraMaxCols || N >= kLoraMaxCols || R >= kLoraMaxCols) { + throw std::invalid_argument( + op + " K, N and R must each be below " + std::to_string(kLoraMaxCols) + + " (2^25): K " + std::to_string(K) + ", N " + std::to_string(N) + + ", R " + std::to_string(R) + "."); + } + const int64_t rows = K > 0 ? static_cast(x.size()) / K : 0; + if (rows >= kLoraMaxRows) { + throw std::invalid_argument( + op + " x must hold fewer than " + std::to_string(kLoraMaxRows) + + " (2^31) rows: it holds " + std::to_string(rows) + "."); + } + if (strength.has_value()) { + const auto& st = *strength; + // 0-D, [1] or [R]. Any other shape would broadcast z to more dims. + const bool shape_ok = st.ndim() == 0 || + (st.ndim() == 1 && (st.shape(0) == 1 || st.shape(0) == R)); + if (!shape_ok || !mx::issubdtype(st.dtype(), mx::floating)) { + throw std::invalid_argument( + op + + " strength must be a float array of no dim, of shape [1] or of " + "shape [R] with R " + + std::to_string(R) + "."); + } + } + if (bands.size() % 4 != 0) { + throw std::invalid_argument( + op + " bands must hold rows of (col_start, col_end, z_start, rank)."); + } + int prev_end = 0; + for (size_t i = 0; i < bands.size(); i += 4) { + const int cs = bands[i]; + const int ce = bands[i + 1]; + const int z0 = bands[i + 2]; + const int rb = bands[i + 3]; + if (cs < prev_end || ce <= cs || ce > N) { + throw std::invalid_argument( + op + + " bands must be in column order, disjoint, not empty and " + "within the " + + std::to_string(N) + " columns: band " + std::to_string(i / 4) + + " covers columns " + std::to_string(cs) + " to " + + std::to_string(ce) + "."); + } + if (z0 < 0 || rb < 0 || rb > R - z0) { + throw std::invalid_argument( + op + " band " + std::to_string(i / 4) + " takes ranks " + + std::to_string(z0) + " to " + std::to_string(z0) + " + " + + std::to_string(rb) + " of " + std::to_string(R) + "."); + } + prev_end = ce; + } + auto s = mx::to_stream(s_); + + // float32 x runs in bfloat16, as in quantized_matmul. + const mx::Dtype out_type = + x.dtype() == mx::float32 ? mx::bfloat16 : x.dtype(); + auto scales = mx::zeros({1}, mx::uint8, s); + if (rows == 0 || N == 0 || R == 0 || no_band) { + // Nothing to add: the base product, with its shape and dtype. + return quantized_matmul( + x, + w, + scales, + kquant_type, + true, + std::nullopt, + std::nullopt, + std::nullopt, + s); + } + + // The kernel walks K in whole K tiles, so a K off the tile runs the + // composition (q8_0 only: the K-quant codecs hold K in blocks of 256). + const int bk = lora_kernel_bk(kquant_type, out_type); + const int mode = lora_mode(); + // The most ranks a column takes: every rank with no table, the rank of + // the largest band with one. + int col_rank = bands.empty() ? R : 0; + for (size_t i = 0; i < bands.size(); i += 4) { + col_rank = std::max(col_rank, bands[i + 3]); + } + bool fused = s.device == mx::Device::gpu && nax_available() && mode != 0 && + bk > 0 && K % bk == 0 && rows >= lora_min_rows() && + (mode == 2 || lora_shape_rule(K, N, rows, col_rank)); + if (bands.empty()) { + fused = fused && R <= kLoraMaxRank; + } else { + fused = fused && bands.size() / 4 <= kLoraMaxBands; + for (size_t i = 0; i < bands.size(); i += 4) { + // A band that ended inside a simdgroup's columns would put z times + // zero on the columns after it, which is not zero for a z that is + // not finite. + const bool end_ok = bands[i + 1] % kLoraSn == 0 || bands[i + 1] == N; + fused = fused && bands[i] % kLoraBn == 0 && end_ok && + bands[i + 3] <= kLoraMaxRank; + } + } + + // x gets a Contiguous node: the strides of an unevaluated array mean + // nothing at build time. The weight and B go in as they are and the eval + // looks at their strides. + auto x_c = fused ? mx::contiguous(mx::astype(x, out_type, s), false, s) + : mx::astype(x, out_type, s); + auto z = mx::matmul(x_c, mx::astype(lora_a, out_type, s), s); + if (strength.has_value()) { + z = mx::multiply(z, mx::astype(*strength, out_type, s), s); + } + auto b = mx::astype(lora_b, out_type, s); + if (!fused) { + return lora_composed(x, w, z, b, kquant_type, bands, s); + } + + kq_check_weight_base_at_build(w, kquant_type, s, "quantized_matmul_lora"); + auto out_shape = x.shape(); + out_shape.back() = N; + return mx::array( + std::move(out_shape), + out_type, + std::make_shared(s, kquant_type, bands), + {std::move(x_c), std::move(w), std::move(z), std::move(b)}); +} + +} // namespace mlx_kquant diff --git a/src/kquant_matmul_resid.cpp b/src/kquant_matmul_resid.cpp new file mode 100644 index 00000000..0d91fae7 --- /dev/null +++ b/src/kquant_matmul_resid.cpp @@ -0,0 +1,411 @@ +// Quantized matmul with a residual epilogue (see kq_quantized_nax_resid.h +// for the kernel). +// +// y = residual + gate * (x @ W^T) +// +// On a NAX GPU, for the codecs with a kernel (q4_k, q5_k, q6_k and q8_0 on +// bfloat16), from kResidMinRows rows, one dispatch computes the product +// and stores residual + gate * product, rounded once. Every other case +// builds the same sum from quantized_matmul and stock ops. +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kquant.h" +#include "kquant_codec.h" +#include "kquant_internal.h" + +#include "mlx/allocator.h" +#include "mlx/ops.h" +#include "mlx/transforms.h" +#include "mlx/utils.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +const char* kResidOp = "[mlx_kquant.quantized_matmul_residual]"; + +// Column tile of the kernel. +constexpr int kResidBn = 64; +// Sizes the op refuses. The kernel offsets a row tile by a 32-bit product +// of 64 rows and a row length (K or N), which wraps from 2^25, and holds +// the row count of a batch in an int. +constexpr int64_t kResidMaxCols = int64_t{1} << 25; +constexpr int64_t kResidMaxRows = int64_t{1} << 31; + +// K tile of the kernel for a codec and an activation dtype, 0 where the op +// builds the composition. +int resid_kernel_bk(const std::string& t, mx::Dtype dt) { + if (dt != mx::bfloat16) { + return 0; + } + if (t == "q4_k" || t == "q5_k" || t == "q6_k") { + return 128; + } + if (t == "q8_0") { + return 64; + } + return 0; +} + +bool resid_tune_live() { + static const bool live = kq_env_lever("KQ_TUNE_LIVE", 0) != 0; + return live; +} + +// KQ_QMM_RESID: 0 = the composition on every call, 2 = the kernel wherever +// it can run, unset, 1 or any other value = the kernel from kResidMinRows +// rows. +int resid_mode() { + static const int cached = kq_env_lever("KQ_QMM_RESID", 1); + return resid_tune_live() ? kq_env_lever("KQ_QMM_RESID", 1) : cached; +} + +// Rows from which the kernel runs under the default mode. +constexpr int64_t kResidMinRows = 1024; + +bool resid_log() { + static const bool cached = std::getenv("KQ_QMM_LOG") != nullptr; + return resid_tune_live() ? std::getenv("KQ_QMM_LOG") != nullptr : cached; +} + +// The sum from quantized_matmul and stock ops. The op builds this where it +// has no kernel, and the derivatives of the kernel are the derivatives of +// this. residual and gate are in the output dtype. +mx::array resid_composed( + const mx::array& x, + const mx::array& w, + const mx::array& residual, + const std::optional& gate, + const std::string& kquant_type, + mx::Stream s) { + auto scales = mx::zeros({1}, mx::uint8, s); + auto y = quantized_matmul( + x, + w, + scales, + kquant_type, + true, + std::nullopt, + std::nullopt, + std::nullopt, + s); + if (gate.has_value()) { + y = mx::multiply(*gate, y, s); + } + return mx::add(residual, y, s); +} + +} // namespace + +#ifdef _METAL_ + +void KQuantMatmulResid::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + const mx::array x = kq_row_major_gpu(inputs[0], s); + // A weight whose rows are dense binds as it is, at its own offset. Any + // other view is copied to dense rows first. The residual and the gate + // bind as they are when they are row-major. + const mx::array w = kq_matrix_dense_gpu(inputs[1], s); + const mx::array r = kq_row_major_gpu(inputs[2], s); + auto& out = outputs[0]; + kq_check_weight_base(w, kquant_type_, "quantized_matmul_residual"); + out.set_data(mx::allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + const bool gated = inputs.size() > 3; + const mx::array g = gated ? kq_row_major_gpu(inputs[3], s) : r; + + const int K = x.shape(-1); + const int N = w.shape(0); + // A gate per batch: the batch dims run in the grid and M is the rows of + // one batch. Else every row folds into M. + const int64_t rows = static_cast(x.size() / K); + const bool per_batch = gated && g.size() > static_cast(N); + const int B = per_batch ? static_cast(g.size() / N) : 1; + const int M = static_cast(rows / B); + const int g_stride = per_batch ? N : 0; + + const int max_k = kq_nax_large_max_k(kquant_type_); + const int bm = max_k > 0 && K > max_k ? 64 : 128; + const int bk = resid_kernel_bk(kquant_type_, x.dtype()); + const bool aligned = N % kResidBn == 0; + const std::string kname = kq_kname_prefix(kquant_type_) + "qmm_t_nax_resid_" + + kq_type_string(x.dtype()) + "_bm" + std::to_string(bm) + "_bn64_bk" + + std::to_string(bk) + "_wm2_wn2_alN_" + (aligned ? "true" : "false") + + "_gate_" + (gated ? "true" : "false"); + + MTL::Size group_dims(32, 2, 2); + MTL::Size grid_dims( + (N + kResidBn - 1) / kResidBn, + (static_cast(M) + bm - 1) / bm, + static_cast(B)); + if (resid_log()) { + std::fprintf( + stderr, + "[kq qmm] %s M=%d N=%d K=%d B=%d grid=%lu,%lu,%lu\n", + kname.c_str(), + M, + N, + K, + B, + static_cast(grid_dims.width), + static_cast(grid_dims.height), + static_cast(grid_dims.depth)); + } + + auto kernel = kq_get_kernel(d, kname); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + int c = 0; + ce.set_input_array(w, c++); + ce.set_input_array(x, c++); + ce.set_input_array(r, c++); + ce.set_input_array(g, c++); + ce.set_output_array(out, c++); + ce.set_bytes(K, c++); + ce.set_bytes(N, c++); + ce.set_bytes(M, c++); + ce.set_bytes(g_stride, c++); + ce.dispatch_threadgroups(grid_dims, group_dims); +} + +#else // !_METAL_ + +void KQuantMatmulResid::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error(std::string(kResidOp) + " requires Metal."); +} + +#endif + +void KQuantMatmulResid::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + std::string(kResidOp) + " CPU streams build the composed ops."); +} + +// The kernel computes the function of the composition, so its transforms +// are the composition's, with the quantized weight frozen as in +// quantized_matmul. +namespace { + +std::optional resid_gate_of(const std::vector& in) { + return in.size() > 3 ? std::optional(in[3]) : std::nullopt; +} + +} // namespace + +std::vector KQuantMatmulResid::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector&) { + for (auto arg : argnums) { + if (arg == 1) { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul_residual vjp: no gradient wrt the " + "quantized weight (the kquant base is frozen)."); + } + } + // The inputs other than the weight, in their order. + std::vector at; + for (size_t i = 0; i < primals.size(); i++) { + if (i != 1) { + at.push_back(primals[i]); + } + } + auto fun = [&](const std::vector& in) { + std::optional g = + in.size() > 2 ? std::optional(in[2]) : std::nullopt; + return std::vector{ + resid_composed(in[0], primals[1], in[1], g, kquant_type_, stream())}; + }; + auto vjps = mx::vjp(fun, at, cotangents).second; + std::vector out; + for (auto arg : argnums) { + out.push_back(vjps[arg == 0 ? 0 : arg - 1]); + } + return out; +} + +// jvp and vmap are the transforms of the composition. A tangent or a batch +// of x or of the weight reaches quantized_matmul, which has neither +// transform and refuses by its own name, as it does where the op builds +// the composition. +std::vector KQuantMatmulResid::jvp( + const std::vector& primals, + const std::vector& tangents, + const std::vector& argnums) { + auto fun = [&](const std::vector& in) { + std::vector full = primals; + for (size_t i = 0; i < argnums.size(); i++) { + full[argnums[i]] = in[i]; + } + return std::vector{resid_composed( + full[0], + full[1], + full[2], + resid_gate_of(full), + kquant_type_, + stream())}; + }; + std::vector at; + for (auto arg : argnums) { + at.push_back(primals[arg]); + } + return mx::jvp(fun, at, tangents).second; +} + +std::pair, std::vector> KQuantMatmulResid::vmap( + const std::vector& inputs, + const std::vector& axes) { + auto fun = [&](const std::vector& in) { + return std::vector{resid_composed( + in[0], in[1], in[2], resid_gate_of(in), kquant_type_, stream())}; + }; + return {mx::vmap(fun, axes, {0})(inputs), {0}}; +} + +bool KQuantMatmulResid::is_equivalent(const mx::Primitive& other) const { + const auto& o = static_cast(other); + return kquant_type_ == o.kquant_type_; +} + +std::vector KQuantMatmulResid::output_shapes( + const std::vector& inputs) { + return {inputs[2].shape()}; +} + +mx::array quantized_matmul_residual( + mx::array x, + mx::array w, + const std::string& kquant_type, + mx::array residual, + const std::optional& gate, + mx::StreamOrDevice s_) { + const std::string op = kResidOp; + const KQuantCodec* codec = codec_by_name(kquant_type); + if (codec == nullptr) { + throw std::invalid_argument( + op + " Unknown kquant_type: '" + kquant_type + "'."); + } + if (w.dtype() != mx::uint8 || w.ndim() != 2) { + throw std::invalid_argument( + op + " w must be 2-D uint8 (raw GGUF wire bytes)."); + } + if (w.shape(-1) % codec->bytes_per_block != 0) { + throw std::invalid_argument( + op + " weight last dim (" + std::to_string(w.shape(-1)) + + " bytes) is not a whole number of " + + std::to_string(codec->bytes_per_block) + "-byte " + codec->name + + " blocks."); + } + const int K = + (w.shape(-1) / codec->bytes_per_block) * codec->weights_per_block; + const int N = w.shape(0); + if (x.ndim() < 1 || x.shape(-1) != K) { + throw std::invalid_argument( + op + " x last dim does not match the weight inner dim (" + + std::to_string(K) + ")."); + } + if (!mx::issubdtype(x.dtype(), mx::floating)) { + throw std::invalid_argument( + op + " Only real floating x dtypes are supported."); + } + auto out_shape = x.shape(); + out_shape.back() = N; + if (residual.shape() != out_shape || + !mx::issubdtype(residual.dtype(), mx::floating)) { + throw std::invalid_argument( + op + + " residual must be a float array of the output shape, x's " + "shape with N " + + std::to_string(N) + " last."); + } + // A gate of [N], or of x's batch dims, then 1, then N. + if (gate.has_value()) { + const auto& g = *gate; + mx::Shape batch_shape = out_shape; + bool ok = mx::issubdtype(g.dtype(), mx::floating); + if (g.ndim() == 1) { + ok = ok && g.shape(0) == N; + } else if (x.ndim() >= 2) { + batch_shape[batch_shape.size() - 2] = 1; + ok = ok && g.shape() == batch_shape; + } else { + ok = false; + } + if (!ok) { + throw std::invalid_argument( + op + + " gate must be a float array of shape [N] or of x's batch " + "dims, then 1, then N, with N " + + std::to_string(N) + "."); + } + } + if (K >= kResidMaxCols || N >= kResidMaxCols) { + throw std::invalid_argument( + op + " K and N must each be below " + std::to_string(kResidMaxCols) + + " (2^25): K " + std::to_string(K) + ", N " + std::to_string(N) + "."); + } + const int64_t rows = K > 0 ? static_cast(x.size()) / K : 0; + if (rows >= kResidMaxRows) { + throw std::invalid_argument( + op + " x must hold fewer than " + std::to_string(kResidMaxRows) + + " (2^31) rows: it holds " + std::to_string(rows) + "."); + } + auto s = mx::to_stream(s_); + + // float32 x runs in bfloat16, as in quantized_matmul. The residual and + // the gate run in the output dtype. + const mx::Dtype out_type = + x.dtype() == mx::float32 ? mx::bfloat16 : x.dtype(); + auto r = mx::astype(residual, out_type, s); + std::optional g; + if (gate.has_value()) { + g = mx::astype(*gate, out_type, s); + } + + const int bk = resid_kernel_bk(kquant_type, out_type); + const int mode = resid_mode(); + const bool fused = s.device == mx::Device::gpu && nax_available() && + mode != 0 && bk > 0 && K > 0 && N > 0 && K % bk == 0 && + (mode == 2 || rows >= kResidMinRows); + if (!fused) { + return resid_composed(x, w, r, g, kquant_type, s); + } + + kq_check_weight_base_at_build(w, kquant_type, s, "quantized_matmul_residual"); + std::vector in = { + mx::astype(x, out_type, s), std::move(w), std::move(r)}; + if (g.has_value()) { + in.push_back(*g); + } + return mx::array( + std::move(out_shape), + out_type, + std::make_shared(s, kquant_type), + std::move(in)); +} + +} // namespace mlx_kquant diff --git a/src/kquant_metal_internal.h b/src/kquant_metal_internal.h index ec5a4d1e..b238b3f8 100644 --- a/src/kquant_metal_internal.h +++ b/src/kquant_metal_internal.h @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -198,14 +199,15 @@ inline int verify_qmv_max_rows() { } // Codecs with a register-resident MMA verify kernel (kq_verify_mma.h) and -// the output rows one of its threadgroups covers (8 simdgroups of 8 rows -// per row tile); 0 for codecs without the kernel. M <= 8 only. +// the output rows one of its threadgroups covers: 8 simdgroups of 8-row +// tiles, two tiles each for q4_0 and q8_0 and one for pq2_0, q2_0 and +// ptq1_0; 0 for codecs without the kernel. M <= 8 only. inline int codec_verify_mma_rows(const std::string& kquant_type) { - if (kquant_type == "pq2_0" || kquant_type == "q2_0" || - kquant_type == "q4_0" || kquant_type == "q8_0") { + if (kquant_type == "q4_0" || kquant_type == "q8_0") { return 128; } - if (kquant_type == "ptq1_0") { + if (kquant_type == "pq2_0" || kquant_type == "q2_0" || + kquant_type == "ptq1_0") { return 64; } return 0; @@ -401,6 +403,30 @@ inline MTL::ComputePipelineState* kq_get_kernel( return d.get_kernel(kname, lib); } +// True when the bundled metallib carries a function named kname. Looks the +// name up without building a pipeline, so a tile override can raise with +// the kernel name before any dispatch. +inline bool kq_has_kernel(Device& d, const std::string& kname) { + static MTL::Library* lib = d.get_library("mlx_kquant", metallib_dir()); + auto pool = NS::TransferPtr(NS::AutoreleasePool::alloc()->init()); + auto fn = NS::TransferPtr(lib->newFunction( + NS::String::string(kname.c_str(), NS::UTF8StringEncoding))); + return fn.get() != nullptr; +} + +// Raises when a lever setting selects a kernel the metallib does not carry, +// naming the setting and the kernel. +inline void kq_require_kernel( + Device& d, + const std::string& setting, + const std::string& kname) { + if (!kq_has_kernel(d, kname)) { + throw std::invalid_argument( + "[mlx_kquant] " + setting + " selects kernel " + kname + + ", which is not instantiated."); + } +} + // Func-constant variant for kernels specialized via an MTLFCList. // hash_name must encode the func-const values so each specialization gets a // distinct pipeline-state cache entry (d.get_kernel(name, lib, hash, consts)). diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index e21e8027..e1e34e52 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -108,13 +108,7 @@ inline bool kq_moe_dedup(int T, int max_t) { return false; } static const bool has_env = std::getenv("KQ_MOE_DEDUP") != nullptr; - if (has_env) { - const char* e = std::getenv("KQ_MOE_DEDUP"); - if (e != nullptr && std::atoi(e) == 0) { - return false; - } - } - return true; + return !has_env || kq_env_flag("KQ_MOE_DEDUP", 1) != 0; } constexpr int KQ_DD_MAX_T = 8; @@ -183,11 +177,8 @@ inline bool kq_moe_sp(int64_t coarse_tgs, int S) { return false; // threadgroup is 64 * S threads; 1024 cap => S <= 16 } static const bool has_env = std::getenv("KQ_MOE_SP") != nullptr; - if (has_env) { - const char* e = std::getenv("KQ_MOE_SP"); - return e != nullptr && std::atoi(e) != 0; - } - return coarse_tgs < 2048; + const int v = has_env ? kq_env_flag("KQ_MOE_SP", -1) : -1; + return v == -1 ? coarse_tgs < 2048 : v != 0; } // Form of a score-mixed down gather (mix and mix_ns) where more than one @@ -304,11 +295,12 @@ inline std::string kq_gather_stem_nx(const std::string& t, int K, int nx) { // float decode is ALU-bound at decode widths (iq2_xs, iq2_xxs, iq3_xxs) run // the "_h" kernels, which stage the grid as half4 pairs, stage the // activation row as half in dynamic threadgroup memory (K halves), and dot -// in half with a float sum per 16-weight chunk. Expert slots only; shared -// experts keep the float path. Read live once set (KQ_MOE_NX pattern) so -// interleaved in-process A/Bs can flip arms. K is capped so the staged row -// plus the largest half table (iq2_xs, 8.3 KB) stays inside the 32 KB -// threadgroup budget. +// each 16-weight chunk in float. Expert slots only; shared experts keep the +// float path. Read live once set (KQ_MOE_NX pattern) so interleaved +// in-process A/Bs can flip arms. K is capped so the staged row plus the +// largest half table (iq2_xs, 8.3 KB) stays inside the 32 KB threadgroup +// budget. An activation past the half range (65504) turns to inf in the +// staged row. inline bool kq_moe_half(const std::string& codec, int K) { static const bool has_env = std::getenv("KQ_MOE_HALF") != nullptr; if (!has_env) { @@ -346,17 +338,38 @@ inline int kq_moe_pick_sg_h(int N, int nx) { inline bool kq_gather_fine(int64_t coarse_tgs) { static const bool has_env = std::getenv("KQ_GATHER_FINE") != nullptr; - if (has_env) { - const char* e = std::getenv("KQ_GATHER_FINE"); - const int v = e == nullptr ? -1 : std::atoi(e); - if (v == 0) { - return false; - } - if (v == 1) { - return true; - } + const int v = has_env ? kq_env_flag("KQ_GATHER_FINE", -1) : -1; + return v == -1 ? coarse_tgs < 2048 : v != 0; +} + +// Layout of the operands the MoE kernels read without strides. An +// unevaluated view reports the default dense strides to an op builder, so +// the builders decide only for an operand that has memory, and each eval +// decides for the rest before it binds a buffer. +// +// The kernels find expert e of a stack at e * N * bytes_per_row and read a +// shared-expert tensor as dense rows. A stack is too large to copy on a +// call, so a stack or a shared-expert tensor that is not row-major is +// refused, with the builder's message. +inline std::string moe_op(const char* name) { + return std::string("[mlx_kquant.") + name + "]"; +} + +inline void moe_check_stack(const char* name, const mx::array& w) { + kq_check_row_major(w, moe_op(name) + " weights"); +} + +inline void moe_check_shexp(const char* name, const mx::array& w) { + kq_check_row_major(w, moe_op(name) + " shared-expert weights"); +} + +// The packed-mxfp4 stacks carry a companion scales tensor. +inline void +moe_check_stack(const char* name, const mx::array& w, const mx::array& sc) { + if (!w.flags().row_contiguous || !sc.flags().row_contiguous) { + throw std::invalid_argument( + moe_op(name) + " weights/scales must be row-contiguous."); } - return coarse_tgs < 2048; } } // namespace @@ -369,7 +382,6 @@ void KQuantMoEGLU::eval_gpu( auto& s = stream(); auto& d = mx::metal::device(s.device); auto& out = outputs[0]; - out.set_data(mx::allocator::malloc(out.nbytes())); const auto& gw = inputs[0]; const auto& gs = inputs[1]; @@ -377,8 +389,12 @@ void KQuantMoEGLU::eval_gpu( const auto& uw = inputs[3]; const auto& us = inputs[4]; const auto& ub = inputs[5]; - const auto& x = inputs[6]; const auto& indices = inputs[7]; + moe_check_stack("moe_glu_gather", gw, gs); + moe_check_stack("moe_glu_gather", uw, us); + // x is read as dense rows: a view with other strides is copied first. + const mx::array x = kq_row_major_gpu(inputs[6], s); + out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); int R = indices.shape(1); @@ -416,13 +432,14 @@ void KQuantGatherQMVBias::eval_gpu( auto& s = stream(); auto& d = mx::metal::device(s.device); auto& out = outputs[0]; - out.set_data(mx::allocator::malloc(out.nbytes())); const auto& w = inputs[0]; const auto& sc = inputs[1]; const auto& b = inputs[2]; - const auto& x = inputs[3]; const auto& indices = inputs[4]; + moe_check_stack("gather_qmv_bias", w, sc); + const mx::array x = kq_row_major_gpu(inputs[3], s); + out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); int R = indices.shape(1); @@ -455,14 +472,15 @@ void KQuantGatherQMVMixBias::eval_gpu( auto& s = stream(); auto& d = mx::metal::device(s.device); auto& out = outputs[0]; - out.set_data(mx::allocator::malloc(out.nbytes())); const auto& w = inputs[0]; const auto& sc = inputs[1]; const auto& b = inputs[2]; - const auto& x = inputs[3]; const auto& indices = inputs[4]; const auto& scores = inputs[5]; + moe_check_stack("gather_qmv_mix_bias", w, sc); + const mx::array x = kq_row_major_gpu(inputs[3], s); + out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); int S = indices.shape(1); @@ -502,10 +520,12 @@ void KQuantMoEGLUKQ::eval_gpu( const bool biased = inputs.size() == 6; const auto& gw = inputs[0]; const auto& uw = inputs[1]; - const auto& x = inputs[biased ? 4 : 2]; const auto& indices = inputs.back(); + moe_check_stack("moe_glu_gather_kq", gw); + moe_check_stack("moe_glu_gather_kq", uw); kq_check_weight_base(gw, kquant_type_, "moe_glu_gather_kq"); kq_check_weight_base(uw, kquant_type_, "moe_glu_gather_kq"); + const mx::array x = kq_row_major_gpu(inputs[biased ? 4 : 2], s); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -574,9 +594,10 @@ void KQuantGatherQMVKQ::eval_gpu( const size_t n_base = inputs.size() - kq_lora_input_count(lora_flags_); const bool biased = n_base == 4; const auto& w = inputs[0]; - const auto& x = inputs[biased ? 2 : 1]; const auto& indices = inputs[n_base - 1]; + moe_check_stack("gather_qmv_kq", w); kq_check_weight_base(w, kquant_type_, "gather_qmv_kq"); + const mx::array x = kq_row_major_gpu(inputs[biased ? 2 : 1], s); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -649,12 +670,16 @@ void KQuantMoEGLUShexpKQ::eval_gpu( const auto& uw = inputs[1]; const auto& sgw = inputs[2]; const auto& suw = inputs[3]; - const auto& x = inputs[4]; const auto& indices = inputs[5]; + moe_check_stack("moe_glu_gather_shexp_kq", gw); + moe_check_stack("moe_glu_gather_shexp_kq", uw); + moe_check_shexp("moe_glu_gather_shexp_kq", sgw); + moe_check_shexp("moe_glu_gather_shexp_kq", suw); kq_check_weight_base(gw, kquant_type_, "moe_glu_gather_shexp_kq"); kq_check_weight_base(uw, kquant_type_, "moe_glu_gather_shexp_kq"); kq_check_weight_base(sgw, shexp_type_, "moe_glu_gather_shexp_kq"); kq_check_weight_base(suw, shexp_up_type_, "moe_glu_gather_shexp_kq"); + const mx::array x = kq_row_major_gpu(inputs[4], s); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -749,11 +774,13 @@ void KQuantGatherQMVMixKQ::eval_gpu( const auto& w = inputs[0]; const auto& sw = inputs[1]; - const auto& x = inputs[2]; const auto& indices = inputs[3]; const auto& scores = inputs[4]; + moe_check_stack("gather_qmv_mix_kq", w); + moe_check_shexp("gather_qmv_mix_kq", sw); kq_check_weight_base(w, kquant_type_, "gather_qmv_mix_kq"); kq_check_weight_base(sw, shexp_type_, "gather_qmv_mix_kq"); + const mx::array x = kq_row_major_gpu(inputs[2], s); out.set_data(mx::allocator::malloc(out.nbytes())); int T = x.shape(0); @@ -882,10 +909,11 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( auto& out = outputs[0]; const auto& w = inputs[0]; - const auto& x = inputs[1]; const auto& indices = inputs[2]; const auto& scores = inputs[3]; + moe_check_stack("gather_qmv_mix_ns_kq", w); kq_check_weight_base(w, kquant_type_, "gather_qmv_mix_ns_kq"); + const mx::array x = kq_row_major_gpu(inputs[1], s); out.set_data(mx::allocator::malloc(out.nbytes())); int T = x.shape(0); @@ -978,10 +1006,12 @@ void KQuantMoERouterTopK::eval_gpu( auto& d = mx::metal::device(s.device); auto& indices = outputs[0]; auto& scores = outputs[1]; + // The kernel reads logits as dense rows: a view with other strides is + // copied first. + const mx::array logits = kq_row_major_gpu(inputs[0], s); indices.set_data(mx::allocator::malloc(indices.nbytes())); scores.set_data(mx::allocator::malloc(scores.nbytes())); - const auto& logits = inputs[0]; int SHARED = shared_ ? 1 : 0; int HAS_PES = has_pes_ ? 1 : 0; int HAS_BIAS = has_bias_ ? 1 : 0; @@ -992,8 +1022,8 @@ void KQuantMoERouterTopK::eval_gpu( int R = top_k_; int NORM = norm_ ? 1 : 0; // pes/bias must be bound buffers even when unused; logits stands in. - const auto& pes = has_pes_ ? inputs[1] : inputs[0]; - const auto& bias = has_bias_ ? inputs[1 + (has_pes_ ? 1 : 0)] : inputs[0]; + const auto& pes = has_pes_ ? inputs[1] : logits; + const auto& bias = has_bias_ ? inputs[1 + (has_pes_ ? 1 : 0)] : logits; std::string kname = "kq_moe_router_topk_" + kq_type_string(logits.dtype()); auto kernel = kq_get_kernel(d, kname); @@ -1338,7 +1368,9 @@ void check_expert_stack( throw std::invalid_argument( std::string(op) + " K must be a multiple of 32."); } - if (!w.flags().row_contiguous || !sc.flags().row_contiguous) { + // An operand without memory is checked at eval. + if ((kq_has_layout(w) && !w.flags().row_contiguous) || + (kq_has_layout(sc) && !sc.flags().row_contiguous)) { throw std::invalid_argument( std::string(op) + " weights/scales must be row-contiguous."); } @@ -1348,6 +1380,18 @@ mx::array prep_bias(const mx::array& b, mx::StreamOrDevice s) { return mx::contiguous(mx::astype(b, mx::float32, s), false, s); } +// x or router logits for a kernel that reads dense rows. An array with +// memory is copied here when it is not row-major. An array without memory +// reports the default dense strides, so it goes in as it is and the eval +// decides (kq_row_major_gpu): x is unevaluated on nearly every call, and a +// node here would cost every one of them. +mx::array moe_dense_x(const mx::array& x, mx::StreamOrDevice s) { + if (kq_has_layout(x) && !x.flags().row_contiguous) { + return mx::contiguous(x, false, s); + } + return x; +} + mx::array prep_indices(const mx::array& idx, mx::StreamOrDevice s) { return mx::contiguous(mx::astype(idx, mx::uint32, s), false, s); } @@ -1384,9 +1428,9 @@ void check_kq_expert_stack( throw std::invalid_argument( std::string(op) + " N must be a multiple of 8."); } - if (!w.flags().row_contiguous) { - throw std::invalid_argument( - std::string(op) + " weights must be row-contiguous."); + // A stack without memory is checked at eval. + if (kq_has_layout(w)) { + kq_check_row_major(w, std::string(op) + " weights"); } } @@ -1416,9 +1460,9 @@ void check_kq_shexp_row( std::string(op) + " shared-expert weights must match the expert stack row shape."); } - if (!w.flags().row_contiguous) { - throw std::invalid_argument( - std::string(op) + " shared-expert weights must be row-contiguous."); + // A tensor without memory is checked at eval. + if (kq_has_layout(w)) { + kq_check_row_major(w, std::string(op) + " shared-expert weights"); } } @@ -1464,7 +1508,7 @@ mx::array moe_glu_gather( "[mlx_kquant.moe_glu_gather] gate/up expert shapes must match."); } - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); mx::Shape out_shape = {x.shape(0), indices.shape(1), gate_w.shape(1)}; return mx::array( std::move(out_shape), @@ -1505,7 +1549,7 @@ mx::array gather_qmv_bias( int K = x.shape(2); check_expert_stack("[mlx_kquant.gather_qmv_bias]", w, scales, bias, K); - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); mx::Shape out_shape = {x.shape(0), x.shape(1), w.shape(1)}; return mx::array( std::move(out_shape), @@ -1548,7 +1592,7 @@ mx::array gather_qmv_mix_bias( int K = x.shape(2); check_expert_stack(op, w, scales, bias, K); - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); mx::Shape out_shape = {x.shape(0), w.shape(1)}; return mx::array( std::move(out_shape), @@ -1637,7 +1681,7 @@ mx::array moe_glu_gather_kq( } } - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); mx::Shape out_shape = {x.shape(0), indices.shape(1), gate_w.shape(1)}; std::vector op_inputs = {std::move(gate_w), std::move(up_w)}; if (biased) { @@ -1702,7 +1746,7 @@ mx::array gather_qmv_kq( } } - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); const int E = w.shape(0); const int N = w.shape(1); mx::Shape out_shape = {x.shape(0), x.shape(1), N}; @@ -1797,7 +1841,7 @@ mx::array moe_glu_gather_shexp_kq( check_kq_shexp_row(op, shexp_gate_w, shexp_type, K, N); check_kq_shexp_row(op, shexp_up_w, shexp_up_type, K, N); - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); mx::Shape out_shape = {x.shape(0), indices.shape(1) + 1, N}; kq_check_weight_base_at_build( gate_w, kquant_type, s, "moe_glu_gather_shexp_kq"); @@ -1866,7 +1910,7 @@ mx::array gather_qmv_mix_kq( check_kq_expert_stack(op, w, kquant_type, K); check_kq_shexp_row(op, shexp_w, shexp_type, K, w.shape(1)); - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); mx::Shape out_shape = {x.shape(0), w.shape(1)}; kq_check_weight_base_at_build(w, kquant_type, s, "gather_qmv_mix_kq"); kq_check_weight_base_at_build(shexp_w, shexp_type, s, "gather_qmv_mix_kq"); @@ -1916,7 +1960,7 @@ mx::array gather_qmv_mix_ns_kq( check_kq_expert_stack(op, w, kquant_type, K); kq_check_weight_base_at_build(w, kquant_type, s, "gather_qmv_mix_ns_kq"); - auto x_c = x.flags().row_contiguous ? x : mx::contiguous(x, false, s); + auto x_c = moe_dense_x(x, s); const int E = w.shape(0); const int N = w.shape(1); mx::Shape out_shape = {x.shape(0), N}; @@ -1997,8 +2041,7 @@ std::vector moe_router_topk( throw std::invalid_argument(std::string(op) + " requires E <= 1024."); } - auto l_c = - logits.flags().row_contiguous ? logits : mx::contiguous(logits, false, s); + auto l_c = moe_dense_x(logits, s); int T = l_c.shape(0); std::vector inputs = {std::move(l_c)}; if (per_expert_scale.has_value()) { diff --git a/src/kquant_ops.cpp b/src/kquant_ops.cpp index 4e6c2e0b..b7767a68 100644 --- a/src/kquant_ops.cpp +++ b/src/kquant_ops.cpp @@ -1,6 +1,8 @@ #include #include +#include +#include #include #include #include @@ -23,48 +25,18 @@ namespace mlx_kquant { namespace { -// Matrix-contiguity helper. A "matrix-contiguous" tensor only needs -// its LAST TWO dims densely packed (stride(-2)==shape(-1) && stride(-1)==1); a -// strided LEADING/batch dim is fine because the qmv/qmm/gather kernels walk it -// via the strides passed in add_strides_and_shapes. Using full -// mx::contiguous() here instead would force a copy of the whole tensor whenever -// the leading dim is strided - e.g. the gate/up halves of a fused MoE -// gate_up_exps tensor are slices [..., :H, :] / [..., H:, :] whose expert -// (leading) stride is 2*H rows: matrix-contiguous but NOT row_contiguous, so -// full contiguify copies ~142 MB PER gather call (the decode MoE 5x -// regression). -// True when the elements of x are laid out row-major in memory (size-1 -// dims ignored), so a kernel can walk them as dense rows without a copy. -inline bool kq_rows_dense(const mx::array& x) { - int64_t expect = 1; - for (int i = x.ndim() - 1; i >= 0; --i) { - if (x.shape(i) == 1) { - continue; - } - if (x.strides()[i] != expect) { - return false; - } - expect *= x.shape(i); - } - return true; -} - -inline mx::array kq_ensure_row_contiguous_matrix( - const mx::array& x, - mx::StreamOrDevice s) { - if (x.ndim() < 2) { - if (x.strides()[0] == 1) { - return x; - } - } else { - auto stride_0 = x.strides()[x.ndim() - 2]; - auto stride_1 = x.strides()[x.ndim() - 1]; - if (stride_0 == static_cast(x.shape(-1)) && stride_1 == 1) { - return x; - } - } - return mx::contiguous(x, false, s); -} +// Matrix-contiguity. A "matrix-contiguous" tensor only needs its LAST TWO +// dims densely packed (stride(-2)==shape(-1) && stride(-1)==1); a strided +// LEADING/batch dim is fine because the qmv/qmm/gather kernels walk it via +// the strides passed in add_strides_and_shapes. Using full mx::contiguous() +// instead would force a copy of the whole tensor whenever the leading dim is +// strided - e.g. the gate/up halves of a fused MoE gate_up_exps tensor are +// slices [..., :H, :] / [..., H:, :] whose expert (leading) stride is 2*H +// rows: matrix-contiguous but NOT row_contiguous, so full contiguify copies +// ~142 MB PER gather call (the decode MoE 5x regression). w and scales go +// through kq_matrix_dense (kquant_dense.cpp), which reads the strides of an +// array that has memory and leaves any other array to a node that decides +// at eval. // x always gets an unconditional Contiguous node: build-time strides of an // unevaluated array are meaningless, so a stride check here can hand the @@ -75,8 +47,162 @@ inline mx::array kq_dense_x(const mx::array& x, mx::StreamOrDevice s) { return mx::contiguous(x, false, s); } +// Codecs whose plain transposed quantized_matmul runs as dequantize then the +// stock matmul on a GPU without NAX before M5, from a floor in the rows of +// x. The qmm_t kernel decodes each weight tile once per 64-row tile of x; +// the dequantize pass decodes the weight once, so it pays from the floor, +// where the stock matmul outruns the in-loop decode. Floors measured on M3 +// Max, and applied to the whole class. +struct KqDequantizeRow { + const char* codec; + int min_m; +}; + +constexpr KqDequantizeRow kDequantizeRows[] = { + {"nvfp4", 2048}, + {"ptq1_0", 2048}, + {"q3_k", 2048}, + {"q4_k", 16384}, + {"q5_k", 2048}, + {"q6_k", 2048}, +}; + +// The measured range of the route. The dense copy of the weight lives while +// the call runs, so N x K caps it at 24576 x 12288 x 2 bytes; a vocabulary +// head keeps the kernel. +constexpr int kDequantizeMinK = 3840; +constexpr int kDequantizeMaxK = 12288; +constexpr int kDequantizeMinN = 3840; +constexpr int kDequantizeMaxN = 24576; + +// The row for a codec and shape, the chip class aside: min_m, min_k, max_k, +// min_n and max_n, or empty. +std::vector +kq_dequantize_shape_row(const std::string& kquant_type, int M, int N, int K) { + if (K < kDequantizeMinK || K > kDequantizeMaxK || N < kDequantizeMinN || + N > kDequantizeMaxN) { + return {}; + } + for (const auto& r : kDequantizeRows) { + if (kquant_type == r.codec && M >= r.min_m) { + return { + r.min_m, + kDequantizeMinK, + kDequantizeMaxK, + kDequantizeMinN, + kDequantizeMaxN}; + } + } + return {}; +} + +// True for the chip classes the route runs on by default: an architecture +// string applegpu_g with a two digit generation from 13 (M1) +// to 16 (M4), a one letter suffix other than p (phones), and no NAX. Any +// other string keeps the kernel. +bool kq_dequantize_class(const std::string& arch, bool nax) { + if (nax) { + return false; + } + const std::string head = "applegpu_g"; + if (arch.size() != head.size() + 3 || arch.compare(0, head.size(), head)) { + return false; + } + const char d0 = arch[head.size()]; + const char d1 = arch[head.size() + 1]; + const char suffix = arch[head.size() + 2]; + if (!std::isdigit(static_cast(d0)) || + !std::isdigit(static_cast(d1)) || + !std::islower(static_cast(suffix)) || suffix == 'p') { + return false; + } + const int gen = (d0 - '0') * 10 + (d1 - '0'); + return gen >= 13 && gen <= 16; +} + +// Whether a plain transposed call on a 2-D weight on a GPU stream takes the +// route: never under off or while KQ_QMM_TILE pins the tile, on any GPU +// under force, on the class under auto, and then only for a codec row and +// shape of the table. +bool kq_dequantize_route_takes( + bool on_class, + int mode, + bool pinned, + const std::string& kquant_type, + int M, + int N, + int K) { + if (mode == kDequantOff || pinned) { + return false; + } + if (mode != kDequantForce && !on_class) { + return false; + } + return !kq_dequantize_shape_row(kquant_type, M, N, K).empty(); +} + +#ifdef _METAL_ +// This GPU's class verdict, read once per process. NAX first. +bool kq_dequantize_class_here() { + static const bool here = []() { + if (kq_is_nax_available()) { + return false; + } + auto& d = mx::metal::device(mx::Device::gpu); + return kq_dequantize_class(d.get_architecture(), false); + }(); + return here; +} +#endif + +// Dequantize then the stock matmul, carried by KQuantDequantRoute for the +// vjp of the kernel path. +mx::array kq_dequantize_matmul( + const mx::array& x, + const mx::array& w, + const mx::array& scales, + const std::string& kquant_type, + mx::Dtype out_type, + mx::Stream s); + } // namespace +std::vector alu_dequantize_row( + const std::string& arch, + bool nax, + const std::string& kquant_type, + int M, + int N, + int K) { + if (!kq_dequantize_class(arch, nax)) { + return {}; + } + return kq_dequantize_shape_row(kquant_type, M, N, K); +} + +bool dequantize_route_takes( + const std::string& arch, + bool nax, + const std::string& mode, + bool pinned, + const std::string& kquant_type, + int M, + int N, + int K) { + int m = kDequantAuto; + if (mode == "off") { + m = kDequantOff; + } else if (mode == "force") { + m = kDequantForce; + } else if (mode != "auto") { + throw std::invalid_argument( + "[mlx_kquant] dequantize route mode must be auto, off or force, got '" + + mode + "'."); + } + return kq_dequantize_route_takes( + kq_dequantize_class(arch, nax), m, pinned, kquant_type, M, N, K); +} + // ----------------------------- ops ----------------------------- mx::array dequantize( @@ -124,8 +250,8 @@ mx::array dequantize( out_shape.back() = (w.shape(-1) / codec->bytes_per_block) * codec->weights_per_block; - // Row-contiguize at the op level so eval_gpu can assume dense inputs. - auto w_c = w.flags().row_contiguous ? w : mx::contiguous(w, false, s); + // Row-major bytes for the kernel. A w without memory is decided at eval. + auto w_c = kq_row_major(w, s); kq_check_weight_base_at_build(w, kquant_type, s, "dequantize"); @@ -196,30 +322,84 @@ mx::array quantized_matmul( } auto s = mx::to_stream(s_); + const bool lora = + lora_a.has_value() || lora_b.has_value() || lora_rows.has_value(); + if (lora && w.ndim() != 2) { + throw std::invalid_argument( + "[mlx_kquant.quantized_matmul] the LoRA epilogue needs a 2-D w."); + } + // Rows of x, which the leading dims give for any inner dim. + int rows = 1; + for (int i = 0; i + 1 < static_cast(x.ndim()); i++) { + rows *= x.shape(i); + } + // An inner dim of 0 makes every sum empty, so the result is zeros. The + // kernels divide by the inner dim. The LoRA arguments pass the checks of + // a full call first. + if (w_inner_dims == 0) { + if (lora) { + std::vector checked; + kq_lora_prep( + "[mlx_kquant.quantized_matmul]", + lora_a, + lora_b, + std::nullopt, + std::nullopt, + lora_rows, + out_type, + w_inner_dims, + w_outer_dims, + 0, + {rows}, + 1, + s, + checked); + } + auto shape = x.shape(); + shape.back() = w_outer_dims; + return mx::zeros(shape, out_type, s); + } kq_check_weight_base_at_build(w, kquant_type, s, "quantized_matmul"); +#ifdef _METAL_ + if (transpose && !lora && w.ndim() == 2 && s.device == mx::Device::gpu) { + const bool here = kq_dequantize_class_here(); + if (kq_dequantize_route_takes( + here, + kq_dequant_route_mode(), + kq_qmm_tile_pinned(), + kquant_type, + rows, + w_outer_dims, + w_inner_dims)) { + if (kq_qmm_log_on()) { + std::fprintf( + stderr, + "[kq dequant] %s M=%d N=%d K=%d\n", + kquant_type.c_str(), + rows, + w_outer_dims, + w_inner_dims); + } + return kq_dequantize_matmul(x, w, scales, kquant_type, out_type, s); + } + } +#endif + // Cast x to the output dtype, then matrix-row-contiguize x / w / scales at // the op level so eval_gpu can assume dense inputs. auto x_c = kq_dense_x(mx::astype(x, out_type, s), s); - auto w_c = kq_ensure_row_contiguous_matrix(w, s); - auto scales_c = kq_ensure_row_contiguous_matrix(scales, s); + auto w_c = kq_matrix_dense(w, s); + auto scales_c = kq_matrix_dense(scales, s); auto out_shape = x_c.shape(); out_shape.back() = w_outer_dims; std::vector op_inputs = {x_c, w_c, scales_c}; int lora_flags = 0; - if (lora_a.has_value() || lora_b.has_value() || lora_rows.has_value()) { - if (w.ndim() != 2) { - throw std::invalid_argument( - "[mlx_kquant.quantized_matmul] the LoRA epilogue needs a 2-D w."); - } - // The epilogue walks x as dense [rows, K] rows; views that are only - // matrix-contiguous with a leading stride gap get copied. - if (!kq_rows_dense(op_inputs[0])) { - op_inputs[0] = mx::contiguous(op_inputs[0], false, s); - } - const int rows = static_cast(x_c.size() / w_inner_dims); + if (lora) { + // The epilogue walks x as dense [rows, K] rows, which the Contiguous + // node of x gives it. lora_flags = kq_lora_prep( "[mlx_kquant.quantized_matmul]", lora_a, @@ -249,6 +429,90 @@ mx::array quantized_matmul( std::move(op_inputs)); } +namespace { + +mx::array kq_dequantize_matmul( + const mx::array& x, + const mx::array& w, + const mx::array& scales, + const std::string& kquant_type, + mx::Dtype out_type, + mx::Stream s) { + auto w_f = dequantize(w, scales, kquant_type, out_type, s); + auto y = mx::stop_gradient( + mx::matmul(mx::astype(x, out_type, s), mx::transpose(w_f, s), s), s); + auto shape = y.shape(); + return mx::array( + std::move(shape), + out_type, + std::make_shared(s, kquant_type), + {y, x, w, scales}); +} + +} // namespace + +void KQuantDequantRoute::eval_cpu( + const std::vector& inputs, + std::vector& outputs) { + outputs[0].copy_shared_buffer(inputs[0]); +} + +void KQuantDequantRoute::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + outputs[0].copy_shared_buffer(inputs[0]); +} + +std::vector KQuantDequantRoute::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector&) { + // primals = {y (stop_gradient), x, w, scales}, as KQuantMatmul::vjp: x + // gets the quantized matmul the other way, w and scales throw. + std::vector vjps; + for (auto arg : argnums) { + if (arg == 0) { + vjps.push_back(cotangents[0]); + } else if (arg == 1) { + auto dx = quantized_matmul( + cotangents[0], + primals[2], + primals[3], + kquant_type_, + false, + std::nullopt, + std::nullopt, + std::nullopt, + stream()); + vjps.push_back(mx::astype(dx, primals[1].dtype(), stream())); + } else if (arg == 2) { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul vjp: no gradient wrt the quantized " + "weights (the kquant base is frozen)."); + } else { + throw std::invalid_argument( + "[mlx_kquant] quantized_matmul vjp: no gradient wrt scales."); + } + } + return vjps; +} + +std::vector KQuantDequantRoute::jvp( + const std::vector&, + const std::vector&, + const std::vector&) { + throw std::invalid_argument( + "[Primitive::jvp] Not implemented for KQuantMatmul."); +} + +std::pair, std::vector> KQuantDequantRoute::vmap( + const std::vector&, + const std::vector&) { + throw std::invalid_argument( + "[Primitive::vmap] Not implemented for KQuantMatmul."); +} + mx::array quantized_matmul_qmv_bias( mx::array x, mx::array w, @@ -323,11 +587,9 @@ mx::array quantized_matmul_qmv_bias( kq_check_weight_base_at_build(w, kquant_type, s, "quantized_matmul_qmv_bias"); auto x_c = kq_dense_x(mx::astype(x, out_type, s), s); - auto w_c = kq_ensure_row_contiguous_matrix(w, s); - auto scales_c = kq_ensure_row_contiguous_matrix(scales, s); - auto bias_c = mx::astype(bias, out_type, s); - bias_c = - bias_c.flags().row_contiguous ? bias_c : mx::contiguous(bias_c, false, s); + auto w_c = kq_matrix_dense(w, s); + auto scales_c = kq_matrix_dense(scales, s); + auto bias_c = kq_row_major(mx::astype(bias, out_type, s), s); auto out_shape = x_c.shape(); out_shape.back() = N; @@ -492,9 +754,9 @@ mx::array gather_qmm( // short-circuit here skipped the compaction on the slice's first eval and // handed the stride-less rhs kernels an expert stride they mis-walk. // Contiguous::eval aliases zero-copy when w turns out packed anyway. - auto w_c = rhs_reachable ? mx::contiguous(w, false, s) - : kq_ensure_row_contiguous_matrix(w, s); - auto scales_c = kq_ensure_row_contiguous_matrix(scales, s); + auto w_c = + rhs_reachable ? mx::contiguous(w, false, s) : kq_matrix_dense(w, s); + auto scales_c = kq_matrix_dense(scales, s); return mx::array( std::move(out_shape), @@ -626,7 +888,7 @@ mx::array gather_qmm_seg( // The kernel computes each expert's base pointer as expert * N * K_w, so w // must be fully row-contiguous (no strided expert dim). auto w_c = mx::contiguous(w, false, s); - auto scales_c = kq_ensure_row_contiguous_matrix(scales, s); + auto scales_c = kq_matrix_dense(scales, s); auto map_c = mx::contiguous(map, false, s); auto cnt_c = mx::contiguous(counts, false, s); @@ -704,14 +966,11 @@ std::vector quantize( (w.shape(-1) / codec->weights_per_block) * codec->bytes_per_block; mx::Shape s_shape = {1}; - // Row-contiguize w (+ imatrix) at the op level; eval_gpu trusts contiguity. - auto w_c = w.flags().row_contiguous ? w : mx::contiguous(w, false, s); - std::vector inputs = {w_c}; + // Row-major w and imatrix for the encoder. An operand without memory is + // decided at eval. + std::vector inputs = {kq_row_major(w, s)}; if (imatrix.has_value()) { - auto im_c = imatrix->flags().row_contiguous - ? *imatrix - : mx::contiguous(*imatrix, false, s); - inputs.push_back(im_c); + inputs.push_back(kq_row_major(*imatrix, s)); } return mx::array::make_arrays( @@ -749,6 +1008,41 @@ bool metallib_loads() { #endif } +std::vector metallib_kernels() { +#ifdef _METAL_ + auto& d = mx::metal::device(mx::Device::gpu); + MTL::Library* lib = d.get_library("mlx_kquant", metallib_dir()); + auto pool = NS::TransferPtr(NS::AutoreleasePool::alloc()->init()); + NS::Array* names = lib->functionNames(); + std::vector out; + out.reserve(names->count()); + for (NS::UInteger i = 0; i < names->count(); ++i) { + out.emplace_back(names->object(i)->utf8String()); + } + std::sort(out.begin(), out.end()); + return out; +#else + return {}; +#endif +} + +int max_threadgroup_memory() { +#ifdef _METAL_ + auto& d = mx::metal::device(mx::Device::gpu); + return static_cast(d.mtl_device()->maxThreadgroupMemoryLength()); +#else + return 0; +#endif +} + +bool probe_kernels() { +#ifdef KQ_PROBE_KERNELS + return true; +#else + return false; +#endif +} + bool nax_available() { #ifdef _METAL_ return kq_is_nax_available(); diff --git a/src/kquant_qk_norm_rope.cpp b/src/kquant_qk_norm_rope.cpp new file mode 100644 index 00000000..46efda1e --- /dev/null +++ b/src/kquant_qk_norm_rope.cpp @@ -0,0 +1,282 @@ +// Fused per-head RMSNorm and RoPE of the DiT attention q and k (see +// kq_qk_norm_rope.h for the kernel). On the GPU one dispatch reads q and k +// through their strides and writes both in the [B, H, L, D] attention +// layout. CPU streams, and head dims without a kernel, run the same math as +// a composition of stock ops: f32 throughout, one round at the end. +#include +#include +#include +#include +#include +#include + +#include "kquant.h" +#include "kquant_internal.h" // kq_type_string + +#include "mlx/ops.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" // kq_get_kernel +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +// Head dims with a kernel; each is a multiple of 32 so a 16-lane group +// holds an even number of elements per lane. +bool qknr_kernel_dim(int D) { + return D == 64 || D == 96 || D == 128 || D == 256; +} + +// f32 rms_norm over the last axis, then the pair rotation, rounded once to +// the input dtype and written as [B, H, L, D]. +mx::array qknr_compose( + const mx::array& x, + const mx::array& w, + const mx::array& cos, + const mx::array& sin, + float eps, + mx::Stream s) { + const int B = x.shape(0); + const int L = x.shape(1); + const int H = x.shape(2); + const int D = x.shape(3); + const int P = D / 2; + auto xf = mx::astype(x, mx::float32, s); + auto ms = mx::mean(mx::square(xf, s), -1, true, s); + auto inv = mx::rsqrt(mx::add(ms, mx::array(eps, mx::float32), s), s); + auto y = mx::multiply(mx::multiply(xf, inv, s), w, s); + y = mx::reshape(y, {B, L, H, P, 2}, s); + auto y0 = mx::reshape( + mx::slice(y, {0, 0, 0, 0, 0}, {B, L, H, P, 1}, s), {B, L, H, P}, s); + auto y1 = mx::reshape( + mx::slice(y, {0, 0, 0, 0, 1}, {B, L, H, P, 2}, s), {B, L, H, P}, s); + auto c = mx::reshape(cos, {1, L, 1, P}, s); + auto sn = mx::reshape(sin, {1, L, 1, P}, s); + auto o0 = mx::subtract(mx::multiply(y0, c, s), mx::multiply(y1, sn, s), s); + auto o1 = mx::add(mx::multiply(y0, sn, s), mx::multiply(y1, c, s), s); + auto o = mx::reshape(mx::stack({o0, o1}, -1, s), {B, L, H, D}, s); + o = mx::astype(o, x.dtype(), s); + return mx::contiguous(mx::transpose(o, {0, 2, 1, 3}, s), false, s); +} + +} // namespace + +#ifdef _METAL_ + +namespace { + +// Mirror of the Metal-side KQQkNormRopeParams (kq_qk_norm_rope.h); the +// layouts must match field for field. +struct KQQkNormRopeParams { + int64_t q_b; + int64_t q_l; + int64_t q_h; + int64_t q_d; + int64_t k_b; + int64_t k_l; + int64_t k_h; + int64_t k_d; + int32_t L; + int32_t H; + float eps; +}; + +// Elements per vector load: a lane holds D / 16 elements and loads them 4 +// at a time when that count is a multiple of 4, else 2 (head dim 96). +int qknr_vec_width(int D) { + return (D / 16) % 4 == 0 ? 4 : 2; +} + +// True when x can take vector loads: a unit d stride, and every head row +// start aligned to the vector width. Size-1 axes never move the start. +bool qknr_vec_ok(const mx::array& x, int vw) { + if (x.strides(3) != 1) { + return false; + } + for (int ax = 0; ax < 3; ++ax) { + if (x.shape(ax) > 1 && x.strides(ax) % vw != 0) { + return false; + } + } + return x.offset() % (int64_t(vw) * x.itemsize()) == 0; +} + +// Simdgroups per threadgroup: a (b, l) row holds 2 * H head rows, two per +// simdgroup pass. Prefer the count that splits the passes evenly. +int qknr_simdgroups(int H) { + const int passes = H; + for (int n : {5, 6, 4, 3, 8, 2}) { + if (passes % n == 0) { + return n; + } + } + return 4; +} + +} // namespace + +void KQuantQkNormRope::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + for (auto& out : outputs) { + out.set_data(mx::allocator::malloc(out.nbytes())); + } + const auto& q = inputs[0]; + const auto& k = inputs[1]; + const int B = q.shape(0); + const int L = q.shape(1); + const int H = q.shape(2); + const int D = q.shape(3); + if (q.size() == 0) { + return; + } + + const int vw = qknr_vec_width(D); + bool vec = qknr_vec_ok(q, vw) && qknr_vec_ok(k, vw); + KQQkNormRopeParams p{ + q.strides(0), + q.strides(1), + q.strides(2), + q.strides(3), + k.strides(0), + k.strides(1), + k.strides(2), + k.strides(3), + L, + H, + eps_}; + + const std::string kname = + "kq_qk_norm_rope_" + kq_type_string(q.dtype()) + "_" + std::to_string(D); + mx::metal::MTLFCList func_consts = { + {&vec, MTL::DataType::DataTypeBool, 350}, + }; + auto kernel = + kq_get_kernel(d, kname, kname + (vec ? "_vec" : "_strided"), func_consts); + const int cap = std::min(int(kernel->maxTotalThreadsPerThreadgroup()), 1024); + const int nsg = std::max(1, std::min(qknr_simdgroups(H), cap / 32)); + + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(q, 0); + ce.set_input_array(k, 1); + ce.set_input_array(inputs[2], 2); + ce.set_input_array(inputs[3], 3); + ce.set_input_array(inputs[4], 4); + ce.set_input_array(inputs[5], 5); + ce.set_output_array(outputs[0], 6); + ce.set_output_array(outputs[1], 7); + ce.set_bytes(p, 8); + ce.dispatch_threadgroups( + MTL::Size(size_t(B) * L, 1, 1), MTL::Size(nsg * 32, 1, 1)); +} + +#else // !_METAL_ + +void KQuantQkNormRope::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error("[mlx_kquant.qk_norm_rope] requires Metal."); +} + +#endif + +void KQuantQkNormRope::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.qk_norm_rope] CPU streams build the composed ops."); +} + +bool KQuantQkNormRope::is_equivalent(const mx::Primitive& other) const { + return eps_ == static_cast(other).eps_; +} + +std::vector KQuantQkNormRope::output_shapes( + const std::vector& inputs) { + const auto& q = inputs[0]; + mx::Shape shape = {q.shape(0), q.shape(2), q.shape(1), q.shape(3)}; + return {shape, shape}; +} + +std::pair qk_norm_rope( + mx::array q, + mx::array k, + mx::array q_w, + mx::array k_w, + mx::array cos, + mx::array sin, + float eps, + mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const std::string op = "[mlx_kquant.qk_norm_rope]"; + if (q.ndim() != 4) { + throw std::invalid_argument(op + " q must be [B, L, H, D]."); + } + if (k.shape() != q.shape() || k.dtype() != q.dtype()) { + throw std::invalid_argument(op + " k must match q in shape and dtype."); + } + const auto dt = q.dtype(); + if (dt != mx::float16 && dt != mx::bfloat16) { + throw std::invalid_argument(op + " q and k must be float16 or bfloat16."); + } + const int B = q.shape(0); + const int L = q.shape(1); + const int H = q.shape(2); + const int D = q.shape(3); + if (D < 2 || D % 2 != 0) { + throw std::invalid_argument(op + " head dim must be even."); + } + for (const auto* w : {&q_w, &k_w}) { + if (w->ndim() != 1 || w->shape(0) != D || + !mx::issubdtype(w->dtype(), mx::floating)) { + throw std::invalid_argument( + op + " q_w and k_w must be floating [" + std::to_string(D) + "]."); + } + } + for (const auto* t : {&cos, &sin}) { + if (t->ndim() != 2 || t->shape(0) != L || t->shape(1) != D / 2 || + !mx::issubdtype(t->dtype(), mx::floating)) { + throw std::invalid_argument( + op + " cos and sin must be floating [L, D / 2] = [" + + std::to_string(L) + ", " + std::to_string(D / 2) + "]."); + } + } + + // The weights and tables are small: cast to f32 and pack them. q and k + // stay as they are; the kernel reads them through their strides. + auto f32 = [&](const mx::array& a) { + return mx::contiguous(mx::astype(a, mx::float32, s), false, s); + }; + auto qw_c = f32(q_w); + auto kw_c = f32(k_w); + auto cos_c = f32(cos); + auto sin_c = f32(sin); + + if (s.device.type == mx::Device::DeviceType::cpu || !qknr_kernel_dim(D)) { + return { + qknr_compose(q, qw_c, cos_c, sin_c, eps, s), + qknr_compose(k, kw_c, cos_c, sin_c, eps, s)}; + } + mx::Shape shape = {B, H, L, D}; + auto outs = mx::array::make_arrays( + {shape, shape}, + {dt, dt}, + std::make_shared(s, eps), + {std::move(q), + std::move(k), + std::move(qw_c), + std::move(kw_c), + std::move(cos_c), + std::move(sin_c)}); + return {outs[0], outs[1]}; +} + +} // namespace mlx_kquant diff --git a/src/kquant_route_shed.cpp b/src/kquant_route_shed.cpp index 1bbab882..ece35446 100644 --- a/src/kquant_route_shed.cpp +++ b/src/kquant_route_shed.cpp @@ -8,6 +8,7 @@ // ordering exactly (parity harness in tests/test_route_shed.py). #include #include +#include #include "kquant.h" #include "kquant_internal.h" @@ -38,14 +39,15 @@ void KQuantRouteShed::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); + // The kernel reads the three operands as dense memory: a view with other + // strides is copied first. + const mx::array indices = kq_row_major_gpu(inputs[0], s); + const mx::array scores = kq_row_major_gpu(inputs[1], s); + const mx::array table = kq_row_major_gpu(inputs[2], s); for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& indices = inputs[0]; - const auto& scores = inputs[1]; - const auto& table = inputs[2]; - int R = indices.shape(-1); int T = int(indices.size() / R); int E = table.shape(0); @@ -107,9 +109,13 @@ void KQuantRouteShed::eval_cpu( const int R = indices.shape(-1); const int64_t T = indices.size() / R; const int E = table.shape(0); - const uint32_t* idp = indices.data(); - const float* scp = scores.data(); - const int32_t* tbp = table.data(); + // Dense memory for the loop: a view with other strides is copied. + std::vector id_store; + std::vector sc_store; + std::vector tb_store; + const uint32_t* idp = kq_row_major_cpu(indices, id_store); + const float* scp = kq_row_major_cpu(scores, sc_store); + const int32_t* tbp = kq_row_major_cpu(table, tb_store); uint32_t* slp = slots.data(); float* mxp = mix.data(); int32_t* mip = miss_ids.data(); @@ -222,15 +228,15 @@ std::vector route_shed( throw std::invalid_argument(std::string(op) + " slot_table must be int32."); } - auto contig = [&](const mx::array& a) { - return a.flags().row_contiguous ? a : mx::contiguous(a, false, s); - }; + // The operands go in as they are. indices and scores are unevaluated on + // most calls, so their strides are not known here, and each eval reads a + // view that is not row-major through a row-major copy. const auto& shape = indices.shape(); return mx::array::make_arrays( {shape, shape, shape, shape}, {mx::uint32, mx::float32, mx::int32, mx::float32}, std::make_shared(s), - {contig(indices), contig(scores), contig(slot_table)}); + {indices, scores, slot_table}); } } // namespace mlx_kquant diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 30f7e736..1e07d039 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -94,6 +94,26 @@ void kq_sdpa_check_batch_heads(const char* op, const array& a) { } } +// The scales and biases of one side of a quantized cache. The kernel reads +// both with the head and sequence strides of the scales and with a unit +// stride in the last dim, so the two need one layout. An unevaluated view +// reports dense strides to the builder, so the check runs here. +void kq_sdpa_check_q8_pair(const char* op, const array& sc, const array& bi) { + if (sc.size() == 0) { + return; + } + bool ok = sc.shape(3) == 1 || (sc.strides(3) == 1 && bi.strides(3) == 1); + for (int i = 0; i < 3 && ok; i++) { + ok = sc.shape(i) == 1 || sc.strides(i) == bi.strides(i); + } + if (!ok) { + throw std::runtime_error( + std::string("[mlx_kquant.") + op + + "] quantized-KV scales and biases need one layout with a " + "contiguous last dim."); + } +} + // Number of key-blocks to split the reduction across. Mirrors MLX's own // sdpa_vector_2pass heuristic: more blocks only when there are enough // simdgroups per kv-head (n_simds = gqa_factor * qL) to justify the extra @@ -366,6 +386,12 @@ void KQuantSDPAGQA::eval_gpu( kq_sdpa_check_batch_heads("sdpa_decode_gqa", inputs[qkv_idx + i]); } } + if (kv_q8) { + for (int i = 0; i < 4; i += 2) { + kq_sdpa_check_q8_pair( + "sdpa_decode_gqa", inputs[qkv_idx + i], inputs[qkv_idx + i + 1]); + } + } int B = q.shape(0); int n_q_heads = q.shape(1); @@ -395,14 +421,13 @@ void KQuantSDPAGQA::eval_gpu( float scale = scale_; // One query at head_dim 512 on NAX GPUs runs pass 1 on the matrix units - // (kq_sdpa_nax.h) with the same partials and merge, from kQGqaNaxMinRows. - // Its rounding differs from the simdgroup kernel's, and its split count - // follows the batch size. KQ_GQA_NAX=0, read at each call, keeps the - // simdgroup kernel. - const char* nax_e = std::getenv("KQ_GQA_NAX"); - const bool use_nax = (!nax_e || std::atoi(nax_e) != 0) && D == 512 && - qL == 1 && gqa_factor <= 8 && !kv_q8 && !paged && !kvarn && !starts && - !ends && size_t(n_kv_heads) * B * kL >= kQGqaNaxMinRows && + // (kq_sdpa_nax.h) with the same partials and merge, from kQGqaNaxMinRows, + // with or without starts. Its rounding differs from the simdgroup + // kernel's, and its split count follows the batch size. KQ_GQA_NAX=0, + // read at each call, keeps the simdgroup kernel. + const bool use_nax = kq_env_flag("KQ_GQA_NAX", 1) != 0 && D == 512 && + qL == 1 && gqa_factor <= 8 && !kv_q8 && !paged && !kvarn && !ends && + size_t(n_kv_heads) * B * kL >= kQGqaNaxMinRows && k_head_stride % 8 == 0 && k_seq_stride % 8 == 0 && v_head_stride % 8 == 0 && v_seq_stride % 8 == 0 && (B == 1 || (k.strides(0) % 8 == 0 && v.strides(0) % 8 == 0)) && @@ -460,7 +485,8 @@ void KQuantSDPAGQA::eval_gpu( // (kv-head, batch, split). if (use_nax) { std::string kname = "kq_sdpa_gqa_nax_2pass_1_" + ts; - std::string hash = kname + "_s" + std::to_string(splits); + std::string hash = + kname + "_s" + std::to_string(splits) + (has_starts ? "_st1" : "_st0"); auto kernel = kq_get_kernel(d, kname, hash, fc); if (kernel->maxTotalThreadsPerThreadgroup() < 256) { throw std::runtime_error( @@ -485,6 +511,8 @@ void KQuantSDPAGQA::eval_gpu( const KQKvBatch kvb = kq_kv_batch(k, v); ce.set_bytes(kvb.k, 13); ce.set_bytes(kvb.v, 14); + // Without starts, sums is a dummy binding (the read is compiled out). + ce.set_input_array(starts ? inputs[starts_idx] : sums, 15); ce.dispatch_threadgroups( MTL::Size(n_kv_heads, B, splits), MTL::Size(32, 8, 1)); } @@ -857,10 +885,7 @@ void KQuantSDPAFAIndexed::eval_gpu( // strip); otherwise the simdgroup kernel, one threadgroup per (32-head // strip, query, split). KQ_SDPA_IDX_NAX=0 forces the simdgroup kernel on // any GPU. - static const bool nax_env = [] { - const char* e = std::getenv("KQ_SDPA_IDX_NAX"); - return !e || std::atoi(e) != 0; - }(); + static const bool nax_env = kq_env_flag("KQ_SDPA_IDX_NAX", 1) != 0; const bool use_nax = nax_env && kq_is_nax_available(); int splits = splits_; @@ -997,6 +1022,12 @@ void KQuantSDPACascade::eval_gpu( for (int i = 4; i < 8; i++) { kq_sdpa_check_batch_heads("sdpa_decode_gqa_cascade", inputs[q8_base + i]); } + for (int i = 0; i < 8; i += 2) { + kq_sdpa_check_q8_pair( + "sdpa_decode_gqa_cascade", + inputs[q8_base + i], + inputs[q8_base + i + 1]); + } } int B = q.shape(0); @@ -1259,6 +1290,7 @@ void KQuantSDPABSPrefill::eval_gpu( const auto& pmask = inputs[4]; const auto& counts = inputs[5]; auto& out = outputs[0]; + kq_sdpa_check_layout("sdpa_prefill_block_sparse", q, k, v); out.set_data(mx::allocator::malloc(out.nbytes())); int n_kv_heads = q.shape(1); diff --git a/src/kquant_sdpa_sparse_decode.cpp b/src/kquant_sdpa_sparse_decode.cpp index 8c1e527a..9ab3c2cd 100644 --- a/src/kquant_sdpa_sparse_decode.cpp +++ b/src/kquant_sdpa_sparse_decode.cpp @@ -47,6 +47,37 @@ mx::array rows_flat(const mx::array& a, const char* op, mx::StreamOrDevice s) { } return mx::reshape(a, {a.shape(0), static_cast(flat)}, s); } + +#ifdef _METAL_ +// The window, pool and pool scales as the kernels read them. A reshape to +// [B, rows * D] keeps a view whenever the dims of its input fold into the +// two of its output, so a row can come out with an element stride other +// than 1, which is copied here. The kernels step the batches of a packed +// pool by the stride of its codes and take an eighth of it for the scales, +// so codes and scales that step differently are copied too. +struct SparseRows { + mx::array win; + mx::array pool; + mx::array scales; // the pool again when it is not packed +}; + +SparseRows sparse_rows( + const mx::array& win, + const mx::array& pool, + const mx::array* scales, + const mx::Stream& s) { + SparseRows r{ + kq_unit_stride_gpu(win, s), + kq_unit_stride_gpu(pool, s), + scales ? kq_unit_stride_gpu(*scales, s) : pool}; + if (scales && r.pool.size() > 0 && r.pool.shape(0) != 1 && + r.pool.strides(0) != 8 * r.scales.strides(0)) { + r.pool = kq_dense_copy_gpu(r.pool, s); + r.scales = kq_dense_copy_gpu(r.scales, s); + } + return r; +} +#endif } // namespace #ifdef _METAL_ @@ -59,11 +90,13 @@ void KQSdpaSparseDecode::eval_gpu( auto& ce = mx::metal::get_command_encoder(s); const auto& q = inputs[0]; - const auto& win = inputs[1]; // [B|1, W * D], rows D apart - const auto& pool = inputs[2]; // [B|1, P * D] (packed: [B|1, P * D / 2] codes) const auto& idx = inputs[3]; size_t next = 4; - const mx::array* pscales = packed_ ? &inputs[next++] : nullptr; + const SparseRows rows = + sparse_rows(inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s); + const auto& win = rows.win; // [B|1, W * D], rows D apart + const auto& pool = rows.pool; // [B|1, P * D] (packed: [B|1, P * D / 2] codes) + const mx::array* pscales = packed_ ? &rows.scales : nullptr; const mx::array* sinks = has_sinks_ ? &inputs[next++] : nullptr; const mx::array* wmask = has_win_mask_ ? &inputs[next++] : nullptr; const mx::array* smask = has_sel_mask_ ? &inputs[next++] : nullptr; @@ -344,8 +377,9 @@ mx::array sdpa_sparse_decode( // (Contiguous donates when the input already is). The window and pool // go through a reshape to [B, rows * D] instead: a cache's prefix slice // is a row-contiguous view, which Reshape keeps as a view at eval while - // Contiguous would copy it (its buffer is larger than the view), and - // any other layout Reshape copies. + // Contiguous would copy it (its buffer is larger than the view). Reshape + // keeps other views too, and the eval copies one whose rows are not + // dense. auto q_c = mx::contiguous(q, false, s); auto win_c = rows_flat(mx::astype(window, dt, s), op, s); auto pool_c = rows_flat( @@ -448,11 +482,13 @@ void KQSdpaSparsePrefill::eval_gpu( auto& ce = mx::metal::get_command_encoder(s); const auto& q = inputs[0]; - const auto& win = inputs[1]; // [B|1, S * D], rows D apart - const auto& pool = inputs[2]; // [B|1, P * D] (packed: [B|1, P * D / 2] codes) const auto& idx = inputs[3]; size_t next = 4; - const mx::array* pscales = packed_ ? &inputs[next++] : nullptr; + const SparseRows rows = + sparse_rows(inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s); + const auto& win = rows.win; // [B|1, S * D], rows D apart + const auto& pool = rows.pool; // [B|1, P * D] (packed: [B|1, P * D / 2] codes) + const mx::array* pscales = packed_ ? &rows.scales : nullptr; const mx::array* sinks = has_sinks_ ? &inputs[next++] : nullptr; const mx::array* smask = has_sel_mask_ ? &inputs[next++] : nullptr; auto& o = outputs[0]; diff --git a/tests/conftest.py b/tests/conftest.py index 5da39db8..cc51381b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,49 @@ import os +import pytest + + +@pytest.fixture(autouse=True) +def _clear_mlx_cache(): + # MLX keeps freed buffers in its cache, and its default limit is near the + # device memory. Clear it after each test, so the process footprint over a + # session stays at the largest single test. + yield + import mlx.core as mx + + mx.clear_cache() + + +def _compile_switch_on(): + # While mx's compile switch is on, a compiled function traces once over + # two calls. While it is off, the function runs on each call. + import mlx.core as mx + + calls = [] + + def f(t): + calls.append(1) + return t + + g = mx.compile(f) + t = mx.zeros((1,)) + g(t) + g(t) + return len(calls) == 1 + + +@pytest.fixture +def compile_on(): + # mx's compile switch on for a test of mx.compile, and back as found after + # it. MLX_DISABLE_COMPILE sets the switch off for the process. + import mlx.core as mx + + found = _compile_switch_on() + mx.enable_compile() + yield + (mx.enable_compile if found else mx.disable_compile)() + def pytest_configure(config): # The logic tests touch only small mx array ops and never dispatch a kquant diff --git a/tests/qmm_probe.py b/tests/qmm_probe.py new file mode 100644 index 00000000..d4e108be --- /dev/null +++ b/tests/qmm_probe.py @@ -0,0 +1,368 @@ +"""Child-process probe for the quantized matmul tile levers. + +The tile levers are read once per process unless ``KQ_TUNE_LIVE=1``, so a +test that sets them needs a fresh interpreter. ``run_probe`` starts one with +a clean ``KQ_*`` environment, runs a list of ``quantized_matmul`` calls, +and returns per call the ``KQ_QMM_LOG`` lines it printed (kernel name and +grid), its error if it raised, and its relative Frobenius error against a +float64 reference built from ``kq.dequantize`` weights. A call with +``"r2": true`` also gets ``rel2`` against those weights rounded to the +activation dtype, and ``floor1``/``floor2``, the error of rounding each +reference to the output dtype. A call with ``"ulp_vs": i`` gets ``ulp``, +the largest ulp distance to call ``i``'s output, and ``rdiff``, the +relative Frobenius distance to it, and with the check ``ulp_big``, the +largest ulp distance on the elements whose reference is above 1% of its +root mean square. ``"check_block": B`` runs the check and the ulp +comparison over every row of x's flattened leading dims, B rows at a +time, so a call with many rows holds float64 copies of B rows; it reports +``checked``, the number of rows it read. Every call gets +``digest``, a hash of its output bits, to compare calls across probe +processes. + +``"design": "lead"`` draws nonnegative weights and activations whose first +K column is 1024. ``"design": "big"`` draws small weights and dense +activations from 2e4 up, with bfloat16 column 0 at 1e5, past the half +range. + +Under ``KQUANT_FORCE_CPU`` the child runs on the CPU device and never +touches Metal, so the same probe covers both streams. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys + +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) + +_CHILD = r""" +import hashlib, json, os, sys +import numpy as np +spec = json.loads(sys.argv[1]) +sys.path.insert(0, spec["tests_dir"]) +import mlx.core as mx +if os.environ.get("KQUANT_FORCE_CPU"): + mx.set_default_device(mx.cpu) +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY, DECODE_ONLY_CODECS +from kqref import is_synth, synth_wire + +DT = {"bfloat16": mx.bfloat16, "float16": mx.float16} +# float32 activations too: the op casts them to bfloat16. +X_DT = {**DT, "float32": mx.float32} +_weights = {} +_dequant = {} + + +def weights(codec, n, k, design): + # design "lead": nonnegative weights, so every partial sum only grows. + # design "tiled": 64 quantized rows repeated to n, for wide N at little + # quantize cost. design "big": small weights, nonnegative where the + # codec quantizes, so large activations give a large sum per block and + # a moderate output. design "cancel": weights near +0.05 on the first + # half of K and -0.05 on the second, so the halves cancel. + key = (codec, n, k, design) + if key not in _weights: + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + if design == "lead": + wf = mx.abs(mx.random.normal((n, k), key=mx.random.key(2))) * 0.05 + w, s = kq.quantize(wf, codec) + elif design == "big" and is_synth(codec): + rng = np.random.default_rng(n * 131 + k) + wire = synth_wire(rng, codec, bpb, n * (k // wpb), d_range=(1e-4, 4e-4)) + w = mx.array(wire.reshape(n, (k // wpb) * bpb)) + s = mx.zeros((1,), dtype=mx.uint8) + elif design == "big": + wf = mx.abs(mx.random.normal((n, k), key=mx.random.key(2))) * 1e-3 + w, s = kq.quantize(wf, codec) + elif design == "cancel": + sign = mx.where(mx.arange(k) < k // 2, 1.0, -1.0) + noise = mx.random.normal((n, k), key=mx.random.key(3)) * 0.002 + w, s = kq.quantize(0.05 * sign[None, :] + noise, codec) + elif design == "tiled": + wf = mx.random.normal((64, k), key=mx.random.key(4)) * 0.1 + w1, s = kq.quantize(wf, codec) + w = mx.contiguous(mx.tile(w1, (n // 64, 1))) + elif is_synth(codec) or codec in DECODE_ONLY_CODECS: + rng = np.random.default_rng(n * 131 + k) + wire = synth_wire(rng, codec, bpb, n * (k // wpb)) + w = mx.array(wire.reshape(n, (k // wpb) * bpb)) + s = mx.zeros((1,), dtype=mx.uint8) + else: + wf = mx.random.normal((n, k), key=mx.random.key(n * 7 + k)) * 0.1 + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + _weights[key] = (w, s) + return _weights[key] + + +def dequantized(codec, n, k, design): + # The float64 weights of weights(), built on the first check. + key = (codec, n, k, design) + if key not in _dequant: + w, s = weights(codec, n, k, design) + wd = kq.dequantize(w, s, codec, dtype=mx.float32) + _dequant[key] = np.array(wd).astype(np.float64) + return _dequant[key] + + +def activations(shape, dtype, layout, seed, design): + # design "lead": nonnegative activations whose first K column is 1024, + # so the first K slice of a split holds most of each output. design + # "cancel": activations near 3e4. + key = mx.random.key(seed) + lead = design == "lead" + + def draw(s): + if lead: + return mx.abs(mx.random.normal(s, key=key)) * 0.01 + if design == "big": + # 2e4 and up, inside the float16 range. + return mx.minimum(mx.abs(mx.random.normal(s, key=key)) * 1e4 + 2e4, 6e4) + if design == "cancel": + return mx.random.uniform(-100.0, 100.0, s, key=key) + 3e4 + return mx.random.normal(s, key=key) * 0.5 + + if layout == "dense": + x = draw(shape) + if lead: + x[..., 0] = 1024.0 + if design == "big" and dtype == mx.bfloat16: + # One column past the half range. + x[..., 0] = 1e5 + x = x.astype(dtype) + mx.eval(x) + return x + if layout == "transposed": + # A lazy transposed view: the op sees an unevaluated swapaxes. + t = shape[:-2] + [shape[-1], shape[-2]] + xt = draw(t) + if lead: + xt[..., 0, :] = 1024.0 + xt = xt.astype(dtype) + mx.eval(xt) + return mx.swapaxes(xt, -1, -2) + if layout == "strided": + # Every other column of a wider array. + wide = shape[:-1] + [2 * shape[-1]] + xw = draw(wide) + if lead: + xw[..., 0] = 1024.0 + xw = xw.astype(dtype) + mx.eval(xw) + return xw[..., ::2] + raise ValueError(layout) + + +def rel(a, b): + return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-30)) + + +def rounded(a, dtype): + # a rounded to dtype, back in float64. + r = mx.array(a.astype(np.float32)).astype(dtype).astype(mx.float32) + return np.array(r).astype(np.float64) + + +def ordered(bits): + u = bits.astype(np.int64) + return np.where(u & 0x8000, 0x8000 - (u & 0x7FFF), u + 0x8000) + + +outs = [] +results = [] +_refs = {} +def as64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + +def blocked_check(c, x, y, other, wd, refs, res): + # The check and the ulp comparison to other over every row, check_block + # rows at a time. Norms add up over the blocks; ulp counts take the max. + # ulp_big needs the root mean square of the whole reference first: the + # first call of a reference keeps its sum of squares. + blk = c["check_block"] + xf = x.reshape(-1, x.shape[-1]) + yf = y.reshape(-1, y.shape[-1]) + of = other.reshape(-1, other.shape[-1]) if other is not None else None + mx.eval(xf, yf) + rows = xf.shape[0] + starts = range(0, rows, blk) + + def ref_of(lo): + return as64(xf[lo : lo + blk]) @ wd.T + + thr = 0.0 + if of is not None: + mx.eval(of) + if "sumsq" not in refs: + refs["sumsq"] = sum(float((ref_of(lo) ** 2).sum()) for lo in starts) + thr = 1e-2 * np.sqrt(refs["sumsq"] / (rows * yf.shape[-1])) + d2 = r2 = dd = b2 = 0.0 + ulp = ulp_big = checked = 0 + for lo in starts: + ref = ref_of(lo) + got = as64(yf[lo : lo + blk]) + d2 += float(((got - ref) ** 2).sum()) + r2 += float((ref * ref).sum()) + checked += ref.shape[0] + if of is not None: + a = ordered(np.array(yf[lo : lo + blk].view(mx.uint16))) + b = ordered(np.array(of[lo : lo + blk].view(mx.uint16))) + d = np.abs(a - b) + ulp = max(ulp, int(d.max())) + big = np.abs(ref) > thr + if big.any(): + ulp_big = max(ulp_big, int(d[big].max())) + yb = as64(of[lo : lo + blk]) + dd += float(((got - yb) ** 2).sum()) + b2 += float((yb * yb).sum()) + refs["sumsq"] = r2 + res["rel"] = float(np.sqrt(d2) / (np.sqrt(r2) + 1e-30)) + res["checked"] = checked + if of is not None: + res["ulp"] = ulp + res["ulp_big"] = ulp_big + res["rdiff"] = float(np.sqrt(dd) / (np.sqrt(b2) + 1e-30)) + + +_ulp_bases = {c["ulp_vs"] for c in spec["calls"] if c.get("ulp_vs") is not None} +for i, c in enumerate(spec["calls"]): + for name, value in c.get("env", {}).items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = str(value) + dtype = X_DT[c.get("dtype", "bfloat16")] + design = c.get("design", "normal") + w, s = weights(c["codec"], c["n"], c["shape"][-1], design) + x = activations( + list(c["shape"]), dtype, c.get("layout", "dense"), c.get("seed", i), design + ) + os.write(2, f"[probe] {i}\n".encode()) + res = {"i": i, "error": None, "rel": None, "ulp": None} + try: + y = kq.quantized_matmul(x, w, s, c["codec"], transpose=True) + mx.eval(y) + except Exception as e: + res["error"] = f"{type(e).__name__}: {e}" + outs.append(None) + results.append(res) + continue + # Keep the output only for the calls that compare against it. + outs.append(y if i in _ulp_bases else None) + # Digest of the output bits, to compare calls across probe processes. + bits = y.view(mx.uint16) if y.dtype in DT.values() else y + res["digest"] = hashlib.sha1(memoryview(np.array(bits))).hexdigest() + base = c.get("ulp_vs") + blocked = spec.get("check", True) and c.get("check_block") + # Calls that differ only in env share their references. + rkey = ( + c["codec"], c["n"], tuple(c["shape"]), str(dtype), + c.get("layout", "dense"), c.get("seed", i), design, + c.get("check_block"), + ) + if blocked: + assert not c.get("r2"), "r2 takes no check_block" + wd = dequantized(c["codec"], c["n"], c["shape"][-1], design) + other = outs[base] if base is not None else None + blocked_check(c, x, y, other, wd, _refs.setdefault(rkey, {}), res) + res["nonfinite"] = int((~mx.isfinite(y)).sum().item()) + elif spec.get("check", True): + if rkey not in _refs: + wd = dequantized(c["codec"], c["n"], c["shape"][-1], design) + xn = as64(x) + _refs[rkey] = {"xn": xn, "ref": xn @ wd.T} + refs = _refs[rkey] + ref = refs["ref"] + got = as64(y) + res["rel"] = float(np.linalg.norm(got - ref) / (np.linalg.norm(ref) + 1e-30)) + res["nonfinite"] = int((~mx.isfinite(y)).sum().item()) + if c.get("r2"): + # R2: weights rounded to the activation dtype. floor1 and floor2: + # the error of rounding each reference to the output dtype. + if "ref2" not in refs: + wd = dequantized(c["codec"], c["n"], c["shape"][-1], design) + refs["ref2"] = refs["xn"] @ rounded(wd, dtype).T + refs["floor1"] = rel(rounded(ref, dtype), ref) + refs["floor2"] = rel(rounded(refs["ref2"], dtype), refs["ref2"]) + res["rel2"] = rel(got, refs["ref2"]) + res["floor1"] = refs["floor1"] + res["floor2"] = refs["floor2"] + if base is not None and outs[base] is not None and not blocked: + a = np.array(y.view(mx.uint16)) + b = np.array(outs[base].view(mx.uint16)) + res["ulp"] = int(np.abs(ordered(a) - ordered(b)).max()) + # Relative Frobenius distance to that call's output. + yb = as64(outs[base]) + ya = as64(y) + res["rdiff"] = rel(ya, yb) + if spec.get("check", True): + # The largest ulp distance on the elements above 1% of the + # reference's root mean square, away from the cancellations + # near zero where any change of summation order moves many ulp. + ref = _refs[rkey]["ref"] + big = np.abs(ref) > 1e-2 * np.sqrt(np.mean(ref * ref)) + d = np.abs(ordered(a) - ordered(b)) + res["ulp_big"] = int(d[big].max()) if big.any() else 0 + results.append(res) + # Return freed buffers, so the child holds about one call at a time. + mx.clear_cache() +print(json.dumps(results)) +""" + +_LOG = re.compile(r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+)$") + + +def run_probe(calls, env=None, check=True, timeout=900): + """Run ``calls`` in a fresh interpreter whose environment keeps no + ``KQ_*`` variable but those in ``env``. Returns one dict per call with + ``kernels`` (list of (name, grid) from KQ_QMM_LOG), ``stderr`` (its + other stderr lines), ``error``, ``rel`` and ``ulp``.""" + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update(env or {}) + spec = {"tests_dir": TESTS_DIR, "calls": calls, "check": check} + r = subprocess.run( + [sys.executable, "-c", _CHILD, json.dumps(spec)], + env=child_env, + capture_output=True, + text=True, + timeout=timeout, + ) + assert r.returncode == 0, r.stderr[-6000:] + results = json.loads(r.stdout.strip().splitlines()[-1]) + logs = {} + other = {} + cur = None + for line in r.stderr.splitlines(): + if line.startswith("[probe] "): + cur = int(line.split()[1]) + logs[cur] = [] + other[cur] = [] + continue + if cur is None: + continue + m = _LOG.match(line) + if m: + grid = tuple(int(g) for g in m.group(5, 6, 7)) + logs[cur].append((m.group(1), grid)) + else: + other[cur].append(line) + for res in results: + res["kernels"] = logs.get(res["i"], []) + res["stderr"] = other.get(res["i"], []) + return results + + +def tile_of(kname): + """(bm, bn, wm, wn) parsed from a NAX qmm_t kernel name.""" + m = re.search(r"_bm(\d+)_bn(\d+)_bk\d+_wm(\d+)_wn(\d+)", kname) + assert m, kname + return tuple(int(g) for g in m.groups()) + + +def force_cpu() -> bool: + return bool(os.environ.get("KQUANT_FORCE_CPU")) diff --git a/tests/test_alu_dequantize_rows.py b/tests/test_alu_dequantize_rows.py new file mode 100644 index 00000000..cc5013de --- /dev/null +++ b/tests/test_alu_dequantize_rows.py @@ -0,0 +1,758 @@ +"""The route that runs a plain quantized_matmul as dequantize then the stock +matmul on GPUs without NAX before M5. + +``_alu_dequantize_row`` answers for any architecture string, so the policy +cases run on every machine. The GPU cases run in a fresh interpreter +(tests/qmm_probe.py and the children below) under ``KQ_QMM_LOG``: a routed +call logs a ``[kq dequant]`` line and no qmm kernel. Where this GPU's class +has no row, ``KQ_DEQUANT_ROUTE=force`` runs the route, so every machine runs +its code. Each numeric pair holds the route against the qmm_t kernel under +``KQ_DEQUANT_ROUTE=off`` and ``KQ_DISABLE_NAX=1``, and both against float64 +with exact ``kq.dequantize`` weights. Under ``KQUANT_FORCE_CPU`` the GPU +cases skip. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys + +import mlx.core as mx +import pytest +from kqref import is_synth +from mlx_kquant._ext import _alu_dequantize_row, _dequantize_route_takes +from qmm_probe import TESTS_DIR, force_cpu, run_probe + +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY + +R1 = {"bfloat16": 6e-3, "float16": 1e-3, "float32": 6e-3} +# Largest ulp distance to the kernel on the elements above 1% of the +# output's root mean square, by output dtype. +ULP = {"bfloat16": 1, "float16": 2} +LIVE = {"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"} + +# Codec -> floor in the rows of x (measured on M3 Max). +ROWS = { + "nvfp4": 2048, + "ptq1_0": 2048, + "q3_k": 2048, + "q4_k": 16384, + "q5_k": 2048, + "q6_k": 2048, +} +MIN_K, MAX_K = 3840, 12288 +MIN_N, MAX_N = 3840, 24576 + +# Architecture strings with a row (generation 13 to 16, any suffix letter +# but p) and without one (phones, M5 and later, earlier generations, +# malformed). +ROW_CLASS = [f"applegpu_g{g}{c}" for g in (13, 14, 15, 16) for c in "gsdc"] +NO_ROW_CLASS = [f"applegpu_g{g}p" for g in (13, 14, 15, 16, 17, 18)] + [ + "applegpu_g17g", + "applegpu_g17s", + "applegpu_g17d", + "applegpu_g18s", + "applegpu_g12s", + "applegpu_g99s", + "applegpu_g15", + "applegpu_g15sx", + "applegpu_g15S", + "applegpu_g015s", + "applegpu_g5s", + "applegpu_gxxs", + "applegpu_15s", + "", + "unknown", +] + +MATMUL_CODECS = sorted(c for c in CODEC_GEOMETRY if kq.codec_has_matmul(c)) + + +def _metal(): + if force_cpu(): + pytest.skip("Metal kernel logs only") + + +def _device_class(): + """(architecture string with its suffix, nax) of this GPU.""" + return mx.device_info().get("architecture", ""), kq.nax_available() + + +def _row_here(codec): + arch, nax = _device_class() + return bool(_alu_dequantize_row(arch, nax, codec, ROWS[codec], MIN_N, MIN_K)) + + +def _route_env(codec): + # The route: by default where this GPU's class has the row, else forced. + return { + "KQ_QMM_TILE": None, + "KQ_DISABLE_NAX": None, + "KQ_DEQUANT_ROUTE": None if _row_here(codec) else "force", + } + + +# The qmm_t kernel without NAX. +KERNEL_ENV = {"KQ_QMM_TILE": None, "KQ_DISABLE_NAX": "1", "KQ_DEQUANT_ROUTE": "off"} + + +def _routed(r): + lines = [x for x in r["stderr"] if x.startswith("[kq dequant] ")] + assert len(lines) <= 1, r["stderr"] + assert bool(lines) != bool(r["kernels"]), (lines, r["kernels"]) + return bool(lines) + + +@pytest.mark.parametrize("codec", MATMUL_CODECS) +def test_rows_by_chip_class(codec): + # A row runs on the classes with one, from its floor in M, inside the K + # and N range, and nowhere else: not on NAX GPUs, phones, M5 and later, + # earlier or malformed strings, or for a codec without a row. + if codec in ROWS: + min_m = ROWS[codec] + row = [min_m, MIN_K, MAX_K, MIN_N, MAX_N] + outside = [ + (1, MIN_N, MIN_K), + (512, MIN_N, MIN_K), + (min_m - 1, MIN_N, MIN_K), + (min_m, MIN_N - 64, MIN_K), + (min_m, MAX_N + 64, MIN_K), + (min_m, 248320, MIN_K), + (min_m, MIN_N, MIN_K - 256), + (min_m, MIN_N, MAX_K + 256), + (min_m, MIN_N, 256), + ] + for arch in ROW_CLASS: + for m in (min_m, min_m + 1, 32769): + for n in (MIN_N, 4096, MAX_N): + for k in (MIN_K, 4096, MAX_K): + got = _alu_dequantize_row(arch, False, codec, m, n, k) + assert got == row, (arch, m, n, k, got) + assert _alu_dequantize_row(arch, True, codec, min_m, MIN_N, MIN_K) == [] + for m, n, k in outside: + got = _alu_dequantize_row(arch, False, codec, m, n, k) + assert got == [], (arch, m, n, k, got) + else: + for arch in ROW_CLASS: + for m in (2048, 16384, 32769): + got = _alu_dequantize_row(arch, False, codec, m, MIN_N, MIN_K) + assert got == [], (arch, m, got) + for arch in NO_ROW_CLASS: + for nax in (False, True): + for m in (2048, 16384, 32769): + got = _alu_dequantize_row(arch, nax, codec, m, MIN_N, MIN_K) + assert got == [], (arch, nax, got) + + +MODES = ["auto", "off", "force"] + + +@pytest.mark.parametrize("codec", sorted(ROWS) + ["q8_0"]) +def test_route_decision_by_mode(codec): + # The op's decision: off and a pinned KQ_QMM_TILE never route, force + # routes on any GPU, auto only on the classes with a row, and each only + # for a codec row at a shape inside its floor and range. + min_m = ROWS.get(codec, 2048) + inside = (min_m, MIN_N, MIN_K) + outside = [ + (min_m - 1, MIN_N, MIN_K), + (min_m, MAX_N + 64, MIN_K), + (min_m, MIN_N, MAX_K + 256), + ] + classes = [(a, False, True) for a in ROW_CLASS] + [ + ("applegpu_g15s", True, False), + ("applegpu_g17s", True, False), + ("applegpu_g17s", False, False), + ("applegpu_g16p", False, False), + ("unknown", False, False), + ] + for arch, nax, on_class in classes: + assert bool(_alu_dequantize_row(arch, nax, codec, *inside)) == ( + on_class and codec in ROWS + ), arch + for mode in MODES: + for pinned in (False, True): + want = ( + codec in ROWS + and not pinned + and (mode == "force" or (mode == "auto" and on_class)) + ) + got = _dequantize_route_takes(arch, nax, mode, pinned, codec, *inside) + assert got == want, (arch, nax, mode, pinned, got) + for shape in outside: + got = _dequantize_route_takes( + arch, nax, mode, pinned, codec, *shape + ) + assert not got, (arch, nax, mode, pinned, shape) + with pytest.raises(ValueError, match="auto, off or force"): + _dequantize_route_takes("applegpu_g15s", False, "on", False, codec, *inside) + + +def _call(codec, shape, n, env, **kw): + return dict({"codec": codec, "shape": shape, "n": n, "env": env}, **kw) + + +@pytest.mark.parametrize("codec", sorted(ROWS)) +def test_dispatch_on_this_gpu(codec): + # Under the route's env a call takes the route from the floor and inside + # the K and N range, and a batch of two counts its rows together. Below + # the floor, outside the range, under off and with KQ_QMM_TILE set the + # kernel runs. By default the route runs where this GPU's class has the + # row. A value the lever does not take warns and reads as auto. + _metal() + min_m = ROWS[codec] + here = _row_here(codec) + route = _route_env(codec) + pin = {"KQ_QMM_TILE": "64,64,2,2"} + cases = [ + ([1, min_m - 32, MIN_K], MIN_N, route, False), + ([1, min_m, MIN_K], MIN_N, route, True), + ([2, min_m // 2, MIN_K], MIN_N, route, True), + ([1, min_m, MIN_K], MIN_N - 64, route, False), + ([1, min_m, MIN_K], MAX_N, route, True), + ([1, min_m, MIN_K], MAX_N + 64, route, False), + ([1, min_m, MIN_K - 256], MIN_N, route, False), + ([1, min_m, MAX_K], MIN_N, route, True), + ([1, min_m, MAX_K + 256], MIN_N, route, False), + ([1, min_m, MIN_K], MIN_N, pin, False), + ([1, min_m, MIN_K], MIN_N, dict(route, KQ_DEQUANT_ROUTE="off"), False), + ([1, min_m, MIN_K], MIN_N, dict(route, KQ_DEQUANT_ROUTE="Force"), True), + ([1, min_m, MIN_K], MIN_N, dict(route, KQ_DEQUANT_ROUTE=None), here), + ([1, min_m, MIN_K], MIN_N, dict(route, KQ_DEQUANT_ROUTE="auto"), here), + ([1, min_m, MIN_K], MIN_N, dict(route, KQ_DEQUANT_ROUTE="on"), here), + ] + # The calls check no values, so the quantized codecs take 64 quantized + # rows tiled to N, which spares quantizing the widest weights. + design = "normal" if is_synth(codec) else "tiled" + calls = [_call(codec, shape, n, env, design=design) for shape, n, env, _ in cases] + res = run_probe(calls, env=LIVE, check=False) + for c, (*_, want), r in zip(calls, cases, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert _routed(r) == want, (c, r["kernels"], r["stderr"]) + for name, _ in r["kernels"]: + assert "_qmm_t_" in name, (c, name) + warned = [x for x in res[-1]["stderr"] if "KQ_DEQUANT_ROUTE=on is not" in x] + assert len(warned) == 1, res[-1]["stderr"] + + +@pytest.mark.parametrize("codec", sorted(ROWS)) +def test_nax_route_unchanged(codec): + # On NAX GPUs a call at the route's sizes runs the kernel it runs under + # KQ_DEQUANT_ROUTE=off, bit for bit, by default and under + # KQ_DISABLE_NAX=1. + _metal() + if not kq.nax_available(): + pytest.skip("NAX GPU") + calls = [] + for rows in (ROWS[codec], 32769): + for nax_off in (None, "1"): + for mode in (None, "off"): + env = {"KQ_DISABLE_NAX": nax_off, "KQ_DEQUANT_ROUTE": mode} + calls.append(_call(codec, [1, rows, MIN_K], MIN_N, env, seed=rows)) + res = run_probe(calls, env=LIVE, check=False) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert not _routed(r), (c, r["stderr"]) + ((name, _),) = r["kernels"] + assert "_qmm_t_" in name, name + if c["env"]["KQ_DISABLE_NAX"]: + assert "_nax_" not in name, name + for i in range(0, len(res), 2): + assert res[i]["kernels"] == res[i + 1]["kernels"], (calls[i], calls[i + 1]) + assert res[i]["digest"] == res[i + 1]["digest"], (calls[i], calls[i + 1]) + + +# The float64 and ulp checks read every row, 1024 rows at a time, so a +# check holds float64 copies of 1024 rows. +CHECK_BLOCK = 1024 + + +def _pair(codec, shape, dtype, layout, seed): + # The route, then the kernel on the same inputs, its ulp held to the + # route's. + route = _call(codec, shape, MIN_N, _route_env(codec), dtype=dtype, layout=layout) + route.update(seed=seed, check_block=CHECK_BLOCK) + return [route, dict(route, env=KERNEL_ENV)] + + +def _with_ulp(calls): + for i in range(1, len(calls), 2): + calls[i]["ulp_vs"] = i - 1 + return calls + + +def _check_pairs(calls, res, dtype): + out_dtype = "bfloat16" if dtype == "float32" else dtype + for i in range(0, len(calls), 2): + c, r, k = calls[i], res[i], res[i + 1] + assert r["error"] is None and k["error"] is None, (c, r["error"], k["error"]) + assert _routed(r), (c, r["kernels"]) + assert not _routed(k), (c, k["stderr"]) + ((name, _),) = k["kernels"] + assert "_qmm_t_" in name and "_nax_" not in name, name + for x in (r, k): + assert x["checked"] == c["shape"][0] * c["shape"][1], (c, x["checked"]) + assert x["rel"] < R1[dtype], (c, x["rel"]) + assert x["nonfinite"] == 0, c + # The route's error against float64 is the kernel's. + assert abs(r["rel"] - k["rel"]) <= 0.01 * k["rel"], (c, r["rel"], k["rel"]) + assert k["ulp_big"] <= ULP[out_dtype], (c, k["ulp_big"]) + assert k["rdiff"] < 1e-4, (c, k["rdiff"]) + + +@pytest.mark.parametrize("layout", ["dense", "transposed", "strided"]) +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", sorted(ROWS)) +def test_route_numerics(codec, dtype, layout): + # The floor in one batch and in two, every row checked: for q4_k all + # 16384 rows, row 8192 where the two batches meet among them. + _metal() + min_m = ROWS[codec] + calls = _with_ulp( + _pair(codec, [1, min_m, MIN_K], dtype, layout, 0) + + _pair(codec, [2, min_m // 2, MIN_K], dtype, layout, 1) + ) + _check_pairs(calls, run_probe(calls, env=LIVE), dtype) + + +@pytest.mark.parametrize("codec", sorted(ROWS)) +def test_route_float32_x(codec): + # float32 x runs as bfloat16, on the route as on the kernel. + _metal() + calls = _with_ulp(_pair(codec, [1, ROWS[codec], MIN_K], "float32", "dense", 0)) + _check_pairs(calls, run_probe(calls, env=LIVE), "float32") + + +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", sorted(ROWS)) +def test_route_k_max(codec, dtype): + _metal() + calls = _with_ulp(_pair(codec, [1, ROWS[codec], MAX_K], dtype, "dense", 0)) + _check_pairs(calls, run_probe(calls, env=LIVE), dtype) + + +@pytest.mark.parametrize("codec", sorted(ROWS)) +def test_route_rows_above_int16(codec): + # 32769 rows on the route against the kernel and float64, every row + # checked, so rows 32767 and 32768 past the int16 range among them. + _metal() + calls = _with_ulp(_pair(codec, [1, 32769, MIN_K], "bfloat16", "dense", 0)) + _check_pairs(calls, run_probe(calls, env=LIVE), "bfloat16") + + +# Children for the calls run_probe does not make. Each prints one JSON line; +# stderr carries the KQ_QMM_LOG lines after a "[case] " marker. +_PRELUDE = r""" +import json, os, sys +import numpy as np +spec = json.loads(sys.argv[1]) +sys.path.insert(0, spec["tests_dir"]) +import mlx.core as mx +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY +from kqref import is_synth, synth_wire + + +def weights(codec, n, k, seed=0): + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + if is_synth(codec): + rng = np.random.default_rng(n * 131 + k + seed) + wire = synth_wire(rng, codec, bpb, n * (k // wpb)) + w = mx.array(wire.reshape(n, (k // wpb) * bpb)) + s = mx.zeros((1,), dtype=mx.uint8) + else: + wf = mx.random.normal((n, k), key=mx.random.key(n * 7 + k + seed)) * 0.1 + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + return w, s + + +def case(name, env): + for key, value in env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + os.write(2, f"[case] {name}\n".encode()) + + +def rel(a, b): + a = np.array(a.astype(mx.float32)).astype(np.float64) + b = np.array(b.astype(mx.float32)).astype(np.float64) + return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-30)) +""" + +_GRAD_CHILD = ( + _PRELUDE + + r""" +import time + +codec, m, n, k = spec["codec"], spec["m"], spec["n"], spec["k"] +w, s = weights(codec, n, k) +dt = getattr(mx, spec["dtype"]) +key = mx.random.key(1) +layout = spec["layout"] +if layout == "batched": + x = (mx.random.normal((2, m // 2, k), key=key) * 0.5).astype(dt) +elif layout == "transposed": + # A lazy transposed view. + x = mx.swapaxes((mx.random.normal((k, m), key=key) * 0.5).astype(dt), 0, 1) +elif layout == "strided": + # Every other column of a wider array. + x = (mx.random.normal((m, 2 * k), key=key) * 0.5).astype(dt)[:, ::2] +else: + x = (mx.random.normal((m, k), key=key) * 0.5).astype(dt) +mx.eval(x) + + +def f(x): + y = kq.quantized_matmul(x, w, s, codec, transpose=True) + return y.astype(mx.float32).sum() + + +def composed(x): + # Dequantize then matmul in stock ops, whose vjp keeps the dense copy. + wf = kq.dequantize(w, s, codec, dtype=mx.bfloat16) + return (x.astype(mx.bfloat16) @ wf.T).astype(mx.float32).sum() + + +def active(window): + # The buffers a finished command buffer frees go back in its completion + # handler, after the wait returns. Nothing allocates meanwhile, so the + # least of the reads over the window is what stays held. + mx.synchronize() + least = mx.get_active_memory() + end = time.monotonic() + window + while time.monotonic() < end: + time.sleep(0.01) + least = min(least, mx.get_active_memory()) + return least + + +def settled(): + # Reads until two windows agree, so the buffers an earlier arm freed + # are back before the base is taken. + prev = active(0.2) + for _ in range(5): + cur = active(0.2) + if cur == prev: + break + prev = cur + return cur + + +def forward(fn): + # What the pending backward graph holds after the forward. + mx.clear_cache() + base = settled() + loss, g = mx.value_and_grad(fn)(x) + mx.eval(loss) + return base, loss, g, active(0.2) - base + + +out = {} +grads = {} +for arm, fn in (("route", f), ("kernel", f), ("composed", composed)): + case(arm, spec.get(arm + "_env", {})) + base, loss, g, held = forward(fn) + # A late completion handler moves a reading. On the route it leaves + # buffers not yet freed in the reading, which then gets up to five + # longer windows. On the stock product it leaves them in the base, + # and the arm runs again, up to five times. + for _ in range(5): + if arm == "route" and held >= spec["held_below"]: + held = active(1.0) - base + elif arm == "composed" and held < spec["dense"]: + mx.eval(g) + del loss, g + base, loss, g, held = forward(fn) + else: + break + mx.eval(g) + grads[arm] = g + out[arm] = {"held": held, "loss": loss.item(), "dtype": str(g.dtype)} +gr, gk = grads["route"], grads["kernel"] +out["grad_equal"] = bool(mx.array_equal(gr, gk).item()) +out["dense"] = n * k * 2 +print(json.dumps(out)) +""" +) + +_WGRAD_CHILD = ( + _PRELUDE + + r""" +codec, m, n, k = spec["codec"], spec["m"], spec["n"], spec["k"] +w, s = weights(codec, n, k) +x = (mx.random.normal((m, k), key=mx.random.key(1)) * 0.5).astype(mx.bfloat16) +mx.eval(x) + + +def loss(x, w, s): + y = kq.quantized_matmul(x, w, s, codec, transpose=True) + return y.astype(mx.float32).sum() + + +out = {} +for arm in ("route", "kernel"): + for argnum in (1, 2): + name = f"{arm}_{argnum}" + case(name, spec[arm + "_env"]) + try: + mx.eval(mx.grad(loss, argnums=argnum)(x, w, s)) + out[name] = None + except Exception as e: + out[name] = f"{type(e).__name__}: {e}" +print(json.dumps(out)) +""" +) + +_TRANSFORM_CHILD = ( + _PRELUDE + + r""" +codec, m, n, k = spec["codec"], spec["m"], spec["n"], spec["k"] +w, s = weights(codec, n, k) +w2, _ = weights(codec, n, k, seed=1) +x = (mx.random.normal((m, k), key=mx.random.key(1)) * 0.5).astype(mx.bfloat16) +a = (mx.random.normal((k, 16), key=mx.random.key(5)) * 0.05).astype(mx.bfloat16) +b = (mx.random.normal((16, n), key=mx.random.key(6)) * 0.05).astype(mx.bfloat16) +r = (mx.random.normal((m, n), key=mx.random.key(7)) * 0.5).astype(mx.bfloat16) +mx.eval(x, a, b, r) +ops = { + "matmul": lambda t: kq.quantized_matmul(t, w, s, codec, transpose=True), + "glu": lambda t: kq.quantized_matmul_glu(t, w, w2, codec), + "lora": lambda t: kq.quantized_matmul_lora(t, w, codec, a, b), + "residual": lambda t: kq.quantized_matmul_residual(t, w, codec, r), +} +transforms = { + "jvp": lambda f: mx.jvp(f, [x], [x])[1], + "vmap": lambda f: mx.vmap(f)(mx.stack([x, x * 0.5])), +} +out = {"nax": kq.nax_available()} +for arm in ("route", "kernel"): + # A forward glu call first, so the log shows the node this env builds. + case(f"{arm}_glu_forward", spec[arm + "_env"]) + mx.eval(ops["glu"](x)) + for op, f in ops.items(): + for tr, run in transforms.items(): + name = f"{arm}_{op}_{tr}" + case(name, spec[arm + "_env"]) + try: + mx.eval(run(f)) + out[name] = None + except Exception as e: + out[name] = f"{type(e).__name__}: {e}" +print(json.dumps(out)) +""" +) + +_EXCLUDE_CHILD = ( + _PRELUDE + + r""" +codec, m, n, k = spec["codec"], spec["m"], spec["n"], spec["k"] +env = spec["env"] +w, s = weights(codec, n, k) +wd = kq.dequantize(w, s, codec, dtype=mx.float32) +x = (mx.random.normal((m, k), key=mx.random.key(1)) * 0.5).astype(mx.bfloat16) +mx.eval(x, wd) +xf = x.astype(mx.float32) +ref = xf @ wd.T +out = {} + + +def run(name, fn, want): + case(name, env) + try: + y = fn() + mx.eval(y) + out[name] = {"rel": rel(y, want), "error": None} + except Exception as e: + out[name] = {"rel": None, "error": f"{type(e).__name__}: {e}"} + + +# The plain call takes the route under this env. +run("plain", lambda: kq.quantized_matmul(x, w, s, codec, transpose=True), ref) +# LoRA arguments keep the fused kernel. +a = (mx.random.normal((k, 16), key=mx.random.key(5)) * 0.05).astype(mx.bfloat16) +b = (mx.random.normal((16, n), key=mx.random.key(6)) * 0.05).astype(mx.bfloat16) +mx.eval(a, b) +lora_ref = ref + (xf @ a.astype(mx.float32)) @ b.astype(mx.float32) +run( + "lora", + lambda: kq.quantized_matmul(x, w, s, codec, transpose=True, lora_a=a, lora_b=b), + lora_ref, +) +# A CPU stream under the GPU default device. +run( + "cpu_stream", + lambda: kq.quantized_matmul(x, w, s, codec, transpose=True, stream=mx.cpu), + ref, +) +# transpose=False: w holds [K, N] rows. +wt, st = weights(codec, k, n, seed=1) +wtd = kq.dequantize(wt, st, codec, dtype=mx.float32) +mx.eval(wtd) +run( + "no_transpose", + lambda: kq.quantized_matmul(x, wt, st, codec, transpose=False), + xf @ wtd, +) +# A 3-D w, from quantized_matmul and from gather_qmm without indices. +w3 = mx.stack([w, w]) +s3 = s if s.size == 1 else mx.stack([s, s]) +mx.eval(w3, s3) +ref3 = mx.stack([ref, ref]) +run("w3d", lambda: kq.quantized_matmul(x, w3, s3, codec, transpose=True), ref3) +run("gather_w3d", lambda: kq.gather_qmm(x, w3, s3, codec, transpose=True), ref3) +# gather_qmm without indices on a 2-D w is the plain call. +run("gather_w2d", lambda: kq.gather_qmm(x, w, s, codec, transpose=True), ref) +print(json.dumps(out)) +""" +) + + +def _run_child(code, spec, env): + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update(env) + spec = dict(spec, tests_dir=TESTS_DIR) + r = subprocess.run( + [sys.executable, "-c", code, json.dumps(spec)], + env=child_env, + capture_output=True, + text=True, + timeout=900, + ) + assert r.returncode == 0, r.stderr[-6000:] + logs = {} + cur = None + for line in r.stderr.splitlines(): + if line.startswith("[case] "): + cur = line.split()[1] + logs[cur] = [] + elif cur is not None and line.startswith("[kq "): + logs[cur].append(line) + return json.loads(r.stdout.strip().splitlines()[-1]), logs + + +def _child_routed(lines): + return any(x.startswith("[kq dequant] ") for x in lines) + + +VJP_CASES = [(c, "bfloat16", "dense") for c in sorted(ROWS)] + [ + ("q6_k", "float32", "dense"), + ("q6_k", "float16", "dense"), + ("nvfp4", "float16", "dense"), + ("q6_k", "bfloat16", "batched"), + ("q6_k", "bfloat16", "transposed"), + ("q6_k", "float16", "strided"), +] + + +@pytest.mark.parametrize("codec,dtype,layout", VJP_CASES) +def test_route_vjp(codec, dtype, layout): + # value_and_grad through a routed call: after the forward the pending + # backward holds no dense copy of the weight, where the same product in + # stock ops holds one, and the gradient wrt x is the kernel path's bit + # for bit (both run the quantized matmul the other way). + _metal() + spec = { + "codec": codec, + "m": ROWS[codec], + "n": MIN_N, + "k": MAX_K, + "dtype": dtype, + "layout": layout, + "route_env": _route_env(codec), + "kernel_env": KERNEL_ENV, + "held_below": MIN_N * MAX_K * 2 // 4, + "dense": MIN_N * MAX_K * 2, + } + out, logs = _run_child(_GRAD_CHILD, spec, LIVE) + assert _child_routed(logs["route"]), logs["route"] + assert not _child_routed(logs["kernel"]), logs["kernel"] + assert out["route"]["held"] < out["dense"] // 4, out + # The reading sees a dense copy where one is held. + assert out["composed"]["held"] >= out["dense"], out + assert out["route"]["dtype"] == out["kernel"]["dtype"] == f"mlx.core.{dtype}", out + assert out["grad_equal"], out + lr, lk = out["route"]["loss"], out["kernel"]["loss"] + assert abs(lr - lk) <= 1e-3 * abs(lk), out + + +@pytest.mark.parametrize("codec", ["nvfp4", "q6_k"]) +def test_route_weight_grad_raises(codec): + # The gradient wrt the wire bytes or the scales raises on the route as + # on the kernel, with the same message. + _metal() + spec = {"codec": codec, "m": ROWS[codec], "n": MIN_N, "k": MIN_K} + spec.update(route_env=_route_env(codec), kernel_env=KERNEL_ENV) + out, logs = _run_child(_WGRAD_CHILD, spec, LIVE) + for argnum, msg in ((1, "no gradient wrt the quantized weights"), (2, "scales")): + r, k = out[f"route_{argnum}"], out[f"kernel_{argnum}"] + assert r is not None and msg in r, out + assert r == k, out + assert _child_routed(logs[f"route_{argnum}"]), logs + assert not _child_routed(logs[f"kernel_{argnum}"]), logs + + +@pytest.mark.parametrize("fused_ops", ["default", "off"]) +def test_route_jvp_vmap_refuse_as_kernel(fused_ops): + # jvp and vmap wrt x of quantized_matmul, glu, lora and residual at a + # routed shape raise by the name KQuantMatmul, as on the kernel, so the + # error text does not change at the row floor. Where glu builds its + # fused node (a NAX GPU with KQ_GLU on, KQ_DISABLE_NAX or not), vmap of + # glu refuses on that node, before any quantized_matmul is built. + # Elsewhere glu builds two quantized_matmul calls, which route and + # refuse as KQuantMatmul. off turns the fused glu, lora and residual + # nodes off, so any GPU builds what a GPU without NAX builds. + _metal() + codec = "q6_k" + spec = {"codec": codec, "m": ROWS[codec], "n": MIN_N, "k": MIN_K} + off = "0" if fused_ops == "off" else None + lever = {"KQ_GLU": off, "KQ_QMM_LORA": off, "KQ_QMM_RESID": off} + force = {"KQ_QMM_TILE": None, "KQ_DISABLE_NAX": None, "KQ_DEQUANT_ROUTE": "force"} + spec.update(route_env=dict(force, **lever), kernel_env=dict(KERNEL_ENV, **lever)) + out, logs = _run_child(_TRANSFORM_CHILD, spec, LIVE) + fused = out["nax"] and fused_ops == "default" + for arm in ("route", "kernel"): + # The forward call runs the fused kernel exactly where the node is. + lines = logs[f"{arm}_glu_forward"] + assert any("_glu_" in x for x in lines) == fused, (arm, out["nax"], lines) + for op in ("matmul", "glu", "lora", "residual"): + for tr in ("jvp", "vmap"): + r, k = out[f"route_{op}_{tr}"], out[f"kernel_{op}_{tr}"] + on_node = fused and (op, tr) == ("glu", "vmap") + node = "KQuantMatmulFused" if on_node else "KQuantMatmul" + want = f"ValueError: [Primitive::{tr}] Not implemented for {node}." + assert r == want, (op, tr, out) + assert k == want, (op, tr, out) + routed = _child_routed(logs[f"route_{op}_{tr}"]) + assert routed != on_node, (op, tr, logs) + assert not _child_routed(logs[f"kernel_{op}_{tr}"]), (op, tr, logs) + + +@pytest.mark.parametrize("codec", ["nvfp4", "q6_k"]) +def test_route_exclusions(codec): + # Under the route's env the plain call routes; LoRA arguments, a CPU + # stream under the GPU default device, transpose=False and a 3-D w + # (quantized_matmul or gather_qmm) keep the kernel. gather_qmm without + # indices on a 2-D w is the plain call and routes. + _metal() + spec = {"codec": codec, "m": ROWS[codec], "n": MIN_N, "k": MIN_K} + spec["env"] = _route_env(codec) + out, logs = _run_child(_EXCLUDE_CHILD, spec, LIVE) + want = { + "plain": True, + "lora": False, + "cpu_stream": False, + "no_transpose": False, + "w3d": False, + "gather_w3d": False, + "gather_w2d": True, + } + for name, routed in want.items(): + assert out[name]["error"] is None, (name, out[name]) + assert out[name]["rel"] < R1["bfloat16"], (name, out[name]) + assert _child_routed(logs[name]) == routed, (name, logs[name]) diff --git a/tests/test_conv_nax.py b/tests/test_conv_nax.py new file mode 100644 index 00000000..54e5eb3d --- /dev/null +++ b/tests/test_conv_nax.py @@ -0,0 +1,448 @@ +"""conv2d_3x3: y = conv(x, w) + bias, 3x3, stride 1, padding 1, channels last. + +On a NAX GPU the op runs one implicit GEMM kernel for bfloat16 with C a +multiple of 16 from 1024 output pixels: f32 sums over the 9 taps and C +channels, the bias added in f32, one rounding. Everywhere else it builds +MLX's convolution on the flipped kernel (its implicit GEMM, no Winograd +transform) and a stock add of the bias. + +- R1: relative Frobenius error against the float64 convolution, at most + 6e-3 on bfloat16 and 1e-3 on float16. +- R2: on the kernel, the error sits at the one rounding of y: at most 1.05 + times the error of the float64 result rounded to bfloat16, and each + element within ULPS ulp of bfloat16 at the float64 result plus 2^-20 + times the sum of |x w| over its taps and channels. +- Where the op has no kernel, y is the composition a caller writes, bit for + bit. + +Under KQUANT_FORCE_CPU the op runs the composition on the CPU stream, and +the tests that name a Metal kernel skip. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys + +import mlx.core as mx +import numpy as np +import pytest +from qmm_probe import force_cpu +from test_operand_views import _dense + +import mlx_kquant as kq + +R1 = {mx.bfloat16: 6e-3, mx.float16: 1e-3} +CPU_REL = 5e-2 +MIN_ROWS = 1024 +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) +# Tiles of the metallib: the two the op picks, and two for KQ_CONV_TILE. +TILES = [(128, 128, 4, 4), (128, 64, 4, 2), (64, 128, 2, 4), (64, 64, 2, 2)] +# R2 per element: the ulp count and the allowance for the f32 sums per unit +# of sum |x w|. Bfloat16 holds 8 significant bits, its smallest normal has +# np.frexp exponent -125. +ULPS = 2 +SUM_TOL = 2.0**-20 + + +def _nax(): + return not force_cpu() and kq.nax_available() + + +def _fused(pixels, c, dtype=mx.bfloat16): + return _nax() and dtype == mx.bfloat16 and c % 16 == 0 and pixels >= MIN_ROWS + + +def _f64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + +def _rel(a, b): + return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-30)) + + +def _inputs(b, h, w, c, o, dtype=mx.bfloat16, bias=True, seed=0): + x = mx.random.normal((b, h, w, c), key=mx.random.key(seed + 1)).astype(dtype) + wt = mx.random.normal((o, 3, 3, c), key=mx.random.key(seed + 2)) / np.sqrt(9 * c) + wt = wt.astype(dtype) + bs = ( + mx.random.normal((o,), key=mx.random.key(seed + 3)).astype(dtype) + if bias + else None + ) + mx.eval(x, wt, *(() if bs is None else (bs,))) + return x, wt, bs + + +def _ref(x, w, bias, mag=False): + """The float64 convolution, pixels flattened. mag=True gives sum |x w| + per element, without the bias.""" + xd, wd = _f64(x), _f64(w) + if mag: + xd, wd, bias = np.abs(xd), np.abs(wd), None + b, h, wid, c = xd.shape + o = wd.shape[0] + xp = np.pad(xd, ((0, 0), (1, 1), (1, 1), (0, 0))) + taps = [xp[:, dy : dy + h, dx : dx + wid, :] for dy in range(3) for dx in range(3)] + y = np.concatenate(taps, axis=-1).reshape(-1, 9 * c) @ wd.reshape(o, 9 * c).T + if bias is not None: + y = y + _f64(bias) + return y + + +def _ulps(got, ref, mag): + """The largest error in ulp of bfloat16 at |ref|, after the allowance of + SUM_TOL * mag for the f32 sums, and its pixel and column.""" + e = np.frexp(np.abs(ref))[1] + e = np.where(ref == 0, -125, np.maximum(e, -125)) + u = (np.abs(got - ref) - SUM_TOL * mag) / np.ldexp(1.0, e - 8) + i = np.unravel_index(int(np.argmax(u)), u.shape) + return float(u[i]), tuple(int(v) for v in i) + + +def _composed(x, w, bias): + rev = mx.array([2, 1, 0]) + wf = mx.take(mx.take(w.astype(x.dtype), rev, axis=1), rev, axis=2) + y = mx.conv_general(x, wf, stride=1, padding=1, flip=True) + if bias is not None: + y = y + bias.astype(x.dtype) + return y + + +def _check(x, w, bias): + """R1, and R2 on the kernel. Returns (R1 error, R2 ratio or None).""" + b, h, wid, c = x.shape + o = w.shape[0] + y = kq.conv2d_3x3(x, w, bias=bias) + mx.eval(y) + assert y.shape == (b, h, wid, o) and y.dtype == x.dtype + got = _f64(y).reshape(-1, o) + assert np.isfinite(got).all() + ref = _ref(x, w, bias) + err = _rel(got, ref) + if not _fused(b * h * wid, c, x.dtype): + comp = _composed(x, w, bias) + mx.eval(comp) + assert np.array_equal(_f64(comp).reshape(-1, o), got) + assert err <= (R1[x.dtype] if _nax() else CPU_REL), err + return err, None + assert err <= R1[x.dtype], err + floor = _rel(_f64(mx.array(ref.astype(np.float32)).astype(x.dtype)), ref) + assert err <= 1.05 * floor + 1e-5, (err, floor) + worst = _ulps(got, ref, _ref(x, w, bias, mag=True)) + assert worst[0] <= ULPS, worst + return err, err / floor + + +# The channel counts of the Z-Image (128, 256, 512) and Qwen-Image 2.1 +# (144, 288, 576, 1152) VAE decoders, and the small and ragged ones. +SHAPES = [ + (1, 37, 53, 128, 128), + (1, 32, 40, 256, 256), + (1, 24, 24, 512, 512), + (1, 33, 29, 144, 144), + (1, 30, 31, 288, 144), + (1, 16, 18, 576, 288), + (1, 12, 12, 1152, 1152), + (1, 40, 26, 32, 64), + (1, 64, 64, 128, 3), + (1, 33, 31, 16, 100), +] + + +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.parametrize("shape", SHAPES, ids=lambda s: "x".join(map(str, s))) +def test_conv_shapes(shape, bias): + _check(*_inputs(*shape, bias=bias)) + + +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("hw", [(31, 33), (16, 32), (1, 1100), (1100, 1)]) +def test_conv_batch_and_thin_images(hw, batch): + # One row or one column of pixels: every tap but the middle one falls + # outside the image on one axis. + _check(*_inputs(batch, *hw, 64, 96, seed=batch)) + + +@pytest.mark.parametrize("pixels", [1023, 1024, 1025]) +def test_conv_pixel_bound(pixels): + _check(*_inputs(1, 1, pixels, 32, 64)) + + +def test_conv_pixels_above_32767(): + # Row offsets pass a 16-bit index: every pixel is compared. + _check(*_inputs(1, 200, 180, 32, 64)) + _check(*_inputs(2, 129, 130, 16, 32)) + + +@pytest.mark.parametrize("c", [8, 24, 40]) +def test_conv_channels_off_16_compose(c): + err, r2 = _check(*_inputs(1, 40, 40, c, 64)) + assert r2 is None + + +def test_conv_float16_and_float32_compose(): + for dt in (mx.float16, mx.float32): + x, w, b = _inputs(1, 40, 40, 64, 64, dtype=dt) + y = kq.conv2d_3x3(x, w, bias=b) + want = _composed(x, w, b) + mx.eval(y, want) + assert y.dtype == dt + assert mx.array_equal(y, want).item() + + +def test_conv_weight_and_bias_run_in_the_dtype_of_x(): + x, w, b = _inputs(1, 40, 40, 64, 64) + y = kq.conv2d_3x3(x, w.astype(mx.float32), bias=b.astype(mx.float32)) + want = kq.conv2d_3x3(x, w, bias=b) + mx.eval(y, want) + assert y.dtype == mx.bfloat16 + assert mx.array_equal(y, want).item() + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +def test_conv_operand_views(state): + """A view of x, w or bias gives the bits of its dense copy, lazy or + evaluated (one rounding apart on the CPU stream).""" + x, w, b = _inputs(1, 40, 36, 64, 96) + x2 = mx.random.normal((1, 36, 40, 128), key=mx.random.key(9)).astype(mx.bfloat16) + w2 = ( + mx.random.normal((96, 3, 3, 128), key=mx.random.key(10)).astype(mx.bfloat16) + * 0.05 + ) + b2 = mx.random.normal((192,), key=mx.random.key(11)).astype(mx.bfloat16) + mx.eval(x2, w2, b2) + views = { + "x_transposed": lambda: (mx.swapaxes(x2[..., :64], 1, 2), w, b), + "x_strided": lambda: (mx.swapaxes(x2, 1, 2)[..., 1::2], w, b), + "w_strided": lambda: (x, w2[..., 0::2], b), + "bias_strided": lambda: (x, w, b2[1::2]), + } + for name, make in views.items(): + xv, wv, bv = make() + want = kq.conv2d_3x3(_dense(xv), _dense(wv), bias=_dense(bv)) + if state == "evaluated": + mx.eval(xv, wv, bv) + got = kq.conv2d_3x3(xv, wv, bias=bv) + mx.eval(got, want) + if force_cpu(): + assert _rel(_f64(got), _f64(want)) <= R1[mx.bfloat16], name + else: + assert mx.array_equal(got, want).item(), name + xv = mx.swapaxes(x2[..., :64], 1, 2) + _check(xv, w, b) + + +@pytest.mark.parametrize("bias", [True, False]) +def test_conv_gradient_is_the_gradient_of_the_composition(bias): + x, w, b = _inputs(1, 40, 36, 32, 48, bias=bias) + cot = mx.random.normal((1, 40, 36, 48), key=mx.random.key(5)) + + def op(*a): + return ( + kq.conv2d_3x3(a[0], a[1], bias=a[2] if bias else None).astype(mx.float32) + * cot + ).sum() + + def comp(*a): + return ( + _composed(a[0], a[1], a[2] if bias else None).astype(mx.float32) * cot + ).sum() + + args = (x, w, b) if bias else (x, w) + nums = tuple(range(len(args))) + got = mx.grad(op, argnums=nums)(*args) + want = mx.grad(comp, argnums=nums)(*args) + mx.eval(got, want) + for g, v in zip(got, want, strict=True): + assert mx.array_equal(g, v).item() + + +def _outcome(fn): + try: + out = fn() + mx.eval(out) + return out, None + except Exception as e: # noqa: BLE001 + return None, type(e) + + +def test_conv_jvp_and_vmap_follow_the_composition(): + """jvp and vmap give the composition's result, or raise where it does.""" + x, w, b = _inputs(1, 40, 36, 32, 48) + tx = mx.random.normal(x.shape, key=mx.random.key(6)).astype(x.dtype) + got, gerr = _outcome( + lambda: mx.jvp(lambda x_: kq.conv2d_3x3(x_, w, bias=b), [x], [tx])[1] + ) + want, werr = _outcome(lambda: mx.jvp(lambda x_: _composed(x_, w, b), [x], [tx])[1]) + assert gerr == werr + if got is not None: + assert mx.array_equal(got[0], want[0]).item() + xs = mx.random.normal((2, 1, 40, 36, 32), key=mx.random.key(7)).astype(x.dtype) + got, gerr = _outcome(lambda: mx.vmap(lambda x_: kq.conv2d_3x3(x_, w, bias=b))(xs)) + want, werr = _outcome(lambda: mx.vmap(lambda x_: _composed(x_, w, b))(xs)) + assert gerr == werr + if got is not None: + assert mx.array_equal(got, want).item() + + +def test_conv_compiles_and_repeats(): + x, w, b = _inputs(1, 40, 36, 64, 64) + f = mx.compile(lambda x_, w_, b_: kq.conv2d_3x3(x_, w_, bias=b_)) + want = kq.conv2d_3x3(x, w, bias=b) + for _ in range(2): + got = f(x, w, b) + mx.eval(got, want) + assert mx.array_equal(got, want).item() + + +def test_conv_empty_batch(): + x, w, b = _inputs(1, 4, 4, 32, 32) + y = kq.conv2d_3x3(mx.zeros((0, 4, 4, 32), dtype=mx.bfloat16), w, bias=b) + mx.eval(y) + assert y.shape == (0, 4, 4, 32) + + +def test_conv_argument_errors(): + x, w, b = _inputs(1, 8, 8, 32, 16) + + def bad(match, x_=x, w_=w, **kw): + with pytest.raises(ValueError, match=r"conv2d_3x3\] .*" + match): + kq.conv2d_3x3(x_, w_, **kw) + + bad("x must be", x_=x[0]) + bad("x must be", x_=x.astype(mx.int32)) + bad("w must be", w_=w[..., :-1]) + bad("w must be", w_=w[:, :2]) + bad("w must be", w_=w.reshape(16, 9, 32)) + bad("w must be", w_=w.astype(mx.int32)) + bad("bias must be", bias=b[:-1]) + bad("bias must be", bias=b.reshape(1, 16)) + bad("bias must be", bias=b.astype(mx.int32)) + + +_PROBE = r""" +import json, os, sys +spec = json.loads(sys.argv[1]) +sys.path.insert(0, spec["tests_dir"]) +import mlx.core as mx +import mlx_kquant as kq +import test_conv_nax as t + +out = [] +for i, c in enumerate(spec["calls"]): + os.environ.update(c.get("env", {})) + x, w, b = t._inputs(*c["shape"], bias=c.get("bias", True)) + print(f"[probe] {i}", file=sys.stderr, flush=True) + res = {"error": None} + try: + y = kq.conv2d_3x3(x, w, bias=b) + mx.eval(y) + got = t._f64(y).reshape(-1, c["shape"][-1]) + ref = t._ref(x, w, b) + res["rel"] = t._rel(got, ref) + res["floor"] = t._rel( + t._f64(mx.array(ref.astype(t.np.float32)).astype(mx.bfloat16)), ref) + res["ulps"] = t._ulps(got, ref, t._ref(x, w, b, mag=True)) + except Exception as e: + res["error"] = f"{type(e).__name__}: {e}" + out.append(res) +print("RESULT " + json.dumps(out)) +""" + + +def _probe(calls, env=None): + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update({"KQ_QMM_LOG": "1", **(env or {})}) + p = subprocess.run( + [ + sys.executable, + "-c", + _PROBE, + json.dumps({"tests_dir": TESTS_DIR, "calls": calls}), + ], + env=child_env, + capture_output=True, + text=True, + timeout=900, + ) + assert p.returncode == 0, p.stderr[-3000:] + line = [ln for ln in p.stdout.splitlines() if ln.startswith("RESULT ")][-1] + results = json.loads(line[len("RESULT ") :]) + cur = None + for ln in p.stderr.splitlines(): + if ln.startswith("[probe] "): + cur = results[int(ln.split()[1])] + cur["kernels"] = [] + elif cur is not None and ln.startswith("[kq qmm] "): + cur["kernels"].append(ln.split()[2]) + return results + + +def _name(tile): + bm, bn, wm, wn = tile + return f"kq_conv3x3_nax_bfloat16_t_bm{bm}_bn{bn}_wm{wm}_wn{wn}" + + +def test_conv_kernels_in_metallib(): + if force_cpu(): + pytest.skip("Metal kernel names only") + names = kq.metallib_kernels() + if not names: + pytest.skip("no metallib loaded") + got = {n for n in names if n.startswith("kq_conv3x3_nax_")} + assert got == {_name(t) for t in TILES}, sorted(got) + + +def test_conv_known_answer_of_every_kernel(): + """Each kernel of the metallib, under KQ_CONV_TILE, against the float64 + convolution and the one-rounding floor, with O on and off its column + tile and C off a multiple of 32.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + calls = [] + for tile in TILES: + env = {"KQ_CONV_TILE": ",".join(map(str, tile))} + for shape in ((1, 37, 41, 144, 200), (2, 24, 26, 64, 128)): + calls.append({"shape": list(shape), "env": env}) + res = _probe(calls, env={"KQ_TUNE_LIVE": "1"}) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + tile = tuple(int(v) for v in c["env"]["KQ_CONV_TILE"].split(",")) + assert r["kernels"] == [_name(tile)], (c, r["kernels"]) + assert r["rel"] <= 1.05 * r["floor"] + 1e-5, (c, r["rel"], r["floor"]) + assert r["ulps"][0] <= ULPS, (c, r["ulps"]) + + +def test_conv_tile_choice_and_levers(): + """By default the 128 column tile runs, or the 64 one where it pads O + less; KQ_CONV_NAX=0 composes, 2 runs the kernel below 1024 pixels, and + a tile the metallib lacks raises with its name.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + calls = [ + ({"shape": [1, 40, 40, 64, 128]}, _name((128, 128, 4, 4))), + ({"shape": [1, 40, 40, 64, 144]}, _name((128, 64, 4, 2))), + ({"shape": [1, 40, 40, 64, 256]}, _name((128, 128, 4, 4))), + ({"shape": [1, 40, 40, 64, 288]}, _name((128, 64, 4, 2))), + ({"shape": [1, 10, 10, 64, 128]}, None), + ({"shape": [1, 10, 10, 64, 128], "env": {"KQ_CONV_NAX": "2"}}, _name(TILES[0])), + ({"shape": [1, 40, 40, 64, 128], "env": {"KQ_CONV_NAX": "0"}}, None), + ({"shape": [1, 40, 40, 64, 128], "env": {"KQ_CONV_NAX": "1"}}, _name(TILES[0])), + ( + {"shape": [1, 40, 40, 64, 128], "env": {"KQ_CONV_TILE": "32,32,1,1"}}, + "raise", + ), + ] + res = _probe([c for c, _ in calls], env={"KQ_TUNE_LIVE": "1"}) + for (c, want), r in zip(calls, res, strict=True): + if want == "raise": + assert ( + r["error"] is not None + and "kq_conv3x3_nax_bfloat16_t_bm32" in r["error"] + ) + continue + assert r["error"] is None, (c, r["error"]) + assert r["kernels"] == ([] if want is None else [want]), (c, r["kernels"]) diff --git a/tests/test_deps.py b/tests/test_deps.py index 6d8f5390..b0f8d508 100644 --- a/tests/test_deps.py +++ b/tests/test_deps.py @@ -39,7 +39,7 @@ def test_require_tools_missing(monkeypatch, _no_mlx_lm): def test_require_tools_broken(monkeypatch, _no_mlx_lm): - # e.g. transformers 5.13 rejecting mlx-lm's tokenizer registration + # e.g. transformers 5.13.0 rejecting mlx-lm's tokenizer registration _install_finder( monkeypatch, AttributeError("'str' object has no attribute '__module__'") ) diff --git a/tests/test_device.py b/tests/test_device.py index b237f539..c5fde52a 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -60,3 +60,12 @@ def test_fine_ceiling_scales_with_cores(monkeypatch): monkeypatch.setenv("KQ_GPU_CORES", "80") assert kq.qmv_fine_max_n("q6_k") == 32768 assert kq.qmv_fine_max_n("nvfp4") == 0 + + +def test_max_threadgroup_memory(): + n = kq.max_threadgroup_memory() + if sys.platform == "darwin": + # Every Apple GPU with Metal 3 has at least 32 KiB. + assert n >= 32768 + else: + assert n == 0 diff --git a/tests/test_dit_bench.py b/tests/test_dit_bench.py new file mode 100644 index 00000000..4fdbbb4c --- /dev/null +++ b/tests/test_dit_bench.py @@ -0,0 +1,594 @@ +"""The DiT benchmark harness: the median CI at any sample count, rebuilding +reports from a raw log, and, at tiny shapes in a child process, the cache +limit, the per-visit memory records and resuming a stopped run.""" + +from __future__ import annotations + +import json +import math +import os +import subprocess +import sys +import textwrap +from fractions import Fraction + +import pytest + +BENCH_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "benchmarks") +sys.path.insert(0, BENCH_DIR) + +import bench_stats as bs # noqa: E402 + + +def _tail(n, k): + """P(X < k) for X ~ Binomial(n, 1/2), exact.""" + return Fraction(sum(math.comb(n, i) for i in range(k)), 2**n) + + +@pytest.mark.parametrize("n", [3, 24, 1100, 5000]) +def test_median_ci_order_statistics(n): + s = list(range(n)) + ci = bs.median_ci(s) + if n == 3: + # Even the widest interval [s(1), s(n)] covers under 95 percent. + assert ci is None + return + lo, hi = ci + k = lo + 1 + assert hi == n - k + # Largest k with P(X < k) <= 0.025. + assert _tail(n, k) <= Fraction(1, 40) < _tail(n, k + 1) + if n == 24: + # Standard table value: the 95 percent CI of the median of 24 is + # [x(7), x(18)]. + assert (lo, hi) == (6, 17) + else: + assert abs(k - (n - 1.96 * math.sqrt(n)) / 2) < 1.5 + + +def test_summarize_ci_finite_for_thousands_of_pairs(): + s = [1.0 + 0.001 * math.sin(i) for i in range(5000)] + out = bs.summarize(s) + assert out["n"] == 5000 + assert out["ci95"] is not None + assert 0 < out["ci95_half"] < 0.001 + + +_WRITER = textwrap.dedent( + """ + import json, sys + sys.path.insert(0, sys.argv[1]) + import bench_dit_shapes as b + raw, json_out, md_out = sys.argv[2:5] + g = {"family": "zimage", "op": "qkv", "K": 256, "N": 704, "L": 64, "B": 1, + "M": 64, "L_model": 4096, "vector": False, "heavy": False, + "codecs": ["q4_k"]} + args = {"rounds": 3, "slots": 8, "inner": 3, "max_inner": 64, + "min_sample_ms": 10.0, "shrink": 16, "arms": ["kq", "kq_aa"]} + machine = {"device": "gpu", "architecture": "applegpu_g17s", "arch_gen": 17, + "arch_letter": "s", "gpu_cores": 40, "nax": True, + "max_threadgroup_memory": 32768, "os": "26", "mlx": "0.32.2", + "mlx_kquant": "0.5.0", "metallib_sha256": None, "tf32": True} + log = b.RawLog(raw) + log.write({"kind": "meta", "machine": machine, "args": args, + "heavy_codecs": [], "groups": [g], "mlx_kquant_file": "kq"}) + log.write({"kind": "warmup", "dtype": "bfloat16", "inner": [3], + "kernels": [[0, "kq/q4_k", "k"], [0, "kq_aa/q4_k", "k"]], + "errors": []}) + visit = 0 + for r in range(2): + for s in range(8): + ms = 1.0 + 0.01 * s + rows = [[0, "kq/q4_k", visit, r, s, 1.0, ms], + [0, "kq_aa/q4_k", visit, r, s, 1.0, ms * 1.01]] + log.write({"kind": "slot", "dtype": "bfloat16", "round": r, + "slot": s, "samples": rows}) + visit += 1 + log.close() + with open(raw, "a") as f: + f.write('{"kind": "slot", "dtype": "bfl') + sys.exit(b.main(["--reaggregate", raw, "--json-out", json_out, + "--md-out", md_out])) + """ +) + + +def test_reaggregate_rebuilds_a_stopped_run(tmp_path): + raw = tmp_path / "run.json.raw.jsonl" + json_out = tmp_path / "run.json" + md_out = tmp_path / "run.md" + script = tmp_path / "writer.py" + script.write_text(_WRITER) + # A child process, since dit_common sets KQ_TUNE_LIVE on import. + subprocess.run( + [sys.executable, str(script), BENCH_DIR, str(raw), str(json_out), str(md_out)], + check=True, + timeout=120, + ) + result = json.loads(json_out.read_text()) + assert result["complete"] is False + (p,) = result["passes"] + assert (p["slots_done"], p["slots_planned"]) == (16, 24) + (ag,) = [a for a in p["aggregate"] if a["ratio"] == "aa/kq"] + assert ag["op"] == "zimage.qkv" + # Round 0 is left out of the pooled figure. + assert ag["pooled"]["n"] == 8 + assert ag["pooled"]["median"] == pytest.approx(1.01) + # A log without a sample mode is from the independent-call bench. + assert result["sample_mode"] == "independent" + md = md_out.read_text() + assert "do not compare with chained runs" in md + assert "Partial run" in md + assert "16 of 24 slots ran." in md + + +_ORDER_WRITER = textwrap.dedent( + """ + import json, sys + sys.path.insert(0, sys.argv[1]) + import bench_dit_shapes as b + import dit_common as dc + raw, json_out, md_out = sys.argv[2:5] + g = {"family": "zimage", "op": "qkv", "K": 256, "N": 704, "L": 64, "B": 1, + "M": 64, "L_model": 4096, "vector": False, "heavy": False, + "codecs": ["q4_k"]} + args = {"rounds": 3, "slots": 8, "inner": 3, "max_inner": 64, + "min_sample_ms": 10.0, "shrink": 16, "arms": ["kq", "kq_aa"]} + machine = {"device": "gpu", "architecture": "applegpu_g17s", "arch_gen": 17, + "arch_letter": "s", "gpu_cores": 40, "nax": True, + "max_threadgroup_memory": 32768, "os": "26", "mlx": "0.32.2", + "mlx_kquant": "0.5.0", "metallib_sha256": None, "tf32": True} + log = b.RawLog(raw) + log.write({"kind": "meta", "machine": machine, "args": args, + "heavy_codecs": [], "groups": [g], "mlx_kquant_file": "kq", + "sample_mode": "chain"}) + log.write({"kind": "warmup", "dtype": "bfloat16", "inner": [3], + "kernels": [[0, "kq/q4_k", "k"], [0, "kq_aa/q4_k", "k"]], + "errors": []}) + # kq_aa takes 1.02x kq, and the second call of a visit 1.1x the first. + visit, t = 0, 0.0 + for r in range(3): + for s in range(8): + arms = [("kq/q4_k", 1.0), ("kq_aa/q4_k", 1.02)] + if dc.THUE_MORSE[s]: + arms.reverse() + rows = [] + for i, (key, ms) in enumerate(arms): + t += 1.0 + rows.append([0, key, visit, r, s, t, ms * (1.1 if i else 1.0)]) + log.write({"kind": "slot", "dtype": "bfloat16", "round": r, + "slot": s, "samples": rows}) + visit += 1 + log.write({"kind": "end", "wall_s": 1.0}) + log.close() + sys.exit(b.main(["--reaggregate", raw, "--json-out", json_out, + "--md-out", md_out])) + """ +) + + +def test_ratio_balances_call_order(tmp_path): + raw = tmp_path / "run.json.raw.jsonl" + json_out = tmp_path / "run.json" + script = tmp_path / "writer.py" + script.write_text(_ORDER_WRITER) + subprocess.run( + [ + sys.executable, + str(script), + BENCH_DIR, + str(raw), + str(json_out), + str(tmp_path / "run.md"), + ], + check=True, + timeout=120, + ) + result = json.loads(json_out.read_text()) + (p,) = result["passes"] + (g,) = p["groups"] + (c,) = g["cells"] + rt = c["ratios"]["aa/kq"] + # The pairs of the two orders, 1.122 and 0.927, meet at 1.02. + assert rt["order"] == pytest.approx(1.1**2) + assert sorted({round(v, 9) for _r, v, _f in rt["pairs_raw"]}) == [ + round(1.02 / 1.1, 9), + round(1.02 * 1.1, 9), + ] + st = rt["settled"] + assert st["median"] == pytest.approx(1.02) + assert st["p10"] == pytest.approx(st["p90"]) + (ag,) = [a for a in p["aggregate"] if a["ratio"] == "aa/kq"] + assert ag["pooled"]["median"] == pytest.approx(1.02) + assert ag["cell_order"]["median"] == pytest.approx(1.1**2) + + +# Two tiny shapes, two rounds of two slots, one call per sample. +_TINY = [ + "--families", + "zimage", + "--ops", + "qkv,out", + "--rows", + "64", + "--batches", + "1", + "--codecs", + "q4_k", + "--heavy-codecs", + "none", + "--arms", + "kq,kq_aa", + "--dtypes", + "bfloat16", + "--rounds", + "2", + "--slots", + "2", + "--inner", + "1", + "--max-inner", + "1", + "--min-sample-ms", + "0.001", + "--shrink", + "16", +] + + +_ZERO_WRITER = textwrap.dedent( + """ + import json, sys + sys.path.insert(0, sys.argv[1]) + import bench_dit_shapes as b + raw, json_out, md_out = sys.argv[2:5] + g = {"family": "zimage", "op": "qkv", "K": 256, "N": 704, "L": 64, "B": 1, + "M": 64, "L_model": 4096, "vector": False, "heavy": False, + "codecs": ["q4_k"]} + args = {"rounds": 2, "slots": 2, "inner": 1, "max_inner": 1, + "min_sample_ms": 0.001, "shrink": 16, "arms": ["kq", "kq_aa"]} + machine = {"device": "gpu", "architecture": "applegpu_g17s", "arch_gen": 17, + "arch_letter": "s", "gpu_cores": 40, "nax": True, + "max_threadgroup_memory": 32768, "os": "26", "mlx": "0.32.2", + "mlx_kquant": "0.5.0", "metallib_sha256": None, "tf32": True} + log = b.RawLog(raw) + log.write({"kind": "meta", "machine": machine, "args": args, + "heavy_codecs": [], "groups": [g], "mlx_kquant_file": "kq", + "sample_mode": "chain"}) + log.write({"kind": "warmup", "dtype": "bfloat16", "inner": [1], + "kernels": [[0, "kq/q4_k", "k"], [0, "kq_aa/q4_k", "k"]], + "errors": []}) + # At tiny shapes a sample (time less glue time) sits near 0. One kq + # sample is exactly 0 and one is below it. + kq_ms = {(0, 0): 0.01, (0, 1): -0.004, (1, 0): 0.0, (1, 1): 0.01} + visit, t = 0, 0.0 + for r in range(2): + for s in range(2): + rows = [] + for key, ms in (("kq/q4_k", kq_ms[(r, s)]), ("kq_aa/q4_k", 0.012)): + t += 1.0 + rows.append([0, key, visit, r, s, t, ms]) + log.write({"kind": "slot", "dtype": "bfloat16", "round": r, + "slot": s, "samples": rows}) + visit += 1 + log.write({"kind": "end", "wall_s": 1.0}) + log.close() + sys.exit(b.main(["--reaggregate", raw, "--json-out", json_out, + "--md-out", md_out])) + """ +) + + +def test_report_builds_with_a_zero_sample(tmp_path): + """A reference sample at or below 0 ms (a time at or under its glue + time) gives its visit no ratio. The report counts the visits left out + and marks the ratio.""" + json_out = tmp_path / "run.json" + script = tmp_path / "writer.py" + script.write_text(_ZERO_WRITER) + out = subprocess.run( + [ + sys.executable, + str(script), + BENCH_DIR, + str(tmp_path / "run.json.raw.jsonl"), + str(json_out), + str(tmp_path / "run.md"), + ], + timeout=120, + capture_output=True, + text=True, + ) + assert out.returncode == 0, out.stderr[-3000:] + result = json.loads(json_out.read_text()) + assert result["complete"] + (p,) = result["passes"] + (g,) = p["groups"] + (c,) = g["cells"] + rt = c["ratios"]["aa/kq"] + # Four visits, two ratios: the kq samples of 0 and of -0.004 ms have + # none. + assert sorted(round(v, 9) for _r, v, _f in rt["pairs_raw"]) == [1.2, 1.2] + assert rt["dropped"] == 2 + assert rt["ref_floor"] is True + assert len(c["arms"]["kq"]["samples"]) == 4 + md = (tmp_path / "run.md").read_text() + assert "zimage.qkv 64 x 1 q4_k aa/kq: 2 of 4" in md + assert "1.200* |" in md + + +def _bench(tmp_path, *extra): + json_out = tmp_path / "run.json" + out = subprocess.run( + [ + sys.executable, + os.path.join(BENCH_DIR, "bench_dit_shapes.py"), + *_TINY, + "--json-out", + str(json_out), + "--md-out", + str(tmp_path / "run.md"), + *extra, + ], + timeout=300, + capture_output=True, + text=True, + ) + if out.returncode != 0: + # pytest prints this with a failure. The exit status alone names no + # cause. + sys.stderr.write(out.stderr[-3000:]) + raise subprocess.CalledProcessError( + out.returncode, out.args, out.stdout, out.stderr + ) + raw = tmp_path / "run.json.raw.jsonl" + recs = [json.loads(line) for line in raw.read_text().splitlines()] + return json.loads(json_out.read_text()), raw, recs, out.stdout + + +def test_cache_limit_and_memory_per_visit(tmp_path): + limit = 1 << 20 + result, _, recs, out = _bench( + tmp_path, + "--cache-limit-gib", + str(limit / 2**30), + "--mem-trace", + "--visit-warmup-calls", + "1,full", + ) + assert result["complete"] + assert result["sample_mode"] == recs[0]["sample_mode"] == "chain" + assert "dependent call chains" in (tmp_path / "run.md").read_text() + assert recs[0]["args"]["cache_limit_gib"] == limit / 2**30 + slots = [r for r in recs if r["kind"] == "slot"] + (warm,) = [r for r in recs if r["kind"] == "warmup"] + assert len(slots) == 4 + assert [r["visit_warmup_calls"] for r in slots] == ["1", "full"] * 2 + for rec in slots: + # One glue sample per visit, beside each arm's. + keys = [row[1] for row in rec["samples"]] + assert keys.count("glue") == 2 + assert len(keys) == 6 + (p,) = result["passes"] + assert all("glue" in g["shared"] for g in p["groups"]) + for rec in [warm, *slots]: + assert sorted(m[0] for m in rec["mem"]) == [0, 1] + for _, active, cache, fp in rec["mem"]: + assert active > 0 + assert cache <= limit + assert (fp is None) == (sys.platform != "darwin") + # The trace prints every visit of the warmup and the four slots. + traced = [line.split()[1] for line in out.splitlines() if line.startswith("mem ")] + assert sorted(traced) == ["0"] * 5 + ["1"] * 5 + + +def test_resume_runs_only_the_slots_left(tmp_path): + _, raw, recs, _ = _bench(tmp_path) + lines = raw.read_text().splitlines() + # Meta, warmup and two slots, then a record cut short. + assert [r["kind"] for r in recs[:4]] == ["meta", "warmup", "slot", "slot"] + raw.write_text("\n".join(lines[:4]) + "\n" + lines[4][:20]) + result, _, after, _ = _bench(tmp_path, "--resume") + assert after[:4] == recs[:4] + assert [r["kind"] for r in after] == ["meta", "warmup"] + ["slot"] * 4 + ["end"] + done = [(r["round"], r["slot"]) for r in after if r["kind"] == "slot"] + assert sorted(done) == [(0, 0), (0, 1), (1, 0), (1, 1)] + visits = [row[2] for r in after if r["kind"] == "slot" for row in r["samples"]] + assert len(set(visits)) == 8 + assert result["complete"] + (p,) = result["passes"] + assert p["slots_done"] == p["slots_planned"] == 4 + # A complete log only rebuilds the reports. + before = raw.read_text() + _bench(tmp_path, "--resume") + assert raw.read_text() == before + + +_CHAIN = textwrap.dedent( + """ + import json, sys + sys.path.insert(0, sys.argv[1]) + import mlx.core as mx + import dit_common as dc + + seen = [] + + def f(a): + seen.append(a) + # The second call's output trips the tie, so the third call is fed + # zeros: only a chained call sees it. + return a + (1e31 if len(seen) == 2 else 0.0) + + x0 = mx.ones((2, 3), dtype=mx.float32) + y = dc.chain(f, x0, 4) + mx.eval(y) + b = dc.tie(mx.ones((4,), dtype=mx.bfloat16), x0.astype(mx.bfloat16)) + ms = dc.chain_ms(lambda a: a * 2, x0, 3) + print(json.dumps({ + "first": [float(a[0, 0]) for a in seen], + "out": float(y[0, 0]), + "tie_dtype": str(b.dtype), + "ms": ms, + })) + """ +) + + +def test_sample_is_one_dependent_chain(tmp_path): + script = tmp_path / "chain.py" + script.write_text(_CHAIN) + # A child process, since dit_common sets KQ_TUNE_LIVE on import. + out = subprocess.run( + [sys.executable, str(script), BENCH_DIR], + check=True, + timeout=120, + capture_output=True, + text=True, + ) + r = json.loads(out.stdout.splitlines()[-1]) + assert r["first"] == [1.0, 1.0, 0.0, 1.0] + assert r["out"] == 1.0 + assert r["tie_dtype"] == "mlx.core.bfloat16" + assert r["ms"] > 0 + + +def test_lever_arm_pairs_with_kq(tmp_path): + result, _, recs, _ = _bench( + tmp_path, "--env-arm", "sw1:KQ_NAX_SWIZZLE=1,KQ_QMM_MIDM=0" + ) + assert recs[0]["args"]["arms"] == ["kq", "kq_aa", "sw1"] + (p,) = result["passes"] + for g in p["groups"]: + (c,) = g["cells"] + assert set(c["arms"]) == {"kq", "kq_aa", "sw1"} + # Every visit pairs the lever arm with kq. A visit whose kq sample + # is at or below 0 has no ratio and is counted as dropped, and at a + # shape this small any visit can be one. + rt = c["ratios"]["sw1/kq"] + visits = len(c["arms"]["kq"]["samples"]) + assert visits == len(c["arms"]["sw1"]["samples"]) > 0 + assert len(rt["pairs_raw"]) + rt["dropped"] == visits + assert "sw1/kq" in (tmp_path / "run.md").read_text() + + +@pytest.mark.parametrize("spec", ["sw1", "sw1:", ":KQ_NAX_SWIZZLE=1", "kq:A=1"]) +def test_lever_arm_malformed(tmp_path, spec): + with pytest.raises(subprocess.CalledProcessError): + _bench(tmp_path, "--env-arm", spec) + + +def test_time_budget_splits_a_run(tmp_path): + # A budget that every slot passes: each process runs one slot and stops, + # and --resume continues until the run is complete. + runs = 0 + while True: + result, _, recs, out = _bench(tmp_path, "--time-budget-min", "1e-9", "--resume") + runs += 1 + if result["complete"]: + break + assert "--resume continues" in out + assert runs < 6 + assert runs == 4 + assert [r["kind"] for r in recs].count("slot") == 4 + assert recs[-1]["kind"] == "end" + (p,) = result["passes"] + assert p["slots_done"] == p["slots_planned"] == 4 + + +# In a fresh process on the bench's device, the MLX peak of nn.silu on an +# input made in the same eval, in outputs, taken as the bench takes its +# peaks: one where silu runs compiled and writes into its input, two where +# it runs as separate ops (MLX_DISABLE_COMPILE). +_SILU_PEAK = textwrap.dedent( + """ + import sys + sys.path.insert(0, sys.argv[1]) + import mlx.core as mx + import mlx.nn as nn + from bench_glu import settle + m, n = int(sys.argv[2]), int(sys.argv[3]) + a = (mx.random.normal((m, n), key=mx.random.key(3)) * 0.5).astype(mx.bfloat16) + mx.eval(a) + for _ in range(2): + base = settle(mx) + mx.reset_peak_memory() + y = nn.silu(a * 2) + mx.eval(y) + peak = (mx.get_peak_memory() - base) / (m * n * 2) + del y + print(peak) + """ +) + + +def test_bench_glu_runs_a_tiny_cell(tmp_path): + """bench_glu.py end to end on one tiny cell: the arms, the report and + the peak of one call per arm. An arm writes at least its output. + two_mm holds two outputs, and the composed arms the peak of silu on + one output plus the other output.""" + out = subprocess.run( + [ + sys.executable, + os.path.join(BENCH_DIR, "bench_glu.py"), + "--codec", + "q4_k", + "--cells", + "bfloat16:256:512:1100", + "--rounds", + "2", + "--calls", + "1", + "--out", + str(tmp_path / "glu"), + ], + timeout=300, + capture_output=True, + text=True, + ) + assert out.returncode == 0, out.stderr[-3000:] + res = json.loads((tmp_path / "glu.json").read_text()) + (cell,) = res["cells"] + assert cell["arms"] == ["composed", "composed_aa", "two_mm", "glu"] + assert len(cell["samples"]) == 2 * 4 + probe = subprocess.run( + [sys.executable, "-c", _SILU_PEAK, BENCH_DIR, "1100", "512"], + check=True, + timeout=120, + capture_output=True, + text=True, + ) + silu = float(probe.stdout.split()[-1]) + pk = cell["peak_outputs"] + assert 1.9 <= pk["two_mm"] <= 2.2, pk + for arm in ("composed", "composed_aa"): + assert silu + 0.9 <= pk[arm] <= silu + 1.2, (pk, silu) + # One output where the fused kernel ran, the composed peak where the op + # composed. + assert 0.95 <= pk["glu"] <= silu + 1.2, (pk, silu) + assert cell["check"]["glu"] < 2e-2 + assert "| glu |" in (tmp_path / "glu.md").read_text() + + +def test_bench_glu_tiles_need_a_probe_build(tmp_path): + """A release build reads no tile lever, so --tiles is refused there + before anything runs.""" + from mlx_kquant._ext import _probe_kernels + + if _probe_kernels(): + pytest.skip("a probe build runs --tiles") + out = subprocess.run( + [ + sys.executable, + os.path.join(BENCH_DIR, "bench_glu.py"), + "--tiles", + "64x64", + "--out", + str(tmp_path / "glu"), + ], + timeout=120, + capture_output=True, + text=True, + ) + assert out.returncode == 2, out.stderr[-2000:] + assert "KQ_PROBE_KERNELS=ON" in out.stderr + assert not (tmp_path / "glu.json").exists() diff --git a/tests/test_env_levers.py b/tests/test_env_levers.py new file mode 100644 index 00000000..8dee9aa3 --- /dev/null +++ b/tests/test_env_levers.py @@ -0,0 +1,88 @@ +"""Environment lever values. + +An empty value reads as unset. 0, false, off and no read as 0, and 1, true, +on and yes as 1, in any case. Any other value reads as unset, with one +warning on stderr the first time the lever holds it. So on a lever whose +unset default differs from 0, an empty or mistyped value keeps the default. +The kernel names of KQ_QMM_LOG show which value each call read. +""" + +from __future__ import annotations + +import pytest +from qmm_probe import force_cpu, run_probe + +import mlx_kquant as kq + +pytestmark = pytest.mark.skipif(force_cpu(), reason="GPU kernel names") + +K, N = 12288, 256 +LOG = { + "KQ_QMM_LOG": "1", + "KQ_TUNE_LIVE": "1", + "KQ_QMM_ROUTE_STRICT": "1", +} + +# KQ_SPLITK_F32 on float16, where unset means float partials: (value, _pf32 +# kernel, warning). Its set values are 0 and 1, so 2 warns too. The second +# "x" call finds the warning already given. +F32_VALUES = [ + ("", True, False), + ("true", True, False), + ("On", True, False), + ("yes", True, False), + ("1", True, False), + ("false", False, False), + ("off", False, False), + ("NO", False, False), + ("0", False, False), + ("2", True, True), + ("x", True, True), + ("x", True, False), +] + + +def _warnings(r, lever, value): + return [line for line in r["stderr"] if f"{lever}={value} is not" in line] + + +def test_splitk_f32_values(): + call = {"codec": "q4_k", "shape": [8, K], "n": N, "dtype": "float16"} + calls = [dict(call, env={"KQ_SPLITK_F32": v}) for v, _, _ in F32_VALUES] + res = run_probe(calls, env={**LOG, "KQ_QMM_ROUTE": "splitk"}) + for (v, pf32, warns), r in zip(F32_VALUES, res, strict=True): + assert r["error"] is None, (v, r["error"]) + ((name, _),) = r["kernels"] + assert "_qmm_t_splitk_" in name and name.endswith("_pf32") == pf32, (v, name) + assert len(_warnings(r, "KQ_SPLITK_F32", v)) == warns, (v, r["stderr"]) + + +# KQ_QMM_SPLITK_BM8 is on unless 0 and read once per process, so each value +# runs in its own probe. +@pytest.mark.parametrize( + "value,bm8", [(None, True), ("", True), ("off", False), ("0", False)] +) +def test_splitk_bm8_values(value, bm8): + call = {"codec": "q4_k", "shape": [8, K], "n": N, "dtype": "bfloat16"} + env = {**LOG, "KQ_QMM_ROUTE": "splitk"} + if value is not None: + env["KQ_QMM_SPLITK_BM8"] = value + (r,) = run_probe([call], env=env) + assert r["error"] is None, r["error"] + ((name, _),) = r["kernels"] + assert ("_qmm_t_splitk_bm8_" in name) == bm8, (value, name) + + +# KQ_VERIFY_NAX_Q4_0_SB is on unless 0 and read live: q4_0 at a K that is a +# multiple of 256 takes the eight-block kernel unless it is off. +@pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") +def test_verify_nax_q4_0_sb_values(): + values = [(None, True), ("", True), ("typo", True), ("no", False)] + call = {"codec": "q4_0", "shape": [8, 3840], "n": 1000, "dtype": "bfloat16"} + calls = [dict(call, env={"KQ_VERIFY_NAX_Q4_0_SB": v}) for v, _ in values] + res = run_probe(calls, env={**LOG, "KQ_QMM_ROUTE": "verify_nax"}) + for (v, sb), r in zip(values, res, strict=True): + assert r["error"] is None, (v, r["error"]) + ((name, _),) = r["kernels"] + assert name.startswith("kquant_q4_0_sb_verify_nax_") == sb, (v, name) + assert _warnings(res[2], "KQ_VERIFY_NAX_Q4_0_SB", "typo"), res[2]["stderr"] diff --git a/tests/test_float32_activations.py b/tests/test_float32_activations.py new file mode 100644 index 00000000..9cc17ed9 --- /dev/null +++ b/tests/test_float32_activations.py @@ -0,0 +1,154 @@ +"""float32 activations run the bfloat16 kernels. + +The ops cast float32 activations to bfloat16 before dispatch, so the +metallib has no quantized matmul or mat-vec kernel for float activations. +On every route of every op, a float32 call returns bfloat16, bit-identical +to the same call on the activations cast to bfloat16. +""" + +from __future__ import annotations + +import functools +import re + +import mlx.core as mx +import numpy as np +import pytest +from qmm_probe import force_cpu, run_probe + +import mlx_kquant as kq + +pytestmark = pytest.mark.skipif(force_cpu(), reason="GPU kernels") + +K, N, E = 1024, 256, 4 +NAX_ROUTES = {"nax", "nax_splitk", "verify_nax"} +# quantized_matmul, transpose=True: (codec, M, forced route or None for the +# default route of that shape). +ROUTES = [ + ("q4_k", 1, None), + ("q6_k", 1, None), + ("q4_k", 4, "qmv"), + ("q4_0", 4, "verify_qmv"), + ("q4_k", 4, "mv_ext"), + ("pq2_0", 4, "verify_mma"), + ("q2_0", 4, "verify_mma"), + ("q4_0", 4, "verify_nax"), + ("q4_k", 8, "splitk"), + ("q4_k", 16, "nax"), + ("q4_k", 16, "nax_splitk"), + ("q4_k", 64, None), +] + + +# The float kernels left are the encode and dequantize kernels, which take +# weights. +def test_no_float_activation_kernels(): + lib = kq.metallib_kernels() + floats = [ + n + for n in lib + if n.startswith("kquant_") + and "_float_gs_" in n + and not re.search(r"_(de)?quantize_", n) + ] + assert floats == [], (len(floats), floats[:4]) + assert "kquant_q4_k_dequantize_float_gs_256_b_4" in lib + assert "kquant_q4_k_quantize_float_gs_256_b_4" in lib + + +@functools.cache +def _probe(): + calls = [] + for codec, m, route in ROUTES: + env = {"KQ_QMM_ROUTE": route, "KQ_QMM_ROUTE_STRICT": "1" if route else None} + for dtype in ("float32", "bfloat16"): + calls.append( + { + "codec": codec, + "shape": [m, K], + "n": N, + "dtype": dtype, + "seed": 7, + "env": env, + } + ) + res = run_probe(calls, env={"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"}) + return {r: (res[2 * i], res[2 * i + 1]) for i, r in enumerate(ROUTES)} + + +@pytest.mark.parametrize("codec,m,route", ROUTES) +def test_quantized_matmul_float32(codec, m, route): + if route in NAX_ROUTES and not kq.nax_available(): + pytest.skip("NAX route") + f32, bf16 = _probe()[(codec, m, route)] + assert f32["error"] is None and bf16["error"] is None, (f32, bf16) + assert f32["digest"] == bf16["digest"] + assert f32["kernels"] == bf16["kernels"] + for name, _ in f32["kernels"]: + assert "_bfloat16_t_" in name, name + + +def _weights(codec, rows, cols, experts=None): + shape = (rows, cols) if experts is None else (experts, rows, cols) + wf = mx.random.normal(shape, key=mx.random.key(1)) * 0.1 + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + return w, s + + +def _x(shape): + return mx.random.normal(shape, key=mx.random.key(2)) * 0.5 + + +def _check(f, x): + got = f(x) + want = f(x.astype(mx.bfloat16)) + mx.eval(got, want) + assert x.dtype == mx.float32 + assert got.dtype == mx.bfloat16, got.dtype + assert mx.array_equal(got, want).item() + + +# transpose=False: qvm at M=1, qmm_n above it. +@pytest.mark.parametrize("m", [1, 64]) +def test_quantized_matmul_n_float32(m): + w, s = _weights("q4_k", K, N) + _check(lambda x: kq.quantized_matmul(x, w, s, "q4_k", transpose=False), _x((m, K))) + + +def test_qmv_bias_float32(): + w, s = _weights("q8_0", N, K) + bias = mx.random.normal((N,), key=mx.random.key(3)).astype(mx.float32) + _check(lambda x: kq.quantized_matmul_qmv_bias(x, w, s, bias, "q8_0"), _x((1, K))) + + +# gather_qmm: the mat-vec (one row per expert pick), the tiled gather, the +# sorted per-expert GEMM and the transpose=False gather. +@pytest.mark.parametrize( + "rows,m,transpose,sort", + [ + (8, 1, True, False), + (4, 32, True, False), + (256, 1, True, True), + (8, 1, False, False), + ], +) +def test_gather_qmm_float32(rows, m, transpose, sort): + w, s = _weights("q4_k", N, K, E) if transpose else _weights("q4_k", K, N, E) + idx = np.random.default_rng(4).integers(0, E, rows).astype(np.uint32) + if sort: + idx = np.sort(idx) + rhs = mx.array(idx) + _check( + lambda x: kq.gather_qmm( + x, w, s, "q4_k", rhs_indices=rhs, transpose=transpose, sorted_indices=sort + ), + _x((rows, m, K)), + ) + + +def test_gather_qmm_seg_float32(): + w, s = _weights("q4_k", N, K, E) + idx = np.sort(np.random.default_rng(5).integers(0, E, 200).astype(np.uint32)) + maps = kq.expert_tile_map(mx.array(idx), E) + _check(lambda x: kq.gather_qmm_seg(x, w, s, "q4_k", *maps), _x((len(idx), K))) diff --git a/tests/test_loader.py b/tests/test_loader.py index c0fd0fec..ff1c207e 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -217,6 +217,63 @@ def test_patched_load_model_file_trusted(tmp_path): mlx_lm_patch._trust_remote_code = False +def _resolver_and_calls(): + # A caller-owned class resolver: a llama subclass, so the checkpoint's + # weights load into it unchanged. + llama_model, llama_args = _get_classes(_tiny_llama_cfg()) + owned = type("OwnedModel", (llama_model,), {}) + calls = [] + + def get_model_classes(config): + calls.append(config["model_type"]) + return owned, llama_args + + return owned, get_model_classes, calls + + +def _saved_tiny_llama(tmp_path): + from mlx_kquant.convert import quantize_model, save + + cfg = _tiny_llama_cfg() + qmodel, qconfig = quantize_model(_build(cfg), cfg, preset="q4_k_m") + save(tmp_path, qmodel, qconfig) + return qmodel + + +def test_load_honors_get_model_classes(tmp_path): + # The caller's resolver picks the class; the weights and the KQuant swap + # are the same as with the default resolver. + qmodel = _saved_tiny_llama(tmp_path) + owned, get_model_classes, calls = _resolver_and_calls() + model, _ = load(tmp_path, get_model_classes=get_model_classes) + assert type(model) is owned + assert calls == ["llama"] + x = mx.array([[1, 2, 3, 4, 5]]) + out, ref = model(x), qmodel(x) + mx.eval(out, ref) + assert bool(mx.array_equal(out, ref).item()) + + +def test_patched_load_model_honors_get_model_classes(tmp_path): + # Through the mlx-lm seam, a kquant checkpoint must build the class the + # caller's get_model_classes returns, as stock load_model does for every + # other checkpoint. + import mlx_lm.utils + + from mlx_kquant.mlx_lm_patch import patch_mlx_lm_load + + qmodel = _saved_tiny_llama(tmp_path) + owned, get_model_classes, calls = _resolver_and_calls() + patch_mlx_lm_load() + model, _ = mlx_lm.utils.load_model(tmp_path, get_model_classes=get_model_classes) + assert type(model) is owned + assert calls == ["llama"] + x = mx.array([[1, 2, 3, 4, 5]]) + out, ref = model(x), qmodel(x) + mx.eval(out, ref) + assert bool(mx.array_equal(out, ref).item()) + + def test_resolve_path_typo_is_not_a_repo_id(): # A nonexistent path-shaped input must say "no such file", not bounce off # the HF Hub validator's repo-id naming-rules message. No network needed. diff --git a/tests/test_matmul_empty.py b/tests/test_matmul_empty.py new file mode 100644 index 00000000..6ddc4ef6 --- /dev/null +++ b/tests/test_matmul_empty.py @@ -0,0 +1,124 @@ +"""quantized_matmul and quantized_matmul_glu on empty inputs. + +No rows or no output columns give an empty array of the product's shape, +and an inner dim of 0 gives zeros. The cases run on the default device, so +the suite covers the GPU and, under KQUANT_FORCE_CPU, the CPU stream, at a +few rows (the mat-vec routes) and at 1100 rows (the matmul routes). +""" + +from __future__ import annotations + +import mlx.core as mx +import pytest + +import mlx_kquant as kq + +N, K = 96, 768 +SCALES = mx.zeros((1,), dtype=mx.uint8) + + +def _wire(codec): + w, _ = kq.quantize(mx.random.normal((N, K), key=mx.random.key(1)) * 0.1, codec) + mx.eval(w) + return w + + +def _ops(codec): + return { + "matmul": lambda x, w: kq.quantized_matmul(x, w, SCALES, codec, True), + "glu": lambda x, w: kq.quantized_matmul_glu(x, w, w, codec), + } + + +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) +@pytest.mark.parametrize("op", ["matmul", "glu"]) +@pytest.mark.parametrize("codec", ["q4_k", "q8_0", "q3_k"]) +def test_empty_inputs(codec, op, dtype): + f = _ops(codec)[op] + w = _wire(codec) + # No rows, alone and under a batch dim. + for shape in ([0, K], [2, 0, K], [0, 3, K]): + y = f(mx.zeros(shape, dtype=dtype), w) + mx.eval(y) + assert y.shape == (*shape[:-1], N) and y.dtype == dtype + for m in (1, 4, 1100): + x = (mx.random.normal((m, K), key=mx.random.key(m)) * 0.5).astype(dtype) + # No output columns. + y = f(x, w[:0]) + mx.eval(y) + assert y.shape == (m, 0) and y.dtype == dtype + # An inner dim of 0: every sum is empty. + y = f(mx.zeros((m, 0), dtype=dtype), mx.zeros((N, 0), dtype=mx.uint8)) + mx.eval(y) + assert y.shape == (m, N) and y.dtype == dtype + assert not y.astype(mx.float32).any().item() + # The stream still runs a full call. + y = f(x, w) + mx.eval(y) + assert y.shape == (m, N) and mx.isfinite(y.astype(mx.float32)).all().item() + + +@pytest.mark.parametrize( + "dtype", [mx.float32, mx.bfloat16, mx.float16], ids=["f32", "bf16", "f16"] +) +@pytest.mark.parametrize("m", [4, 1100]) +def test_inner_dim_of_zero_has_the_dtype_of_a_full_call(m, dtype): + """The zeros for an inner dim of 0 have the dtype the op computes in, + which for float32 x is not the dtype of x.""" + full = kq.quantized_matmul( + mx.zeros((m, K), dtype=dtype), _wire("q4_k"), SCALES, "q4_k", True + ) + y = kq.quantized_matmul( + mx.zeros((m, 0), dtype=dtype), + mx.zeros((N, 0), dtype=mx.uint8), + SCALES, + "q4_k", + True, + ) + mx.eval(y) + assert y.shape == (m, N) and y.dtype == full.dtype + assert not y.astype(mx.float32).any().item() + if dtype == mx.float32: + assert y.dtype == mx.bfloat16 + + +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) +def test_inner_dim_of_zero_checks_the_lora_arguments(dtype): + """An inner dim of 0 returns zeros only after the LoRA arguments pass + the checks of a full call.""" + m, r = 4, 8 + x = mx.zeros((m, 0), dtype=dtype) + w = mx.zeros((N, 0), dtype=mx.uint8) + a = mx.zeros((0, r), dtype=dtype) + b = mx.zeros((r, N), dtype=dtype) + rows = mx.ones((m,)) + + def call(**lora): + return kq.quantized_matmul(x, w, SCALES, "q4_k", True, **lora) + + y = call(lora_a=a, lora_b=b, lora_rows=rows) + mx.eval(y) + assert y.shape == (m, N) and y.dtype == dtype + assert not y.astype(mx.float32).any().item() + other = mx.float16 if dtype == mx.bfloat16 else mx.bfloat16 + for bad in ( + {"lora_a": a}, + {"lora_rows": rows}, + {"lora_a": a.astype(other), "lora_b": b}, + {"lora_a": mx.zeros((3, r), dtype=dtype), "lora_b": b}, + {"lora_a": a, "lora_b": mx.zeros((r, N + 1), dtype=dtype)}, + {"lora_a": a, "lora_b": mx.zeros((r + 1, N), dtype=dtype)}, + {"lora_a": a, "lora_b": b, "lora_rows": mx.ones((m + 1,))}, + ): + with pytest.raises(ValueError): + call(**bad) + with pytest.raises(ValueError, match="2-D w"): + kq.quantized_matmul( + x, + mx.zeros((2, N, 0), dtype=mx.uint8), + SCALES, + "q4_k", + True, + lora_a=a, + lora_b=b, + ) diff --git a/tests/test_matmul_glu.py b/tests/test_matmul_glu.py new file mode 100644 index 00000000..e40e6271 --- /dev/null +++ b/tests/test_matmul_glu.py @@ -0,0 +1,1025 @@ +"""quantized_matmul_glu against a float64 reference. + +The reference is act(x @ Wg^T) * (x @ Wu^T) in float64 on the CPU, with the +exact kq.dequantize weights and x as the op receives it. The fused kernel +runs the activation and the product in float32 and rounds the activation +once and the product once, so it should be no less accurate than the +composition, which rounds each matmul to the activation dtype first. + +The op runs the kernel from MIN_ROWS rows, for the codecs of GLU_CODECS +(q8_0 on bfloat16 only) and an inner dim that is a whole number of the +codec's K tiles, and builds the composition otherwise. In-process cases +cover both sides of the row bound, each codec, partial row tiles, both +parities of the row-tile count and counts above 32767. Inputs come +dense, as a lazy transposed view, as a strided view and with one and two +batch dims. The kernel below MIN_ROWS and the probe tiles run in a fresh +interpreter, since the levers are read once per process. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +import time + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +import pytest +from mlx_kquant._ext import _probe_kernels +from qmm_probe import force_cpu + +import mlx_kquant as kq + +R1 = {mx.bfloat16: 6e-3, mx.float16: 1e-3} +DTYPES = [mx.bfloat16, mx.float16] +# Gate of the composition on the CPU stream, as in test_cpu_neon. +CPU_REL = 5e-2 +K = 768 +# Rows from which the op runs the fused kernel (kGluMinRows). +MIN_ROWS = 1024 +# Codecs with a fused kernel, and the K tile each runs. +GLU_CODECS = {"q4_k": 128, "q5_k": 128, "q6_k": 128, "q8_0": 64} +OTHER_CODECS = ["q5_k", "q6_k", "q8_0"] +# Codecs whose kernel runs on bfloat16 only. On float16 the op builds the +# composition, which measured faster there. +BF16_ONLY = {"q8_0"} +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) + +_W = {} + + +def _nax(): + return not force_cpu() and kq.nax_available() + + +def _fused(rows=MIN_ROWS, codec="q4_k", k=K, dtype=mx.bfloat16): + tile = GLU_CODECS.get(codec, 0) + if codec in BF16_ONLY and dtype != mx.bfloat16: + tile = 0 + return _nax() and rows >= MIN_ROWS and tile > 0 and k % tile == 0 + + +def _weights(codec, n, k, seed=0): + """Two packed weights [n, k] of one codec and their float64 values.""" + key = (codec, n, k, seed) + if key not in _W: + out = [] + for i in range(2): + wf = ( + mx.random.normal( + (n, k), key=mx.random.key(1000 * seed + 17 * n + k + i) + ) + * 0.1 + ) + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + wd = np.array(kq.dequantize(w, s, codec, dtype=mx.float32)).astype( + np.float64 + ) + out.append((w, wd)) + _W[key] = out + return _W[key] + + +def _x(shape, dtype, seed, scale=0.5): + x = (mx.random.normal(shape, key=mx.random.key(seed)) * scale).astype(dtype) + mx.eval(x) + return x + + +def _f64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + +def _act(g, activation): + if activation == "gelu_tanh": + return 0.5 * g * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (g + 0.044715 * g**3))) + with np.errstate(over="ignore"): + return g / (1.0 + np.exp(-g)) + + +def _ref(x, wg, wu, activation="silu"): + x64 = _f64(x) + return _act(x64 @ wg.T, activation) * (x64 @ wu.T) + + +def _rel(a, b): + return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-30)) + + +def _composed(x, wg, wu, codec, activation="silu"): + """The form a caller writes without the op. Where the op has no fused + kernel it builds this, bit for bit.""" + s = mx.zeros((1,), dtype=mx.uint8) + g = kq.quantized_matmul(x, wg, s, codec, True) + u = kq.quantized_matmul(x, wu, s, codec, True) + act = nn.gelu_approx if activation == "gelu_tanh" else nn.silu + return act(g) * u + + +def _check(x, n, dtype, activation="silu", codec="q4_k", k=K): + (wg, wg64), (wu, wu64) = _weights(codec, n, k) + y = kq.quantized_matmul_glu(x, wg, wu, codec, activation) + mx.eval(y) + assert y.dtype == dtype + assert y.shape == tuple(x.shape[:-1]) + (n,) + ref = _ref(x, wg64, wu64, activation) + got = _f64(y) + assert np.isfinite(got).all() + err = _rel(got, ref) + comp = _composed(x, wg, wu, codec, activation) + mx.eval(comp) + err_c = _rel(_f64(comp), ref) + if not _fused(x.size // k, codec, k, dtype): + # The op is the composition here. On the CPU a few rows run the int8 + # GEMV path, which is good to about 1e-2 per matmul. + assert mx.array_equal(y, comp).item() + assert err <= (R1[dtype] if _nax() else CPU_REL), err + return err, err_c + assert err <= R1[dtype], (err, R1[dtype]) + # No less accurate than the composition, which runs the same NAX tiles. + assert err <= 1.1 * err_c + 1e-7, (err, err_c) + return err, err_c + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize( + "m", [1, 17, 64, 129, 577, 1023, 1024, 1025, 1039, 1040, 1100, 1151, 1152] +) +@pytest.mark.parametrize("n", [64, 96, 200]) +def test_glu_rows_and_columns(m, n, dtype): + # Below MIN_ROWS the op is the composition. From there up: a last row + # tile of 1 row, 15 and 16 rows (a partial and a whole simdgroup + # fragment), 76 rows, 127 rows and none. + _check(_x([m, K], dtype, m * 7 + n), n, dtype) + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("m", [4127, 4128, 4160, 4224]) +def test_glu_image_rows(m, dtype): + # 4127, 4128 and 4160 leave a partial last row tile of 128, 4224 none. + _check(_x([m, K], dtype, m), 128, dtype) + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("m", [32768, 32769, 33792]) +def test_glu_rows_above_32767(m, dtype): + (wg, wg64), (wu, wu64) = _weights("q4_k", 96, 256) + x = _x([m, 256], dtype, m) + y = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y) + ref = _ref(x, wg64, wu64) + got = _f64(y) + assert _rel(got, ref) <= R1[dtype] + # The first and last row tiles, where a 16-bit row span would wrap. + for rows in (slice(0, 256), slice(m - 300, m)): + assert _rel(got[rows], ref[rows]) <= R1[dtype] + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("activation", ["silu", "gelu_tanh"]) +def test_glu_activations(activation, dtype): + _check(_x([1100, K], dtype, 5, scale=2.0), 96, dtype, activation) + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("layout", ["batch1", "batch2", "transposed", "strided"]) +def test_glu_layouts(layout, dtype): + m, n = 1030, 96 + if layout == "batch1": + x = _x([1, m, K], dtype, 3) + elif layout == "batch2": + x = _x([2, m, K], dtype, 4) + elif layout == "transposed": + # A lazy transposed view: the op sees an unevaluated swapaxes. + x = mx.swapaxes(_x([K, m], dtype, 5), 0, 1) + else: + x = _x([m, 2 * K], dtype, 6)[:, ::2] + _check(x, n, dtype) + + +def test_glu_float32_x_runs_in_bfloat16(): + (wg, wg64), (wu, wu64) = _weights("q4_k", 96, K) + x = _x([1030, K], mx.float32, 9) + y = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y) + assert y.dtype == mx.bfloat16 + ref = _ref(x.astype(mx.bfloat16), wg64, wu64) + assert _rel(_f64(y), ref) <= R1[mx.bfloat16] + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +def test_glu_row_slices_of_one_weight(dtype): + """Gate and up as the two row bands of one fused weight.""" + n = 96 + (wg, wg64), (wu, wu64) = _weights("q4_k", n, K) + both = mx.concatenate([wg, wu], axis=0) + mx.eval(both) + x = _x([1100, K], dtype, 11) + y = kq.quantized_matmul_glu(x, both[:n], both[n:], "q4_k") + y2 = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y, y2) + assert _rel(_f64(y), _ref(x, wg64, wu64)) <= R1[dtype] + assert mx.array_equal(y, y2).item() + + +def _settle(): + """Active memory once the buffers of finished evals are released.""" + mx.synchronize() + last, same = -1, 0 + for _ in range(400): + mx.clear_cache() + cur = mx.get_active_memory() + same = same + 1 if cur == last else 0 + last = cur + if same >= 4: + break + time.sleep(0.005) + return last + + +def _peak_outputs(build, out_bytes): + """MLX peak of one evaluated call over the memory live before it, in + output buffers. The call runs once first, so pipelines and constants + are in place.""" + y = build() + mx.eval(y) + del y + base = _settle() + mx.reset_peak_memory() + y = build() + mx.eval(y) + peak = mx.get_peak_memory() - base + del y + return peak / out_bytes + + +def test_glu_row_bands_are_read_in_place(): + """Two row bands of one weight on the fused path: the call allocates + its output and no copy of a band. The two bands are half an output.""" + if not _fused(): + pytest.skip("the fused kernel needs a NAX GPU") + n, k, m = 2048, 1024, 1100 + (wg, _), (wu, _) = _weights("q4_k", n, k) + both = mx.concatenate([wg, wu], axis=0) + mx.eval(both) + x = _x([m, k], mx.bfloat16, 41) + out_bytes = m * n * 2 + assert 2 * wg.nbytes > 0.5 * out_bytes + for lazy in (True, False): + bands = (both[:n], both[n:]) + if not lazy: + mx.eval(*bands) + + def build(bands=bands, lazy=lazy): + g, u = (both[:n], both[n:]) if lazy else bands + return kq.quantized_matmul_glu(x, g, u, "q4_k") + + peak = _peak_outputs(build, out_bytes) + assert 0.99 <= peak <= 1.05, (lazy, peak) + y = kq.quantized_matmul_glu(x, both[:n], both[n:], "q4_k") + ref = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y, ref) + assert mx.array_equal(y, ref).item() + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("activation", ["silu", "gelu_tanh"]) +def test_glu_composition_is_the_stock_form(activation, dtype): + """A codec with no fused kernel: the op has the bits and the peak of + the stock activation on one matmul times the other. That is the peak + of the activation on an input made in the same eval, plus the other + output. Where the activation runs compiled, that is two outputs with + nn.silu, which writes into its input, and three with nn.gelu_approx, + which does not. As separate ops (MLX_DISABLE_COMPILE) it is one more.""" + n, k, m = 2048, 1024, 1100 + (wg, _), (wu, _) = _weights("q3_k", n, k) + x = _x([m, k], dtype, 47) + y = kq.quantized_matmul_glu(x, wg, wu, "q3_k", activation) + comp = _composed(x, wg, wu, "q3_k", activation) + mx.eval(y, comp) + assert mx.array_equal(y, comp).item() + out_bytes = m * n * 2 + op = _peak_outputs( + lambda: kq.quantized_matmul_glu(x, wg, wu, "q3_k", activation), out_bytes + ) + stock = _peak_outputs(lambda: _composed(x, wg, wu, "q3_k", activation), out_bytes) + assert abs(op - stock) <= 0.05, (op, stock) + # On a materialized input nn.silu peaks at one output compiled or not, + # so the input is made in the measured eval. + act = nn.gelu_approx if activation == "gelu_tanh" else nn.silu + a = _x([m, n], dtype, 53) + act_peak = _peak_outputs(lambda: act(a * 2), out_bytes) + assert op <= act_peak + 1.05, (op, act_peak) + + +def _node_kinds(*outputs): + """Kinds of the primitives in the graph of the outputs. A compiled + function is one node, named after the ops it holds.""" + import tempfile + + with tempfile.NamedTemporaryFile("r", suffix=".dot") as f: + mx.export_to_dot(f.name, *outputs) + text = f.read() + names = re.findall(r'label ="([A-Za-z0-9_]+)"', text) + return sorted("Compiled" if n.startswith("Compiled") else n for n in names) + + +@pytest.mark.usefixtures("compile_on") +@pytest.mark.parametrize("activation", ["silu", "gelu_tanh"]) +def test_glu_unfused_ops_follow_the_compile_switch(activation): + """The unfused path builds its activation as mlx.nn does in either + state of mx's compile switch: one compiled function while the switch is + on and separate ops while it is off, whatever the state was at an + earlier call. 64 rows build the unfused ops on every device.""" + (wg, _), (wu, _) = _weights("q4_k", 96, K) + x = _x([64, K], mx.bfloat16, 71) + for on in (True, False, True, False): + (mx.enable_compile if on else mx.disable_compile)() + y = kq.quantized_matmul_glu(x, wg, wu, "q4_k", activation) + ref = _composed(x, wg, wu, "q4_k", activation) + kinds = _node_kinds(y) + assert kinds == _node_kinds(ref), (on, kinds) + assert ("Compiled" in kinds) == on, (on, kinds) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), on + + +@pytest.mark.parametrize("activation", ["silu", "gelu_tanh"]) +def test_glu_unfused_ops_run_on_the_stream_of_the_call(activation): + """An unfused call on a stream that is not the default one has the bits + of the stock form built on that stream. For silu the two devices round + the activation differently at this input, so the bits tell where it ran. + For gelu_tanh they agree here, and that case checks the value alone.""" + if force_cpu() or mx.default_device() != mx.gpu: + pytest.skip("the CPU is the default device") + (wg, _), (wu, _) = _weights("q3_k", 96, K) + x = _x([64, K], mx.bfloat16, 61) + y = kq.quantized_matmul_glu(x, wg, wu, "q3_k", activation, stream=mx.cpu) + with mx.stream(mx.cpu): + ref = _composed(x, wg, wu, "q3_k", activation) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item() + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +def test_glu_strided_weight_views(state, dtype): + """Weights whose rows are not dense: every other row of one weight, and + a column band of a wider one. The fused path copies them at eval.""" + n = 96 + (wg, wg64), (wu, wu64) = _weights("q4_k", n, K) + rows = mx.stack([wg, wu], axis=1).reshape(2 * n, -1) + cols = mx.concatenate([wg, wu], axis=1) + mx.eval(rows, cols) + kb = wg.shape[1] + x = _x([1100, K], dtype, 43) + ref = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + for make in ( + lambda: (rows[0::2], rows[1::2]), + lambda: (cols[:, :kb], cols[:, kb:]), + ): + g, u = make() + if state == "evaluated": + mx.eval(g, u) + y = kq.quantized_matmul_glu(x, g, u, "q4_k") + mx.eval(y, ref) + assert mx.array_equal(y, ref).item() + assert _rel(_f64(ref), _ref(x, wg64, wu64)) <= R1[dtype] + + +def test_glu_float16_store_saturates(): + """Products past the half range store 65504 with the sign kept.""" + n = 64 + (wg, wg64), (wu, wu64) = _weights("q4_k", n, K) + x = _x([1088, K], mx.float16, 13, scale=200.0) + ref = _ref(x, wg64, wu64) + assert (np.abs(ref) > 65504.0).mean() > 0.05 + # The activation alone stays in range, so the clipped product is exact. + assert np.abs(_act(_f64(x) @ wg64.T, "silu")).max() < 65504.0 + y = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y) + got = _f64(y) + if not _fused(): + # The composition overflows to inf where the fused store clamps. + ok = np.isfinite(got) + assert _rel(got[ok], ref[ok]) <= 2e-2 + return + assert np.isfinite(got).all() + clipped = np.clip(ref, -65504.0, 65504.0) + assert _rel(got, clipped) <= R1[mx.float16] + over = np.abs(ref) > 66000.0 + assert (np.abs(got[over]) == 65504.0).all() + assert (np.sign(got[over]) == np.sign(ref[over])).all() + + +def test_glu_bfloat16_result_above_the_float16_range(): + """The float16 store saturates at 65504. bfloat16 holds the products + past it, and the op returns them.""" + n = 64 + (wg, wg64), (wu, wu64) = _weights("q4_k", n, K) + x = _x([1088, K], mx.bfloat16, 13, scale=200.0) + ref = _ref(x, wg64, wu64) + assert (np.abs(ref) > 65504.0).mean() > 0.05 + y = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y) + got = _f64(y) + assert np.isfinite(got).all() + assert _rel(got, ref) <= (R1[mx.bfloat16] if _nax() else CPU_REL) + over = np.abs(ref) > 70000.0 + assert over.any() and (np.abs(got[over]) > 65536.0).all() + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("activation", ["silu", "gelu_tanh"]) +@pytest.mark.parametrize("codec", ["q4_k", "q6_k"]) +def test_glu_nan_and_inf_pass(codec, activation, dtype): + """A NaN or an infinity in x gives NaN and inf where the unfused ops + give them. The float16 store saturates finite values only.""" + n, m = 96, 1100 + (wg, _), (wu, _) = _weights(codec, n, K) + xn = np.array(_x([m, K], dtype, 31).astype(mx.float32)) + # Rows of the first, a middle and the partial last row tile. + xn[3, 5] = np.nan + xn[200, 700] = np.inf + xn[511, 0] = -np.inf + xn[640, 100] = np.inf + xn[640, 101] = -np.inf + xn[1090, 767] = -np.inf + xn[1099, 300] = np.nan + x = mx.array(xn).astype(dtype) + mx.eval(x) + y = kq.quantized_matmul_glu(x, wg, wu, codec, activation) + comp = _composed(x, wg, wu, codec, activation) + mx.eval(y, comp) + got = np.array(y.astype(mx.float32)) + ref = np.array(comp.astype(mx.float32)) + assert np.isnan(ref).any() and np.isinf(ref).any() + assert (np.isnan(got) == np.isnan(ref)).all() + assert (np.isinf(got) == np.isinf(ref)).all() + inf = np.isinf(ref) + assert (got[inf] == ref[inf]).all() + ok = np.isfinite(ref) + assert _rel(got[ok].astype(np.float64), ref[ok].astype(np.float64)) <= 2 * R1[dtype] + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("m", [1023, 1024, 1039, 1100, 1152]) +@pytest.mark.parametrize("n", [64, 200]) +@pytest.mark.parametrize("codec", OTHER_CODECS) +def test_glu_codecs_rows_and_columns(codec, n, m, dtype): + _check(_x([m, K], dtype, m * 7 + n), n, dtype, codec=codec) + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("codec", OTHER_CODECS) +def test_glu_codecs_image_rows(codec, dtype): + _check(_x([4128, K], dtype, 4128), 128, dtype, codec=codec) + + +@pytest.mark.parametrize("codec", OTHER_CODECS) +def test_glu_codecs_gelu(codec): + _check( + _x([1100, K], mx.bfloat16, 5, scale=2.0), 96, mx.bfloat16, "gelu_tanh", codec + ) + + +@pytest.mark.parametrize("codec", OTHER_CODECS) +def test_glu_codecs_rows_above_32767(codec): + m = 32769 + (wg, wg64), (wu, wu64) = _weights(codec, 96, 256) + x = _x([m, 256], mx.bfloat16, m) + y = kq.quantized_matmul_glu(x, wg, wu, codec) + mx.eval(y) + ref = _ref(x, wg64, wu64) + got = _f64(y) + for rows in (slice(0, m), slice(0, 256), slice(m - 300, m)): + assert _rel(got[rows], ref[rows]) <= R1[mx.bfloat16] + + +@pytest.mark.parametrize("k", [832, 800]) +def test_glu_q8_0_inner_dim_rule(k): + """q8_0 has a K tile of 64: 832 is 13 tiles and runs the kernel, 800 is + not a whole number of tiles and builds the composition.""" + assert _fused(1100, "q8_0", k) == (_nax() and k == 832) + _check(_x([1100, k], mx.bfloat16, k), 96, mx.bfloat16, codec="q8_0", k=k) + + +def test_glu_q8_0_float16_composes(): + """q8_0 on float16 is the composition at every row count.""" + assert not _fused(4128, "q8_0", K, mx.float16) + _check(_x([4128, K], mx.float16, 23), 96, mx.float16, codec="q8_0") + + +@pytest.mark.parametrize("codec", ["q3_k", "q4_0", "q2_0", "pq2_0"]) +def test_glu_codec_without_kernel_composes(codec): + (wg, wg64), (wu, wu64) = _weights(codec, 96, K) + x = _x([1100, K], mx.bfloat16, 15) + y = kq.quantized_matmul_glu(x, wg, wu, codec) + mx.eval(y) + assert _rel(_f64(y), _ref(x, wg64, wu64)) <= R1[mx.bfloat16] + comp = _composed(x, wg, wu, codec) + assert mx.array_equal(y, comp).item() + + +@pytest.mark.parametrize("activation", ["silu", "gelu_tanh"]) +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize( + "dtype", [mx.bfloat16, mx.float16, mx.float32], ids=["bf16", "f16", "f32"] +) +def test_glu_gradient_is_the_gradient_of_the_composition(dtype, m, activation): + """mx.grad wrt x on both sides of the row bound. The fused kernel has + the vjp of the unfused ops, bit for bit on float16 and bfloat16 x. + + float32 x is cast to bfloat16 before the fused node, so the node sums + the gate term and the up term of the gradient in bfloat16 and the + unfused ops sum them in float32. The two differ by the rounding of that + one sum. bfloat16 keeps 8 significant bits, so one rounding moves a + value by at most 2^-8 = 3.9e-3 of itself and by about 1.7e-3 in the + root mean square, and the difference measures 1.8e-3. The bound is + 6e-3, the bound on a bfloat16 result against float64.""" + n = 96 + (wg, _), (wu, _) = _weights("q4_k", n, K) + x = _x([m, K], dtype, 53) + c = mx.random.normal((m, n), key=mx.random.key(59)) + mx.eval(c) + + def loss(f): + return lambda t: (f(t).astype(mx.float32) * c).sum() + + g_op = mx.grad( + loss(lambda t: kq.quantized_matmul_glu(t, wg, wu, "q4_k", activation)) + )(x) + g_ref = mx.grad(loss(lambda t: _composed(t, wg, wu, "q4_k", activation)))(x) + mx.eval(g_op, g_ref) + assert g_op.shape == x.shape and g_op.dtype == x.dtype + ref = _f64(g_ref) + assert np.isfinite(ref).all() and np.abs(ref).max() > 0 + if dtype == mx.float32 and _fused(m): + rel = _rel(_f64(g_op), ref) + assert rel <= R1[mx.bfloat16], rel + else: + assert mx.array_equal(g_op, g_ref).item() + + +def test_glu_jvp_and_vmap_do_not_depend_on_the_row_count(): + """quantized_matmul has no jvp and no vmap, so both raise for the + composition. The fused kernel raises the same way.""" + (wg, _), (wu, _) = _weights("q4_k", 96, K) + seen = {"jvp": [], "vmap": []} + for m in (1023, 1100): + x = _x([m, K], mx.bfloat16, 61) + + def f(t): + return kq.quantized_matmul_glu(t, wg, wu, "q4_k") + + with pytest.raises(ValueError) as e: + mx.eval(mx.jvp(f, [x], [x])[1]) + seen["jvp"].append(str(e.value)) + with pytest.raises(ValueError) as e: + mx.eval(mx.vmap(f)(x)) + seen["vmap"].append(type(e.value)) + assert seen["jvp"][0] == seen["jvp"][1], seen["jvp"] + assert "KQuantMatmul" in seen["jvp"][0] + + +def test_glu_compiled_function_with_both_activations(): + """Two calls that differ in the activation alone stay two nodes when a + compiled function holds both.""" + (wg, _), (wu, _) = _weights("q4_k", 96, K) + x = _x([1100, K], mx.bfloat16, 67) + + def both(t): + a = kq.quantized_matmul_glu(t, wg, wu, "q4_k", "silu") + b = kq.quantized_matmul_glu(t, wg, wu, "q4_k", "gelu_tanh") + # mx.compile merges equivalent nodes but never an output of the + # function, so the two calls feed the outputs. The casts are exact. + return a.astype(mx.float32), b.astype(mx.float32) + + a, b = mx.compile(both)(x) + ra, rb = both(x) + mx.eval(a, b, ra, rb) + assert mx.array_equal(a, ra).item() + assert mx.array_equal(b, rb).item() + assert not mx.array_equal(ra, rb).item() + + +def test_glu_compiles_and_repeats(): + (wg, _), (wu, _) = _weights("q4_k", 96, K) + x = _x([1100, K], mx.bfloat16, 17) + f = mx.compile(lambda t: kq.quantized_matmul_glu(t, wg, wu, "q4_k")) + a = f(x) + b = f(x) + c = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(a, b, c) + assert mx.array_equal(a, c).item() + assert mx.array_equal(b, c).item() + + +def test_glu_argument_errors(): + (wg, _), (wu, _) = _weights("q4_k", 96, K) + x = _x([4, K], mx.bfloat16, 19) + with pytest.raises(ValueError, match="activation"): + kq.quantized_matmul_glu(x, wg, wu, "q4_k", "relu") + with pytest.raises(ValueError, match="one shape"): + kq.quantized_matmul_glu(x, wg, wu[:64], "q4_k") + with pytest.raises(ValueError, match="inner dim"): + kq.quantized_matmul_glu(x[:, :256], wg, wu, "q4_k") + with pytest.raises(ValueError, match="Unknown kquant_type"): + kq.quantized_matmul_glu(x, wg, wu, "q4_z") + with pytest.raises(ValueError, match="uint8"): + kq.quantized_matmul_glu(x, wg.astype(mx.float32), wu, "q4_k") + + +def test_glu_kernels_in_metallib(): + if force_cpu(): + pytest.skip("Metal kernel names only") + names = kq.metallib_kernels() + if not names: + pytest.skip("no metallib loaded") + glu = {n for n in names if "_qmm_t_nax_glu_" in n} + # The tile the op runs for each codec. + want = { + f"kquant_{c}_qmm_t_nax_glu_{t}_bm128_bn64_bk{bk}w_wm2_wn2_alN_{al}" + for c, bk in GLU_CODECS.items() + for t in ("float16_t", "bfloat16_t") + for al in ("true", "false") + if t == "bfloat16_t" or c not in BF16_ONLY + } + assert want <= glu, sorted(want - glu) + # A release build carries those and no other GLU kernel. A build with + # KQ_PROBE_KERNELS=ON carries the probe tiles too, and the module that + # reads their levers reports the option. + if _probe_kernels(): + assert glu > want + else: + assert glu == want, sorted(glu - want) + # Every GLU kernel is a tile of a known form of one of those codecs: K + # tile 64 or 128, with d (two unpadded Ws buffers) or w (two K walks). + pat = re.compile( + r"^kquant_(" + "|".join(GLU_CODECS) + r")_qmm_t_nax_glu_" + r"(bfloat16_t|float16_t)_bm(64|128)_bn64" + r"_bk(64|128|128d|64w|128w)_wm2_wn2_alN_(true|false)$" + ) + bad = sorted(n for n in glu if not pat.match(n)) + assert not bad, bad[:5] + + +def _built_tiles(codec="q4_k"): + """(BM, K tile) of every GLU kernel of a codec in the metallib.""" + tiles = set() + head = f"kquant_{codec}_qmm_t_nax_glu_bfloat16_t_" + for n in kq.metallib_kernels(): + if n.startswith(head) and n.endswith("alN_true"): + f = n.split("_") + bm = next(p for p in f if p.startswith("bm"))[2:] + bk = next(p for p in f if p.startswith("bk"))[2:] + tiles.add((int(bm), bk)) + return sorted(tiles) + + +_CHILD = r""" +import hashlib, json, os, sys +import numpy as np +spec = json.loads(sys.argv[1]) +sys.path.insert(0, spec["tests_dir"]) +import mlx.core as mx +import mlx_kquant as kq +import test_matmul_glu as t + +out = [] +for m, n, dt in spec["cells"]: + dtype = getattr(mx, dt) + (wg, wg64), (wu, wu64) = t._weights(spec["codec"], n, t.K) + x = t._x([m, t.K], dtype, m * 7 + n) + row = {"m": m, "n": n, "dtype": dt} + for bm, bk in spec["tiles"]: + os.environ["KQ_GLU_BM"] = str(bm) + os.environ["KQ_GLU_BK"] = bk.rstrip("dw") + os.environ["KQ_GLU_DUAL128"] = "1" if bk.endswith("d") else "0" + os.environ["KQ_GLU_WALKS"] = "1" if bk.endswith("w") else "0" + y = kq.quantized_matmul_glu(x, wg, wu, spec["codec"]) + mx.eval(y) + got = t._f64(y) + row[f"{bm},{bk}"] = { + "rel": t._rel(got, t._ref(x, wg64, wu64)), + "hash": hashlib.sha256(np.array(y.astype(mx.float32))).hexdigest(), + } + out.append(row) +print("RESULT " + json.dumps(out)) +""" + + +@pytest.mark.parametrize("codec", list(GLU_CODECS)) +def test_glu_kernel_below_min_rows_and_tiles_agree(codec): + """The kernel at any row count, and one set of bits per kernel form.""" + if not _nax(): + pytest.skip("the fused kernel needs a NAX GPU") + # Every tile in the metallib. One K walk: K tile 64 on two Ws buffers, + # 128 on one, 128d on two unpadded ones. Two K walks (64w, 128w) round + # the gate activation to the output dtype before the product, so those + # bits form their own group. + tiles = _built_tiles(codec) + assert (128, f"{GLU_CODECS[codec]}w") in tiles + rows = (1, 16, 17, 63, 65, 127, 129, 300, 449, 577, 4128) + cells = [ + (m, n, dt) + for m in (rows if codec == "q4_k" else rows[::3]) + for n in (64, 96, 200) + for dt in ("bfloat16", "float16") + ] + spec = {"tests_dir": TESTS_DIR, "tiles": tiles, "cells": cells, "codec": codec} + env = dict(os.environ, KQ_TUNE_LIVE="1", KQ_QMM_LOG="1", KQ_GLU_MIN_ROWS="1") + p = subprocess.run( + [sys.executable, "-c", _CHILD, json.dumps(spec)], + env=env, + capture_output=True, + text=True, + timeout=900, + ) + assert p.returncode == 0, p.stderr[-2000:] + line = [ln for ln in p.stdout.splitlines() if ln.startswith("RESULT ")][-1] + rows = json.loads(line[len("RESULT ") :]) + logged = { + ln.split()[2] + for ln in p.stderr.splitlines() + if ln.startswith("[kq qmm] ") and "_glu_" in ln + } + for bm, bk in tiles: + assert any(f"_bm{bm}_bn64_bk{bk}_" in n for n in logged), ( + bm, + bk, + sorted(logged), + ) + for row in rows: + tol = R1[getattr(mx, row["dtype"])] + hashes = {False: set(), True: set()} + for bm, bk in tiles: + r = row[f"{bm},{bk}"] + assert r["rel"] <= tol, (row["m"], row["n"], bm, bk, r["rel"]) + hashes[bk.endswith("w")].add(r["hash"]) + for walks, group in hashes.items(): + # A release build has no one-walk tile, so that group is empty. + assert len(group) <= 1, (row["m"], row["n"], row["dtype"], walks) + + +_PROBE = r""" +import json, os, sys +spec = json.loads(sys.argv[1]) +sys.path.insert(0, spec["tests_dir"]) +import mlx.core as mx +import mlx_kquant as kq +import test_matmul_glu as t + +out = [] +for i, c in enumerate(spec["calls"]): + dtype = getattr(mx, c["dtype"]) + act = c.get("activation", "silu") + k = c.get("k", t.K) + (wg, wg64), (wu, wu64) = t._weights(c["codec"], c["n"], k) + x = t._x([c["m"], k], dtype, 7 * c["m"] + c["n"]) + os.environ.update(c.get("env", {})) + print(f"[probe] {i}", file=sys.stderr, flush=True) + res = {"error": None} + try: + y = kq.quantized_matmul_glu(x, wg, wu, c["codec"], act) + mx.eval(y) + res["rel"] = t._rel(t._f64(y), t._ref(x, wg64, wu64, act)) + except Exception as e: + res["error"] = f"{type(e).__name__}: {e}" + out.append(res) +print("RESULT " + json.dumps(out)) +""" + + +def _glu_probe(calls, env=None): + """Run GLU calls in a fresh interpreter that keeps no KQ_ variable but + those of env and of each call. Per call: glu (the GLU kernels that + KQ_QMM_LOG named), plain (the other kernels), stderr (its other + lines), error and rel.""" + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update({"KQ_QMM_LOG": "1", **(env or {})}) + spec = {"tests_dir": TESTS_DIR, "calls": calls} + p = subprocess.run( + [sys.executable, "-c", _PROBE, json.dumps(spec)], + env=child_env, + capture_output=True, + text=True, + timeout=900, + ) + assert p.returncode == 0, p.stderr[-3000:] + line = [ln for ln in p.stdout.splitlines() if ln.startswith("RESULT ")][-1] + results = json.loads(line[len("RESULT ") :]) + cur = None + for ln in p.stderr.splitlines(): + if ln.startswith("[probe] "): + cur = results[int(ln.split()[1])] + cur.update(glu=[], plain=[], stderr=[]) + elif cur is None: + continue + elif ln.startswith("[kq qmm] "): + name = ln.split()[2] + cur["glu" if "_glu_" in name else "plain"].append(name) + else: + cur["stderr"].append(ln) + return results + + +def test_glu_min_rows_lever_starts_at_one(): + """KQ_GLU_MIN_ROWS moves the row bound. A value below 1 keeps the + bound of 1024 rows, with one warning per value.""" + if not _nax(): + pytest.skip("the fused kernel needs a NAX GPU") + call = {"codec": "q4_k", "m": 600, "n": 96, "dtype": "bfloat16"} + values = [ + ("600", True, False), + ("601", False, False), + ("1", True, False), + ("0", False, True), + ("0", False, False), + ("-4", False, True), + ("false", False, True), + ] + calls = [dict(call, env={"KQ_GLU_MIN_ROWS": v}) for v, _, _ in values] + res = _glu_probe(calls, env={"KQ_TUNE_LIVE": "1"}) + for (v, fused, warns), r in zip(values, res, strict=True): + assert r["error"] is None, (v, r["error"]) + assert r["rel"] <= R1[mx.bfloat16], (v, r["rel"]) + assert len(r["glu"]) == (1 if fused else 0), (v, r["glu"]) + assert len(r["plain"]) == (0 if fused else 2), (v, r["plain"]) + warned = [ln for ln in r["stderr"] if f"KQ_GLU_MIN_ROWS={v} is not" in ln] + assert len(warned) == warns, (v, r["stderr"]) + if warns: + assert "at least 1" in warned[0] + + +def test_glu_tile_levers_belong_to_a_probe_build(): + """A release build reads no tile lever. A probe build raises for a tile + the codec has no kernel for, with the levers and the kernel name.""" + if not _nax(): + pytest.skip("the fused kernel needs a NAX GPU") + call = {"codec": "q5_k", "m": 1100, "n": 96, "dtype": "bfloat16"} + env = {"KQ_GLU_BM": "64", "KQ_GLU_WALKS": "0"} + if not _probe_kernels(): + # A release build reads none of the four. + env.update({"KQ_GLU_BK": "64", "KQ_GLU_DUAL128": "1"}) + (r,) = _glu_probe([call], env=env) + head = "kquant_q5_k_qmm_t_nax_glu_bfloat16_t_" + if _probe_kernels(): + # q5_k has no one-walk kernel in any build. + assert r["error"] is not None and r["error"].startswith("ValueError") + assert "KQ_GLU_BM=64" in r["error"] and "KQ_GLU_WALKS=0" in r["error"] + assert head + "bm64_bn64_bk128_wm2_wn2_alN_false" in r["error"] + assert "not instantiated" in r["error"] + else: + assert r["error"] is None, r["error"] + assert r["glu"] == [head + "bm128_bn64_bk128w_wm2_wn2_alN_false"] + assert r["rel"] <= R1[mx.bfloat16] + + +def _glu_name(codec, dtype, n): + """The kernel the op runs for a codec, an activation dtype and N.""" + aligned = "true" if n % 64 == 0 else "false" + return ( + f"kquant_{codec}_qmm_t_nax_glu_{dtype}_t_bm128_bn64" + f"_bk{GLU_CODECS[codec]}w_wm2_wn2_alN_{aligned}" + ) + + +@pytest.mark.parametrize("codec", list(GLU_CODECS)) +def test_glu_kernel_that_runs(codec): + """KQ_QMM_LOG names the kernel of each call. From 1024 rows a codec + with a kernel for the dtype runs it, once, and no matmul tile: for both + dtypes, N on and off the column tile and both activations. At 1023 + rows, and for q8_0 on float16, the two matmuls run and no GLU kernel.""" + if not _nax(): + pytest.skip("the fused kernel needs a NAX GPU") + calls = [ + {"codec": codec, "m": m, "n": n, "dtype": dt, "activation": act} + for dt in ("bfloat16", "float16") + for m, n, act in ( + (1100, 64, "silu"), + (1100, 96, "silu"), + (1100, 64, "gelu_tanh"), + (1100, 96, "gelu_tanh"), + (1024, 96, "silu"), + (1023, 96, "silu"), + ) + ] + res = _glu_probe(calls) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] <= R1[getattr(mx, c["dtype"])], (c, r["rel"]) + has_kernel = c["dtype"] == "bfloat16" or codec not in BF16_ONLY + if c["m"] >= MIN_ROWS and has_kernel: + assert r["glu"] == [_glu_name(codec, c["dtype"], c["n"])], (c, r["glu"]) + assert r["plain"] == [], (c, r["plain"]) + else: + assert r["glu"] == [], (c, r["glu"]) + assert len(r["plain"]) == 2, (c, r["plain"]) + + +def test_glu_lever_off_runs_the_unfused_ops(): + """KQ_GLU=0: no GLU kernel at any row count, for any codec.""" + if not _nax(): + pytest.skip("the fused kernel needs a NAX GPU") + calls = [ + {"codec": codec, "m": 1100, "n": 96, "dtype": "bfloat16"} + for codec in GLU_CODECS + ] + off = _glu_probe(calls, env={"KQ_GLU": "0"}) + on = _glu_probe(calls, env={"KQ_GLU": "1"}) + for c, r0, r1 in zip(calls, off, on, strict=True): + assert r0["error"] is None and r1["error"] is None + assert r0["glu"] == [] and len(r0["plain"]) == 2, (c, r0) + assert r1["glu"] == [_glu_name(c["codec"], "bfloat16", 96)], (c, r1) + assert r0["rel"] <= R1[mx.bfloat16] + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("codec", ["q4_k", "q6_k"]) +def test_glu_gelu_below_saturation(codec, dtype): + """gelu_tanh on gate sums of about 1.4, where the tanh is far from 1 + and the cubic term of its argument shows in the result.""" + _check(_x([1100, K], dtype, 71, scale=0.5), 96, dtype, "gelu_tanh", codec) + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("n", [1, 17, 63]) +def test_glu_columns_below_one_tile(n, dtype): + """N below the column tile of 64: one partial tile.""" + _check(_x([1100, K], dtype, 73 + n), n, dtype) + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +def test_glu_columns_above_32767(dtype): + """N past 32767 with a partial last row tile and N off the column + tile, where a 16-bit column span would wrap. The reference covers the + first and the last columns.""" + n, k, m = 32833, 256, 1100 + (wg, wg64), (wu, wu64) = _weights("q4_k", n, k) + x = _x([m, k], dtype, 77) + y = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y) + assert y.shape == (m, n) + x64 = _f64(x) + for cols in (slice(0, 256), slice(n - 300, n)): + ref = _act(x64 @ wg64[cols].T, "silu") * (x64 @ wu64[cols].T) + got = _f64(y[:, cols]) + # All rows, and the partial last row tile alone. + for rows in (slice(0, m), slice(1024, m)): + assert _rel(got[rows], ref[rows]) <= R1[dtype], (cols, rows) + del y + + +def test_glu_float16_activation_saturates_beside_a_small_up(): + """A float16 gate activation past the half range is stored as 65504, so + the product with a small up sum is finite: 65504 times the up sum. The + gate weight reads the first two K blocks and the up weight the third, + so x sets the size of each sum.""" + m, n = 1100, 96 + if not _fused(m, dtype=mx.float16): + pytest.skip("the unfused ops return inf here") + kg = 512 + wf = mx.random.normal((2, n, K), key=mx.random.key(83)) * 0.1 + cols = mx.arange(K) + wg, _ = kq.quantize(mx.where(cols < kg, wf[0], 0.0), "q4_k") + wu, _ = kq.quantize(mx.where(cols < kg, 0.0, wf[1]), "q4_k") + s = mx.zeros((1,), dtype=mx.uint8) + wg64, wu64 = (_f64(kq.dequantize(w, s, "q4_k", dtype=mx.float32)) for w in (wg, wu)) + big = mx.random.uniform(-3e4, 3e4, (m, K), key=mx.random.key(85)) + small = mx.random.normal((m, K), key=mx.random.key(87)) * 0.1 + x = mx.where(cols < kg, big, small).astype(mx.float16) + mx.eval(x, wg, wu) + g = _f64(x) @ wg64.T + u = _f64(x) @ wu64.T + over = (g > 66000.0) & (np.abs(u) < 0.9) & (np.abs(u) > 1e-3) + assert over.mean() > 0.02 + y = kq.quantized_matmul_glu(x, wg, wu, "q4_k") + mx.eval(y) + got = _f64(y) + assert np.isfinite(got[over]).all() + assert _rel(got[over], 65504.0 * u[over]) <= 2 * R1[mx.float16] + + +def test_glu_zero_size_inputs(): + """No rows, no output columns and an inner dim of 0 give an array of + the composed shape.""" + (wg, _), (wu, _) = _weights("q4_k", 96, K) + for shape in ([0, K], [2, 0, K], [0, 3, K]): + y = kq.quantized_matmul_glu(mx.zeros(shape, dtype=mx.bfloat16), wg, wu, "q4_k") + mx.eval(y) + assert y.shape == (*shape[:-1], 96) and y.dtype == mx.bfloat16 + x = _x([1100, K], mx.bfloat16, 89) + y = kq.quantized_matmul_glu(x, wg[:0], wu[:0], "q4_k") + mx.eval(y) + assert y.shape == (1100, 0) and y.dtype == mx.bfloat16 + w0 = mx.zeros((96, 0), dtype=mx.uint8) + y = kq.quantized_matmul_glu(mx.zeros((1100, 0), dtype=mx.float16), w0, w0, "q4_k") + mx.eval(y) + assert y.shape == (1100, 96) and y.dtype == mx.float16 + assert not y.astype(mx.float32).any().item() diff --git a/tests/test_matmul_lora.py b/tests/test_matmul_lora.py new file mode 100644 index 00000000..2619e299 --- /dev/null +++ b/tests/test_matmul_lora.py @@ -0,0 +1,1441 @@ +"""quantized_matmul_lora: y = x @ W^T + (strength * (x @ A)) @ B. + +On a NAX GPU the op runs one kernel for q4_k, q5_k, q6_k and q8_0 weights +on bfloat16 from 1024 rows: stock ops compute z = strength * (x @ A), and +the kernel walks the rank for z and B and then K for the weight in the same +f32 sums. Everywhere else it builds the sum from quantized_matmul and stock +ops. By itself the op gives the kernel only the calls of its shape rule +(K up to 4096, N of at least K and, below 4096 rows, N of at least 8192). +The weights here are narrow, so the tests of this module run under +KQ_QMM_LORA=2, which runs the kernel wherever it can run, and the rule is +tested in fresh interpreters. + +- R1: relative Frobenius error against the float64 sum, at most 6e-3 on + bfloat16 and 1e-3 on float16. +- R2: against the model of the kernel, with the weight and z rounded to + the activation dtype, exact sums and one rounding of y. The kernel sits + at the error of that one rounding. +- Where the op has no kernel, y is the composition a caller writes, bit + for bit, and every column of it is held against the float64 sum (COL): + the composition is stock ops, and a stock op that misreads one column + moves the whole-output error of R1 by too little to see. + +Under KQUANT_FORCE_CPU the op runs the composition on the CPU stream, and +the tests that name a Metal kernel skip. +""" + +from __future__ import annotations + +import itertools +import json +import os +import subprocess +import sys + +import mlx.core as mx +import numpy as np +import pytest +from qmm_probe import force_cpu +from test_operand_views import _dense, _nodes, _strided +from test_weight_base import BASE_ALIGN, _shift + +import mlx_kquant as kq + +R1 = {mx.bfloat16: 6e-3, mx.float16: 1e-3} +DTYPES = [mx.bfloat16, mx.float16] +# Gate of the composition on the CPU stream, as in test_cpu_neon. +CPU_REL = 5e-2 +# Gate of one column of the composition: the error of the column against +# the size of a mean column. A column read from the wrong place is off by +# about its own size. +COL = {mx.bfloat16: 5e-2, mx.float16: 2e-2} +CPU_COL = 0.2 +K = 768 +# Rows from which the op runs the kernel (kLoraMinRows). +MIN_ROWS = 1024 +# Column tile: a band starts on a multiple of it. A band ends on the column +# tile of a simdgroup, or at N. +BN = 64 +SN = 32 +# Largest rank of a band, and most bands, the kernel runs. +MAX_RANK = 256 +MAX_BANDS = 16 +# The shape rule: by itself the op gives the kernel a weight with K up to +# RULE_K and N of at least K, from RULE_ROWS rows where no column takes more +# than RULE_RANK ranks, and below that row count or above that rank only +# with N of at least RULE_WIDE_N. +RULE_K = 4096 +RULE_ROWS = 4096 +RULE_WIDE_N = 8192 +RULE_RANK = 128 +# Codecs with a kernel on bfloat16, and the K tile of each kernel. K is a +# whole number of K tiles where the kernel runs. +KERNEL_BK = {"q4_k": 128, "q5_k": 128, "q6_k": 128, "q8_0": 64} +CODECS = list(KERNEL_BK) +KERNELS = {(c, mx.bfloat16) for c in CODECS} +# Every kernel, and float16 once: float16 has no kernel. +PAIRS = [(c, mx.bfloat16) for c in CODECS] + [("q4_k", mx.float16)] +PAIR_IDS = [f"{c}-bf16" for c in CODECS] + ["q4_k-f16"] +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) + +_W = {} + + +@pytest.fixture(autouse=True, scope="module") +def _kernel_wherever_it_can_run(): + """KQ_QMM_LORA=2 for the calls of this module. The op reads the lever + at its first call in the process and keeps it, and no other module + calls the op.""" + old = os.environ.get("KQ_QMM_LORA") + os.environ["KQ_QMM_LORA"] = "2" + yield + if old is None: + del os.environ["KQ_QMM_LORA"] + else: + os.environ["KQ_QMM_LORA"] = old + + +def _nax(): + return not force_cpu() and kq.nax_available() + + +def _shape_rule(k, n, rows, rank): + """The calls on which the op picks the kernel by itself. rank is the + most ranks a column takes.""" + near = rows >= RULE_ROWS and rank <= RULE_RANK + return k <= RULE_K and n >= k and (n >= RULE_WIDE_N or near) + + +def _fused(rows, rank, n, codec="q4_k", dtype=mx.bfloat16, bands=None, k=K): + """Where the kernel can run: the dispatch rule of the op under + KQ_QMM_LORA=2.""" + if not _nax() or (codec, dtype) not in KERNELS: + return False + if k % KERNEL_BK[codec]: + return False + if rows < MIN_ROWS or rank == 0: + return False + if bands is None: + return rank <= MAX_RANK + return len(bands) <= MAX_BANDS and all( + cs % BN == 0 and (ce % SN == 0 or ce == n) and rb <= MAX_RANK + for cs, ce, _, rb in bands + ) + + +def _weight(codec, n, k, seed=0): + """A packed weight [n, k] and its float64 values.""" + key = (codec, n, k, seed) + if key not in _W: + wf = mx.random.normal((n, k), key=mx.random.key(1000 * seed + 17 * n + k)) * 0.1 + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + _W[key] = (w, _f64(kq.dequantize(w, s, codec, dtype=mx.float32))) + return _W[key] + + +def _x(shape, dtype, seed, scale=0.5): + x = (mx.random.normal(shape, key=mx.random.key(seed)) * scale).astype(dtype) + mx.eval(x) + return x + + +def _ab(k, n, r, dtype, seed=0): + """A [k, r] and B [r, n], sized so that the low-rank term is near the + size of the product.""" + a = mx.random.normal((k, r), key=mx.random.key(31 * seed + r)) / np.sqrt(k) + b = mx.random.normal((r, n), key=mx.random.key(37 * seed + r + n)) * 0.3 + a, b = a.astype(dtype), b.astype(dtype) + mx.eval(a, b) + return a, b + + +def _f64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + +def _rounded(v, dtype): + return _f64(mx.array(v.astype(np.float32)).astype(dtype)) + + +def _rel(a, b): + return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-30)) + + +def _col(a, b): + """The largest error of a column of a against b, over the size of a + mean column of b. Rows are the leading axis.""" + mean = np.linalg.norm(b) / np.sqrt(b.shape[1]) + 1e-30 + return float(np.linalg.norm(a - b, axis=0).max() / mean) + + +def _delta(z, b, n, bands): + if bands is None: + return z @ b + d = np.zeros((z.shape[0], n)) + for cs, ce, z0, rb in bands: + d[:, cs:ce] += z[:, z0 : z0 + rb] @ b[z0 : z0 + rb, cs:ce] + return d + + +def _ref(x, wd, a, b, strength=None, bands=None, model=None): + """The float64 sum, rows flattened. model=dtype rounds the weight and z + to that dtype, as the kernel reads them.""" + out = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + x64 = _f64(x.astype(out)).reshape(-1, x.shape[-1]) + z = x64 @ _f64(a.astype(out)) + if model is not None: + z = _rounded(z, model) + wd = _rounded(wd, model) + if strength is not None: + z = z * _f64(strength.astype(out)).reshape(-1) + if model is not None: + z = _rounded(z, model) + with np.errstate(invalid="ignore", over="ignore"): + return x64 @ wd.T + _delta(z, _f64(b.astype(out)), wd.shape[0], bands) + + +def _composed(x, w, a, b, codec="q4_k", strength=None, bands=None): + """The form a caller writes without the op. Where the op has no kernel + it builds this, bit for bit. A band of one column goes to mx.addmm as a + dense copy: on the GPU the stock op misreads a first argument of one + column that is a slice of a wider array (docs/upstream-defects.md).""" + out = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + base = kq.quantized_matmul(x, w, mx.zeros((1,), dtype=mx.uint8), codec, True) + z = x.astype(out) @ a.astype(out) + if strength is not None: + z = z * strength.astype(out) + b = b.astype(out) + if bands is None: + return mx.addmm(base, z, b) + if len(bands) == 0: + return base + parts, col = [], 0 + for cs, ce, z0, rb in bands: + if cs > col: + parts.append(base[..., col:cs]) + y = base[..., cs:ce] + if rb > 0: + if ce - cs == 1: + y = mx.contiguous(y) + y = mx.addmm(y, z[..., z0 : z0 + rb], b[z0 : z0 + rb, cs:ce]) + parts.append(y) + col = ce + if col < base.shape[-1]: + parts.append(base[..., col:]) + return parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=-1) + + +def _check(x, w, wd, a, b, codec="q4_k", strength=None, bands=None, cols=None): + """R1 and, on the kernel, R2. cols limits the comparison to a column + range. Returns (R1 error, R2 error or None).""" + n, k = wd.shape + dtype = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + y = kq.quantized_matmul_lora(x, w, codec, a, b, strength=strength, bands=bands) + mx.eval(y) + assert y.dtype == dtype + assert y.shape == tuple(x.shape[:-1]) + (n,) + sel = slice(*cols) if cols else slice(None) + got = _f64(y).reshape(-1, n)[:, sel] + assert np.isfinite(got).all() + ref = _ref(x, wd, a, b, strength, bands)[:, sel] + err = _rel(got, ref) + if not _fused(x.size // k, a.shape[1], n, codec, dtype, bands, k): + comp = _composed(x, w, a, b, codec, strength, bands) + mx.eval(comp) + assert np.array_equal(_f64(comp).reshape(-1, n)[:, sel], got) + assert err <= (R1[dtype] if _nax() else CPU_REL), err + # The composition is stock ops: no column of it is taken on trust. + col = _col(got, ref) + assert col <= (COL[dtype] if _nax() else CPU_COL), col + return err, None + assert err <= R1[dtype], (err, R1[dtype]) + ref2 = _ref(x, wd, a, b, strength, bands, model=dtype)[:, sel] + floor2 = _rel(_rounded(ref2, dtype), ref2) + err2 = _rel(got, ref2) + assert err2 <= 1.05 * floor2 + 2e-5, (err2, floor2) + return err, err2 + + +def _case(m, n, r, dtype=mx.bfloat16, codec="q4_k", k=K, seed=0, **kw): + w, wd = _weight(codec, n, k) + a, b = _ab(k, n, r, dtype, seed) + return _check(_x([m, k], dtype, 7 * m + n + r), w, wd, a, b, codec, **kw) + + +@pytest.mark.parametrize(("codec", "dtype"), PAIRS, ids=PAIR_IDS) +@pytest.mark.parametrize( + "m", [1, 17, 577, 1023, 1024, 1025, 1039, 1040, 1100, 1151, 1152] +) +@pytest.mark.parametrize("n", [64, 96, 200]) +def test_lora_rows_and_columns(m, n, codec, dtype): + # Below MIN_ROWS the op is the composition. From there up: a last row + # tile of 1 row, 15 and 16 rows, 76 rows, 127 rows and none, with N on + # and off the column tile. + _case(m, n, 16, dtype, codec) + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("r", [1, 4, 16, 31, 32, 33, 64, 80, 96, 130, 192, 256, 257]) +@pytest.mark.parametrize("n", [128, 200]) +def test_lora_ranks(r, n, codec): + # A rank below, at and above the step of 32, no multiple of 64, several + # steps, the largest rank of the kernel and one above it. + _case(1100, n, r, codec=codec) + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("m", [4127, 4128, 4608]) +def test_lora_image_rows(m, codec): + n = 192 + bands = [(0, 64, 0, 64), (64, 128, 64, 64), (128, 192, 128, 64)] + _case(m, n, 192, codec=codec, bands=bands) + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("m", [32768, 32769, 33792]) +def test_lora_rows_above_32767(m, codec): + # Row offsets pass a 16-bit index: every row is compared. + _case(m, 64, 16, codec=codec, k=256) + + +@pytest.mark.parametrize("k", [832, 800, 96, 32]) +def test_lora_q8_0_inner_dim_off_the_k_tile(k): + # q8_0 holds K in blocks of 32 and its kernel walks K in tiles of 64: a + # K off the tile is the composition. + err, err2 = _case(1100, 96, 16, codec="q8_0", k=k) + assert (err2 is not None) == (_nax() and k % 64 == 0) + + +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "f16"]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_lora_batch(batch, dtype): + n, r, m = 96, 16, 1100 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, dtype) + x = _x([batch, m, K], dtype, 11 + batch) + _check(x, w, wd, a, b) + y = kq.quantized_matmul_lora(x, w, "q4_k", a, b) + flat = kq.quantized_matmul_lora(x.reshape(batch * m, K), w, "q4_k", a, b) + mx.eval(y, flat) + if _fused(batch * m, r, n, dtype=dtype) == _fused(m, r, n, dtype=dtype): + assert mx.array_equal(y.reshape(batch * m, n), flat).item() + + +def test_lora_one_dim_x(): + n, r = 96, 16 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([K], mx.bfloat16, 5) + y = kq.quantized_matmul_lora(x, w, "q4_k", a, b) + two = kq.quantized_matmul_lora(x.reshape(1, K), w, "q4_k", a, b) + mx.eval(y, two) + assert y.shape == (n,) + assert mx.array_equal(y, two[0]).item() + + +def test_lora_float32_x_runs_in_bfloat16(): + n, r = 96, 16 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.float32) + x = _x([1100, K], mx.float32, 13) + _check(x, w, wd, a, b) + y = kq.quantized_matmul_lora(x, w, "q4_k", a, b) + y16 = kq.quantized_matmul_lora( + x.astype(mx.bfloat16), + w, + "q4_k", + a.astype(mx.bfloat16), + b.astype(mx.bfloat16), + ) + mx.eval(y, y16) + assert y.dtype == mx.bfloat16 + assert mx.array_equal(y, y16).item() + + +BANDS = { + # Three bands over every column, as a fused q, k and v weight. + "three": (192, 48, [(0, 64, 0, 16), (64, 128, 16, 16), (128, 192, 32, 16)]), + # Columns before, between and after the bands take no low-rank term. + "gaps": (320, 40, [(64, 128, 0, 24), (192, 256, 24, 16)]), + # A band that ends on the column tile of a simdgroup, and a band that + # ends at N off the tile. + "end_on_half_tile": (300, 40, [(0, 96, 0, 24), (128, 300, 24, 16)]), + # A band of the last 44 columns, after columns in no band. + "end_at_n": (300, 40, [(64, 128, 0, 24), (256, 300, 24, 16)]), + # A band that ends inside the columns of a simdgroup: the composition. + "end_in_tile": (300, 40, [(0, 100, 0, 24), (128, 250, 24, 16)]), + # A band of rank 0, and a band that takes no rank from the first columns + # of z. + "rank_zero": (192, 32, [(0, 64, 5, 0), (64, 192, 12, 20)]), + # Two bands that share z columns, and z ranges out of column order. + "shared_z": (192, 40, [(0, 64, 20, 20), (64, 128, 0, 40), (128, 192, 10, 7)]), + # Ranks that are no multiple of the step of 32, over several steps. + "odd_ranks": (192, 133, [(0, 64, 0, 33), (64, 192, 33, 100)]), + # A start off the column tile: the op builds the composition. + "start_off_tile": (192, 32, [(0, 32, 0, 16), (32, 192, 16, 16)]), + # A band above the largest rank of the kernel: the composition. + "rank_above_max": (128, 300, [(0, 64, 0, 257), (64, 128, 257, 43)]), + # More bands than the kernel looks through: the composition. + "seventeen": (17 * 64, 34, [(64 * i, 64 * i + 64, 2 * i, 2) for i in range(17)]), + "sixteen": (16 * 64, 32, [(64 * i, 64 * i + 64, 2 * i, 2) for i in range(16)]), +} + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize("name", list(BANDS)) +def test_lora_bands(name, m, codec): + n, r, bands = BANDS[name] + _case(m, n, r, codec=codec, bands=bands) + + +ONE_COLUMN = { + # A band of one column at the start, inside and at the end of a column + # tile, and at the last column. + "first": (200, [(0, 1, 0, 8)]), + "inside": (200, [(33, 34, 0, 8)]), + "tile_start": (200, [(64, 65, 0, 8), (128, 192, 8, 8)]), + "last": (200, [(199, 200, 0, 8)]), + # One-column bands side by side, then a wide band. + "side_by_side": (200, [(0, 1, 0, 4), (1, 2, 4, 4), (2, 200, 8, 8)]), + # A band of the one column past the last column tile: it starts on a + # column tile and ends at N, so the kernel runs it. + "past_last_tile": (193, [(192, 193, 0, 8)]), + # A weight of one row, with a band and with none. + "n_of_one": (1, [(0, 1, 0, 8)]), + "n_of_one_no_table": (1, None), +} + + +@pytest.mark.parametrize( + "dtype", [mx.bfloat16, mx.float16, mx.float32], ids=["bf16", "f16", "f32"] +) +@pytest.mark.parametrize("shape", [[2], [1023], [1100], [2, 600], [3, 2, 200]]) +@pytest.mark.parametrize("name", list(ONE_COLUMN)) +def test_lora_band_of_one_column(name, shape, dtype): + """A band of one column is right in every column, on both paths, for + every dtype of x and with batch dims. The composition adds each band + with the stock mx.addmm, which on the GPU misreads a first argument of + one column that is a slice of a wider array: the op passes it a dense + copy of that column.""" + n, bands = ONE_COLUMN[name] + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, 16, dtype) + x = _x([*shape, K], dtype, 131 + len(shape) + shape[-1]) + _check(x, w, wd, a, b, bands=bands) + + +def test_lora_no_band_and_no_table(): + """bands=None is no table: every column takes every rank. bands=[] is a + table with no band: no column takes a rank, and y is the product alone, + bit for bit, at every row count and whatever the strength.""" + n, r = 128, 16 + w, wd = _weight("q4_k", n, K) + s8 = mx.zeros((1,), dtype=mx.uint8) + for m, dtype in itertools.product((1023, 1100), (mx.bfloat16, mx.float32)): + a, b = _ab(K, n, r, dtype) + x = _x([m, K], dtype, 137) + base = kq.quantized_matmul(x, w, s8, "q4_k", True) + for bands in ([], ()): + for strength in (None, mx.array(2.0), mx.ones((r,))): + y = kq.quantized_matmul_lora( + x, w, "q4_k", a, b, strength=strength, bands=bands + ) + mx.eval(y, base) + assert y.dtype == base.dtype and y.shape == base.shape + assert mx.array_equal(y, base).item(), (m, dtype, bands) + none = kq.quantized_matmul_lora(x, w, "q4_k", a, b, bands=None) + default = kq.quantized_matmul_lora(x, w, "q4_k", a, b) + mx.eval(none, default) + assert mx.array_equal(none, default).item() + assert not mx.array_equal(none, base).item() + _check(x, w, wd, a, b, bands=None) + + +@pytest.mark.parametrize("end", [96, 100]) +@pytest.mark.parametrize("codec", CODECS) +def test_lora_band_columns_are_isolated(codec, end): + """A band reads its own ranks of z and its own block of B, whatever the + rest holds: an infinite strength on the ranks of the other band and + infinite values in B outside the two blocks leave the columns of the + first band finite and right. The first band ends on the column tile of + a simdgroup (the kernel) and inside one (the composition).""" + n, r, m = 192, 32, 1100 + bands = [(0, end, 0, 20), (128, 192, 20, 12)] + w, wd = _weight(codec, n, K) + a, b = _ab(K, n, r, mx.bfloat16) + mask = np.ones((r, n), dtype=bool) + for cs, ce, z0, rb in bands: + mask[z0 : z0 + rb, cs:ce] = False + b = mx.where(mx.array(mask), mx.array(float("inf")), b).astype(mx.bfloat16) + strength = mx.array([1.0] * 20 + [float("inf")] * 12) + x = _x([m, K], mx.bfloat16, 23) + # The columns of the first band, then the columns outside every band. + kw = dict(codec=codec, strength=strength, bands=bands) + _check(x, w, wd, a, b, cols=(0, end), **kw) + _check(x, w, wd, a, b, cols=(end, 128), **kw) + + +@pytest.mark.parametrize("end", [96, 100]) +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize("codec", CODECS) +def test_lora_z_that_is_not_finite_stays_in_its_band(codec, m, end): + """A column outside every band takes no term of z. With an infinite + strength every column of a band is not finite, and every other column + is finite and right: the columns before the first band, the columns + after it in its column tile, a column tile between two bands, and the + columns after the last band. The first band ends on the column tile of + a simdgroup (the kernel) and inside one (the composition).""" + n, r = 448, 32 + bands = [(64, end, 0, 20), (192, 256, 20, 12)] + w, wd = _weight(codec, n, K) + a, b = _ab(K, n, r, mx.bfloat16) + strength = mx.array(float("inf")) + x = _x([m, K], mx.bfloat16, 113) + kw = dict(codec=codec, strength=strength, bands=bands) + for cols in ((0, 64), (end, 128), (128, 192), (256, 448)): + _check(x, w, wd, a, b, cols=cols, **kw) + y = _f64( + kq.quantized_matmul_lora(x, w, codec, a, b, strength=strength, bands=bands) + ) + for cs, ce, _, _ in bands: + assert not np.isfinite(y[:, cs:ce]).any() + + +@pytest.mark.parametrize("m", [1023, 1100]) +def test_lora_strength(m): + n, r = 128, 24 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([m, K], mx.bfloat16, 29) + per_col = mx.array(np.linspace(-1.5, 1.5, r).astype(np.float32)) + for strength in (mx.array([0.5]), mx.array(-2.0), per_col): + _check(x, w, wd, a, b, strength=strength) + # A strength of one is no strength, and a strength of zero is the + # product alone. + y = kq.quantized_matmul_lora(x, w, "q4_k", a, b) + one = kq.quantized_matmul_lora(x, w, "q4_k", a, b, strength=mx.ones((r,))) + zero = kq.quantized_matmul_lora(x, w, "q4_k", a, b, strength=mx.zeros((1,))) + base = kq.quantized_matmul(x, w, mx.zeros((1,), dtype=mx.uint8), "q4_k", True) + mx.eval(y, one, zero, base) + assert mx.array_equal(y, one).item() + if _fused(m, r, n): + # The K walk of the kernel may not be the tile of quantized_matmul. + assert _rel(_f64(zero), _f64(base)) <= 1e-3 + else: + assert mx.array_equal(zero, base).item() + + +@pytest.mark.usefixtures("compile_on") +def test_lora_strength_change_needs_no_new_trace(): + n, r, m = 128, 16, 1100 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([m, K], mx.bfloat16, 31) + traces = [] + + def f(t, s): + traces.append(1) + return kq.quantized_matmul_lora(t, w, "q4_k", a, b, strength=s) + + cf = mx.compile(f) + for value in (1.0, -0.25, 0.0, 3.0): + s = mx.array([value]) + got = cf(x, s) + want = kq.quantized_matmul_lora(x, w, "q4_k", a, b, strength=s) + mx.eval(got, want) + assert mx.array_equal(got, want).item(), value + assert len(traces) == 1 + + +@pytest.mark.parametrize("m", [1023, 1100]) +def test_lora_stacked_pair(m): + """Two adapters in one call: their A columns and B rows joined, one + strength per column. Equal to the sum of the two terms.""" + n = 128 + w, wd = _weight("q4_k", n, K) + a1, b1 = _ab(K, n, 16, mx.bfloat16, seed=1) + a2, b2 = _ab(K, n, 64, mx.bfloat16, seed=2) + a = mx.concatenate([a1, a2], axis=1) + b = mx.concatenate([b1, b2], axis=0) + strength = mx.array([0.75] * 16 + [-0.5] * 64) + x = _x([m, K], mx.bfloat16, 37) + _check(x, w, wd, a, b, strength=strength) + y = kq.quantized_matmul_lora(x, w, "q4_k", a, b, strength=strength) + x64 = _f64(x) + want = ( + x64 @ wd.T + + 0.75 * (x64 @ _f64(a1)) @ _f64(b1) + - 0.5 * (x64 @ _f64(a2)) @ _f64(b2) + ) + assert _rel(_f64(y), want) <= (R1[mx.bfloat16] if _nax() else CPU_REL) + + +def test_lora_rank_zero_and_empty_inputs(): + n = 96 + w, wd = _weight("q4_k", n, K) + s8 = mx.zeros((1,), dtype=mx.uint8) + for m in (1023, 1100): + x = _x([m, K], mx.bfloat16, 41) + # Rank 0: the product alone, bit for bit. + y = kq.quantized_matmul_lora(x, w, "q4_k", mx.zeros((K, 0)), mx.zeros((0, n))) + base = kq.quantized_matmul(x, w, s8, "q4_k", True) + mx.eval(y, base) + assert y.dtype == base.dtype and mx.array_equal(y, base).item() + a, b = _ab(K, n, 16, mx.bfloat16) + for shape in ([0, K], [2, 0, K], [0, 3, K]): + y = kq.quantized_matmul_lora( + mx.zeros(shape, dtype=mx.bfloat16), w, "q4_k", a, b + ) + mx.eval(y) + assert y.shape == tuple(shape[:-1]) + (n,) and y.dtype == mx.bfloat16 + # No output column. + w0 = w[:0] + y = kq.quantized_matmul_lora( + _x([1100, K], mx.bfloat16, 43), w0, "q4_k", a, mx.zeros((16, 0)) + ) + mx.eval(y) + assert y.shape == (1100, 0) + # No inner dim: every sum is empty, so y is zero. + y = kq.quantized_matmul_lora( + mx.zeros((5, 0), dtype=mx.bfloat16), + mx.zeros((n, 0), dtype=mx.uint8), + "q4_k", + mx.zeros((0, 16)), + b, + ) + mx.eval(y) + assert y.shape == (5, n) and y.dtype == mx.bfloat16 + assert not _f64(y).any() + + +def _check_views(views, call, state): + """call(view) against call(dense copy of the view). On the GPU the two + have the same bits. On the CPU stream the bound is one rounding of the + activation dtype (R1), not bit equality: the stock matmul there sums a + product with a transposed second operand and a product of dense + operands in different orders. Both are right to about one rounding + against the exact sum, the dense order with about 1.3 times the error + of the other. That is the accuracy of the stock op on bfloat16 and + float16, so the two calls may differ by a rounding. A view read with + the wrong strides differs by its whole size.""" + for name, make in views.items(): + want = call(_dense(make())) + view = make() + if state == "evaluated": + mx.eval(view) + got = call(view) + mx.eval(got, want) + if force_cpu(): + assert _rel(_f64(got), _f64(want)) <= R1[mx.bfloat16], (name, state) + else: + assert mx.array_equal(got, want).item(), (name, state) + + +def _weight_views(w): + """Views of one weight whose rows are not dense.""" + n, kb = w.shape + other = mx.array(np.random.default_rng(3).integers(0, 255, (n, kb), dtype=np.uint8)) + rows = mx.stack([w, other], axis=1).reshape(2 * n, kb) + cols = mx.concatenate([other, w], axis=1) + rev = mx.array(np.ascontiguousarray(np.array(w)[::-1])) + mx.eval(rows, cols, rev) + return { + "even_rows": lambda: rows[0::2], + "col_band": lambda: cols[:, kb:], + "reversed_rows": lambda: rev[::-1], + } + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +def test_lora_operand_views(state, m, codec): + """A view of any operand gives the bits of its dense copy, lazy or + evaluated, on both sides of the row bound.""" + n, r = 96, 16 + w, wd = _weight(codec, n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([m, K], mx.bfloat16, 47) + s = mx.array(np.linspace(0.5, 1.5, r).astype(np.float32)).astype(mx.bfloat16) + bands = [(0, 64, 0, 16), (64, 96, 4, 12)] + + def call(x_=x, w_=w, a_=a, b_=b, s_=s): + return kq.quantized_matmul_lora(x_, w_, codec, a_, b_, strength=s_, bands=bands) + + _check_views( + _strided(_x([2 * m, 2 * K], mx.bfloat16, 53)), lambda v: call(x_=v), state + ) + _check_views( + _strided(_x([2 * K, 2 * r], mx.bfloat16, 59)), lambda v: call(a_=v), state + ) + _check_views( + _strided(_x([2 * r, 2 * n], mx.bfloat16, 61)), lambda v: call(b_=v), state + ) + _check_views(_weight_views(w), lambda v: call(w_=v), state) + s2 = _x([2 * r], mx.bfloat16, 67) + views = {"even": lambda: s2[0::2], "reversed": lambda: s2[::-1][:r]} + _check_views(views, lambda v: call(s_=v), state) + # The views of x keep the result right, not only equal to the copy. + xv = _strided(_x([2 * m, 2 * K], mx.bfloat16, 53))["odd_cols"]() + _check(xv, w, wd, a, b, codec, strength=s, bands=bands) + + +@pytest.mark.parametrize( + ("codec", "dtype"), + [ + ("q3_k", mx.bfloat16), + ("q4_0", mx.bfloat16), + ("q5_k", mx.float16), + ("q6_k", mx.float16), + ("q8_0", mx.float16), + ("q2_0", mx.bfloat16), + ("pq2_0", mx.bfloat16), + ], + ids=[ + "q3_k-bf16", + "q4_0-bf16", + "q5_k-f16", + "q6_k-f16", + "q8_0-f16", + "q2_0-bf16", + "pq2_0-bf16", + ], +) +def test_lora_codec_without_kernel_composes(codec, dtype): + n, r = 96, 16 + w, wd = _weight(codec, n, K) + a, b = _ab(K, n, r, dtype) + err, err2 = _check(_x([1100, K], dtype, 71), w, wd, a, b, codec) + assert err2 is None + + +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize( + "dtype", [mx.bfloat16, mx.float16, mx.float32], ids=["bf16", "f16", "f32"] +) +@pytest.mark.parametrize("banded", [False, True]) +def test_lora_gradient_is_the_gradient_of_the_composition(banded, dtype, m): + """mx.grad wrt x, A, B and the strength on both sides of the row bound. + The kernel has the vjp of the unfused ops.""" + n, r = 128, 16 + w, _ = _weight("q4_k", n, K) + a, b = _ab(K, n, r, dtype) + x = _x([m, K], dtype, 73) + s = mx.array(np.linspace(0.5, 1.5, r).astype(np.float32)).astype(dtype) + c = mx.random.normal((m, n), key=mx.random.key(79)) + mx.eval(c) + bands = [(0, 64, 0, 10), (64, 128, 10, 6)] if banded else None + + def loss(f): + def run(x_, a_, b_, s_): + y = f(x_, w, a_, b_, "q4_k", strength=s_, bands=bands) + return (y.astype(mx.float32) * c).sum() + + return run + + def op(x_, w_, a_, b_, codec, **kw): + return kq.quantized_matmul_lora(x_, w_, codec, a_, b_, **kw) + + g_op = mx.grad(loss(op), argnums=[0, 1, 2, 3])(x, a, b, s) + g_ref = mx.grad(loss(_composed), argnums=[0, 1, 2, 3])(x, a, b, s) + mx.eval(g_op, g_ref) + for got, want, arg in zip(g_op, g_ref, (x, a, b, s), strict=True): + assert got.shape == arg.shape and got.dtype == arg.dtype + ref = _f64(want) + assert np.isfinite(ref).all() and np.abs(ref).max() > 0 + # bfloat16 keeps 8 significant bits, and the terms of a gradient + # may be summed in another order than in the composition. + tol = 1e-3 if dtype == mx.float16 else 6e-3 + assert _rel(_f64(got), ref) <= tol + + +ARGS = ("x", "a", "b", "s") +SUBSETS = [c for k in range(1, len(ARGS) + 1) for c in itertools.combinations(ARGS, k)] + + +def _args_case(m, banded, seed, dtype=mx.bfloat16): + """A weight, the four float arguments by name, the float64 weight and + a band table or None.""" + n, r = 128, 16 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, dtype) + s = mx.array(np.linspace(0.5, 1.5, r).astype(np.float32)).astype(dtype) + args = {"x": _x([m, K], dtype, seed), "a": a, "b": b, "s": s} + bands = [(0, 64, 0, 10), (64, 128, 10, 6)] if banded else None + return w, wd, args, bands + + +def _of(call, w, args, bands, names): + """call as a function of the arguments in names, the rest fixed.""" + + def f(*vals): + kw = {**args, **dict(zip(names, vals, strict=True))} + return call(kw["x"], w, kw["a"], kw["b"], strength=kw["s"], bands=bands) + + return f + + +def _op(x, w, a, b, **kw): + return kq.quantized_matmul_lora(x, w, "q4_k", a, b, **kw) + + +def _unfused(x, w, a, b, strength=None, bands=None): + return _composed(x, w, a, b, "q4_k", strength, bands) + + +@pytest.mark.parametrize("banded", [False, True], ids=["no_table", "bands"]) +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize("names", SUBSETS, ids=["".join(c) for c in SUBSETS]) +def test_lora_gradient_of_every_subset_of_the_arguments(names, m, banded): + """mx.grad with respect to each of the 15 subsets of x, A, B and the + strength, on both sides of the row bound: each gradient has the shape + of its argument and is the gradient of the composition.""" + w, _, args, bands = _args_case(m, banded, 139) + c = mx.random.normal((m, 128), key=mx.random.key(149)) + + def loss(call): + f = _of(call, w, args, bands, names) + return lambda *vals: (f(*vals).astype(mx.float32) * c).sum() + + at = [args[k] for k in names] + argnums = list(range(len(names))) + got = mx.grad(loss(_op), argnums=argnums)(*at) + want = mx.grad(loss(_unfused), argnums=argnums)(*at) + mx.eval(got, want) + # mx.grad returns one array for one argument and a tuple for more. + got = [got] if isinstance(got, mx.array) else list(got) + want = [want] if isinstance(want, mx.array) else list(want) + assert len(got) == len(names) + for g, r, name in zip(got, want, names, strict=True): + assert g.shape == args[name].shape and g.dtype == args[name].dtype, name + ref = _f64(r) + assert np.isfinite(ref).all() and np.abs(ref).max() > 0 + assert _rel(_f64(g), ref) <= 6e-3, name + + +def test_lora_compiles_and_repeats(): + n, r, m = 128, 16, 1100 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([m, K], mx.bfloat16, 83) + f = mx.compile(lambda t: kq.quantized_matmul_lora(t, w, "q4_k", a, b)) + first = f(x) + second = f(x) + plain = kq.quantized_matmul_lora(x, w, "q4_k", a, b) + mx.eval(first, second, plain) + assert mx.array_equal(first, second).item() + assert mx.array_equal(first, plain).item() + + +def test_lora_compiled_function_with_two_band_tables(): + """Two calls that differ in the band table alone stay two nodes when a + compiled function holds both.""" + n, r, m = 128, 32, 1100 + w, _ = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([m, K], mx.bfloat16, 103) + one = [(0, 64, 0, 16), (64, 128, 16, 16)] + two = [(0, 64, 16, 16), (64, 128, 0, 16)] + + def both(t): + p = kq.quantized_matmul_lora(t, w, "q4_k", a, b, bands=one) + q = kq.quantized_matmul_lora(t, w, "q4_k", a, b, bands=two) + # mx.compile merges equivalent nodes but never an output of the + # function, so the two calls feed the outputs. The casts are exact. + return p.astype(mx.float32), q.astype(mx.float32) + + p, q = mx.compile(both)(x) + rp, rq = both(x) + mx.eval(p, q, rp, rq) + assert mx.array_equal(p, rp).item() + assert mx.array_equal(q, rq).item() + assert not mx.array_equal(rp, rq).item() + + +def _kq_nodes(y): + """The mlx_kquant nodes in the graph of an output that is not yet + evaluated.""" + return [n for n in _nodes(y) if n.startswith("KQuant")] + + +@pytest.mark.usefixtures("compile_on") +@pytest.mark.parametrize("banded", [False, True], ids=["no_table", "bands"]) +def test_lora_shapeless_compile_keeps_the_path_of_its_trace(banded): + """A function compiled with shapeless=True and traced on a call that + runs the kernel runs the kernel at every later row count, below the + row bound too, with the shape and the values of each call. A call with + another number of dimensions is a new trace and is not part of this.""" + n, r = 128, 16 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + bands = [(0, 64, 0, 10), (64, 128, 10, 6)] if banded else None + f = mx.compile( + lambda t: kq.quantized_matmul_lora(t, w, "q4_k", a, b, bands=bands), + shapeless=True, + ) + if not _fused(1100, r, n, bands=bands): + pytest.skip("the kernel needs a NAX GPU") + for shape in ([1100, K], [600, K], [1023, K], [2200, K], [1, K]): + x = _x(shape, mx.bfloat16, 151 + shape[0]) + y = f(x) + assert _kq_nodes(y) == ["KQuantMatmulLora"], (shape, _kq_nodes(y)) + mx.eval(y) + assert y.shape == (*shape[:-1], n), (shape, y.shape) + got = _f64(y).reshape(-1, n) + err = _rel(got, _ref(x, wd, a, b, None, bands)) + assert err <= R1[mx.bfloat16], (shape, err) + + +@pytest.mark.usefixtures("compile_on") +def test_lora_shapeless_compile_of_the_unfused_ops(): + """Traced below the row bound, a shapeless function builds the unfused + ops and keeps them above the bound. With bands the unfused ops hold + stock Slice ops, which have no shape rule: the call raises, by the name + of this op.""" + n, r = 128, 16 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + f = mx.compile( + lambda t: kq.quantized_matmul_lora(t, w, "q4_k", a, b), shapeless=True + ) + for m in (600, 1023, 2200): + x = _x([m, K], mx.bfloat16, 157 + m) + y = f(x) + assert _kq_nodes(y) == ["KQuantMatmul"], m + mx.eval(y) + assert y.shape == (m, n) + err = _rel(_f64(y), _ref(x, wd, a, b)) + assert err <= (R1[mx.bfloat16] if _nax() else CPU_REL), (m, err) + bands = [(0, 64, 0, 10), (64, 128, 10, 6)] + g = mx.compile( + lambda t: kq.quantized_matmul_lora(t, w, "q4_k", a, b, bands=bands), + shapeless=True, + ) + for m in (600, 2200): + with pytest.raises(ValueError, match=r"quantized_matmul_lora\] A shapeless"): + mx.eval(g(_x([m, K], mx.bfloat16, 163 + m))) + + +def test_lora_runs_on_the_stream_of_the_call(): + """A call on a stream that is not the default one builds the + composition on that stream, at a row count that takes the kernel on + the default one.""" + if force_cpu() or mx.default_device() != mx.gpu: + pytest.skip("the CPU is the default device") + n, r, m = 96, 16, 1100 + w, wd = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([m, K], mx.bfloat16, 97) + y = kq.quantized_matmul_lora(x, w, "q4_k", a, b, stream=mx.cpu) + with mx.stream(mx.cpu): + ref = _composed(x, w, a, b) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item() + assert _rel(_f64(y), _ref(x, wd, a, b)) <= CPU_REL + + +FLOAT_SUBSETS = [c for c in SUBSETS if "x" not in c] +X_SUBSETS = [("x",), ("x", "a"), ("x", "a", "b", "s")] +# The text of a transform that quantized_matmul does not have. The period +# ends the name: KQuantMatmulLora is another node. +NO_TRANSFORM = r"\[Primitive::{}\] Not implemented for KQuantMatmul\." + + +@pytest.mark.parametrize("banded", [False, True], ids=["no_table", "bands"]) +@pytest.mark.parametrize("m", [1023, 1100]) +def test_lora_jvp_does_not_depend_on_the_row_count(m, banded): + """mx.jvp with respect to A, B and the strength, alone and together, + gives the tangent of the composition on both sides of the row bound. + With respect to x it raises on both sides, by the name of + quantized_matmul, which has no jvp.""" + w, wd, args, bands = _args_case(m, banded, 167) + tans = { + k: _x(list(v.shape), mx.bfloat16, 173 + i, scale=0.1) + for i, (k, v) in enumerate(args.items()) + } + for names in FLOAT_SUBSETS: + at = [args[k] for k in names] + tan = [tans[k] for k in names] + out, got = mx.jvp(_of(_op, w, args, bands, names), at, tan) + _, want = mx.jvp(_of(_unfused, w, args, bands, names), at, tan) + mx.eval(out, got, want) + assert got[0].shape == (m, 128) and got[0].dtype == mx.bfloat16 + # y is linear in each of A, B and the strength: the float64 tangent + # is the sum of the low-rank term with one factor replaced. + ref = np.zeros((m, 128)) + for k in names: + kw = {**args, k: tans[k]} + ref += _ref(kw["x"], 0 * wd, kw["a"], kw["b"], kw["s"], bands) + assert np.abs(ref).max() > 0 + tol = 3 * R1[mx.bfloat16] if _nax() else CPU_REL + assert _rel(_f64(got[0]), ref) <= tol, names + assert _rel(_f64(got[0]), _f64(want[0])) <= R1[mx.bfloat16], names + for names in X_SUBSETS: + at = [args[k] for k in names] + tan = [tans[k] for k in names] + with pytest.raises(ValueError, match=NO_TRANSFORM.format("jvp")): + mx.eval(mx.jvp(_of(_op, w, args, bands, names), at, tan)[1]) + + +@pytest.mark.parametrize("banded", [False, True], ids=["no_table", "bands"]) +@pytest.mark.parametrize("m", [1023, 1100]) +def test_lora_vmap_does_not_depend_on_the_row_count(m, banded): + """mx.vmap over A, B and the strength, alone and together, gives the + call of each element on both sides of the row bound. Over x it raises + on both sides, by the name of quantized_matmul, which has no vmap.""" + w, wd, args, bands = _args_case(m, banded, 179) + count = 3 + stacks = { + k: _x([count, *v.shape], mx.bfloat16, 181 + i, scale=0.3) + for i, (k, v) in enumerate(args.items()) + } + for names in FLOAT_SUBSETS: + y = mx.vmap(_of(_op, w, args, bands, names))(*[stacks[k] for k in names]) + mx.eval(y) + assert y.shape == (count, m, 128) and y.dtype == mx.bfloat16 + for i in range(count): + kw = {**args, **{k: stacks[k][i] for k in names}} + ref = _ref(kw["x"], wd, kw["a"], kw["b"], kw["s"], bands) + err = _rel(_f64(y[i]), ref) + assert err <= (R1[mx.bfloat16] if _nax() else CPU_REL), (names, i) + for names in X_SUBSETS: + with pytest.raises(ValueError, match=NO_TRANSFORM.format("vmap")): + f = mx.vmap(_of(_op, w, args, bands, names)) + mx.eval(f(*[stacks[k] for k in names])) + + +def test_lora_weight_is_frozen(): + """A gradient with respect to the weight raises at every row count, by + the name of the op that holds the weight: this op where the kernel + runs, quantized_matmul where the op builds the composition.""" + n, r = 96, 16 + w, _ = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + for m in (1023, 1100): + x = _x([m, K], mx.bfloat16, 107) + + def loss(w_, x_=x): + y = kq.quantized_matmul_lora(x_, w_, "q4_k", a, b) + return y.astype(mx.float32).sum() + + op = "quantized_matmul_lora" if _fused(m, r, n) else "quantized_matmul" + text = rf"\[mlx_kquant\] {op} vjp: no gradient wrt the quantized weight" + with pytest.raises(ValueError, match=text): + mx.eval(mx.grad(loss)(w)) + + +@pytest.mark.parametrize("codec", CODECS) +def test_lora_weight_off_its_byte_boundary_raises(codec): + """The kernel reads the packed weight with vector loads, as the kernels + of quantized_matmul do: a weight that starts off the byte boundary of + its codec raises by the name of the op that reads it. An evaluated + weight raises from the call and a lazy view of one from the eval.""" + if force_cpu() or mx.default_device() != mx.gpu: + pytest.skip("the start check guards the Metal kernels") + n, r = 96, 16 + w, _ = _weight(codec, n, K) + a, b = _ab(K, n, r, mx.bfloat16) + align = BASE_ALIGN.get(codec, 2) + for m in (1023, 1100): + x = _x([m, K], mx.bfloat16, 191) + op = "quantized_matmul_lora" if _fused(m, r, n, codec) else "quantized_matmul" + text = ( + rf"\[mlx_kquant\] {op}: the {codec} weight starts 1 byte past a " + rf"{align}-byte boundary" + ) + with pytest.raises(ValueError, match=text): + kq.quantized_matmul_lora(x, _shift(w, 1), codec, a, b) + y = kq.quantized_matmul_lora(x, _shift(w, 1, evaluate=False), codec, a, b) + with pytest.raises(ValueError, match=text): + mx.eval(y) + # A start on the boundary runs, with the bits of the weight itself. + good = kq.quantized_matmul_lora(x, _shift(w, 16), codec, a, b) + ref = kq.quantized_matmul_lora(x, w, codec, a, b) + mx.eval(good, ref) + assert mx.array_equal(good, ref).item() + + +def test_lora_argument_errors(): + n, r = 96, 16 + w, _ = _weight("q4_k", n, K) + a, b = _ab(K, n, r, mx.bfloat16) + x = _x([4, K], mx.bfloat16, 89) + + def bad(match, x_=x, w_=w, a_=a, b_=b, codec="q4_k", **kw): + # The error names this op, not an op it builds from. + with pytest.raises(ValueError, match=r"quantized_matmul_lora\] .*" + match): + kq.quantized_matmul_lora(x_, w_, codec, a_, b_, **kw) + + bad("Unknown kquant_type", codec="q9_z") + bad("2-D uint8", w_=w.astype(mx.float32)) + bad("2-D uint8", w_=w.reshape(1, n, -1)) + bad("whole number", w_=w[:, :-1]) + bad("x last dim", x_=mx.array(1.0, dtype=mx.bfloat16)) + bad("x last dim", x_=x[:, :-1]) + bad("floating x", x_=x.astype(mx.int32)) + bad("lora_a must be", a_=a[:-1]) + bad("lora_a must be", a_=a.astype(mx.int32)) + bad("lora_a must be", a_=a.reshape(K, r, 1)) + bad("lora_b must be", b_=b[:-1]) + bad("lora_b must be", b_=b[:, :-1]) + bad("lora_b must be", b_=b.astype(mx.int32)) + bad("lora_b must be", b_=b.reshape(r, n, 1)) + # A strength is 0-D, [1] or [R]. Every other shape raises, at a row + # count of the composition and at one of the kernel. + for x_ in (x, _x([1100, K], mx.bfloat16, 91)): + for shape in ((r + 1,), (2,), (1, 1), (1, 1, 1), (1, r), (r, 1), (0,)): + bad("strength must be", x_=x_, strength=mx.ones(shape)) + bad("strength must be", x_=x_, strength=mx.ones((r,), dtype=mx.int32)) + bad("strength must be", x_=x_, strength=mx.array(1)) + for shape in ((), (1,), (r,)): + kq.quantized_matmul_lora(x_, w, "q4_k", a, b, strength=mx.ones(shape)) + bad("rows of", bands=[(0, 64, 0)]) + bad("rows of", bands=[(0, 64, 0), (64, 96, 0, 8, 8)]) + bad("column order", bands=[(64, 96, 0, 8), (0, 64, 8, 8)]) + bad("column order", bands=[(0, 64, 0, 8), (63, 96, 8, 8)]) + bad("column order", bands=[(32, 32, 0, 8)]) + bad("column order", bands=[(0, n + 1, 0, 8)]) + bad("column order", bands=[(-1, 64, 0, 8)]) + bad("takes ranks", bands=[(0, 64, r - 7, 8)]) + bad("takes ranks", bands=[(0, 64, -1, 8)]) + bad("takes ranks", bands=[(0, 64, 0, -1)]) + + +def test_lora_size_limits(): + """K, N or R of 2^25 and a row count of 2^31 raise by the name of the + op: past them the kernel would compute an address in 32 bits that no + longer fits. The arrays here are lazy and stay so: the op raises from + their shapes, and no test evaluates one.""" + cols, rows = 1 << 25, 1 << 31 + f16 = mx.bfloat16 + + def lazy(*shape, dtype=mx.bfloat16): + return mx.zeros(shape, dtype=dtype) + + def w8(n, k): + # q8_0 holds 32 weights in 34 bytes. + return lazy(n, k // 32 * 34, dtype=mx.uint8) + + cases = { + "K": (lazy(cols, dtype=f16), w8(1, cols), lazy(cols, 1), lazy(1, 1)), + "N": (lazy(4, 32, dtype=f16), w8(cols, 32), lazy(32, 1), lazy(1, cols)), + "R": (lazy(4, 32, dtype=f16), w8(64, 32), lazy(32, cols), lazy(cols, 64)), + } + for name, (x, w, a, b) in cases.items(): + text = ( + rf"quantized_matmul_lora\] K, N and R must each be below {cols} " + rf"\(2\^25\): .*{name} {cols}" + ) + with pytest.raises(ValueError, match=text): + kq.quantized_matmul_lora(x, w, "q8_0", a, b) + x = lazy(2, rows // 2, 32, dtype=f16) + text = rf"quantized_matmul_lora\] x must hold fewer than {rows} \(2\^31\) rows" + with pytest.raises(ValueError, match=text): + kq.quantized_matmul_lora(x, w8(64, 32), "q8_0", lazy(32, 1), lazy(1, 64)) + # One K tile below the bound the op builds its graph (not evaluated). + k = cols - 64 + y = kq.quantized_matmul_lora( + lazy(k, dtype=f16), w8(1, k), "q8_0", lazy(k, 1), lazy(1, 1) + ) + assert y.shape == (1,) + + +def test_lora_kernels_in_metallib(): + if force_cpu(): + pytest.skip("Metal kernel names only") + names = kq.metallib_kernels() + if not names: + pytest.skip("no metallib loaded") + lora = {n for n in names if "_qmm_t_nax_lora_" in n} + want = {_name(n, codec) for n in (64, 96) for codec in CODECS} + assert lora == want, sorted(lora ^ want) + + +_PROBE = r""" +import json, os, sys +spec = json.loads(sys.argv[1]) +sys.path.insert(0, spec["tests_dir"]) +import mlx.core as mx +import mlx_kquant as kq +import test_matmul_lora as t + +out = [] +for i, c in enumerate(spec["calls"]): + dtype = getattr(mx, c.get("dtype", "bfloat16")) + codec = c.get("codec", "q4_k") + k = c.get("k", t.K) + if c.get("tile"): + # A wide weight as copies of a block of rows. + w, wd = t._weight(codec, c["tile"], k) + reps = -(-c["n"] // c["tile"]) + w = mx.contiguous(mx.tile(w, (reps, 1))[: c["n"]]) + wd = t.np.tile(wd, (reps, 1))[: c["n"]] + else: + w, wd = t._weight(codec, c["n"], k) + a, b = t._ab(k, c["n"], c["r"], dtype) + x = t._x(c.get("shape", [c["m"], k]), dtype, 7 * c["m"] + c["n"]) + bands = c.get("bands") + os.environ.update(c.get("env", {})) + print(f"[probe] {i}", file=sys.stderr, flush=True) + res = {"error": None} + try: + y = kq.quantized_matmul_lora(x, w, codec, a, b, bands=bands) + mx.eval(y) + got = t._f64(y).reshape(-1, c["n"]) + ref = t._ref(x, wd, a, b, None, bands) + res["rel"] = t._rel(got, ref) + res["col"] = t._col(got, ref) + except Exception as e: + res["error"] = f"{type(e).__name__}: {e}" + out.append(res) +print("RESULT " + json.dumps(out)) +""" + + +def _probe(calls, env=None, stderr=None): + """Run calls in a fresh interpreter that keeps no KQ_ variable but + those of env and of each call. A call's env is set before it and stays + set for the calls after it. Per call: lora (the low-rank kernels that + KQ_QMM_LOG named), plain (the other kernels), error, rel and col. + stderr, a list, takes the lines the child wrote to stderr.""" + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update({"KQ_QMM_LOG": "1", **(env or {})}) + spec = {"tests_dir": TESTS_DIR, "calls": calls} + p = subprocess.run( + [sys.executable, "-c", _PROBE, json.dumps(spec)], + env=child_env, + capture_output=True, + text=True, + timeout=900, + ) + assert p.returncode == 0, p.stderr[-3000:] + line = [ln for ln in p.stdout.splitlines() if ln.startswith("RESULT ")][-1] + results = json.loads(line[len("RESULT ") :]) + if stderr is not None: + stderr.extend(p.stderr.splitlines()) + cur = None + for ln in p.stderr.splitlines(): + if ln.startswith("[probe] "): + cur = results[int(ln.split()[1])] + cur.update(lora=[], plain=[]) + elif cur is not None and ln.startswith("[kq qmm] "): + name = ln.split()[2] + cur["lora" if "_lora_" in name else "plain"].append(name) + return results + + +def _name(n, codec="q4_k"): + aligned = "true" if n % 64 == 0 else "false" + return ( + f"kquant_{codec}_qmm_t_nax_lora_bfloat16_t_bm128_bn64_bk{KERNEL_BK[codec]}" + f"_wm2_wn2_alN_{aligned}" + ) + + +def _expect(calls, results, kernel): + for c, r, k in zip(calls, results, kernel, strict=True): + assert r["error"] is None, (c, r["error"]) + dtype = getattr(mx, c.get("dtype", "bfloat16")) + assert r["rel"] <= R1[dtype], (c, r["rel"]) + assert r["col"] <= COL[dtype], (c, r["col"]) + if k: + assert r["lora"] == [_name(c["n"], c.get("codec", "q4_k"))], (c, r["lora"]) + assert r["plain"] == [], (c, r["plain"]) + else: + assert r["lora"] == [], (c, r["lora"]) + assert len(r["plain"]) == 1, (c, r["plain"]) + + +def test_lora_kernel_that_runs(): + """KQ_QMM_LOG names the kernel of each call. With KQ_QMM_LORA=2, where + the kernel can run it runs, once, and no matmul tile. Everywhere else + one matmul tile runs and no low-rank kernel.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + sixteen = BANDS["sixteen"] + seventeen = BANDS["seventeen"] + calls = [ + ({"m": 1024, "n": 128, "r": 16}, True), + ({"m": 1023, "n": 128, "r": 16}, False), + ({"m": 1100, "n": 96, "r": 16}, True), + ({"m": 1100, "n": 128, "r": 256}, True), + ({"m": 1100, "n": 128, "r": 257}, False), + ({"m": 1100, "n": sixteen[0], "r": sixteen[1], "bands": sixteen[2]}, True), + ( + {"m": 1100, "n": seventeen[0], "r": seventeen[1], "bands": seventeen[2]}, + False, + ), + ({"m": 1100, "n": 192, "r": 32, "bands": BANDS["start_off_tile"][2]}, False), + ({"m": 1100, "n": 300, "r": 40, "bands": BANDS["end_on_half_tile"][2]}, True), + ({"m": 1100, "n": 300, "r": 40, "bands": BANDS["end_at_n"][2]}, True), + ({"m": 1100, "n": 300, "r": 40, "bands": BANDS["end_in_tile"][2]}, False), + ({"m": 1100, "n": 128, "r": 256, "bands": [(0, 128, 0, 256)]}, True), + ({"m": 1100, "n": 128, "r": 300, "bands": BANDS["rank_above_max"][2]}, False), + ({"m": 2200, "n": 128, "r": 16, "shape": [2, 1100, K]}, True), + # Batch dims fold into rows: 1200 rows of 600 and of 200 pass the + # row bound, 1000 rows of 500 do not. + ({"m": 1200, "n": 128, "r": 16, "shape": [2, 600, K]}, True), + ({"m": 1200, "n": 128, "r": 16, "shape": [3, 2, 200, K]}, True), + ({"m": 1000, "n": 128, "r": 16, "shape": [2, 500, K]}, False), + ({"m": 1100, "n": 128, "r": 16, "dtype": "float16"}, False), + ({"m": 1100, "n": 128, "r": 16, "codec": "q5_k"}, True), + ({"m": 1100, "n": 96, "r": 16, "codec": "q5_k"}, True), + ({"m": 1100, "n": 128, "r": 16, "codec": "q6_k"}, True), + ({"m": 1100, "n": 96, "r": 16, "codec": "q6_k"}, True), + ({"m": 1100, "n": 128, "r": 16, "codec": "q8_0"}, True), + ({"m": 1100, "n": 96, "r": 16, "codec": "q8_0"}, True), + ({"m": 1100, "n": 128, "r": 16, "codec": "q8_0", "k": 832}, True), + ({"m": 1100, "n": 128, "r": 16, "codec": "q8_0", "k": 800}, False), + ({"m": 1100, "n": 128, "r": 16, "codec": "q6_k", "dtype": "float16"}, False), + ({"m": 1100, "n": 128, "r": 16, "codec": "q3_k"}, False), + ({"m": 1100, "n": 128, "r": 16, "codec": "q4_0"}, False), + ] + res = _probe([c for c, _ in calls], env={"KQ_QMM_LORA": "2"}) + _expect([c for c, _ in calls], res, [k for _, k in calls]) + + +def _wide(k, n, **kw): + """A call on a weight [n, k] made of copies of 64 rows.""" + return {"m": 1100, "n": n, "r": 16, "k": k, "tile": 64, **kw} + + +# Both sides of each bound of the shape rule, one row, one column, one rank +# or one K tile apart. From RULE_ROWS rows: N against K, and K against +# RULE_K. Below RULE_ROWS rows: N against RULE_WIDE_N. The next K above +# RULE_K that a kernel can walk is one K tile of q8_0 above it (64); the +# K-quant codecs hold K in blocks of 256. Then the rank against RULE_RANK +# where N is below RULE_WIDE_N: the rank of a call with no table, and the +# rank of the largest band with one. +_HALVES = [(0, 128, 0, RULE_RANK), (128, 256, RULE_RANK, RULE_RANK)] +_HALVES_OVER = [(0, 128, 0, RULE_RANK + 1), (128, 256, RULE_RANK + 1, 1)] +RULE_CASES = [ + (_wide(256, 256, m=RULE_ROWS), True), + (_wide(256, 256, m=RULE_ROWS - 1), False), + (_wide(256, 257, m=RULE_ROWS), True), + (_wide(256, 255, m=RULE_ROWS), False), + (_wide(RULE_K, RULE_K, m=RULE_ROWS), True), + (_wide(RULE_K + 256, RULE_K + 256, m=RULE_ROWS), False), + (_wide(RULE_K, RULE_K, m=RULE_ROWS, codec="q8_0"), True), + (_wide(RULE_K + 64, RULE_K + 64, m=RULE_ROWS, codec="q8_0"), False), + (_wide(64, 64, m=RULE_ROWS, codec="q8_0"), True), + (_wide(64, 63, m=RULE_ROWS, codec="q8_0"), False), + (_wide(256, RULE_WIDE_N), True), + (_wide(256, RULE_WIDE_N - 1), False), + (_wide(256, RULE_WIDE_N + 1), True), + (_wide(RULE_K, RULE_WIDE_N), True), + (_wide(RULE_K + 256, RULE_WIDE_N), False), + (_wide(RULE_K + 64, RULE_WIDE_N, codec="q8_0"), False), + (_wide(256, 256, m=RULE_ROWS, r=RULE_RANK), True), + (_wide(256, 256, m=RULE_ROWS, r=RULE_RANK + 1), False), + (_wide(256, RULE_WIDE_N - 1, m=RULE_ROWS, r=RULE_RANK + 1), False), + (_wide(256, RULE_WIDE_N, m=RULE_ROWS, r=RULE_RANK + 1), True), + (_wide(256, RULE_WIDE_N, r=MAX_RANK), True), + (_wide(256, 256, m=RULE_ROWS, r=2 * RULE_RANK, bands=_HALVES), True), + (_wide(256, 256, m=RULE_ROWS, r=RULE_RANK + 2, bands=_HALVES_OVER), False), +] + + +def test_lora_shape_rule(): + """With KQ_QMM_LORA unset the op gives the kernel the calls of the + shape rule and builds the composition for every other call.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + calls = [c for c, _ in RULE_CASES] + want = [k for _, k in RULE_CASES] + + def rank(c): + return max(b[3] for b in c["bands"]) if c.get("bands") else c["r"] + + assert want == [_shape_rule(c["k"], c["n"], c["m"], rank(c)) for c in calls] + _expect(calls, _probe(calls), want) + + +def test_lora_levers(): + """KQ_QMM_LORA=0 (or off) builds the composition on every call. 2 runs + the kernel outside the shape rule, inside its other bounds. 1, or a + value that is no mode, is the shape rule. KQ_QMM_LORA_MIN_ROWS moves the + row bound, down to one row of a 1-D x.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + inside, outside = _wide(K, RULE_WIDE_N), _wide(K, 128) + calls = [inside, outside, {**inside, "m": 4608, "r": 64}] + for value in ("0", "off"): + res = _probe(calls, env={"KQ_QMM_LORA": value}) + _expect(calls, res, [False, False, False]) + calls = [ + inside, + outside, + _wide(RULE_K + 256, RULE_K + 256), + {**outside, "m": 1023}, + {**outside, "r": 257}, + {**outside, "dtype": "float16"}, + ] + res = _probe(calls, env={"KQ_QMM_LORA": "2"}) + _expect(calls, res, [True, True, True, False, False, False]) + calls = [inside, outside] + for value in ("1", "3", "-1", "on", "x"): + _expect(calls, _probe(calls, env={"KQ_QMM_LORA": value}), [True, False]) + calls = [ + {"m": 600, "n": 128, "r": 16}, + {"m": 511, "n": 128, "r": 16}, + {"m": 512, "n": 128, "r": 16}, + ] + res = _probe(calls, env={"KQ_QMM_LORA": "2", "KQ_QMM_LORA_MIN_ROWS": "512"}) + _expect(calls, res, [True, False, True]) + calls = [{"m": 1, "n": 128, "r": 16, "shape": [K]}] + res = _probe(calls, env={"KQ_QMM_LORA": "2", "KQ_QMM_LORA_MIN_ROWS": "1"}) + _expect(calls, res, [True]) + # The row bound holds inside the shape rule too. + calls = [{**inside, "m": 1023}, {**inside, "m": 1024}] + _expect(calls, _probe(calls), [False, True]) + + +@pytest.mark.parametrize( + ("value", "text"), + [("0", "at least 1"), ("-5", "at least 1"), ("x", "is not an integer")], + ids=["0", "-5", "x"], +) +def test_lora_min_rows_below_one_keeps_the_bound(value, text): + """KQ_QMM_LORA_MIN_ROWS takes an integer of at least 1. Any other value + keeps the bound of 1024 rows and warns on stderr.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + calls = [ + {"m": 600, "n": 128, "r": 16}, + {"m": MIN_ROWS - 1, "n": 128, "r": 16}, + {"m": MIN_ROWS, "n": 128, "r": 16}, + ] + err = [] + env = {"KQ_QMM_LORA": "2", "KQ_QMM_LORA_MIN_ROWS": value} + _expect(calls, _probe(calls, env=env, stderr=err), [False, False, True]) + warned = [ln for ln in err if "KQ_QMM_LORA_MIN_ROWS" in ln] + assert warned and text in warned[0], err[:5] + assert "keeps its default" in warned[0], warned[0] + + +def test_lora_levers_are_read_once_or_live(): + """The op reads its levers at its first call in a process and keeps + them. Under KQ_TUNE_LIVE=1 it reads them at every call.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + call = {"m": 1100, "n": 128, "r": 16} + small = {"m": 600, "n": 128, "r": 16} + calls = [ + {**call, "env": {"KQ_QMM_LORA": "2"}}, + {**call, "env": {"KQ_QMM_LORA": "0"}}, + {**call, "env": {"KQ_QMM_LORA": "2"}}, + {**small, "env": {"KQ_QMM_LORA_MIN_ROWS": "512"}}, + {**small, "env": {"KQ_QMM_LORA_MIN_ROWS": "1024"}}, + ] + res = _probe(calls, env={"KQ_TUNE_LIVE": "1"}) + _expect(calls, res, [True, False, True, True, False]) + # Without KQ_TUNE_LIVE the first read stays: the kernel of the first + # call runs after the lever is set to 0, and the row bound of the first + # call holds after the lever moves it. + _expect(calls, _probe(calls), [True, True, True, False, False]) + # KQ_TUNE_LIVE itself is read once. + late = [dict(c) for c in calls] + late[1] = {**call, "env": {"KQ_QMM_LORA": "0", "KQ_TUNE_LIVE": "1"}} + _expect(late, _probe(late), [True, True, True, False, False]) diff --git a/tests/test_matmul_resid.py b/tests/test_matmul_resid.py new file mode 100644 index 00000000..5a643659 --- /dev/null +++ b/tests/test_matmul_resid.py @@ -0,0 +1,702 @@ +"""quantized_matmul_residual: y = residual + gate * (x @ W^T). + +On a NAX GPU the op runs one kernel for q4_k, q5_k, q6_k and q8_0 weights +on bfloat16 from 1024 rows: it walks K into f32 sums and stores residual + +gate * sum, rounded once. Everywhere else it builds the sum from +quantized_matmul and stock ops. + +- R1: relative Frobenius error against the float64 sum, at most 6e-3 on + bfloat16 and 1e-3 on float16. +- R2: against the model of the kernel, with the weight rounded to the + activation dtype, exact sums and one rounding of y. The kernel sits at + the error of that one rounding, and each element is within ULPS ulp of + the output dtype at the model plus 2^-20 times the sum of |g x_k w_k|. +- Where the op has no kernel, y is the composition a caller writes, bit + for bit, and every column of it is held against the float64 sum (COL). + +Under KQUANT_FORCE_CPU the op runs the composition on the CPU stream, and +the tests that name a Metal kernel skip. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys + +import mlx.core as mx +import numpy as np +import pytest +from qmm_probe import force_cpu +from test_operand_views import _dense, _strided +from test_weight_base import BASE_ALIGN, _shift + +import mlx_kquant as kq + +R1 = {mx.bfloat16: 6e-3, mx.float16: 1e-3} +# Gate of the composition on the CPU stream, as in test_cpu_neon. +CPU_REL = 5e-2 +COL = {mx.bfloat16: 5e-2, mx.float16: 2e-2} +CPU_COL = 0.2 +K = 768 +# Rows from which the op runs the kernel (kResidMinRows). +MIN_ROWS = 1024 +# Codecs with a kernel on bfloat16, the K tile of each kernel, and the K +# above which the kernel takes 64 rows a tile in place of 128 (the K +# ceiling of the codec's large-row tile in quantized_matmul). +KERNEL_BK = {"q4_k": 128, "q5_k": 128, "q6_k": 128, "q8_0": 64} +MAX_K = {"q4_k": 10240, "q5_k": 0, "q6_k": 0, "q8_0": 10240} +CODECS = list(KERNEL_BK) +GATES = ["none", "n", "batch"] +# R2 per element: the ulp count and the allowance for the f32 sums per unit +# of sum |g x_k w_k|. Bfloat16 holds 8 significant bits, its smallest +# normal has np.frexp exponent -125. +ULPS = 2 +SUM_TOL = 2.0**-20 +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) + +_W = {} + + +def _nax(): + return not force_cpu() and kq.nax_available() + + +def _fused(rows, codec="q4_k", dtype=mx.bfloat16, k=K): + """Where the kernel runs under the default lever.""" + if not _nax() or codec not in KERNEL_BK or dtype != mx.bfloat16: + return False + return k % KERNEL_BK[codec] == 0 and rows >= MIN_ROWS + + +def _f64(a): + return np.array(a.astype(mx.float32)).astype(np.float64) + + +def _rounded(v, dtype): + return _f64(mx.array(v.astype(np.float32)).astype(dtype)) + + +def _rel(a, b): + return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-30)) + + +def _ulps(got, ref, mag): + """The largest error in ulp of bfloat16 at |ref|, after the allowance of + SUM_TOL * mag for the f32 sums, and its row and column.""" + e = np.frexp(np.abs(ref))[1] + e = np.where(ref == 0, -125, np.maximum(e, -125)) + u = (np.abs(got - ref) - SUM_TOL * mag) / np.ldexp(1.0, e - 8) + i = np.unravel_index(int(np.argmax(u)), u.shape) + return float(u[i]), tuple(int(v) for v in i) + + +def _col(a, b): + """The largest error of a column of a against b, over the size of a + mean column of b. Rows are the leading axis.""" + mean = np.linalg.norm(b) / np.sqrt(b.shape[1]) + 1e-30 + return float(np.linalg.norm(a - b, axis=0).max() / mean) + + +def _weight(codec, n, k, seed=0): + """A packed weight [n, k] and its float64 values.""" + key = (codec, n, k, seed) + if key not in _W: + wf = mx.random.normal((n, k), key=mx.random.key(1000 * seed + 17 * n + k)) * 0.1 + w, s = kq.quantize(wf, codec) + mx.eval(w, s) + _W[key] = (w, _f64(kq.dequantize(w, s, codec, dtype=mx.float32))) + return _W[key] + + +def _x(shape, dtype, seed, scale=0.5): + x = (mx.random.normal(shape, key=mx.random.key(seed)) * scale).astype(dtype) + mx.eval(x) + return x + + +def _gate(kind, xshape, n, dtype, seed): + """No gate, a gate [n], or one gate row per batch of x.""" + if kind == "none": + return None + shape = (n,) if kind == "n" else tuple(xshape[:-2]) + (1, n) + g = (mx.random.uniform(shape=shape, key=mx.random.key(seed)) * 3 - 1.5).astype( + dtype + ) + mx.eval(g) + return g + + +def _ref(x, wd, r, g, model=None, mag=False): + """The float64 sum, rows flattened. model=dtype rounds the weight to + that dtype, as the kernel reads it. mag=True gives sum |g x_k w_k| per + element in place of the sum.""" + out = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + n, k = wd.shape + x64 = _f64(x.astype(out)).reshape(-1, k) + if model is not None: + wd = _rounded(wd, model) + if mag: + x64, wd = np.abs(x64), np.abs(wd) + prod = (x64 @ wd.T).reshape(tuple(x.shape[:-1]) + (n,)) + if g is not None: + gd = _f64(g.astype(out)) + prod = (np.abs(gd) if mag else gd) * prod + if mag: + return prod.reshape(-1, n) + return (_f64(r.astype(out)) + prod).reshape(-1, n) + + +def _composed(x, w, r, g, codec="q4_k"): + """The form a caller writes without the op. Where the op has no kernel + it builds this, bit for bit.""" + out = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + y = kq.quantized_matmul(x, w, mx.zeros((1,), dtype=mx.uint8), codec, True) + if g is not None: + y = g.astype(out) * y + return r.astype(out) + y + + +def _check(x, w, wd, r, g, codec="q4_k"): + """R1, and R2 on the kernel. Returns (R1 error, R2 error or None).""" + n, k = wd.shape + dtype = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + y = kq.quantized_matmul_residual(x, w, codec, r, gate=g) + mx.eval(y) + assert y.dtype == dtype + assert y.shape == tuple(x.shape[:-1]) + (n,) + got = _f64(y).reshape(-1, n) + assert np.isfinite(got).all() + ref = _ref(x, wd, r, g) + err = _rel(got, ref) + if not _fused(x.size // k, codec, dtype, k): + comp = _composed(x, w, r, g, codec) + mx.eval(comp) + assert np.array_equal(_f64(comp).reshape(-1, n), got) + assert err <= (R1[dtype] if _nax() else CPU_REL), err + col = _col(got, ref) + assert col <= (COL[dtype] if _nax() else CPU_COL), col + return err, None + assert err <= R1[dtype], (err, R1[dtype]) + assert _col(got, ref) <= COL[dtype] + ref2 = _ref(x, wd, r, g, model=dtype) + floor2 = _rel(_rounded(ref2, dtype), ref2) + err2 = _rel(got, ref2) + assert err2 <= 1.05 * floor2 + 2e-5, (err2, floor2) + mag = _ref(x, wd, r, g, model=dtype, mag=True) + worst = _ulps(got, ref2, mag) + assert worst[0] <= ULPS, worst + return err, err2 + + +def _case(shape, n, gate="batch", dtype=mx.bfloat16, codec="q4_k", seed=0): + k = shape[-1] + w, wd = _weight(codec, n, k) + x = _x(shape, dtype, 7 * shape[-2] + n + seed) + r = _x(tuple(shape[:-1]) + (n,), dtype, 11 * shape[-2] + n + seed, scale=1.0) + g = _gate(gate, shape, n, dtype, 13 * shape[-2] + n + seed) + return _check(x, w, wd, r, g, codec) + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("m", [1, 17, 1023, 1024, 1025, 1039, 1100, 1151, 1152]) +@pytest.mark.parametrize("n", [64, 96, 200]) +@pytest.mark.parametrize("gate", ["none", "n"]) +def test_resid_rows_and_columns(gate, n, m, codec): + # Below MIN_ROWS the op is the composition. From there up: a last row + # tile of 1 row, 15 rows, 76 rows, 127 rows and none, with N on and + # off the column tile. + _case([m, K], n, gate, codec=codec) + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("rows", [500, 1100]) +@pytest.mark.parametrize("gate", GATES) +def test_resid_batch(gate, rows, batch, codec): + # A gate row per batch runs the batch in the grid; a gate [N] or none + # folds the batch into the rows. 2 x 500 rows passes the row bound. + _case([batch, rows, K], 96, gate, codec=codec) + + +def test_resid_batch_dims_of_several_axes(): + _case([3, 2, 200, K], 128, "batch") + _case([3, 2, 200, K], 128, "n") + + +@pytest.mark.parametrize("codec", CODECS) +def test_resid_inner_dim_above_the_large_tile_ceiling(codec): + # q4_k and q8_0 take 64 rows a tile above K 10240, q5_k and q6_k do + # not: both tiles hold the same sums. + _case([2, 1100, 10496], 64, "batch", codec=codec) + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("m", [32768, 32769, 33792]) +def test_resid_rows_above_32767(m, codec): + # Row offsets pass a 16-bit index: every row is compared. + _case([1, m, 256], 64, "batch", codec=codec) + + +def test_resid_batch_rows_above_32767(): + _case([2, 33000, 256], 64, "batch") + + +@pytest.mark.parametrize("codec", ["q3_k", "q4_0", "q5_k", "q8_0"]) +def test_resid_codec_or_dtype_without_kernel_composes(codec): + dtype = mx.float16 if codec in KERNEL_BK else mx.bfloat16 + err, err2 = _case([1100, K], 96, "batch", dtype=dtype, codec=codec) + assert err2 is None + + +def test_resid_q8_0_inner_dim_off_the_k_tile(): + # q8_0 blocks are 32 weights: K 800 is no whole number of K tiles. + err, err2 = _case([1100, 800], 96, "n", codec="q8_0") + assert err2 is None + _, err2 = _case([1100, 832], 96, "n", codec="q8_0") + assert (err2 is not None) == _nax() + + +def test_resid_float32_x_runs_in_bfloat16(): + w, wd = _weight("q4_k", 96, K) + x = _x([1100, K], mx.float32, 3) + r = _x([1100, 96], mx.float32, 5) + g = _gate("n", x.shape, 96, mx.float32, 9) + y = kq.quantized_matmul_residual(x, w, "q4_k", r, gate=g) + want = kq.quantized_matmul_residual( + x.astype(mx.bfloat16), + w, + "q4_k", + r.astype(mx.bfloat16), + gate=g.astype(mx.bfloat16), + ) + mx.eval(y, want) + assert y.dtype == mx.bfloat16 + assert mx.array_equal(y, want).item() + + +def test_resid_one_dim_x(): + w, wd = _weight("q4_k", 96, K) + x = _x([K], mx.bfloat16, 3) + r = _x([96], mx.bfloat16, 5) + g = _gate("n", x.shape, 96, mx.bfloat16, 9) + _check(x, w, wd, r, g) + + +def test_resid_empty_inputs(): + w, _ = _weight("q4_k", 96, K) + y = kq.quantized_matmul_residual( + mx.zeros((0, K), dtype=mx.bfloat16), + w, + "q4_k", + mx.zeros((0, 96), dtype=mx.bfloat16), + ) + mx.eval(y) + assert y.shape == (0, 96) + + +def _check_views(views, call, state): + """call(view) against call(dense copy of the view): the same bits on the + GPU, one rounding apart on the CPU stream (see test_matmul_lora).""" + for name, make in views.items(): + want = call(_dense(make())) + view = make() + if state == "evaluated": + mx.eval(view) + got = call(view) + mx.eval(got, want) + if force_cpu(): + assert _rel(_f64(got), _f64(want)) <= R1[mx.bfloat16], (name, state) + else: + assert mx.array_equal(got, want).item(), (name, state) + + +def _weight_views(w): + n, kb = w.shape + other = mx.array(np.random.default_rng(3).integers(0, 255, (n, kb), dtype=np.uint8)) + rows = mx.stack([w, other], axis=1).reshape(2 * n, kb) + cols = mx.concatenate([other, w], axis=1) + rev = mx.array(np.ascontiguousarray(np.array(w)[::-1])) + mx.eval(rows, cols, rev) + return { + "even_rows": lambda: rows[0::2], + "col_band": lambda: cols[:, kb:], + "reversed_rows": lambda: rev[::-1], + } + + +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +def test_resid_operand_views(state, m, codec): + """A view of any operand gives the bits of its dense copy, lazy or + evaluated, on both sides of the row bound.""" + n = 96 + w, wd = _weight(codec, n, K) + x = _x([m, K], mx.bfloat16, 47) + r = _x([m, n], mx.bfloat16, 49) + g = _x([n], mx.bfloat16, 51) + + def call(x_=x, w_=w, r_=r, g_=g): + return kq.quantized_matmul_residual(x_, w_, codec, r_, gate=g_) + + _check_views( + _strided(_x([2 * m, 2 * K], mx.bfloat16, 53)), lambda v: call(x_=v), state + ) + _check_views( + _strided(_x([2 * m, 2 * n], mx.bfloat16, 59)), lambda v: call(r_=v), state + ) + _check_views(_weight_views(w), lambda v: call(w_=v), state) + g2 = _x([2 * n], mx.bfloat16, 67) + views = {"even": lambda: g2[0::2], "reversed": lambda: g2[::-1][:n]} + _check_views(views, lambda v: call(g_=v), state) + # A transposed x keeps the result right, not only equal to the copy. + xv = _strided(_x([2 * m, 2 * K], mx.bfloat16, 53))["transposed"]() + _check(xv, w, wd, r, g, codec) + + +@pytest.mark.parametrize("m", [1023, 1100]) +@pytest.mark.parametrize("gate", GATES) +def test_resid_gradient_is_the_gradient_of_the_composition(gate, m): + """mx.grad wrt x, residual and gate on both sides of the row bound.""" + n = 96 + w, _ = _weight("q4_k", n, K) + x = _x([2, m // 2 + 1, K], mx.bfloat16, 71) + r = _x([2, m // 2 + 1, n], mx.bfloat16, 73) + g = _gate(gate, x.shape, n, mx.bfloat16, 79) + cot = _x([2, m // 2 + 1, n], mx.float32, 83) + + def op(x_, r_, *g_): + y = kq.quantized_matmul_residual(x_, w, "q4_k", r_, gate=g_[0] if g_ else None) + return (y.astype(mx.float32) * cot).sum() + + def comp(x_, r_, *g_): + y = _composed(x_, w, r_, g_[0] if g_ else None) + return (y.astype(mx.float32) * cot).sum() + + args = (x, r) if g is None else (x, r, g) + nums = tuple(range(len(args))) + got = mx.grad(op, argnums=nums)(*args) + want = mx.grad(comp, argnums=nums)(*args) + mx.eval(got, want) + for a, b in zip(got, want, strict=True): + assert mx.array_equal(a, b).item() + + +def test_resid_jvp_and_vmap_over_residual_and_gate(): + n = 96 + w, _ = _weight("q4_k", n, K) + for m in (1023, 1100): + x = _x([m, K], mx.bfloat16, 91) + r = _x([m, n], mx.bfloat16, 93) + g = _x([n], mx.bfloat16, 95) + t = _x([m, n], mx.bfloat16, 97) + + def f(r_, g_, x_=x): + return kq.quantized_matmul_residual(x_, w, "q4_k", r_, gate=g_) + + def fc(r_, g_, x_=x): + return _composed(x_, w, r_, g_) + + (got,), (want,) = ( + mx.jvp(f, [r, g], [t, mx.zeros_like(g)])[1], + mx.jvp(fc, [r, g], [t, mx.zeros_like(g)])[1], + ) + mx.eval(got, want) + assert mx.array_equal(got, want).item() + rs = mx.stack([r, r * 0.5]) + got = mx.vmap(lambda r_, g_=g: f(r_, g_), in_axes=0)(rs) + want = mx.stack([f(r, g), f(r * 0.5, g)]) + mx.eval(got, want) + assert _rel(_f64(got), _f64(want)) <= R1[mx.bfloat16] + + +def test_resid_weight_is_frozen(): + n = 96 + w, _ = _weight("q4_k", n, K) + for m in (1023, 1100): + x = _x([m, K], mx.bfloat16, 107) + r = _x([m, n], mx.bfloat16, 109) + + def loss(w_, x_=x, r_=r): + return ( + kq.quantized_matmul_residual(x_, w_, "q4_k", r_) + .astype(mx.float32) + .sum() + ) + + op = "quantized_matmul_residual" if _fused(m) else "quantized_matmul" + text = rf"\[mlx_kquant\] {op} vjp: no gradient wrt the quantized weight" + with pytest.raises(ValueError, match=text): + mx.eval(mx.grad(loss)(w)) + + +def test_resid_compiles_and_repeats(): + n = 96 + w, _ = _weight("q5_k", n, K) + x = _x([1100, K], mx.bfloat16, 111) + r = _x([1100, n], mx.bfloat16, 113) + g = _x([n], mx.bfloat16, 117) + f = mx.compile( + lambda x_, r_, g_: kq.quantized_matmul_residual(x_, w, "q5_k", r_, gate=g_) + ) + want = kq.quantized_matmul_residual(x, w, "q5_k", r, gate=g) + for _ in range(2): + got = f(x, r, g) + mx.eval(got, want) + assert mx.array_equal(got, want).item() + + +@pytest.mark.parametrize("codec", CODECS) +def test_resid_weight_off_its_byte_boundary_raises(codec): + if force_cpu() or mx.default_device() != mx.gpu: + pytest.skip("the start check guards the Metal kernels") + n = 96 + w, _ = _weight(codec, n, K) + align = BASE_ALIGN.get(codec, 2) + for m in (1023, 1100): + x = _x([m, K], mx.bfloat16, 191) + r = _x([m, n], mx.bfloat16, 193) + op = "quantized_matmul_residual" if _fused(m, codec) else "quantized_matmul" + text = ( + rf"\[mlx_kquant\] {op}: the {codec} weight starts 1 byte past a " + rf"{align}-byte boundary" + ) + with pytest.raises(ValueError, match=text): + kq.quantized_matmul_residual(x, _shift(w, 1), codec, r) + y = kq.quantized_matmul_residual(x, _shift(w, 1, evaluate=False), codec, r) + with pytest.raises(ValueError, match=text): + mx.eval(y) + good = kq.quantized_matmul_residual(x, _shift(w, 16), codec, r) + ref = kq.quantized_matmul_residual(x, w, codec, r) + mx.eval(good, ref) + assert mx.array_equal(good, ref).item() + + +def test_resid_argument_errors(): + n = 96 + w, _ = _weight("q4_k", n, K) + x = _x([2, 4, K], mx.bfloat16, 89) + r = _x([2, 4, n], mx.bfloat16, 90) + + def bad(match, x_=x, w_=w, r_=r, codec="q4_k", **kw): + with pytest.raises(ValueError, match=r"quantized_matmul_residual\] .*" + match): + kq.quantized_matmul_residual(x_, w_, codec, r_, **kw) + + bad("Unknown kquant_type", codec="q9_z") + bad("2-D uint8", w_=w.astype(mx.float32)) + bad("whole number", w_=w[:, :-1]) + bad("x last dim", x_=x[..., :-1]) + bad("floating x", x_=x.astype(mx.int32)) + bad("residual must be", r_=r[..., :-1]) + bad("residual must be", r_=r[:1]) + bad("residual must be", r_=r.reshape(8, n)) + bad("residual must be", r_=r.astype(mx.int32)) + for shape in ((n + 1,), (1, n), (2, 4, n), (2, 1, 1, n), (1, 1, n), (2, 1, n - 1)): + bad("gate must be", gate=mx.ones(shape)) + bad("gate must be", gate=mx.ones((n,), dtype=mx.int32)) + for shape in ((n,), (2, 1, n)): + kq.quantized_matmul_residual(x, w, "q4_k", r, gate=mx.ones(shape)) + + +def test_resid_size_limits(): + """K or N of 2^25 and a row count of 2^31 raise by the name of the op. + The arrays are lazy and stay so.""" + cols, rows = 1 << 25, 1 << 31 + + def lazy(*shape, dtype=mx.bfloat16): + return mx.zeros(shape, dtype=dtype) + + def w8(n, k): + return lazy(n, k // 32 * 34, dtype=mx.uint8) + + cases = { + "K": (lazy(cols), w8(1, cols), lazy(1)), + "N": (lazy(4, 32), w8(cols, 32), lazy(4, cols)), + } + for name, (x, w, r) in cases.items(): + text = ( + rf"quantized_matmul_residual\] K and N must each be below {cols} " + rf".*{name} {cols}" + ) + with pytest.raises(ValueError, match=text): + kq.quantized_matmul_residual(x, w, "q8_0", r) + x = lazy(2, rows // 2, 32) + text = rf"quantized_matmul_residual\] x must hold fewer than {rows} \(2\^31\) rows" + with pytest.raises(ValueError, match=text): + kq.quantized_matmul_residual(x, w8(64, 32), "q8_0", lazy(2, rows // 2, 64)) + + +def _name(codec, n, k, gated): + max_k = MAX_K[codec] + bm = 64 if max_k and k > max_k else 128 + al = "true" if n % 64 == 0 else "false" + gt = "true" if gated else "false" + return ( + f"kquant_{codec}_qmm_t_nax_resid_bfloat16_t_bm{bm}_bn64_bk{KERNEL_BK[codec]}" + f"_wm2_wn2_alN_{al}_gate_{gt}" + ) + + +def test_resid_kernels_in_metallib(): + if force_cpu(): + pytest.skip("Metal kernel names only") + names = kq.metallib_kernels() + if not names: + pytest.skip("no metallib loaded") + got = {n for n in names if "_qmm_t_nax_resid_" in n} + want = { + _name(c, n, k, gt) + for c in CODECS + for n in (64, 96) + for k in (256, 10496) + for gt in (True, False) + } + assert got == want, sorted(got ^ want) + + +_PROBE = r""" +import json, os, sys +spec = json.loads(sys.argv[1]) +sys.path.insert(0, spec["tests_dir"]) +import mlx.core as mx +import mlx_kquant as kq +import test_matmul_resid as t + +out = [] +for i, c in enumerate(spec["calls"]): + codec = c.get("codec", "q4_k") + shape = c["shape"] + n = c["n"] + w, wd = t._weight(codec, n, shape[-1]) + x = t._x(shape, mx.bfloat16, 7 * shape[-2] + n) + r = t._x(shape[:-1] + [n], mx.bfloat16, 11 * shape[-2] + n, scale=1.0) + g = t._gate(c.get("gate", "batch"), shape, n, mx.bfloat16, 13 * shape[-2] + n) + os.environ.update(c.get("env", {})) + print(f"[probe] {i}", file=sys.stderr, flush=True) + res = {"error": None} + try: + y = kq.quantized_matmul_residual(x, w, codec, r, gate=g) + mx.eval(y) + got = t._f64(y).reshape(-1, n) + ref = t._ref(x, wd, r, g) + res["rel"] = t._rel(got, ref) + res["col"] = t._col(got, ref) + ref2 = t._ref(x, wd, r, g, model=mx.bfloat16) + res["rel2"] = t._rel(got, ref2) + res["floor2"] = t._rel(t._rounded(ref2, mx.bfloat16), ref2) + mag = t._ref(x, wd, r, g, model=mx.bfloat16, mag=True) + res["ulps"] = t._ulps(got, ref2, mag) + except Exception as e: + res["error"] = f"{type(e).__name__}: {e}" + out.append(res) +print("RESULT " + json.dumps(out)) +""" + + +def _probe(calls, env=None): + """Run calls in a fresh interpreter that keeps no KQ_ variable but + those of env and of each call. Per call: resid (the residual kernels + that KQ_QMM_LOG named), plain (the other kernels), error and errors.""" + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update({"KQ_QMM_LOG": "1", **(env or {})}) + spec = {"tests_dir": TESTS_DIR, "calls": calls} + p = subprocess.run( + [sys.executable, "-c", _PROBE, json.dumps(spec)], + env=child_env, + capture_output=True, + text=True, + timeout=900, + ) + assert p.returncode == 0, p.stderr[-3000:] + line = [ln for ln in p.stdout.splitlines() if ln.startswith("RESULT ")][-1] + results = json.loads(line[len("RESULT ") :]) + cur = None + for ln in p.stderr.splitlines(): + if ln.startswith("[probe] "): + cur = results[int(ln.split()[1])] + cur.update(resid=[], plain=[]) + elif cur is not None and ln.startswith("[kq qmm] "): + name = ln.split()[2] + cur["resid" if "_resid_" in name else "plain"].append(name) + return results + + +def _expect(calls, results, kernel): + for c, r, k in zip(calls, results, kernel, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] <= R1[mx.bfloat16], (c, r["rel"]) + assert r["col"] <= COL[mx.bfloat16], (c, r["col"]) + if k: + want = _name( + c.get("codec", "q4_k"), c["n"], c["shape"][-1], c.get("gate") != "none" + ) + assert r["resid"] == [want], (c, r["resid"]) + assert r["plain"] == [], (c, r["plain"]) + assert r["rel2"] <= 1.05 * r["floor2"] + 2e-5, (c, r["rel2"], r["floor2"]) + assert r["ulps"][0] <= ULPS, (c, r["ulps"]) + else: + assert r["resid"] == [], (c, r["resid"]) + assert len(r["plain"]) == 1, (c, r["plain"]) + + +def test_resid_known_answer_of_every_kernel(): + """Each kernel of the metallib runs once, named by KQ_QMM_LOG, against + the float64 sum and the one-rounding model: every codec, both row + tiles, N on and off the column tile, with and without a gate.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + calls = [] + for codec in CODECS: + # K above the ceiling takes the BM 64 kernels, where a codec has one. + for k in (256, 10496) if MAX_K[codec] else (256,): + for n in (64, 96): + for gate in ("batch", "none"): + calls.append( + {"codec": codec, "shape": [1, 1100, k], "n": n, "gate": gate} + ) + res = _probe(calls) + _expect(calls, res, [True] * len(calls)) + names = {r["resid"][0] for r in res} + assert len(names) == len(calls) + assert names == {n for n in kq.metallib_kernels() if "_qmm_t_nax_resid_" in n} + + +def test_resid_kernel_that_runs_and_levers(): + """KQ_QMM_LOG names the kernel of each call: from 1024 rows by default, + at every row count under KQ_QMM_RESID=2, never under 0. Batch dims fold + into rows for the row bound.""" + if not _nax(): + pytest.skip("the kernel needs a NAX GPU") + calls = [ + ({"shape": [1024, K], "n": 128, "gate": "n"}, True), + ({"shape": [1023, K], "n": 128, "gate": "n"}, False), + ({"shape": [2, 512, K], "n": 128}, True), + ({"shape": [2, 511, K], "n": 128}, False), + ({"shape": [1100, 800], "n": 128, "codec": "q8_0", "gate": "n"}, False), + ({"shape": [17, K], "n": 96, "gate": "n", "env": {"KQ_QMM_RESID": "2"}}, True), + ( + {"shape": [1100, K], "n": 96, "gate": "n", "env": {"KQ_QMM_RESID": "0"}}, + False, + ), + ( + {"shape": [1100, K], "n": 96, "gate": "n", "env": {"KQ_QMM_RESID": "1"}}, + True, + ), + ] + res = _probe([c for c, _ in calls], env={"KQ_TUNE_LIVE": "1"}) + _expect([c for c, _ in calls], res, [k for _, k in calls]) + # Without KQ_TUNE_LIVE the lever is read once, at the first call. + calls = [ + ( + {"shape": [1100, K], "n": 96, "gate": "n", "env": {"KQ_QMM_RESID": "0"}}, + False, + ), + ( + {"shape": [1100, K], "n": 96, "gate": "n", "env": {"KQ_QMM_RESID": "1"}}, + False, + ), + ] + res = _probe([c for c, _ in calls]) + _expect([c for c, _ in calls], res, [k for _, k in calls]) diff --git a/tests/test_mlx_pin.py b/tests/test_mlx_pin.py new file mode 100644 index 00000000..462e3b6d --- /dev/null +++ b/tests/test_mlx_pin.py @@ -0,0 +1,59 @@ +"""Pinned-mlx behavior this package relies on. + +mlx < 0.32.2 encoded a strided quantized operand's contiguity copy after +binding earlier operands (ml-explore/mlx#4381): dequantize read the copy +kernel's buffers whenever biases arrived strided. Fixed at the 0.32.2 pin; +these tests hold the fix so a venv drifting off the pin fails loudly. +""" + +import itertools + +import mlx.core as mx +import pytest + +GROUP, BITS = 64, 8 + + +def _strided(t): + # row-prefix slice of a taller buffer: lazy, strided when B*heads > 1 + pad = mx.concatenate([t, mx.zeros_like(t)], axis=-2) + return pad[..., : t.shape[-2], :] + + +def _quantized(shape): + mx.random.seed(3) + x = mx.random.normal(shape, dtype=mx.float16) + return x, mx.quantize(x, group_size=GROUP, bits=BITS) + + +def test_mlx_pin_version(): + assert mx.__version__ == "0.32.3", ( + f"mlx {mx.__version__}: package is certified against the 0.32.3 pin" + ) + + +@pytest.mark.parametrize( + "sw,ss,sb", + [c for c in itertools.product([False, True], repeat=3) if any(c)], +) +def test_dequantize_strided_operands(sw, ss, sb): + x, (w, scales, biases) = _quantized((2, 4, 96, 128)) + ref = mx.dequantize(w, scales, biases, group_size=GROUP, bits=BITS) + got = mx.dequantize( + _strided(w) if sw else w, + _strided(scales) if ss else scales, + _strided(biases) if sb else biases, + group_size=GROUP, + bits=BITS, + ) + err = float(mx.abs(got - ref).max().item()) + # fp16 dequant is exact vs itself; any operand misread lands ~1e3 off + assert err == 0.0, f"strided (w={sw},s={ss},b={sb}) diverged: {err}" + + +def test_quantize_strided_input(): + x, (w, scales, biases) = _quantized((2, 4, 96, 128)) + w2, s2, b2 = mx.quantize(_strided(x), group_size=GROUP, bits=BITS) + assert bool(mx.all(w2 == w).item()) + assert bool(mx.all(s2 == scales).item()) + assert bool(mx.all(b2 == biases).item()) diff --git a/tests/test_moe_glu.py b/tests/test_moe_glu.py index 1d15ff26..e4cad94a 100644 --- a/tests/test_moe_glu.py +++ b/tests/test_moe_glu.py @@ -156,9 +156,12 @@ def _rel(got, ref): return float(np.linalg.norm(g - ref) / (np.linalg.norm(ref) + 1e-6)) -def _check_codec(codec, sx=None, act="silu", dtype=mx.float16, k=K, limit=0.0): - """Run all five gather ops for one (expert codec, shexp codec) pair. - Returns list of (name, rel, ok).""" +def _check_codec( + codec, sx=None, act="silu", dtype=mx.float16, k=K, limit=0.0, x_scale=0.1 +): + """Run all five gather ops for one (expert codec, shexp codec) pair on + activations of standard deviation x_scale. Returns list of (name, rel, + ok).""" scodec = sx or codec wire, ref = _wire_and_ref(codec, k=k) if wire is None: @@ -175,7 +178,7 @@ def _check_codec(codec, sx=None, act="silu", dtype=mx.float16, k=K, limit=0.0): sgw, suw, sdw = mx.array(swire[2]), mx.array(swire[3]), mx.array(swire[1]) sg_ref, su_ref, sd_ref = sref[2], sref[3], sref[1] - x_np = (rng.standard_normal((T, k)) * 0.1).astype(np.float16) + x_np = (rng.standard_normal((T, k)) * x_scale).astype(np.float16) x = mx.array(x_np).astype(dtype) xf = np.array(x.astype(mx.float32)) # exactly what the kernel reads inds_np = rng.integers(0, E, size=(T, R)).astype(np.uint32) @@ -206,7 +209,7 @@ def _check_codec(codec, sx=None, act="silu", dtype=mx.float16, k=K, limit=0.0): out.append(("glu", rel, rel < REL_BOUND)) # gather_qmv_kq - h_np = (rng.standard_normal((T, R, k)) * 0.1).astype(np.float16) + h_np = (rng.standard_normal((T, R, k)) * x_scale).astype(np.float16) h = mx.array(h_np).astype(dtype) hf = np.array(h.astype(mx.float32)) got = kq.gather_qmv_kq(h, dw, codec, inds) @@ -259,7 +262,7 @@ def _check_codec(codec, sx=None, act="silu", dtype=mx.float16, k=K, limit=0.0): # gather_qmv_mix_kq (shexp codec = scodec) S = R + 1 - hs_np = (rng.standard_normal((T, S, k)) * 0.1).astype(np.float16) + hs_np = (rng.standard_normal((T, S, k)) * x_scale).astype(np.float16) hs = mx.array(hs_np).astype(dtype) hsf = np.array(hs.astype(mx.float32)) sc_np = rng.uniform(0.05, 0.9, size=(T, S)).astype(np.float32) @@ -792,6 +795,100 @@ def test_swiglu_clamp_validation(): kq.gather_qmv_kq(h, wq8, "q8_0", inds, bias=b) +def _nan_rows(got): + return np.isnan(np.array(got.astype(mx.float32))) + + +@pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16], ids=["f16", "bf16"]) +def test_limited_glu_passes_a_nan(dtype): + """The limited activations clamp the gate from above and the up value + on both sides. A NaN gate or up value gives NaN, as mx.minimum and + mx.clip give, and never the value of a bound. A NaN in x reaches both + values, and a NaN block scale in the up weight or a NaN up bias reaches + the up value alone.""" + rng = np.random.default_rng(3) + inds_np = np.array([[0, 1], [2, 3], [1, 0]], dtype=np.uint32) + inds = mx.array(inds_np) + x_np = (rng.standard_normal((T, K)) * 0.1).astype(np.float16) + x = mx.array(x_np).astype(dtype) + x_bad_np = x_np.copy() + x_bad_np[1, 7] = np.nan + x_bad = mx.array(x_bad_np).astype(dtype) + row = 5 + + def expect_up(e): + # NaN where a slot routes to expert e, at output row `row`. + want = np.zeros((T, R, N), dtype=bool) + want[:, :, row] = inds_np == e + return want + + # silu_limit on q8_0 wire: block 0 of up row [2, row] gets a NaN scale. + wire, _ = _wire_and_ref("q8_0") + up_np = wire[::-1].copy() + up_bad = up_np.copy() + up_bad[2, row, 0:2] = np.frombuffer(np.float16(np.nan).tobytes(), np.uint8) + gw, uw, uwb = mx.array(wire), mx.array(up_np), mx.array(up_bad) + for shexp in (False, True): + + def call(xx, up, shexp=shexp): + if shexp: + return kq.moe_glu_gather_shexp_kq( + xx, gw, up, gw[0], uw[0], "q8_0", inds, act="silu_limit", limit=0.05 + )[:, :R] + return kq.moe_glu_gather_kq( + xx, gw, up, "q8_0", inds, act="silu_limit", limit=0.05 + ) + + clean = call(x, uw) + got_x = call(x_bad, uw) + got_w = call(x, uwb) + mx.eval(clean, got_x, got_w) + assert not _nan_rows(clean).any() + nx, nw = _nan_rows(got_x), _nan_rows(got_w) + assert nx[1].all() and not nx[[0, 2]].any(), shexp + assert (nw == expect_up(2)).all(), shexp + keep = mx.array(~expect_up(2)) + assert mx.array_equal(mx.where(keep, got_w, 0), mx.where(keep, clean, 0)).item() + + # swiglu_clamp with biases on mxfp4 wire, and the packed mxfp4 kernel. + wire, _ = _wire_and_ref("mxfp4") + up_np = wire[::-1].copy() + gw, uw = mx.array(wire), mx.array(up_np) + gb_np = rng.uniform(-0.2, 0.2, (E, N)).astype(np.float32) + ub_np = rng.uniform(-0.2, 0.2, (E, N)).astype(np.float32) + bad_np = ub_np.copy() + bad_np[1, row] = np.nan + gb, ub, ub_bad = mx.array(gb_np), mx.array(ub_np), mx.array(bad_np) + pgw, pgs = (mx.array(a) for a in _mxfp4_deinterleave_np(wire)) + puw, pus = (mx.array(a) for a in _mxfp4_deinterleave_np(up_np)) + calls = { + "wire": lambda xx, b: kq.moe_glu_gather_kq( + xx, + gw, + uw, + "mxfp4", + inds, + act="swiglu_clamp", + limit=0.05, + gate_bias=gb, + up_bias=b, + alpha=1.702, + ), + "packed": lambda xx, b: kq.moe_glu_gather( + xx, pgw, pgs, gb, puw, pus, b, inds, 1.702, 0.05 + ), + } + for name, call in calls.items(): + clean = call(x, ub) + got_x = call(x_bad, ub) + got_b = call(x, ub_bad) + mx.eval(clean, got_x, got_b) + assert not _nan_rows(clean).any(), name + nx, nb = _nan_rows(got_x), _nan_rows(got_b) + assert nx[1].all() and not nx[[0, 2]].any(), name + assert (nb == expect_up(1)).all(), name + + @pytest.mark.parametrize("scoring", ["softmax", "sqrtsoftplus", "sigmoid"]) @pytest.mark.parametrize("e", [16, 64, 250, 512, 1000, 1024]) def test_router_topk_exact_order(scoring, e): @@ -896,6 +993,47 @@ def test_shexp_gate_and_up_in_two_codecs(codec, sg, su): assert _rel(got, want) < REL_BOUND, (dtype, t, r) +@pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16]) +def test_shexp_two_codecs_takes_x_views(dtype): + """The split-codec shared expert reads a strided and a transposed x as + it reads the same values in a dense array.""" + codec, sg, su = "iq2_s", "iq4_xs", "iq3_s" + wire, _ = _wire_and_ref(codec) + gwire, _ = _wire_and_ref(sg, seed=13) + uwire, _ = _wire_and_ref(su, seed=17) + if wire is None or gwire is None or uwire is None: + pytest.skip("K-quant fixture missing") + dw, up_stack = mx.array(wire), mx.array(wire[::-1].copy()) + rng = np.random.default_rng(5) + t, r = 4, 3 + wide = mx.array((rng.standard_normal((t, 2 * K)) * 0.1).astype(np.float16)) + wide = wide.astype(dtype) + inds = mx.array(rng.integers(0, E, size=(t, r)).astype(np.uint32)) + dense = mx.contiguous(wide[:, ::2]) + turned = mx.contiguous(dense.T) + mx.eval(wide, dense, turned) + + def run(x): + y = kq.moe_glu_gather_shexp_kq( + x, + dw, + up_stack, + mx.array(gwire[2]), + mx.array(uwire[3]), + codec, + inds, + shexp_kquant_type=sg, + shexp_up_kquant_type=su, + ) + mx.eval(y) + return y + + want = run(dense) + assert float(mx.abs(want.astype(mx.float32)).max()) > 0 + assert mx.array_equal(run(wide[:, ::2]), want).item() + assert mx.array_equal(run(turned.T), want).item() + + def test_shexp_split_codec_rejects_what_has_no_kernel(): x = mx.zeros((1, K), dtype=mx.float16) inds = mx.zeros((1, 2), dtype=mx.uint32) diff --git a/tests/test_moe_glu_variants.py b/tests/test_moe_glu_variants.py index a47da2a4..039be01f 100644 --- a/tests/test_moe_glu_variants.py +++ b/tests/test_moe_glu_variants.py @@ -182,30 +182,36 @@ def test_sg_bit_identical(codec, sg, tmp_path): import test_moe_glu as tm codec, sx, dtype = sys.argv[2], sys.argv[3] or None, getattr(mx, sys.argv[4]) -res = tm._check_codec(codec, sx=sx, dtype=dtype) +res = tm._check_codec(codec, sx=sx, dtype=dtype, x_scale=float(sys.argv[5])) print("RESULT " + json.dumps( None if res is None else [(n, float(r), bool(o)) for n, r, o in res])) """ -# Half dots round each 4-wide product sum to half before the float chunk -# accumulate; the fused outputs land near 1e-3 relative to the f32 -# reference against the 2e-3 bound of the float kernels. +# Half dots stage the grid and the activations in half and sum each +# chunk's products in float; the fused outputs land near 2e-4 relative to +# the f32 reference on float16 and near 2e-3 on bfloat16, where the float +# kernels are bound at 2e-3. HALF_REL_BOUND = 6e-3 +HALF_OPS = ("glu", "qmv", "shexp", "mix") -@pytest.mark.parametrize("codec", ["iq2_xs", "iq2_xxs", "iq3_xxs"]) -@pytest.mark.parametrize("sx", ["", "q6_k"]) -@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) -def test_half_dot_dispatch_and_accuracy(codec, sx, dtype): - """KQ_MOE_HALF=1 routes the grid codecs' glu / qmv / shexp / mix gathers - to the "_h" kernels (mix_ns has no half form and stays on the float - kernel) and every op stays within the half-dot bound of the f32 - reference.""" +def _half_run(codec, sx, dtype, x_scale): + """KQ_MOE_HALF=1 results of test_moe_glu._check_codec in a child, and + the kernel names it logged.""" tests_dir = os.path.dirname(os.path.abspath(__file__)) r = _run_child( - [sys.executable, "-c", _HALF_SNIPPET, tests_dir, codec, sx, dtype], + [ + sys.executable, + "-c", + _HALF_SNIPPET, + tests_dir, + codec, + sx, + dtype, + str(x_scale), + ], {"KQ_MOE_HALF": "1", "KQ_MOE_NX_LOG": "1"}, - f"half {codec} sx={sx} {dtype}", + f"half {codec} sx={sx} {dtype} x_scale={x_scale}", ) line = [l for l in r.stdout.splitlines() if l.startswith("RESULT ")] assert line, r.stdout + r.stderr @@ -214,7 +220,18 @@ def test_half_dot_dispatch_and_accuracy(codec, sx, dtype): res = json.loads(line[-1][len("RESULT ") :]) if res is None: pytest.skip(f"{sx or codec} fixture missing") - names = r.stderr + return res, r.stderr + + +@pytest.mark.parametrize("codec", ["iq2_xs", "iq2_xxs", "iq3_xxs"]) +@pytest.mark.parametrize("sx", ["", "q6_k"]) +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +def test_half_dot_dispatch_and_accuracy(codec, sx, dtype): + """KQ_MOE_HALF=1 routes the grid codecs' glu / qmv / shexp / mix gathers + to the "_h" kernels (mix_ns has no half form and stays on the float + kernel) and every op stays within the half-dot bound of the f32 + reference.""" + res, names = _half_run(codec, sx, dtype, 0.1) stem = codec if not sx else f"{codec}_sx_{sx}" for op in ( "_moe_glu_gather_h_", @@ -230,11 +247,26 @@ def test_half_dot_dispatch_and_accuracy(codec, sx, dtype): assert want in names, (want, names) assert f"kq_{codec}_gather_qmv_mix_ns" in names and "mix_ns_h" not in names, names for name, rel, ok in res: - bound = HALF_REL_BOUND if name in ("glu", "qmv", "shexp", "mix") else None - if bound is None: + if name in HALF_OPS: + assert rel < HALF_REL_BOUND, (name, rel) + else: assert ok, (name, rel) + + +@pytest.mark.parametrize("codec", ["iq2_xs", "iq2_xxs", "iq3_xxs"]) +def test_half_dot_range(codec): + """Activations of standard deviation 2e3 stay finite under + KQ_MOE_HALF=1: a grid code (up to 62) times such an activation passes + 65504, so each product and the chunk sum run in float. bfloat16 keeps + the GLU outputs in range.""" + res, names = _half_run(codec, "", "bfloat16", 2e3) + for op in ("_moe_glu_gather_h_", "_gather_qmv_h"): + assert f"kq_{codec}{op}" in names, (op, names) + for name, rel, ok in res: + if name in HALF_OPS: + assert rel < HALF_REL_BOUND, (name, rel) else: - assert rel < bound, (name, rel) + assert ok, (name, rel) _DD_SNIPPET = r""" diff --git a/tests/test_nax_bk128.py b/tests/test_nax_bk128.py new file mode 100644 index 00000000..a1c1a6bb --- /dev/null +++ b/tests/test_nax_bk128.py @@ -0,0 +1,296 @@ +"""NAX qmm_t on a K tile of 128 (the half super-block and pair loaders). + +On NAX GPUs the BM=64 and BM=128 tiles of the codecs in CODECS dispatch +the ``_bk128`` kernels by default, the codec's ``KQ_NAX__BK128=0`` +keeps the ``_bk64`` ones, and the two give bit-identical outputs. A 32-weight +block codec takes them only at K % 128 == 0. The other codecs carry no +``_bk128`` kernels. The BM=32 tile and the ``_db`` variant keep BK=64. +Each case runs in a fresh interpreter (tests/qmm_probe.py) and checks the +result against float64 with exact kq.dequantize weights, on the CPU stream +too, where the lever leaves the result alone. +""" + +from __future__ import annotations + +import pytest +from kqref import is_synth +from qmm_probe import force_cpu, run_probe, tile_of + +import mlx_kquant as kq + +R1 = {"bfloat16": 6e-3, "float16": 1e-3} +LIVE = {"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"} +# Three super-blocks, so the loader walks two halves per block across +# block boundaries. +K = 768 +# codec: (group size, bits) +CODECS = { + "q6_k": (256, 6), + "q5_k": (256, 5), + "q3_k": (256, 3), + "q4_k": (256, 4), + "q4_0": (32, 4), + "q4_1": (32, 4), + "q5_0": (32, 5), + "iq4_nl": (32, 4), + "iq3_xxs": (256, 3), + "iq3_s": (256, 3), + "iq2_xs": (256, 2), + "iq2_s": (256, 2), +} +# Codecs whose BM=64 and BM=128 tiles take the K tile of 128 unless the +# lever is 0: all of them. +DEFAULT_ON = set(CODECS) +# Codecs without a BM=32 tile, whose M 16 runs the BM=64 one. +NO_BM32 = {"iq2_xs", "iq2_s"} + + +def _lever(codec): + if codec.endswith("_k"): + return "KQ_NAX_" + codec.replace("_", "").upper() + "_BK128" + return "KQ_NAX_" + codec.upper() + "_BK128" + + +def _designs(codec): + # "lead" quantizes real values, which the synthetic-wire codecs skip. + return ("normal",) if is_synth(codec) else ("normal", "lead") + + +def _nax_gpu(): + return not force_cpu() and kq.nax_available() + + +def _kernel(res): + assert res["error"] is None, res["error"] + assert len(res["kernels"]) == 1, res["kernels"] + return res["kernels"][0] + + +def _bk(name): + return 128 if "_bk128_" in name else 64 + + +def test_bk128_kernels_in_metallib(): + if force_cpu(): + pytest.skip("Metal kernel names only") + names = kq.metallib_kernels() + if not names: + pytest.skip("no metallib loaded") + # Every qmm_t tile kernel on a K tile of 128, of any codec, so the other + # codecs carry none. test_matmul_lora.py holds the low-rank names and + # test_matmul_resid.py the residual ones. + got = sorted( + n + for n in names + if n.startswith("kquant_") + and "_qmm_t_nax_" in n + and "_bk128_" in n + and "_qmm_t_nax_lora_" not in n + and "_qmm_t_nax_resid_" not in n + ) + want = sorted( + f"kquant_{c}_qmm_t_nax_{t}_gs_{gs}_b_{b}_bm{bm}_bn64_bk128_wm2_wn2" + f"_alN_{al}_batch_{bt}" + for c, (gs, b) in CODECS.items() + for t in ("float16_t", "bfloat16_t") + for bm in (64, 128) + for al in ("true", "false") + for bt in (0, 1) + ) + assert got == want + + +@pytest.mark.parametrize("codec", sorted(CODECS)) +def test_bk128_dispatch(codec): + if not _nax_gpu(): + pytest.skip("NAX tile policy") + lever = _lever(codec) + on = "1" + # The probe keeps each call's env for the calls after it. + # KQ_QMM_MIDM=0 keeps M 48 and M 80 on the qmm tile. + base = {"codec": codec, "n": 128, "env": {"KQ_QMM_MIDM": "0"}} + cases = [ + # (call, bm, bk, db) + ( + dict(base, shape=[4096, K], env={"KQ_QMM_MIDM": "0", lever: on}), + 128, + 128, + False, + ), + (dict(base, shape=[4128, K]), 128, 128, False), + (dict(base, shape=[160, K]), 64, 128, False), + (dict(base, shape=[80, K]), 64, 128, False), + (dict(base, shape=[48, K]), 64, 128, False), + # KQ_QMM_ROUTE=nax sends M 16 to the BM=32 tile. + ( + dict(base, shape=[16, K], env={"KQ_QMM_ROUTE": "nax"}), + 64 if codec in NO_BM32 else 32, + 128 if codec in NO_BM32 else 64, + False, + ), + ( + dict( + base, + shape=[4096, K], + env={"KQ_QMM_ROUTE": None, "KQ_NAX_TILE": "64,64,2,2"}, + ), + 64, + 128, + False, + ), + ( + dict(base, shape=[4096, K], env={"KQ_NAX_TILE": None, lever: "0"}), + 128, + 64, + False, + ), + (dict(base, shape=[160, K]), 64, 64, False), + ( + dict(base, shape=[4096, K], env={lever: None}), + 128, + 128 if codec in DEFAULT_ON else 64, + False, + ), + ] + if codec == "q6_k": + # The M 33-64 _db band from N 16384. + cases.append( + (dict(base, shape=[48, 256], n=16384, env={lever: on}), 64, 64, True) + ) + calls = [c for c, *_ in cases] + res = run_probe(calls, env=LIVE, check=False) + for (c, bm, bk, db), r in zip(cases, res, strict=True): + name, _ = _kernel(r) + assert tile_of(name)[0] == bm, (c, name) + assert _bk(name) == bk, (c, name) + assert name.endswith("_db") == db, (c, name) + + +@pytest.mark.parametrize("codec", sorted(CODECS)) +def test_bk128_lever_keeps_first_read(codec): + if not _nax_gpu(): + pytest.skip("NAX tile policy") + lever = _lever(codec) + call = {"codec": codec, "shape": [4096, K], "n": 128} + for first, want in (("0", 64), ("1", 128)): + flip = "1" if first == "0" else "0" + calls = [dict(call), dict(call, env={lever: flip})] + res = run_probe(calls, env={"KQ_QMM_LOG": "1", lever: first}, check=False) + assert [_bk(_kernel(r)[0]) for r in res] == [want, want] + + +# The 2-bit block codecs carry no kernel on a K tile of 128: with a lever of +# their name set they run the BK=64 kernels of the BM=128 and BM=64 tiles. +@pytest.mark.parametrize("codec", ["pq2_0", "q2_0"]) +def test_no_bk128_codec_keeps_bk64(codec): + nax = _nax_gpu() + calls = [ + {"codec": codec, "shape": [rows, K], "n": n, "seed": rows} + for rows, n in ((4096, 128), (4128, 96), (160, 128)) + ] + res = run_probe(calls, env=dict(LIVE, **{_lever(codec): "1"})) + for c, r, bm in zip(calls, res, (128, 128, 64), strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1["bfloat16"], (c, r["rel"]) + assert r["nonfinite"] == 0, c + if nax: + name, _ = _kernel(r) + assert _bk(name) == 64, (c, name) + assert tile_of(name)[0] == bm, (c, name) + + +# (shape, n): BM=128 at an even and an odd row-tile count, BM=64, N not a +# multiple of 64, and batches of two flattened into 512 rows (BM=128) and +# 544 rows (BM=64). +SHAPES = [ + ([4096, K], 128), + ([4128, K], 96), + ([160, K], 128), + ([160, K], 96), + ([2, 256, K], 128), + ([2, 272, K], 96), +] + + +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", sorted(CODECS)) +def test_bk128_bit_identical(codec, dtype): + nax = _nax_gpu() + lever = _lever(codec) + calls = [] + for shape, n in SHAPES: + for design in _designs(codec): + base = len(calls) + call = { + "codec": codec, + "shape": shape, + "n": n, + "dtype": dtype, + "seed": base, + "design": design, + } + calls.append(dict(call, env={lever: "1"})) + calls.append(dict(call, env={lever: "0"}, ulp_vs=base)) + res = run_probe(calls, env=LIVE) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1[dtype], (c, r["rel"]) + assert r["nonfinite"] == 0, c + if not nax: + continue + name, _ = _kernel(r) + want = 64 if c["env"][lever] == "0" else 128 + assert _bk(name) == want, (c, name) + if c.get("ulp_vs") is not None: + assert r["ulp"] == 0, (c, r["ulp"]) + assert r["digest"] == res[c["ulp_vs"]]["digest"], c + + +@pytest.mark.parametrize("codec", sorted(CODECS)) +def test_bk128_rows_above_int16(codec): + # 32768 rows, and a batch of two at 16896 rows flattened to 33792, past + # the int16 row span the tile clamps before narrowing, on the BM=64 and + # BM=128 tiles. + nax = _nax_gpu() + lever = _lever(codec) + calls = [] + for shape in ([1, 32768, K], [2, 16896, K]): + for n in (96, 128): + for tile in ("64,64,2,2", "128,64,2,2"): + base = len(calls) + call = {"codec": codec, "shape": shape, "n": n, "seed": base} + calls.append(dict(call, env={lever: "1", "KQ_NAX_TILE": tile})) + calls.append( + dict(call, env={lever: "0", "KQ_NAX_TILE": tile}, ulp_vs=base) + ) + res = run_probe(calls, env=LIVE) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1["bfloat16"], (c, r["rel"]) + if not nax: + continue + name, _ = _kernel(r) + assert _bk(name) == (64 if c["env"][lever] == "0" else 128), (c, name) + assert tile_of(name)[0] == int(c["env"]["KQ_NAX_TILE"].split(",")[0]) + if c.get("ulp_vs") is not None: + assert r["ulp"] == 0, (c, r["ulp"]) + assert r["digest"] == res[c["ulp_vs"]]["digest"], c + + +@pytest.mark.parametrize( + "codec", sorted(c for c, (gs, _) in CODECS.items() if gs == 32) +) +def test_bk128_k_multiple_of_128(codec): + # A 32-weight block codec takes the K tile of 128 at K 896 and keeps + # BK=64 at K 704, a multiple of 64 and not of 128. + nax = _nax_gpu() + lever = _lever(codec) + calls = [ + {"codec": codec, "shape": [4096, k], "n": 128, "seed": k} for k in (896, 704) + ] + res = run_probe(calls, env=dict(LIVE, **{lever: "1"})) + for c, r, bk in zip(calls, res, (128, 64), strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1["bfloat16"], (c, r["rel"]) + if nax: + assert _bk(_kernel(r)[0]) == bk, c diff --git a/tests/test_nax_bm128_rows.py b/tests/test_nax_bm128_rows.py new file mode 100644 index 00000000..69ccba68 --- /dev/null +++ b/tests/test_nax_bm128_rows.py @@ -0,0 +1,210 @@ +"""BM=128 at large M with an odd ceil(M/64): KQ_NAX_BM128_ODD_M and the +row-live guard of the NAX qmm_t tile. + +Row counts 4128, 4256 and 4640, and a batch of two at 4192 rows flattened +to 8384, leave the last BM=128 row tile at most half live, so the guard +skips the dead simdgroup row there. 32768 rows and a batch of two at 16896 +rows (33792) run past the int16 row span. Each case checks the result +against float64 with exact kq.dequantize weights, on the CPU stream too, +and on NAX GPUs checks the dispatched tile and that BM=128 agrees with the +BM=64 tile within 1 ulp. +""" + +from __future__ import annotations + +import pytest +from qmm_probe import force_cpu, run_probe, tile_of + +import mlx_kquant as kq + +R1 = {"bfloat16": 6e-3, "float16": 1e-3} +K = 256 +LIVE = {"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"} +BM64 = "64,64,2,2" +MAIN_CODECS = ["q4_k", "q5_k", "q6_k", "q8_0"] +# (shape, rows after the batch flattens); every ceil(rows/64) here is odd +# except the last two, which run past 32767 rows. +SHAPES = [ + ([1, 4128, K], 4128), + ([1, 4256, K], 4256), + ([1, 4640, K], 4640), + ([2, 4192, K], 8384), + ([1, 32768, K], 32768), + ([2, 16896, K], 33792), +] + + +# The codecs with NAX kernels; each carries the BM=128 tile. +NAX_CODECS = sorted( + [ + "q2_k", "q3_k", "q4_k", "q5_k", "q6_k", "q8_0", + "q4_0", "q4_1", "q5_0", "q5_1", + "iq1_s", "iq1_m", "iq2_xxs", "iq2_xs", "iq2_s", "iq3_xxs", "iq3_s", + "iq4_nl", "iq4_xs", "stq1_0", "pq2_0", "ptq1_0", "q2_0", + ] +) # fmt: skip + + +def _metallib_bm128_codecs(): + """The BM=128 codecs the metallib carries, or None when there is no + metallib to read: a build without Metal, or KQUANT_FORCE_CPU, under + which this module never touches the GPU.""" + if force_cpu(): + return None + names = kq.metallib_kernels() + if not names: + return None + return sorted( + { + n[len("kquant_") : n.index("_qmm_t_nax_")] + for n in names + if n.startswith("kquant_") + and "_qmm_t_nax_bfloat16_t_" in n + and "_bm128_bn64_" in n + and "_gather_" not in n + and "_splitk_" not in n + } + ) + + +METALLIB_BM128_CODECS = _metallib_bm128_codecs() +BM128_CODECS = NAX_CODECS if METALLIB_BM128_CODECS is None else METALLIB_BM128_CODECS + + +def _nax_gpu(): + return not force_cpu() and kq.nax_available() + + +def _kernel(res): + assert res["error"] is None, res["error"] + assert len(res["kernels"]) == 1, res["kernels"] + return res["kernels"][0] + + +def test_bm128_codecs_found(): + if METALLIB_BM128_CODECS is None: + pytest.skip("no metallib loaded") + assert METALLIB_BM128_CODECS == NAX_CODECS + assert set(MAIN_CODECS) <= set(BM128_CODECS) + + +def test_odd_rows_policy(): + if not _nax_gpu(): + pytest.skip("NAX tile policy") + rows = {"codec": "q4_k", "shape": [1, 4128, K], "n": 128} + below = {"codec": "q4_k", "shape": [1, 4000, K], "n": 128} + batched = {"codec": "q4_k", "shape": [2, 4192, K], "n": 128} + calls = [ + dict(rows), + dict(rows, env={"KQ_NAX_BM128_ODD_M": "0"}), + dict(rows, env={"KQ_NAX_BM128_ODD_M": "8192"}), + dict(batched), + dict(rows, env={"KQ_NAX_BM128_ODD_M": None, "KQ_NAX_BM128": "0"}), + dict(rows, env={"KQ_NAX_BM128": None}), + dict(below), + ] + res = run_probe(calls, env=LIVE) + got = [_kernel(r) for r in res] + assert [tile_of(n)[0] for n, _ in got] == [128, 64, 64, 128, 64, 128, 64] + assert [g for _, g in got] == [ + (2, 33, 1), + (2, 65, 1), + (2, 65, 1), + (2, 66, 1), + (2, 65, 1), + (2, 33, 1), + (2, 63, 1), + ] + for r in res: + assert r["rel"] < R1["bfloat16"], r + + +def test_odd_rows_lever_keeps_first_read(): + if not _nax_gpu(): + pytest.skip("NAX tile policy") + call = {"codec": "q4_k", "shape": [1, 4128, K], "n": 128} + calls = [dict(call), dict(call, env={"KQ_NAX_BM128_ODD_M": None})] + res = run_probe( + calls, env={"KQ_QMM_LOG": "1", "KQ_NAX_BM128_ODD_M": "0"}, check=False + ) + assert [tile_of(_kernel(r)[0])[0] for r in res] == [64, 64] + + +# K ceilings of the large-M rows (kq_largem_max_k): above it, from M 4096, +# the BM=64 tile runs at either parity. +# q2_0 takes the pq2_0 ceiling. +MAX_K = {"q2_k": 4096, "q2_0": 4096, "q8_0": 10240, "q4_0": 10240, "stq1_0": 12288} + + +@pytest.mark.parametrize("codec", sorted(MAX_K)) +def test_large_rows_k_ceiling(codec): + if not _nax_gpu(): + pytest.skip("NAX tile policy") + top = MAX_K[codec] + calls = [] + for k in (top, top + 2048): + for rows in (4096, 4128): + calls.append({"codec": codec, "shape": [1, rows, k], "n": 64}) + # KQ_NAX_BM128_ODD_M=0 turns the large-M rows off: the parity rule + # runs BM=128 at an even ceil(M/64) at any K. + calls.append( + { + "codec": codec, + "shape": [1, 4096, top + 2048], + "n": 64, + "env": {"KQ_NAX_BM128_ODD_M": "0"}, + } + ) + res = run_probe(calls, env=LIVE) + got = [tile_of(_kernel(r)[0])[0] for r in res] + assert got == [128, 128, 64, 64, 128], got + for c, r in zip(calls, res, strict=True): + assert r["rel"] < R1["bfloat16"], (c, r["rel"]) + + +# Codecs whose large-M row has no K ceiling, on BK=64 and on BK=128 tiles. +NO_MAX_K = ["q6_k", "iq2_s", "iq2_xs", "iq3_s", "iq3_xxs", "iq4_xs"] + + +@pytest.mark.parametrize("codec", NO_MAX_K) +def test_large_rows_no_k_ceiling(codec): + if not _nax_gpu(): + pytest.skip("NAX tile policy") + calls = [{"codec": codec, "shape": [1, 4128, 16384], "n": 64}] + res = run_probe(calls, env=LIVE) + assert tile_of(_kernel(res[0])[0])[0] == 128 + assert res[0]["rel"] < R1["bfloat16"], res[0]["rel"] + + +@pytest.mark.parametrize("codec", BM128_CODECS) +def test_bm128_rows_known_answer(codec): + nax = _nax_gpu() + dtypes = ["bfloat16", "float16"] if codec in MAIN_CODECS else ["bfloat16"] + calls = [] + for dtype in dtypes: + for shape, rows in SHAPES: + for n in (128, 96) if rows == 4128 else (128,): + base = len(calls) + call = { + "codec": codec, + "shape": shape, + "n": n, + "dtype": dtype, + "seed": base, + "rows": rows, + } + calls.append(dict(call, env={"KQ_NAX_TILE": None})) + if nax: + calls.append(dict(call, env={"KQ_NAX_TILE": BM64}, ulp_vs=base)) + res = run_probe(calls, env=LIVE) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1[c["dtype"]], (c, r["rel"]) + if not nax: + continue + want = 64 if c["env"]["KQ_NAX_TILE"] else 128 + name, grid = _kernel(r) + assert tile_of(name)[0] == want, (c, name) + assert grid[1] == -(-c["rows"] // want), (c, grid) + if r["ulp"] is not None: + assert r["ulp"] <= 1, (c, r["ulp"]) diff --git a/tests/test_nax_smallm.py b/tests/test_nax_smallm.py index 8eeaa9d6..943e6537 100644 --- a/tests/test_nax_smallm.py +++ b/tests/test_nax_smallm.py @@ -21,6 +21,7 @@ from __future__ import annotations +import functools import json import os import subprocess @@ -32,6 +33,7 @@ sys.path.insert(0, os.path.dirname(__file__)) import mlx.core as mx # noqa: E402 from kqref import synth_wire # noqa: E402 +from qmm_probe import run_probe # noqa: E402 from test_codecs import CODECS # noqa: E402 import mlx_kquant as kq # noqa: E402 @@ -210,12 +212,32 @@ def test_alu_splitk_band_iq(codec, nax_off): # The mid-M rule takes split-K on non-NAX GPUs while the BM64 qmm grid, # ceil(N/64) x ceil(M/64) threadgroups, is under three per core. With 40 -# cores pinned, N 1000 at M 48 (16 threadgroups) matches the forced -# split-K output bit for bit and N 8192 (128) matches plain qmm. The two -# routes differ in their bits, so the comparison tells them apart. +# cores pinned, N 1000 at M 48 (16 threadgroups) runs the split-K kernel +# that KQ_QMM_MIDM=1 forces, bit for bit, and N 8192 (128) runs the plain +# qmm of KQ_QMM_MIDM=0. KQ_QMM_LOG names the kernel of each arm. +ALU_MIDM_CASES = [(c, n) for c in ("q4_k", "q8_0") for n in (1000, 8192)] + + +@functools.cache +def _alu_midm_logged(): + """_logged entries of the default, KQ_QMM_MIDM=1 and =0 arms per + ALU_MIDM_CASES case, with NAX off and 40 cores pinned.""" + calls = [ + {"codec": codec, "shape": [48, K], "n": n_out, "env": {"KQ_QMM_MIDM": v}} + for codec, n_out in ALU_MIDM_CASES + for v in (None, "1", "0") + ] + e = _logged(calls, {"KQ_DISABLE_NAX": "1", "KQ_GPU_CORES": "40"}) + return {case: e[3 * i : 3 * i + 3] for i, case in enumerate(ALU_MIDM_CASES)} + + @pytest.mark.parametrize("n_out,taken", [(1000, "1"), (8192, "0")]) @pytest.mark.parametrize("codec", ["q4_k", "q8_0"]) def test_alu_midm_splitk_rule(codec, n_out, taken, nax_off, monkeypatch): + default, split, tile = (_one_kernel(e) for e in _alu_midm_logged()[(codec, n_out)]) + assert "_qmm_t_splitk_" in split[0] and split[1][2] > 1, split + assert "splitk" not in tile[0] and tile[1][2] == 1, tile + assert default == (split if taken == "1" else tile), default monkeypatch.setenv("KQ_GPU_CORES", "40") w, s, _ = _encodable_setup(codec, n_out) x = (mx.random.normal((48, K)) * 0.5).astype(mx.bfloat16) @@ -229,9 +251,7 @@ def run(route): mx.eval(y) return y - default = run(None) - assert _bits_equal(default, run(taken)) - assert not _bits_equal(default, run("0" if taken == "1" else "1")) + assert _bits_equal(run(None), run(taken)) # Mid-M band on NAX GPUs, where N 1000 keeps the rule on through M 128. @@ -242,10 +262,11 @@ def run(route): NAX_MIDM_MS = [33, 40, 48, 49, 56, 64, 65, 80, 96, 97, 112, 128, 129] +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("codec", ENCODABLE) -def test_nax_midm_band(codec): +def test_nax_midm_band(codec, dtype): w, s, ref_w = _encodable_setup(codec, 1000) - _sweep(codec, w, s, ref_w, 1000, ms=NAX_MIDM_MS) + _sweep(codec, w, s, ref_w, 1000, ms=NAX_MIDM_MS, dtype=dtype) @pytest.mark.parametrize("codec", IQ) @@ -257,22 +278,48 @@ def test_nax_midm_band_iq(codec): # The NAX mid-M rule, with 40 cores pinned. Under three threadgroups per # core the route runs: N 1000 at M 48 and 112 (16 and 32). From M 65 to 96 # it runs up to 16 per core: N 8192 at M 96 (256). N 8192 at M 48 and 112 -# (128 and 256) keep the BM64 qmm. Each default output matches the forced -# arm it names bit for bit and differs from the other, and -# KQ_QMM_SPLITK_NAX=0 turns the route off with the small-M one. +# (128 and 256) keep the BM64 qmm. Each default arm runs the kernel of the +# forced arm it names, bit for bit, and KQ_QMM_SPLITK_NAX=0 turns the route +# off with the small-M one. KQ_QMM_LOG names the kernel of each arm. +NAX_MIDM_CASES = [ + (48, 1000, "1"), + (112, 1000, "1"), + (96, 8192, "1"), + (48, 8192, "0"), + (112, 8192, "0"), +] + + +@functools.cache +def _nax_midm_logged(): + """_logged entries of the default, KQ_QMM_MIDM=1, KQ_QMM_MIDM=0 and + KQ_QMM_SPLITK_NAX=0 arms per (codec, M, N), with 40 cores pinned.""" + cases = [(c, m, n) for c in ("q4_k", "q8_0") for m, n, _ in NAX_MIDM_CASES] + arms = [ + {"KQ_QMM_MIDM": None, "KQ_QMM_SPLITK_NAX": None}, + {"KQ_QMM_MIDM": "1"}, + {"KQ_QMM_MIDM": "0"}, + {"KQ_QMM_MIDM": None, "KQ_QMM_SPLITK_NAX": "0"}, + ] + calls = [ + {"codec": codec, "shape": [m, K], "n": n_out, "env": env} + for codec, m, n_out in cases + for env in arms + ] + e = _logged(calls, {"KQ_GPU_CORES": "40"}) + return {case: e[4 * i : 4 * i + 4] for i, case in enumerate(cases)} + + @pytest.mark.skipif(not kq.nax_available(), reason="NAX tile only") -@pytest.mark.parametrize( - "m,n_out,taken", - [ - (48, 1000, "1"), - (112, 1000, "1"), - (96, 8192, "1"), - (48, 8192, "0"), - (112, 8192, "0"), - ], -) +@pytest.mark.parametrize("m,n_out,taken", NAX_MIDM_CASES) @pytest.mark.parametrize("codec", ["q4_k", "q8_0"]) def test_nax_midm_splitk_rule(codec, m, n_out, taken, monkeypatch): + entries = _nax_midm_logged()[(codec, m, n_out)] + default, split, tile, off = (_one_kernel(e) for e in entries) + assert "_qmm_t_nax_splitk_" in split[0] and split[1][2] > 1, split + assert "splitk" not in tile[0] and tile[1][2] == 1, tile + assert default == (split if taken == "1" else tile), default + assert off == tile, off monkeypatch.setenv("KQ_GPU_CORES", "40") w, s, _ = _encodable_setup(codec, n_out) x = (mx.random.normal((m, K)) * 0.5).astype(mx.bfloat16) @@ -286,12 +333,62 @@ def run(route): mx.eval(y) return y - default = run(None) - tile = run("0") - assert _bits_equal(default, run(taken)) - assert not _bits_equal(default, run("0" if taken == "1" else "1")) + tile_y = run("0") + assert _bits_equal(run(None), run(taken)) monkeypatch.setenv("KQ_QMM_SPLITK_NAX", "0") - assert _bits_equal(run(None), tile) + assert _bits_equal(run(None), tile_y) + + +# Slice targets of the NAX mid-M rule with 40 cores pinned. At K 5120 q4_0 +# and q8_0 cut 80 units of 64 weights, so the slice count is the target. +# With float partials the two codecs target 10 from M 33 to 64 and 5 from +# M 65 to 96, and 8 from M 65 to 96 above K 5120, where 5 would leave 4 +# slices of the 96 or 272 units. With partials in the activation dtype they +# keep 16 under three threadgroups per core and 8 from M 65 to 96 up to 16 +# per core. M 112 keeps 16 either way. +# (M, N, K): (float partials, activation-dtype partials). +NAX_MIDM_TARGETS = { + (48, 1000, 5120): (10, 16), + (80, 1000, 5120): (5, 16), + (80, 5120, 5120): (5, 8), + (80, 5120, 6144): (8, 8), + (96, 5120, 17408): (8, 8), + (112, 1000, 5120): (16, 16), +} +# (activation dtype, KQ_SPLITK_F32): whether the partials are float. +NAX_MIDM_PARTIALS = { + ("float16", None): True, + ("float16", "0"): False, + ("bfloat16", None): False, + ("bfloat16", "1"): True, +} + + +@functools.cache +def _nax_midm_targets_logged(dtype, f32_env): + """_logged entry per (codec, M, N, K) of NAX_MIDM_TARGETS in one probe + process with KQ_SPLITK_F32 as given and 40 cores pinned.""" + cases = [(c, *shape) for c in ("q4_0", "q8_0") for shape in NAX_MIDM_TARGETS] + calls = [ + {"codec": c, "shape": [m, k], "n": n, "dtype": dtype} for c, m, n, k in cases + ] + env = {"KQ_GPU_CORES": "40"} + if f32_env is not None: + env["KQ_SPLITK_F32"] = f32_env + return dict(zip(cases, _logged(calls, env), strict=True)) + + +@pytest.mark.skipif(not kq.nax_available(), reason="NAX tile only") +@pytest.mark.parametrize("dtype,f32_env", list(NAX_MIDM_PARTIALS)) +@pytest.mark.parametrize("codec", ["q4_0", "q8_0"]) +def test_nax_midm_targets(codec, dtype, f32_env): + f32 = NAX_MIDM_PARTIALS[(dtype, f32_env)] + logged = _nax_midm_targets_logged(dtype, f32_env) + for (m, n, k), targets in NAX_MIDM_TARGETS.items(): + name, grid = _one_kernel(logged[(codec, m, n, k)]) + assert "_qmm_t_nax_splitk_" in name, name + assert name.endswith("_pf32") == f32, name + assert grid[2] == targets[0 if f32 else 1], (m, n, k, name, grid) # Register-resident MMA verify band (kq_verify_mma.h): KQ_VERIFY_MMA=2 @@ -333,13 +430,35 @@ def _verify_mma_setup(codec, n_out, k=K): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) -@pytest.mark.parametrize("n_out", [1024, 1000]) +@pytest.mark.parametrize("n_out", [1024, 1000, 40]) @pytest.mark.parametrize("codec", VERIFY_MMA_CODECS) def test_verify_mma_band(codec, n_out, dtype, verify_mma_forced): w, s, ref_w = _verify_mma_setup(codec, n_out) _sweep(codec, w, s, ref_w, n_out, ms=VERIFY_MMA_MS, dtype=dtype) +# Output rows per verify_mma threadgroup: 8 simdgroups of 8-row tiles, one +# tile each on pq2_0, ptq1_0 and q2_0 and two on q4_0 and q8_0. +VERIFY_MMA_ROWS = {"pq2_0": 64, "ptq1_0": 64, "q2_0": 64, "q4_0": 128, "q8_0": 128} + + +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", VERIFY_MMA_CODECS) +def test_verify_mma_grid(codec, dtype): + ns = (1024, 1000) + calls = [{"codec": codec, "shape": [8, K], "n": n, "dtype": dtype} for n in ns] + env = { + "KQ_QMM_ROUTE": "verify_mma", + "KQ_QMM_ROUTE_STRICT": "1", + "KQ_VERIFY_MMA_SPLITS": "1", + } + rows = VERIFY_MMA_ROWS[codec] + for n, entry in zip(ns, _logged(calls, env), strict=True): + name, grid = _one_kernel(entry) + assert "_verify_mma_" in name, name + assert grid == (-(-n // rows), 1, 1), (name, grid) + + @pytest.mark.parametrize("codec", VERIFY_MMA_CODECS) def test_verify_mma_ragged_chunk(codec, verify_mma_forced): # 2176 = 17 blocks of 128: one split of 4 full 512-wide chunks plus a @@ -353,9 +472,16 @@ def test_verify_mma_ragged_chunk(codec, verify_mma_forced): @pytest.mark.parametrize("codec", VERIFY_MMA_CODECS) def test_verify_mma_default_matches_forced(codec, monkeypatch): # The default route at M8 and the forced route agree to the bf16 - # output rounding, whichever kernel the default picks. The route is - # chosen when the graph evaluates, so each output evaluates under its - # own environment. + # output rounding, whichever kernel the default picks. KQ_QMM_LOG shows + # that the forced arm runs verify_mma and the arm with it off does not. + # The route is chosen when the graph evaluates, so each output + # evaluates under its own environment. + calls = [ + {"codec": codec, "shape": [8, K], "n": 1024, "env": {"KQ_VERIFY_MMA": v}} + for v in ("2", "0") + ] + forced, off = (_route_of(e) for e in _logged(calls, {"KQ_VERIFY_NAX": "0"})) + assert forced == "verify_mma" and off != "verify_mma", (forced, off) w, s, ref_w = _verify_mma_setup(codec, 1024) x = (mx.random.normal((8, K)) * 0.5).astype(mx.bfloat16) monkeypatch.setenv("KQ_VERIFY_NAX", "0") @@ -365,7 +491,6 @@ def test_verify_mma_default_matches_forced(codec, monkeypatch): monkeypatch.setenv("KQ_VERIFY_MMA", "0") y_off = kq.quantized_matmul(x, w, s, codec, transpose=True) mx.eval(y_off) - assert not bool(mx.array_equal(y_forced, y_off).item()) err = float( mx.abs(y_forced.astype(mx.float32) - y_off.astype(mx.float32)).max() / (mx.abs(y_off.astype(mx.float32)).max() + 1e-6) @@ -530,24 +655,77 @@ def test_nax_splitk_ragged(codec, dtype, monkeypatch): mx.eval(kq.quantized_matmul(x, w, s, codec, transpose=True)) +def _logged(calls, env): + """(error, kernels) per call from one probe process with KQ_QMM_LOG on. + Two routes can agree in every output bit, the split routes most of all + with float partials, so which route ran, and with how many K slices, + need not show in the output; the log names both. A probe that fails + returns its failure in every entry instead of raising, so the tests + that share it fail on their own entries.""" + try: + res = run_probe(calls, env={**env, "KQ_QMM_LOG": "1"}, check=False) + except (AssertionError, subprocess.TimeoutExpired) as e: + return [(f"probe failed: {e}", [])] * len(calls) + return [(r["error"], r["kernels"]) for r in res] + + +def _one_kernel(entry): + """(name, grid) of the one kernel an entry of _logged dispatched.""" + error, kernels = entry + assert error is None, error + assert len(kernels) == 1, kernels + return kernels[0] + + +def _depth(entry): + return _one_kernel(entry)[1][2] + + +# Kernel name fragment of each route that logs. The mat-vec routes (qmv, +# mv_ext, verify_qmv) log nothing. +ROUTE_KERNEL = { + "verify_nax": "_verify_nax_", + "verify_mma": "_verify_mma_", + "nax_splitk": "_qmm_t_nax_splitk_", + "splitk": "_qmm_t_splitk_", +} + + +def _route_of(entry): + """The route an entry of _logged ran: a ROUTE_KERNEL key, "tile" for + another logged kernel, or "matvec" when nothing logged.""" + error, kernels = entry + assert error is None, error + if not kernels: + return "matvec" + assert len(kernels) == 1, kernels + name = kernels[0][0] + return next((r for r, frag in ROUTE_KERNEL.items() if frag in name), "tile") + + +def _family(route): + """What _route_of reports for a route name.""" + return "matvec" if route in ("qmv", "mv_ext", "verify_qmv") else route + + # Ragged slices replace the divisor count only when they more than double # it. 68 units divide into 4 and run as 14 ragged slices. 35 units divide # into 7, and 12 ragged slices would not double that. @pytest.mark.skipif(not kq.nax_available(), reason="NAX tile only") @pytest.mark.parametrize("units, ragged", [(68, True), (35, False)]) def test_nax_splitk_ragged_rule(units, ragged, monkeypatch): - monkeypatch.delenv("KQ_SPLITK_RAGGED", raising=False) k = 256 * units + calls = [ + {"codec": "q4_k", "shape": [8, k], "n": 256, "env": {"KQ_SPLITK_RAGGED": v}} + for v in ("0", "1") + ] + env = {"KQ_QMM_ROUTE": "nax_splitk", "KQ_QMM_ROUTE_STRICT": "1"} + d0, d1 = (_depth(e) for e in _logged(calls, env)) + assert (d0 != d1) == ragged, (d0, d1) + monkeypatch.setenv("KQ_SPLITK_RAGGED", "1") w, s, ref_w = _verify_mma_setup("q4_k", 256, k=k) monkeypatch.setenv("KQ_QMM_ROUTE", "nax_splitk") monkeypatch.setenv("KQ_QMM_ROUTE_STRICT", "1") - x = (mx.random.normal((8, k)) * 0.5).astype(mx.bfloat16) - ys = {} - for v in ("0", "1"): - monkeypatch.setenv("KQ_SPLITK_RAGGED", v) - ys[v] = kq.quantized_matmul(x, w, s, "q4_k", transpose=True) - mx.eval(ys[v]) - assert _bits_equal(ys["0"], ys["1"]) != ragged _sweep("q4_k", w, s, ref_w, 256, ms=[8], k=k) @@ -561,9 +739,9 @@ def test_verify_mma_entries_f16(codec, monkeypatch): _sweep(codec, w, s, ref_w, 1000, ms=[2, 3, 4, 5, 8, 9], dtype=mx.float16) -# verify_mma sums each block in a half accumulator. One activation channel -# near the top of the half range must not overflow it; q4_0 multiplies it -# by up to 8 before the kernel's 1/16 prescale. +# verify_mma sums each block in float, so one activation channel near the +# top of the half range stays finite and within the contract +# (tests/test_verify_mma_range.py goes past that range). @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("codec", VMMA_CODECS) def test_verify_mma_outlier_channel(codec, dtype, monkeypatch): @@ -699,25 +877,32 @@ def test_verify_nax_word_weight_alignment(codec, offset, monkeypatch): # KQ_VERIFY_NAX_Q4_0_SB=0 keeps q4_0 on the two-block kernel at a K that # is a multiple of 256, where the eight-block kernel runs by default. At K # 3840 a split target of 4 gives the eight-block kernel 3 splits of its 15 -# steps and the two-block kernel 4 of its 60, so a lever that changed -# nothing would return the same bits. Both kernels match the reference. +# steps and the two-block kernel 4 of its 60. KQ_QMM_LOG names the kernel +# and slice count of each arm, and both kernels match the reference. @pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") -@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) def test_verify_nax_q4_0_two_block_lever(dtype, monkeypatch): k = 3840 + base = {"codec": "q4_0", "shape": [8, k], "n": 1000, "dtype": dtype} + calls = [ + dict(base, env={"KQ_VERIFY_NAX_Q4_0_SB": None}), + dict(base, env={"KQ_VERIFY_NAX_Q4_0_SB": "0"}), + ] + env = { + "KQ_QMM_ROUTE": "verify_nax", + "KQ_QMM_ROUTE_STRICT": "1", + "KQ_VERIFY_NAX_SPLITS": "4", + } + sb, two = (_one_kernel(e) for e in _logged(calls, env)) + assert sb[0].startswith("kquant_q4_0_sb_verify_nax_") and sb[1][2] == 3, sb + assert two[0].startswith("kquant_q4_0_verify_nax_") and two[1][2] == 4, two + mdt = getattr(mx, dtype) w, s, ref_w = _verify_mma_setup("q4_0", 1000, k=k) - monkeypatch.setenv("KQ_QMM_ROUTE", "verify_nax") - monkeypatch.setenv("KQ_QMM_ROUTE_STRICT", "1") - monkeypatch.setenv("KQ_VERIFY_NAX_SPLITS", "4") - x = (mx.random.normal((8, k)) * 0.5).astype(dtype) - y_sb = kq.quantized_matmul(x, w, s, "q4_0", transpose=True) - mx.eval(y_sb) - _sweep("q4_0", w, s, ref_w, 1000, ms=[1, 3, 8], dtype=dtype, k=k) + for name, value in env.items(): + monkeypatch.setenv(name, value) + _sweep("q4_0", w, s, ref_w, 1000, ms=[1, 3, 8], dtype=mdt, k=k) monkeypatch.setenv("KQ_VERIFY_NAX_Q4_0_SB", "0") - y_two = kq.quantized_matmul(x, w, s, "q4_0", transpose=True) - mx.eval(y_two) - assert not _bits_equal(y_sb, y_two) - _sweep("q4_0", w, s, ref_w, 1000, ms=[1, 3, 8], dtype=dtype, k=k) + _sweep("q4_0", w, s, ref_w, 1000, ms=[1, 3, 8], dtype=mdt, k=k) # The q4_k and q5_k kernels read their weights as 16-byte words. A base @@ -783,14 +968,14 @@ def test_verify_nax_half_word_weight_alignment(codec, offset, monkeypatch): # or above max(ceil(1280 / simdgroups), ceil(K / 2560)), at most 16, else # the largest divisor under it. N 3072, K 4096 targets 14 and takes 16 of # 32 steps. N 1024, K 5120 targets 16, which does not divide 40 steps, and -# takes 10. N 20480, K 10240 targets 4 from the K walk. The default count -# runs bit for bit like the count forced through the lever, and unlike the -# neighbouring count. +# takes 10. N 20480, K 10240 targets 4 from the K walk. The default runs +# that many slices, as the lever forcing it does, and bit for bit like it; +# the lever also moves the count to the neighbouring one. +Q6_K_SPLITS = [(3072, 4096, 16, 8), (1024, 5120, 10, 8), (20480, 10240, 4, 2)] + + @pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") -@pytest.mark.parametrize( - "n_out, k, count, other", - [(3072, 4096, 16, 8), (1024, 5120, 10, 8), (20480, 10240, 4, 2)], -) +@pytest.mark.parametrize("n_out, k, count, other", Q6_K_SPLITS) def test_verify_nax_q6_k_split_rule(n_out, k, count, other, monkeypatch): wf = mx.random.normal((64, k), key=mx.random.key(4)) * 0.1 w1, s = kq.quantize(wf, "q6_k") @@ -812,7 +997,8 @@ def run(splits=None): y_def = run() assert _bits_equal(y_def, run(count)) - assert not _bits_equal(y_def, run(other)) + arms = _verify_nax_depths()[("q6_k", n_out, k)] + assert tuple(_depth(e) for e in arms) == (count, count, other) ref = x.astype(mx.float32) @ ref_w err = float(mx.abs(y_def[:, :64].astype(mx.float32) - ref).max()) assert err < 2e-2 * float(mx.abs(ref).max()) @@ -827,8 +1013,9 @@ def run(splits=None): # takes 7 over 35. N 2560, K 9728 takes 38 at 2560 and 19 at 640, where 2 # splits ran up to 1.65x slower than 19. N 6144, K 5120 in 256-weight steps # takes 2 over 10. N 1024, K 5120 takes 20, above 16, which on q4_k is -# every K step. The default count runs bit for bit like the count forced -# through the lever, and unlike the neighbouring count. +# every K step. The default runs that many slices, as the lever forcing it +# does, and bit for bit like it; the lever also moves the count to the +# neighbouring one. WAVE_SPLITS = [ *[ (codec, n_out, k, count, other) @@ -849,6 +1036,30 @@ def run(splits=None): ] +@functools.cache +def _verify_nax_depths(): + """_logged entries of the (default, count, other) arms per (codec, N, + K) of the split rule tests, from one probe process.""" + cases = [("q6_k", n, k, c, o) for n, k, c, o in Q6_K_SPLITS] + WAVE_SPLITS + calls = [ + { + "codec": codec, + "shape": [8, k], + "n": n_out, + "design": "tiled", + "env": {"KQ_VERIFY_NAX_SPLITS": splits}, + } + for codec, n_out, k, count, other in cases + for splits in (None, count, other) + ] + env = {"KQ_QMM_ROUTE": "verify_nax", "KQ_QMM_ROUTE_STRICT": "1"} + e = _logged(calls, env) + return { + (codec, n_out, k): tuple(e[3 * i : 3 * i + 3]) + for i, (codec, n_out, k, _, _) in enumerate(cases) + } + + @pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") @pytest.mark.parametrize("codec, n_out, k, count, other", WAVE_SPLITS) def test_verify_nax_wave_split_rule(codec, n_out, k, count, other, monkeypatch): @@ -872,7 +1083,8 @@ def run(splits=None): y_def = run() assert _bits_equal(y_def, run(count)) - assert not _bits_equal(y_def, run(other)) + arms = _verify_nax_depths()[(codec, n_out, k)] + assert tuple(_depth(e) for e in arms) == (count, count, other) ref = x.astype(mx.float32) @ ref_w err = float(mx.abs(y_def[:, :64].astype(mx.float32) - ref).max()) assert err < 2e-2 * float(mx.abs(ref).max()) @@ -1048,8 +1260,10 @@ def test_verify_nax_iq4_nl_weight_alignment(offset, monkeypatch): # and on a vocab head, M 4 at K 4096 to 5120 and at N 10240 and wider on a # narrower K, M 5 on the other shapes from N 2048, none under N 2048. On # float16 an entry under M 5 sits one row higher, except on a vocab head. -# The default output carries the bits of verify_nax from the entry through -# M 8 and other bits below it. +# KQ_QMM_LOG names verify_nax as the default route from the entry through +# M 8 and another route below it, and the default output carries the bits +# of verify_nax from the entry. Below it the bits can agree, since the +# split routes store float partials on float16. VNAX_IQ4_ENTRY = [ (2560, 6144, 3), (2048, 4096, 4), @@ -1077,12 +1291,18 @@ def test_verify_nax_iq4_entry(codec, n_out, k, entry, dtype, monkeypatch): w = mx.contiguous(mx.tile(mx.array(seed), (n_out // 64 + 1, 1))[:n_out]) s = mx.zeros((1,), dtype=mx.uint8) mx.eval(w) - for m in (2, 3, 4, 5, 8): + ms = (2, 3, 4, 5, 8) + name = "float16" if dtype == mx.float16 else "bfloat16" + calls = [{"codec": codec, "shape": [m, k], "n": n_out, "dtype": name} for m in ms] + routes = [_route_of(e) for e in _logged(calls, {})] + for m, route in zip(ms, routes, strict=True): x = (mx.random.normal((m, k)) * 0.5).astype(dtype) y_def = _run_route(monkeypatch, x, w, s, codec) y_nax = _run_route(monkeypatch, x, w, s, codec, "verify_nax") tag = f"{codec} N{n_out} K{k} M{m} entry {entry}" - assert _bits_equal(y_def, y_nax) == (0 < entry <= m), tag + assert (route == "verify_nax") == (0 < entry <= m), (tag, route) + if 0 < entry <= m: + assert _bits_equal(y_def, y_nax), tag # KQ_DISABLE_NAX turns the route off with the other NAX routes. @@ -1128,13 +1348,12 @@ def _bits_equal(a, b): # The default route dispatches verify_nax from its entry, and -# KQ_VERIFY_NAX moves the entry both ways: the default output is -# bit-identical to the route each case expects and differs from the other -# kernels timed, and verify_nax agrees with a second 16-bit kernel -# (verify_mma, or NAX split-K for the codecs without it) to the output -# rounding. Below the entry and with verify_nax off, the case expects the -# table's route: pq2_0 and q4_0 fall back to verify_mma at M 8, the others -# to NAX split-K. +# KQ_VERIFY_NAX moves the entry both ways: KQ_QMM_LOG names the route each +# case expects, the default output is bit-identical to that route forced, +# and verify_nax agrees with a second kernel (verify_mma, or NAX split-K +# for the codecs without it) to the output rounding. Below the entry and +# with verify_nax off, the case expects the table's route: pq2_0 and q4_0 +# fall back to verify_mma at M 8, the others to NAX split-K. # The route is chosen when the graph evaluates, so each output evaluates # under its own environment. Per-row qmv claims these widths at N 1000 by # default, so it is off here. @@ -1164,6 +1383,20 @@ def run(x, route=None): (2, "2", "verify_nax"), (8, "0", _nax_default_route(codec, 1000, 8, qmv=False, vnax=False)), ] + dname = "float16" if dtype == mx.float16 else "bfloat16" + calls = [ + { + "codec": codec, + "shape": [m, K], + "n": 1000, + "dtype": dname, + "env": {"KQ_VERIFY_NAX": entry}, + } + for m, entry, _ in cases + ] + routes = [_route_of(e) for e in _logged(calls, {"KQ_NAX_QMV": "0"})] + for (m, entry, expect), route in zip(cases, routes, strict=True): + assert route == _family(expect), (codec, m, entry, expect, route) for m, entry, expect in cases: if entry is None: monkeypatch.delenv("KQ_VERIFY_NAX", raising=False) @@ -1174,8 +1407,7 @@ def run(x, route=None): alt = "verify_mma" if codec in VMMA_CODECS else "nax_splitk" ys = {r: run(x, r) for r in ("verify_nax", alt, expect)} tag = f"{codec} M{m} KQ_VERIFY_NAX={entry}" - for r, y in ys.items(): - assert _bits_equal(y_def, y) == (r == expect), f"{tag} vs {r}" + assert _bits_equal(y_def, ys[expect]), f"{tag} vs {expect}" y_nax, y_mma = ys["verify_nax"], ys[alt] err = float( mx.abs(y_nax.astype(mx.float32) - y_mma.astype(mx.float32)).max() @@ -1188,13 +1420,28 @@ def run(x, route=None): # default verify_nax entry, so an A/B arm times the route it names. An # explicit KQ_VERIFY_NAX wins back. On a codec verify_mma does not serve, # a forced KQ_VERIFY_MMA leaves verify_nax in place. Per-row qmv is off, -# as above. +# as above. KQ_QMM_LOG names the route of each arm, and each arm matches +# that route forced bit for bit. @pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") @pytest.mark.parametrize("codec", VNAX_CODECS) def test_verify_nax_yields_to_forced_levers(codec, monkeypatch): + m = max(6, VNAX_ENTRY[codec]) + mma = "verify_mma" if codec in VMMA_CODECS else "verify_nax" + arms = [ + ({"KQ_VERIFY_MMA": "2"}, mma), + ({"KQ_VERIFY_MMA": str(m + 1)}, "verify_nax"), + ({"KQ_QMM_SPLITK_NAX": "8"}, "nax_splitk"), + ({"KQ_QMM_SPLITK_NAX": "8", "KQ_VERIFY_NAX": "3"}, "verify_nax"), + ] + clear = {"KQ_VERIFY_MMA": None, "KQ_QMM_SPLITK_NAX": None, "KQ_VERIFY_NAX": None} + calls = [ + {"codec": codec, "shape": [m, K], "n": 1000, "env": {**clear, **env}} + for env, _ in arms + ] + got = [_route_of(e) for e in _logged(calls, {"KQ_NAX_QMV": "0"})] + assert got == [r for _, r in arms], got monkeypatch.setenv("KQ_NAX_QMV", "0") w, s, _ = _setup(codec, 1000) - m = max(6, VNAX_ENTRY[codec]) x = (mx.random.normal((m, K)) * 0.5).astype(mx.bfloat16) def run(route=None, **env): @@ -1212,10 +1459,8 @@ def run(route=None, **env): y_nax = run("verify_nax") y_sk = run("nax_splitk", KQ_QMM_SPLITK_NAX="8") - assert not _bits_equal(y_nax, y_sk) if codec in VMMA_CODECS: y_mma = run("verify_mma") - assert not _bits_equal(y_nax, y_mma) assert _bits_equal(run(KQ_VERIFY_MMA="2"), y_mma) else: assert _bits_equal(run(KQ_VERIFY_MMA="2"), y_nax) @@ -1275,10 +1520,9 @@ def test_verify_nax_head_entry(codec, monkeypatch): QMV_BN = {"q4_k": 4, "q8_0": 8} PRISM_VERIFY_MAX_M = {"pq2_0": 4, "ptq1_0": 4, "q2_0": 4} # Per-row qmv, mv_ext and verify_qmv dot the exact weights in float32 and -# agree bit for bit on most codecs, so a pin cannot tell them apart. The -# NAX tiles and the verify kernels round the weights to a 16-bit type, so -# the pins check every seam against those, whose results differ. -BITSAME = {"qmv", "mv_ext", "verify_qmv"} +# agree bit for bit on most codecs, and log no kernel, so a pin cannot tell +# them apart. verify_mma also sums the exact weights in float32, and +# KQ_QMM_LOG names it and the other logged routes. # Read once per process, so the default routing tests skip when they are set. STATIC_LEVERS = ("KQ_VERIFY_EXT", "KQ_QMM_SPLITK") @@ -1333,41 +1577,76 @@ def _clear_levers(monkeypatch): # The default small-M routing on NAX GPUs runs the kernel the per-codec # table names, bit for bit, on both sides of every N bucket edge and of # the q2_k verify_nax width, and across each seam: the qmv limit, the -# verify routes, the split-K entry and the mat-vec band between. Each pin -# is also checked against the route a limit one step off would give, -# unless both are in BITSAME. +# verify routes, the split-K entry and the mat-vec band between. KQ_QMM_LOG +# names each default route that logs, and shows the mat-vec pins log none. +NAX_SMALL_M_N = [512, 513, 1024, 1025, 2048, 2049, 4096, 4097] + + +def _nax_small_m_ms(codec): + return [2, 3, 4, 5, 6, 7, 8, 9] + ([12, 13] if codec == "iq2_xxs" else []) + + +@functools.cache +def _nax_small_m_routes(): + """_route_of the default call per (codec, N, M), from one probe.""" + cases = [ + (c, n, m) for c in sorted(NAX_SMALL_M) for n in NAX_SMALL_M_N + for m in _nax_small_m_ms(c) + ] # fmt: skip + calls = [{"codec": c, "shape": [m, K], "n": n} for c, n, m in cases] + return dict(zip(cases, map(_route_of, _logged(calls, {})), strict=True)) + + @pytest.mark.skipif(not kq.nax_available(), reason="NAX routing only") -@pytest.mark.parametrize("n_out", [512, 513, 1024, 1025, 2048, 2049, 4096, 4097]) +@pytest.mark.parametrize("n_out", NAX_SMALL_M_N) @pytest.mark.parametrize("codec", sorted(NAX_SMALL_M)) def test_nax_small_m_default_dispatch(codec, n_out, monkeypatch): _clear_levers(monkeypatch) + routes = _nax_small_m_routes() w, s, _ = _setup(codec, n_out) - ms = [2, 3, 4, 5, 6, 7, 8, 9] - if codec == "iq2_xxs": - ms += [12, 13] - for m in ms: + for m in _nax_small_m_ms(codec): x = (mx.random.normal((m, K)) * 0.5).astype(mx.bfloat16) expect = _nax_default_route(codec, n_out, m) + tag = f"{codec} N{n_out} M{m} expect {expect}" + assert routes[(codec, n_out, m)] == _family(expect), tag y_def = _run_route(monkeypatch, x, w, s, codec) y_exp = _run_route(monkeypatch, x, w, s, codec, expect) - tag = f"{codec} N{n_out} M{m} expect {expect}" assert _bits_equal(y_def, y_exp), tag - if expect == "qmv": - alt = _nax_default_route(codec, n_out, m, qmv=False) - else: - alt = "qmv" - if not {expect, alt} <= BITSAME: - y_alt = _run_route(monkeypatch, x, w, s, codec, alt) - assert not _bits_equal(y_def, y_alt), f"{tag} vs {alt}" # KQ_NAX_QMV=0 turns the route off and a value sets the limit at every N. # The default yields to a forced KQ_VERIFY_NAX or KQ_VERIFY_MMA on a codec # they serve and to a forced KQ_QMM_SPLITK_NAX, and an explicit # KQ_NAX_QMV wins back. KQ_DISABLE_NAX turns it off with the NAX routes. +# KQ_QMM_LOG names the route of each arm (codec, N, M, levers, route), and +# each arm matches that route forced bit for bit. +QMV_LEVER_ARMS = [ + ("q4_0", 500, 4, {}, "qmv"), + ("q4_0", 500, 4, {"KQ_NAX_QMV": "0"}, "verify_nax"), + ("q4_0", 500, 4, {"KQ_VERIFY_NAX": "3"}, "verify_nax"), + ("q4_0", 500, 4, {"KQ_VERIFY_NAX": "5"}, "qmv"), + ("q4_0", 500, 4, {"KQ_VERIFY_MMA": "2", "KQ_VERIFY_NAX": "0"}, "verify_mma"), + ("q4_0", 500, 4, {"KQ_QMM_SPLITK_NAX": "8"}, "nax_splitk"), + ("q4_0", 500, 4, {"KQ_QMM_SPLITK_NAX": "8", "KQ_NAX_QMV": "4"}, "qmv"), + ("q4_1", 500, 4, {"KQ_VERIFY_NAX": "2"}, "qmv"), + ("q4_1", 500, 4, {"KQ_VERIFY_MMA": "2"}, "qmv"), + ("q4_1", 4096, 6, {}, "nax_splitk"), + ("q4_1", 4096, 6, {"KQ_NAX_QMV": "6"}, "qmv"), + ("pq2_0", 500, 5, {"KQ_DISABLE_NAX": "1"}, "verify_mma"), + ("pq2_0", 500, 5, {"KQ_DISABLE_NAX": "1", "KQ_NAX_QMV": "8"}, "verify_mma"), +] + + @pytest.mark.skipif(not kq.nax_available(), reason="NAX routing only") def test_nax_qmv_levers(monkeypatch): _clear_levers(monkeypatch) + levers = dict.fromkeys(k for arm in QMV_LEVER_ARMS for k in arm[3]) + calls = [ + {"codec": c, "shape": [m, K], "n": n, "env": {**levers, **env}} + for c, n, m, env, _ in QMV_LEVER_ARMS + ] + got = [_route_of(e) for e in _logged(calls, {})] + assert got == [_family(arm[4]) for arm in QMV_LEVER_ARMS], got def run(codec, w, s, x, route=None, **env): _clear_levers(monkeypatch) @@ -1383,8 +1662,6 @@ def run(codec, w, s, x, route=None, **env): y_nax = run("q4_0", w, s, x, "verify_nax") y_mma = run("q4_0", w, s, x, "verify_mma") y_sk = run("q4_0", w, s, x, "nax_splitk", KQ_QMM_SPLITK_NAX="8") - for y in (y_nax, y_mma, y_sk): - assert not _bits_equal(y, y_qmv) assert _bits_equal(run("q4_0", w, s, x), y_qmv) assert _bits_equal(run("q4_0", w, s, x, KQ_NAX_QMV="0"), y_nax) assert _bits_equal(run("q4_0", w, s, x, KQ_VERIFY_NAX="3"), y_nax) @@ -1400,7 +1677,6 @@ def run(codec, w, s, x, route=None, **env): w, s, _ = _setup("q4_1", 500) x = (mx.random.normal((4, K)) * 0.5).astype(mx.bfloat16) y_qmv = run("q4_1", w, s, x, "qmv") - assert not _bits_equal(y_qmv, run("q4_1", w, s, x, "nax_splitk")) assert _bits_equal(run("q4_1", w, s, x, KQ_VERIFY_NAX="2"), y_qmv) assert _bits_equal(run("q4_1", w, s, x, KQ_VERIFY_MMA="2"), y_qmv) @@ -1409,23 +1685,22 @@ def run(codec, w, s, x, route=None, **env): x = (mx.random.normal((6, K)) * 0.5).astype(mx.bfloat16) y_qmv = run("q4_1", w, s, x, "qmv") y_sk = run("q4_1", w, s, x, "nax_splitk") - assert not _bits_equal(y_qmv, y_sk) assert _bits_equal(run("q4_1", w, s, x), y_sk) assert _bits_equal(run("q4_1", w, s, x, KQ_NAX_QMV="6"), y_qmv) - # KQ_DISABLE_NAX turns the route off even when KQ_NAX_QMV asks for it. - # pq2_0 at M 5 is the check: its non-NAX routes there (mv_ext, - # verify_mma) differ from qmv bit for bit. + # KQ_DISABLE_NAX turns the route off even when KQ_NAX_QMV asks for it: + # pq2_0 at M 5 runs verify_mma either way. w, s, _ = _setup("pq2_0", 500) x = (mx.random.normal((5, K)) * 0.5).astype(mx.bfloat16) y_off = run("pq2_0", w, s, x, KQ_DISABLE_NAX="1") - assert not _bits_equal(y_off, run("pq2_0", w, s, x, "qmv")) + assert _bits_equal(y_off, run("pq2_0", w, s, x, "verify_mma")) y = run("pq2_0", w, s, x, KQ_DISABLE_NAX="1", KQ_NAX_QMV="8") assert _bits_equal(y, y_off) # Probe for the process-static levers: prints, per named route, whether the -# default output equals that route's output bit for bit. +# default output equals that route's output bit for bit, and marks each run +# on stderr ahead of the kernels KQ_QMM_LOG names for it. _STATIC_PROBE = """ import json, os, sys import mlx.core as mx @@ -1438,6 +1713,7 @@ def run(route=None): os.environ.pop("KQ_QMM_ROUTE", None) if route: os.environ["KQ_QMM_ROUTE"] = route + os.write(2, ("[run] " + (route or "default") + chr(10)).encode()) y = kq.quantized_matmul(x, w, s, codec, transpose=True) mx.eval(y) return y @@ -1446,6 +1722,18 @@ def run(route=None): """ +def _logged_kernels(stderr): + """Kernel names KQ_QMM_LOG printed per [run] marker of _STATIC_PROBE.""" + names, cur = {}, None + for line in stderr.splitlines(): + if line.startswith("[run] "): + cur = line.split()[1] + names[cur] = [] + elif line.startswith("[kq qmm] ") and cur is not None: + names[cur].append(line.split()[2]) + return names + + # KQ_VERIFY_EXT and KQ_QMM_SPLITK are read once per process, so each case # runs in a fresh interpreter. A set KQ_VERIFY_EXT keeps per-row qmv and # NAX split-K off the mat-vec band, and a forced KQ_QMM_SPLITK keeps them @@ -1466,6 +1754,7 @@ def test_nax_small_m_yields_to_static_levers(lever, codec, m, expect, other): env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} name, value = lever.split("=") env[name] = value + env["KQ_QMM_LOG"] = "1" others = other.split(",") out = subprocess.run( [sys.executable, "-c", _STATIC_PROBE, codec, str(m), expect, *others], @@ -1474,5 +1763,10 @@ def test_nax_small_m_yields_to_static_levers(lever, codec, m, expect, other): text=True, check=True, ) - res = json.loads(out.stdout.strip().splitlines()[-1]) + bits = json.loads(out.stdout.strip().splitlines()[-1]) + names = _logged_kernels(out.stderr) + # A route matches the default when its output bits and its logged + # kernels both do: two routes can agree in every bit, split routes most + # of all with float partials. + res = {r: bits[r] and names[r] == names["default"] for r in bits} assert res[expect] and not any(res[o] for o in others), (lever, codec, m, res) diff --git a/tests/test_nax_swizzle.py b/tests/test_nax_swizzle.py new file mode 100644 index 00000000..acec77ce --- /dev/null +++ b/tests/test_nax_swizzle.py @@ -0,0 +1,148 @@ +"""KQ_NAX_SWIZZLE and the large-M rows of the NAX qmm_t tile policy. + +The swizzle folds 2^log row tiles into grid.x, so consecutive threadgroups +share one weight column band. The lever sets the fold log on the BM=64 and +BM=128 tiles, capped at the row-tile count; unset, the codec's large-M row +sets it on the row's own tile. Each case runs in a fresh interpreter +(tests/qmm_probe.py), reads the dispatched kernel and grid from +KQ_QMM_LOG, and checks the result against float64 with exact +kq.dequantize weights, on the CPU stream too, where the lever leaves the +result alone. A fold changes only the order of the row tiles, so its +output is bit-identical to the unfolded grid. +""" + +from __future__ import annotations + +import pytest +from qmm_probe import force_cpu, run_probe, tile_of + +import mlx_kquant as kq + +R1 = {"bfloat16": 6e-3, "float16": 1e-3} +K = 512 +LIVE = {"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"} +TILES = {64: "64,64,2,2", 128: "128,64,2,2"} +# The fold log of each codec's large-M row (src/kquant_matmul.cpp, +# kq_largem_row); codecs not listed fold nothing. +POLICY_LOG: dict[str, int] = {} +# Codecs whose large-M row is measured; the rest take their loader class's +# row, the same tile and floor today. +MEASURED = {"q4_k", "q5_k", "q6_k", "q8_0"} +# K ceiling of each row (kq_largem_max_k); codecs not listed have none. +MAX_K = { + "q2_k": 4096, "pq2_0": 4096, "q2_0": 4096, + "q4_k": 10240, "q8_0": 10240, "q4_0": 10240, "q4_1": 10240, + "q5_0": 10240, "q5_1": 10240, + "iq4_nl": 12288, "stq1_0": 12288, +} # fmt: skip +# A k-quant on K tiles of 64 and of 128, a flat block codec, IQ grid codecs +# and a ternary codec. +CODECS = ["q4_k", "q6_k", "q8_0", "q4_0", "iq2_xxs", "iq4_xs", "stq1_0"] + + +def _nax_gpu(): + return not force_cpu() and kq.nax_available() + + +def _kernel(res): + assert res["error"] is None, res["error"] + assert len(res["kernels"]) == 1, res["kernels"] + return res["kernels"][0] + + +def _fold(rows, bm, n, log): + """Grid of a fold of 2^log row tiles, capped at the row-tile count.""" + y = -(-rows // bm) + while log > 0 and (1 << log) > y: + log -= 1 + x = -(-n // 64) + return (x << log, -(-y // (1 << log)), 1) + + +@pytest.mark.parametrize("codec", CODECS) +def test_swizzle_grid(codec): + if not _nax_gpu(): + pytest.skip("NAX grid") + cases = [] + # 4128 rows: 33 row tiles of 128 and 65 of 64. + for bm in (64, 128): + for log in (0, 1, 2, 3): + env = {"KQ_NAX_TILE": TILES[bm], "KQ_NAX_SWIZZLE": str(log)} + cases.append((dict(shape=[4128, K], n=192, env=env), bm, 4128, log)) + # One row tile of 128: nothing to fold. + env = {"KQ_NAX_TILE": TILES[128], "KQ_NAX_SWIZZLE": "2"} + cases.append((dict(shape=[100, K], n=192, env=env), 128, 100, 2)) + # Two row tiles take a fold of 2 at most. + env = {"KQ_NAX_TILE": TILES[64], "KQ_NAX_SWIZZLE": "3"} + cases.append((dict(shape=[2, 64, K], n=192, env=env), 64, 128, 3)) + # Unset: the large-M row's fold on the row's own tile, none under + # KQ_NAX_TILE, and none below the row's floor. + env = {"KQ_NAX_TILE": None, "KQ_NAX_SWIZZLE": None} + log = POLICY_LOG.get(codec, 0) + cases.append((dict(shape=[4128, K], n=192, env=env), 128, 4128, log)) + env = {"KQ_NAX_TILE": TILES[128]} + cases.append((dict(shape=[4128, K], n=192, env=env), 128, 4128, 0)) + env = {"KQ_NAX_TILE": None} + cases.append((dict(shape=[1050, K], n=192, env=env), 64, 1050, 0)) + calls = [dict(c, codec=codec) for c, *_ in cases] + res = run_probe(calls, env=dict(LIVE, KQ_QMM_MIDM="0")) + for (c, bm, rows, log), r in zip(cases, res, strict=True): + name, grid = _kernel(r) + assert tile_of(name)[0] == bm, (c, name) + assert grid == _fold(rows, bm, c["n"], log), (c, grid) + assert r["rel"] < R1["bfloat16"], (c, r["rel"]) + + +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", CODECS) +def test_swizzle_bit_identical(codec, dtype): + # An odd row-tile count with N a multiple of 64, and a batch of two + # flattened into 4160 rows with N not a multiple of 64. + nax = _nax_gpu() + calls = [] + for shape, n in (([4128, K], 192), ([2, 2080, K], 96)): + for bm in (64, 128): + base = len(calls) + for log in (0, 1, 2, 3): + env = {"KQ_NAX_TILE": TILES[bm], "KQ_NAX_SWIZZLE": str(log)} + calls.append( + { + "codec": codec, + "shape": shape, + "n": n, + "dtype": dtype, + "seed": base, + "env": env, + "ulp_vs": base if log else None, + } + ) + res = run_probe(calls, env=LIVE) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1[dtype], (c, r["rel"]) + assert r["nonfinite"] == 0, c + if nax and c["ulp_vs"] is not None: + assert r["ulp"] == 0, (c, r["ulp"]) + assert r["digest"] == res[c["ulp_vs"]]["digest"], c + + +def test_swizzle_lever_keeps_first_read(): + if not _nax_gpu(): + pytest.skip("NAX grid") + call = {"codec": "q4_k", "shape": [4128, K], "n": 192} + env = {"KQ_QMM_LOG": "1", "KQ_NAX_TILE": TILES[128], "KQ_NAX_SWIZZLE": "2"} + calls = [dict(call), dict(call, env={"KQ_NAX_SWIZZLE": "0"})] + res = run_probe(calls, env=env, check=False) + want = _fold(4128, 128, 192, 2) + assert [_kernel(r)[1] for r in res] == [want, want] + + +@pytest.mark.parametrize("codec", kq.codecs()) +def test_large_row(codec): + from mlx_kquant._ext import _nax_large_row + + row = _nax_large_row(codec) + if not row: + pytest.skip("no Metal") + want = [128, 64, 2, 2, POLICY_LOG.get(codec, 0), 4096, 0, MAX_K.get(codec, 0)] + assert row == want + [int(codec not in MEASURED)] diff --git a/tests/test_operand_views.py b/tests/test_operand_views.py new file mode 100644 index 00000000..638587f0 --- /dev/null +++ b/tests/test_operand_views.py @@ -0,0 +1,747 @@ +"""Views of the operands that the MoE, routing and attention kernels read by +layout. + +An op builder sees the strides of an operand only when the operand has +memory: an unevaluated view reports dense strides whatever it becomes. The +ops here decide the layout of an operand in their eval, where every input +has memory, so a view gives the same result whether it is evaluated before +the call or not: + +- an activation, router logits and the three operands of route_shed are + copied when they are not row-major; +- an expert stack, a shared-expert tensor and the packed mxfp4 weights and + scales are refused when they are not row-major, with the message the op + builder gives for a view that has memory; +- the keys and values of sdpa_prefill_block_sparse are refused when their + head dim is not dense, as in the other attention ops; +- the window, pool and pool scales of the sparse attention ops and the + keys of dsa_indexer_score_decode are copied when a row of theirs is not + dense, and a packed pool is copied when its codes and scales step + batches differently; +- the scales and biases of a quantized cache are refused when the two + differ in layout or their last dim is not dense. + +An operand that is dense adds no node to the graph. + +route_shed runs on the default device, so the suite covers the GPU and, +under KQUANT_FORCE_CPU, the CPU stream. The other ops are Metal-only +kernels with no CPU eval, and their cases are skipped there. +""" + +from __future__ import annotations + +import os +import re + +import mlx.core as mx +import numpy as np +import pytest + +import mlx_kquant as kq + +metal_only = pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")), + reason="Metal-only kernel; no CPU path.", +) + +E, N, K = 4, 64, 512 +T, R = 3, 2 +S = R + 1 +_W = {} + + +def _dense(view): + """A copy of a view in its own row-major buffer. numpy has no bfloat16 + and float32 holds it exactly.""" + if view.dtype == mx.bfloat16: + wide = np.ascontiguousarray(np.array(view.astype(mx.float32))) + return mx.array(wide).astype(mx.bfloat16) + return mx.array(np.ascontiguousarray(np.array(view))) + + +def _nodes(*outputs): + """Names of the primitives in the graph of the outputs.""" + import tempfile + + with tempfile.NamedTemporaryFile("r", suffix=".dot") as f: + mx.export_to_dot(f.name, *outputs) + text = f.read() + return sorted(re.findall(r'label ="([A-Za-z0-9_]+)"', text)) + + +def _same(got, want): + got = got if isinstance(got, (list, tuple)) else [got] + want = want if isinstance(want, (list, tuple)) else [want] + mx.eval(*got, *want) + return all(mx.array_equal(g, w).item() for g, w in zip(got, want, strict=True)) + + +def _check_views(views, call, state): + """call(view) equals call(dense copy of the view) for every view.""" + for name, make in views.items(): + want = call(_dense(make())) + view = make() + if state == "evaluated": + mx.eval(view) + assert _same(call(view), want), (name, state) + + +def _strided(base): + """Views of one [rows, cols] array that are not row-major, as builders + of a fresh view. base is [2 * rows, 2 * cols] random values.""" + rows, cols = base.shape[0] // 2, base.shape[1] // 2 + turned = _dense(mx.swapaxes(base[:rows, :cols], 0, 1)) + mx.eval(turned) + return { + "even_cols": lambda: base[:rows, 0::2], + "odd_cols": lambda: base[:rows, 1::2], + "even_rows": lambda: base[0::2, :cols], + "reversed_rows": lambda: base[::-1][:rows, :cols], + "col_band": lambda: base[:rows, cols // 2 : cols // 2 + cols], + "transposed": lambda: mx.swapaxes(turned, 0, 1), + } + + +# ----------------------------- route_shed ----------------------------- + + +def _route_case(seed=0, n_experts=256): + rng = np.random.default_rng(seed) + idx = mx.array(rng.integers(0, n_experts, size=(2 * T, 2 * R)).astype(np.uint32)) + sc = mx.array(rng.uniform(0.05, 0.9, size=(2 * T, 2 * R)).astype(np.float32)) + slots = np.where( + rng.random(2 * n_experts) < 0.7, rng.integers(0, 64, 2 * n_experts), -1 + ) + slot = mx.array(slots.astype(np.int32)) + mx.eval(idx, sc, slot) + return idx, sc, slot + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +def test_route_shed_operand_views(state): + idx2, sc2, slot2 = _route_case() + n_experts = slot2.shape[0] // 2 + idx = _dense(idx2[:T, :R]) + sc = _dense(sc2[:T, :R]) + slot = _dense(slot2[:n_experts]) + _check_views(_strided(idx2), lambda v: kq.route_shed(v, sc, slot), state) + _check_views(_strided(sc2), lambda v: kq.route_shed(idx, v, slot), state) + _check_views( + { + "every_other": lambda: slot2[0::2], + "reversed": lambda: slot2[::-1][:n_experts], + "band": lambda: slot2[n_experts // 2 : n_experts // 2 + n_experts], + }, + lambda v: kq.route_shed(idx, sc, v), + state, + ) + + +def test_route_shed_adds_no_node(): + """The call a streamed decode makes: unevaluated router outputs. The + graph holds the op and the ops that feed it, and no layout node.""" + idx2, sc2, slot2 = _route_case(1) + idx = _dense(idx2[:T, :R]) + sc = _dense(sc2[:T, :R]) + feed = set(_nodes(idx + 0, sc * 1.0)) + assert "Add" in feed and "Multiply" in feed + outs = kq.route_shed(idx + 0, sc * 1.0, slot2) + assert set(_nodes(*outs)) == feed | {"KQuantRouteShed"} + outs = kq.route_shed(idx2[:T, 0::2], sc, slot2[0::2]) + assert set(_nodes(*outs)) == {"KQuantRouteShed", "Slice"} + + +# ------------------------- fused MoE gathers -------------------------- + + +def _stack(codec, seed, n=N, experts=E): + """A K-quant expert stack uint8 [experts, n, bytes_per_row].""" + key = (codec, seed, n, experts) + if key not in _W: + wf = mx.random.normal((experts * n, K), key=mx.random.key(seed)) * 0.1 + w, _ = kq.quantize(wf, codec) + w = _dense(w.reshape(experts, n, -1)) + mx.eval(w) + _W[key] = w + return _W[key] + + +def _shexp(codec, e): + """A shared-expert tensor uint8 [N, bytes_per_row] in its own buffer.""" + key = (codec, "shexp", e) + if key not in _W: + _W[key] = _dense(_stack(codec, 3)[e]) + mx.eval(_W[key]) + return _W[key] + + +def _acts(shape, dtype, seed): + """Activations of twice the row and column count, for _strided.""" + full = (*shape[:-2], 2 * shape[-2], 2 * shape[-1]) + x = (mx.random.normal(full, key=mx.random.key(seed)) * 0.1).astype(dtype) + mx.eval(x) + return x + + +def _x_views(x2): + """_strided over the last two dims of a 2-D or 3-D activation base.""" + if x2.ndim == 2: + return _strided(x2) + t, s2, k2 = x2.shape + s, k = s2 // 2, k2 // 2 + return { + "even_cols": lambda: x2[:, :s, 0::2], + "even_slots": lambda: x2[:, 0::2, :k], + "reversed": lambda: x2[::-1][:, :s, :k], + "col_band": lambda: x2[:, :s, k // 2 : k // 2 + k], + } + + +def _routing(seed=3, slots=R): + rng = np.random.default_rng(seed) + inds = mx.array(rng.integers(0, E, size=(T, slots)).astype(np.uint32)) + sc = mx.array(rng.uniform(0.05, 0.9, size=(T, slots)).astype(np.float32)) + mx.eval(inds, sc) + return inds, sc + + +def _kq_ops(codec): + """name -> (shape of x, call(x)) for the K-quant fused MoE ops.""" + gw, uw = _stack(codec, 1), _stack(codec, 2) + sgw, suw = _shexp(codec, 0), _shexp(codec, 1) + inds, _ = _routing() + inds_s, sc_s = _routing(4, S) + r = 4 + la = (mx.random.normal((E, K, r), key=mx.random.key(8)) * 0.05).astype(mx.float16) + lb = (mx.random.normal((E, r, N), key=mx.random.key(9)) * 0.05).astype(mx.float16) + mx.eval(la, lb) + return { + "moe_glu_gather_kq": ( + (T, K), + lambda x: kq.moe_glu_gather_kq(x, gw, uw, codec, inds), + ), + "gather_qmv_kq": ( + (T, R, K), + lambda x: kq.gather_qmv_kq(x, gw, codec, inds), + ), + "gather_qmv_kq_lora": ( + (T, R, K), + lambda x: kq.gather_qmv_kq(x, gw, codec, inds, lora_a=la, lora_b=lb), + ), + "moe_glu_gather_shexp_kq": ( + (T, K), + lambda x: kq.moe_glu_gather_shexp_kq(x, gw, uw, sgw, suw, codec, inds), + ), + "gather_qmv_mix_kq": ( + (T, S, K), + lambda x: kq.gather_qmv_mix_kq(x, gw, sgw, codec, inds, sc_s), + ), + "gather_qmv_mix_ns_kq": ( + (T, S, K), + lambda x: kq.gather_qmv_mix_ns_kq(x, gw, codec, inds_s, sc_s), + ), + "gather_qmv_mix_ns_kq_lora": ( + (T, S, K), + lambda x: kq.gather_qmv_mix_ns_kq( + x, gw, codec, inds_s, sc_s, lora_a=la, lora_b=lb + ), + ), + } + + +KQ_OPS = [ + "moe_glu_gather_kq", + "gather_qmv_kq", + "gather_qmv_kq_lora", + "moe_glu_gather_shexp_kq", + "gather_qmv_mix_kq", + "gather_qmv_mix_ns_kq", + "gather_qmv_mix_ns_kq_lora", +] + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("codec", ["q4_k", "q8_0"]) +@pytest.mark.parametrize("op", KQ_OPS) +def test_moe_activation_views(op, codec, state): + shape, call = _kq_ops(codec)[op] + _check_views(_x_views(_acts(shape, mx.float16, 21)), call, state) + + +def _mxfp4_stack(seed): + """A packed mxfp4 stack: uint32 [E, N, K / 8] codes, uint8 [E, N, K / 32] + block exponents near 2^-6, float32 [E, N] bias.""" + rng = np.random.default_rng(seed) + w = mx.array(rng.integers(0, 2**32, size=(E, N, K // 8), dtype=np.uint32)) + sc = mx.array(rng.integers(118, 124, size=(E, N, K // 32), dtype=np.uint8)) + b = mx.array(rng.uniform(-0.2, 0.2, (E, N)).astype(np.float32)) + mx.eval(w, sc, b) + return w, sc, b + + +def _mxfp4_ops(): + gw, gs, gb = _mxfp4_stack(1) + uw, us, ub = _mxfp4_stack(2) + inds, sc = _routing() + return { + "moe_glu_gather": ( + (T, K), + lambda x: kq.moe_glu_gather(x, gw, gs, gb, uw, us, ub, inds), + ), + "gather_qmv_bias": ( + (T, R, K), + lambda x: kq.gather_qmv_bias(x, gw, gs, gb, inds), + ), + "gather_qmv_mix_bias": ( + (T, R, K), + lambda x: kq.gather_qmv_mix_bias(x, gw, gs, gb, inds, sc), + ), + } + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize( + "op", ["moe_glu_gather", "gather_qmv_bias", "gather_qmv_mix_bias"] +) +def test_moe_mxfp4_activation_views(op, state): + shape, call = _mxfp4_ops()[op] + _check_views(_x_views(_acts(shape, mx.float16, 23)), call, state) + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize( + "dtype", [mx.float32, mx.float16, mx.bfloat16], ids=["f32", "f16", "bf16"] +) +def test_moe_router_logit_views(dtype, state): + n_experts = 128 + lw = (mx.random.normal((2 * 7, 2 * n_experts), key=mx.random.key(5)) * 2).astype( + dtype + ) + mx.eval(lw) + for shared in (False, True): + _check_views( + _strided(lw), + lambda v, shared=shared: kq.moe_router_topk(v, 4, True, shared_gate=shared), + state, + ) + + +@metal_only +def test_moe_dense_activation_adds_no_node(): + """The call a decode step makes: an unevaluated activation that is + row-major. Its graph is the op, the ops that feed x and the layout node + of the indices.""" + gw, uw = _stack("q4_k", 1), _stack("q4_k", 2) + inds, _ = _routing() + x = _dense(_acts((T, K), mx.float16, 21)[:T, :K]) + feed = _nodes(x * 1.0) + assert "Multiply" in feed and "Contiguous" not in feed + y = kq.moe_glu_gather_kq(x * 1.0, gw, uw, "q4_k", inds) + assert _nodes(y) == sorted([*feed, "Contiguous", "KQuantMoEGLUKQ"]) + lw = mx.random.normal((7, 128), key=mx.random.key(5)) + mx.eval(lw) + outs = kq.moe_router_topk(lw * 1.0, 4, True, shared_gate=False) + assert set(_nodes(*outs)) == {*_nodes(lw * 1.0), "KQuantMoERouterTopK"} + + +def _refused(call, kind, match): + with pytest.raises(kind, match=match): + mx.eval(call()) + + +@metal_only +@pytest.mark.parametrize("codec", ["q4_k", "q8_0"]) +def test_moe_strided_stack_is_refused(codec): + """The gate and up halves of one fused stack are not row-major: the + expert stride is the fused row count. The op refuses them, evaluated + (at the call) or not (at eval), and never reads them as dense.""" + gu = _stack(codec, 5, n=2 * N) + gw, uw = _stack(codec, 1), _stack(codec, 2) + sg2 = _dense(_stack(codec, 3, n=2 * N)[0]) + sgw = _shexp(codec, 0) + mx.eval(sg2) + inds, _ = _routing() + inds_s, sc_s = _routing(4, S) + x = _dense(_acts((T, K), mx.float16, 21)[:T, :K]) + h = _dense(_acts((T, R, K), mx.float16, 22)[:, :R, :K]) + hs = _dense(_acts((T, S, K), mx.float16, 22)[:, :S, :K]) + msg = "weights must be row-contiguous" + calls = { + "glu_gate": lambda g: kq.moe_glu_gather_kq(x, g, uw, codec, inds), + "glu_up": lambda g: kq.moe_glu_gather_kq(x, gw, g, codec, inds), + "qmv": lambda g: kq.gather_qmv_kq(h, g, codec, inds), + "shexp_gate": lambda g: kq.moe_glu_gather_shexp_kq( + x, g, uw, sgw, sgw, codec, inds + ), + "mix": lambda g: kq.gather_qmv_mix_kq(hs, g, sgw, codec, inds, sc_s), + "mix_ns": lambda g: kq.gather_qmv_mix_ns_kq(hs, g, codec, inds_s, sc_s), + } + for name, call in calls.items(): + for half in (lambda: gu[:, :N], lambda: gu[:, N:]): + _refused(lambda call=call, half=half: call(half()), ValueError, msg) + view = half() + mx.eval(view) + with pytest.raises(ValueError, match=msg): + call(view) + # The same rows in their own buffer run. + assert _same(call(_dense(gu[:, :N])), call(_dense(gu[:, :N]))), name + # A shared-expert tensor of every other row. + shexp = { + "shexp_gate": lambda v: kq.moe_glu_gather_shexp_kq( + x, gw, uw, v, sgw, codec, inds + ), + "shexp_up": lambda v: kq.moe_glu_gather_shexp_kq( + x, gw, uw, sgw, v, codec, inds + ), + "mix": lambda v: kq.gather_qmv_mix_kq(hs, gw, v, codec, inds, sc_s), + } + for call in shexp.values(): + _refused(lambda call=call: call(sg2[0::2]), ValueError, "shared-expert " + msg) + view = sg2[0::2] + mx.eval(view) + with pytest.raises(ValueError, match="shared-expert " + msg): + call(view) + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +def test_moe_row_major_stack_views_run(state): + """A band of experts from a larger stack is row-major at an offset. It + is read in place, evaluated or not.""" + codec = "q4_k" + big = _stack(codec, 6, experts=2 * E) + uw = _stack(codec, 2) + inds, _ = _routing() + x = _dense(_acts((T, K), mx.float16, 21)[:T, :K]) + want = kq.moe_glu_gather_kq(x, _dense(big[2 : 2 + E]), uw, codec, inds) + band = big[2 : 2 + E] + if state == "evaluated": + mx.eval(band) + assert _same(kq.moe_glu_gather_kq(x, band, uw, codec, inds), want) + + +@metal_only +def test_moe_mxfp4_strided_stack_is_refused(): + gw, gs, gb = _mxfp4_stack(1) + uw, us, ub = _mxfp4_stack(2) + w2 = mx.concatenate([gw, uw], axis=1) + s2 = mx.concatenate([gs, us], axis=1) + mx.eval(w2, s2) + inds, sc = _routing() + x = _dense(_acts((T, K), mx.float16, 21)[:T, :K]) + h = _dense(_acts((T, R, K), mx.float16, 22)[:, :R, :K]) + msg = "weights/scales must be row-contiguous" + calls = [ + lambda w, s: kq.moe_glu_gather(x, w, s, gb, uw, us, ub, inds), + lambda w, s: kq.moe_glu_gather(x, gw, gs, gb, w, s, ub, inds), + lambda w, s: kq.gather_qmv_bias(h, w, s, gb, inds), + lambda w, s: kq.gather_qmv_mix_bias(h, w, s, gb, inds, sc), + ] + for call in calls: + # A strided weight with dense scales, and dense weights with + # strided scales. + _refused(lambda call=call: call(w2[:, :N], gs), ValueError, msg) + _refused(lambda call=call: call(gw, s2[:, N:]), ValueError, msg) + wv, sv = w2[:, :N], s2[:, N:] + mx.eval(wv, sv) + with pytest.raises(ValueError, match=msg): + call(wv, gs) + with pytest.raises(ValueError, match=msg): + call(gw, sv) + + +# ------------------- sdpa_prefill_block_sparse k and v ------------------- + + +def _bs_case(): + """The block sparse prefill at head dim 256, with keys and values twice + as wide as the op reads.""" + from test_sdpa import _bs_prefill_lists + + hq, hkv, d, s, length, topk, r = 24, 2, 256, 4096, 64, 16, 4 + offset = s - length + q = mx.random.normal((1, hq, length, d), key=mx.random.key(1)).astype(mx.bfloat16) + k2 = mx.random.normal((1, hkv, s, 2 * d), key=mx.random.key(2)).astype(mx.bfloat16) + v2 = mx.random.normal((1, hkv, s, 2 * d), key=mx.random.key(3)).astype(mx.bfloat16) + mx.eval(q, k2, v2) + rng = np.random.default_rng(3) + sel = np.zeros((length, topk), dtype=np.int64) + for t in range(length): + sel[t] = rng.choice((offset + t + 1) // r, size=topk, replace=False) + pages, pmask, counts = (mx.array(a) for a in _bs_prefill_lists(sel, offset, s)) + return q, k2, v2, pages, pmask, counts, offset + + +@metal_only +def test_bs_prefill_strided_head_dim(): + """Keys and values whose head dim is every other element of a wider + buffer. Evaluated, the op copies them. Unevaluated, it cannot see the + strides when it is built, and the eval refuses them as the other + attention ops do.""" + q, k2, v2, pages, pmask, counts, offset = _bs_case() + + def call(k, v): + return kq.sdpa_prefill_block_sparse( + q, k, v, 0.0625, pages, pmask, counts, offset + ) + + want = call(_dense(k2[..., 0::2]), _dense(v2[..., 0::2])) + mx.eval(want) + assert mx.isfinite(want.astype(mx.float32)).all().item() + kv, vv = k2[..., 0::2], v2[..., 0::2] + mx.eval(kv, vv) + assert _same(call(kv, vv), want) + msg = "k/v head_dim must be contiguous" + _refused(lambda: call(k2[..., 0::2], _dense(v2[..., 0::2])), RuntimeError, msg) + _refused(lambda: call(_dense(k2[..., 0::2]), v2[..., 0::2]), RuntimeError, msg) + + +# ----------------- sparse attention window and pool rows ----------------- + + +def _embedded(a, rows_axis): + """Views that hold the values of a and are not row-major, as builders of + a fresh view. The first three have dims that fold into one run, which a + reshape keeps as a view with an element stride other than 1.""" + junk = (a.astype(mx.float32) * 3 + 1).astype(a.dtype) + cols = a.shape[-1] + even = mx.stack([a, junk], axis=-1).reshape(*a.shape[:-1], 2 * cols) + odd = mx.stack([junk, a], axis=-1).reshape(*a.shape[:-1], 2 * cols) + flip = _dense(a[..., ::-1, ::-1]) + longer = mx.concatenate([a, junk], axis=rows_axis) + turned = _dense(mx.swapaxes(a, -1, -2)) + mx.eval(even, odd, flip, longer, turned) + rows = a.shape[rows_axis] + pre = (slice(None),) * rows_axis + return { + "even_cols": lambda: even[..., 0::2], + "odd_cols": lambda: odd[..., 1::2], + "reversed": lambda: flip[..., ::-1, ::-1], + "prefix": lambda: longer[(*pre, slice(0, rows))], + "transposed": lambda: mx.swapaxes(turned, -1, -2), + } + + +def _check_embedded(views, call, want, state): + for name, make in views.items(): + view = make() + if state == "evaluated": + mx.eval(view) + assert _same(call(view), want), (name, state) + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_sparse_decode_row_views(batch, state): + from test_sdpa_sparse_decode import _case + + args, _ = _case( + batch, 16, 2, 128, 24, 300, 40, mx.float16, np.int32, True, True, True + ) + q, win, pool, idx, scale, kw = args + mx.eval(q, win, pool, idx) + want = kq.sdpa_sparse_decode(q, win, pool, idx, scale, **kw) + _check_embedded( + _embedded(win, 2), + lambda v: kq.sdpa_sparse_decode(q, v, pool, idx, scale, **kw), + want, + state, + ) + _check_embedded( + _embedded(pool, 1), + lambda v: kq.sdpa_sparse_decode(q, win, v, idx, scale, **kw), + want, + state, + ) + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_sparse_prefill_row_views(batch, state): + from test_sdpa_sparse_prefill import _case + + args, _ = _case( + batch, 16, 20, 128, 40, 300, 48, 16, mx.float16, np.int32, True, True, True + ) + q, win, pool, idx, scale, band, kw = args + mx.eval(q, win, pool, idx) + want = kq.sdpa_sparse_prefill(q, win, pool, idx, scale, band, **kw) + _check_embedded( + _embedded(win, 2), + lambda v: kq.sdpa_sparse_prefill(q, v, pool, idx, scale, band, **kw), + want, + state, + ) + _check_embedded( + _embedded(pool, 1), + lambda v: kq.sdpa_sparse_prefill(q, win, v, idx, scale, band, **kw), + want, + state, + ) + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("op", ["decode", "prefill"]) +def test_sparse_packed_pool_views(op, batch, state): + """A packed pool is codes and scales. Each is read with a unit stride in + a row, and a batch is stepped by the stride of the codes, with an eighth + of it for the scales. A view of either that steps differently gives the + values of its copy.""" + from test_sdpa_sparse_decode import _on_grid_pool + + if op == "decode": + from test_sdpa_sparse_decode import _case + + args, _ = _case( + batch, 16, 2, 128, 24, 300, 40, mx.float16, np.int32, True, True, True + ) + q, win, pool, idx, scale, kw = args + + def run(codes, scales): + return kq.sdpa_sparse_decode( + q, win, codes, idx, scale, pool_scales=scales, **kw + ) + else: + from test_sdpa_sparse_prefill import _case + + args, _ = _case( + batch, 16, 20, 128, 40, 300, 48, 16, mx.float16, np.int32, True, True, True + ) + q, win, pool, idx, scale, band, kw = args + + def run(codes, scales): + return kq.sdpa_sparse_prefill( + q, win, codes, idx, scale, band, pool_scales=scales, **kw + ) + + codes, scales = kq.latent_fp4_pack(_on_grid_pool(pool)) + mx.eval(q, win, idx, codes, scales) + want = run(codes, scales) + # The prefix view of the codes alone steps batches by the longer + # buffer, and the dense scales do not. + _check_embedded(_embedded(codes, 1), lambda v: run(v, scales), want, state) + _check_embedded(_embedded(scales, 1), lambda v: run(codes, v), want, state) + + +@metal_only +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16], ids=["f16", "bf16"]) +def test_indexer_score_decode_key_views(dtype, state): + """One sequence passes its keys through a reshape that keeps a view.""" + rng = np.random.default_rng(11) + heads, ql, rows, d = 64, 2, 1027, 128 + q = mx.array(rng.standard_normal((1, heads, ql, d)) * 0.4).astype(dtype) + k = mx.array(rng.standard_normal((1, rows, d)) * 0.4).astype(dtype) + w = mx.array(rng.standard_normal((1, ql, heads)) * 0.2).astype(dtype) + mx.eval(q, k, w) + want = kq.dsa_indexer_score_decode(q, k, w, 4106, 4) + _check_embedded( + _embedded(k, 1), + lambda v: kq.dsa_indexer_score_decode(q, v, w, 4106, 4), + want, + state, + ) + + +# ---------------- scales and biases of a quantized cache ---------------- + + +def _q8_case(batch=1, hq=16, hkv=4, rows=256, d=128): + q = mx.random.normal((batch, hq, 1, d), key=mx.random.key(64)).astype(mx.float16) + k = mx.random.normal((batch, hkv, rows, d), key=mx.random.key(65)) + v = mx.random.normal((batch, hkv, rows, d), key=mx.random.key(66)) + kw, ks, kb = mx.quantize(k.astype(mx.float16), group_size=64, bits=8) + vw, vs, vb = mx.quantize(v.astype(mx.float16), group_size=64, bits=8) + mx.eval(q, kw, ks, kb, vw, vs, vb) + return q, (kw, ks, kb), (vw, vs, vb) + + +def _q8_pair_cells(call, ks, kb): + """call(k_scales, k_biases) with one pair of a quantized cache.""" + msg = "scales and biases need one layout" + want = call(ks, kb) + mx.eval(want) + views_s = _embedded(ks, 2) + views_b = _embedded(kb, 2) + # Both as prefix slices of longer buffers, the layout of a cache with + # spare rows: read in place. + for state in ("lazy", "evaluated"): + sv, bv = views_s["prefix"](), views_b["prefix"]() + if state == "evaluated": + mx.eval(sv, bv) + assert _same(call(sv, bv), want), state + # One of the two as a prefix slice: the head strides differ. + for state in ("lazy", "evaluated"): + bv = views_b["prefix"]() + sv = views_s["prefix"]() + if state == "evaluated": + mx.eval(sv, bv) + _refused(lambda bv=bv: call(ks, bv), RuntimeError, msg) + _refused(lambda sv=sv: call(sv, kb), RuntimeError, msg) + # A last dim that is every other element. With memory the builder + # copies it, without memory the eval refuses it. + for views, pick in ((views_s, 0), (views_b, 1)): + for name in ("even_cols", "odd_cols", "reversed"): + view = views[name]() + pair = [ks, kb] + pair[pick] = view + _refused(lambda pair=pair: call(*pair), RuntimeError, msg) + view = views["even_cols"]() + mx.eval(view) + pair = [ks, kb] + pair[pick] = view + assert _same(call(*pair), want) + + +@metal_only +def test_q8_scale_and_bias_layout_decode_gqa(): + q, (kw, ks, kb), (vw, vs, vb) = _q8_case() + scale = 128**-0.5 + + def call_k(s, b): + return kq.sdpa_decode_gqa( + q, kw, vw, scale, k_scales=s, k_biases=b, v_scales=vs, v_biases=vb + ) + + def call_v(s, b): + return kq.sdpa_decode_gqa( + q, kw, vw, scale, k_scales=ks, k_biases=kb, v_scales=s, v_biases=b + ) + + _q8_pair_cells(call_k, ks, kb) + _q8_pair_cells(call_v, vs, vb) + + +@metal_only +def test_q8_scale_and_bias_layout_cascade(): + q, (kw, ks, kb), (vw, vs, vb) = _q8_case(batch=2, rows=64) + _, (sw, ss, sb), (tw, ts, tb) = _q8_case(batch=1, rows=511) + scale = 128**-0.5 + + def call(shared_k, priv_k): + return kq.sdpa_decode_gqa_cascade( + q, + sw, + tw, + kw, + vw, + scale, + k_shared_scales=shared_k[0], + k_shared_biases=shared_k[1], + v_shared_scales=ts, + v_shared_biases=tb, + k_priv_scales=priv_k[0], + k_priv_biases=priv_k[1], + v_priv_scales=vs, + v_priv_biases=vb, + ) + + _q8_pair_cells(lambda s, b: call((s, b), (ks, kb)), ss, sb) + _q8_pair_cells(lambda s, b: call((ss, sb), (s, b)), ks, kb) diff --git a/tests/test_prism_mv_range.py b/tests/test_prism_mv_range.py new file mode 100644 index 00000000..82d1c12f --- /dev/null +++ b/tests/test_prism_mv_range.py @@ -0,0 +1,81 @@ +"""PQ2_0 and Q2_0 mat-vec past the half range. + +The "big" design draws activations from 2e4 up, with one bfloat16 column +at 1e5, against small weights, so the output stays inside the bfloat16 +range. The M=1 mat-vec and the verify-shaped mat-vec stage bfloat16 +activations in float, so the result is finite and sits within a tenth of +an output rounding of float64 over the exact weights. float16 activations +keep the half pipe, whose codes are pre-scaled by 1/16 so a lane's sum +stays below the largest activation; its half sums add about two output +roundings. q2_0 runs the same kernels on a 64-weight block. +""" + +from __future__ import annotations + +import functools +import math +import subprocess + +import pytest +from qmm_probe import force_cpu, run_probe + +pytestmark = pytest.mark.skipif(force_cpu(), reason="GPU mat-vec route") + +# (M, forced route): M=1 takes the mat-vec by default. +ROUTES = [(1, None), (4, "qmv"), (4, "verify_qmv")] +DTYPES = ["bfloat16", "float16"] +DESIGNS = ["big", "normal"] +# Error beyond the rounding of the float64 result, in output roundings. +EXCESS = {"bfloat16": 0.1, "float16": 2.5} + + +def _excess(r, floor): + return math.sqrt(max(0.0, r * r - floor * floor)) / floor + + +@functools.cache +def _probe(codec="pq2_0"): + cases = [(d, t, m, r) for d in DESIGNS for t in DTYPES for m, r in ROUTES] + calls = [ + { + "codec": codec, + "shape": [m, 1024], + "n": 256, + "dtype": dtype, + "seed": 3, + "design": design, + "r2": True, + "env": { + "KQ_QMM_ROUTE": route, + "KQ_QMM_ROUTE_STRICT": "1" if route else None, + }, + } + for design, dtype, m, route in cases + ] + try: + res = run_probe(calls, env={"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"}) + except (AssertionError, subprocess.TimeoutExpired) as e: + res = [{"error": f"probe failed: {e}", "kernels": []}] * len(calls) + return dict(zip(cases, res, strict=True)) + + +@pytest.mark.parametrize("m,route", ROUTES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("design", DESIGNS) +def test_pq2_0_mv_range(design, dtype, m, route): + _check(_probe()[(design, dtype, m, route)], dtype) + + +@pytest.mark.parametrize("m,route", ROUTES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("design", DESIGNS) +def test_q2_0_mv_range(design, dtype, m, route): + _check(_probe("q2_0")[(design, dtype, m, route)], dtype) + + +def _check(r, dtype): + assert r["error"] is None, r["error"] + # The mat-vec routes log no tile kernel. + assert r["kernels"] == [], r["kernels"] + assert r["nonfinite"] == 0, r["nonfinite"] + assert _excess(r["rel"], r["floor1"]) < EXCESS[dtype], (r["rel"], r["floor1"]) diff --git a/tests/test_qk_norm_rope.py b/tests/test_qk_norm_rope.py new file mode 100644 index 00000000..2125feb0 --- /dev/null +++ b/tests/test_qk_norm_rope.py @@ -0,0 +1,365 @@ +"""qk_norm_rope against a float64 reference. + +The reference normalizes and rotates in float64 on the CPU. The op does its +math in float32 and rounds once, so its output should equal the reference +rounded to the output dtype, within 1 ulp on nearly every element, and be no +less accurate than the unfused composition, which rounds the normalized +value to the input dtype before the rotation. + +Inputs come in the layouts a caller passes: packed, sliced from one fused +qkv output, lazily transposed, misaligned for vector loads by the start, the +head stride or the row stride, and with a non-unit head-dim stride. +""" + +from __future__ import annotations + +import os +import re +import tempfile + +import mlx.core as mx +import numpy as np +import pytest +from qmm_probe import force_cpu + +import mlx_kquant as kq + +EPS = 1e-6 +DTYPES = [mx.bfloat16, mx.float16] +KERNEL_DIMS = [64, 96, 128, 256] +# Views that start aligned and have an odd head or row stride, and a packed +# view that starts one element off. Each misses the vector width in one place. +ODD = ["hstride", "lstride", "offset"] +LAYOUTS = ["packed", "qkv", "transposed", "misaligned", "dstride"] + ODD + + +def _metal(): + if force_cpu(): + return False + try: + return kq.metallib_loads() + except Exception: + return False + + +def _tables(L, D, seed): + rng = np.random.default_rng(seed) + ang = rng.uniform(-60.0, 60.0, (L, D // 2)).astype(np.float32) + return mx.array(np.cos(ang)), mx.array(np.sin(ang)) + + +def _weights(D, dtype, seed): + rng = np.random.default_rng(seed) + w = (1.0 + 0.2 * rng.standard_normal((2, D))).astype(np.float32) + return mx.array(w[0]).astype(dtype), mx.array(w[1]).astype(dtype) + + +def _odd_view(x, layout): + """A view of the values of x [B, L, H, D] in one of the ODD layouts.""" + B, L, H, D = x.shape + if layout == "hstride": + return mx.pad(x, [(0, 0), (0, 0), (0, 0), (0, 1)])[..., :D] + if layout == "lstride": + wide = mx.pad(x.reshape(B, L, H * D), [(0, 0), (0, 0), (0, 1)]) + return wide[..., : H * D].reshape(B, L, H, D) + if layout == "offset": + return mx.pad(x.reshape(-1), (1, 0))[1:].reshape(B, L, H, D) + raise ValueError(layout) + + +def _inputs(B, L, H, D, dtype, layout, seed=0, scale=None): + """q and k [B, L, H, D] in the given layout, evaluated. Standard normal + values, or with `scale` magnitudes uniform in [scale / 2, scale] with + random signs.""" + rng = np.random.default_rng(seed) + + def draw(shape): + if scale is None: + x = rng.standard_normal(shape) + else: + x = scale * rng.uniform(0.5, 1.0, shape) * rng.choice([-1.0, 1.0], shape) + return mx.array(x.astype(np.float32)).astype(dtype) + + if layout == "packed": + q, k = draw((B, L, H, D)), draw((B, L, H, D)) + elif layout == "qkv": + qkv = draw((B, L, 3 * H * D)).reshape(B, L, 3, H, D) + q, k = qkv[:, :, 0], qkv[:, :, 1] + elif layout == "transposed": + q = draw((B, H, L, D)).transpose(0, 2, 1, 3) + k = draw((B, H, L, D)).transpose(0, 2, 1, 3) + elif layout == "misaligned": + wide = draw((B, L, 2 * H * D + 2)) + q = wide[..., 1 : 1 + H * D].reshape(B, L, H, D) + k = wide[..., 1 + H * D : 1 + 2 * H * D].reshape(B, L, H, D) + elif layout == "dstride": + q = draw((B, L, D, H)).transpose(0, 1, 3, 2) + k = draw((B, L, D, H)).transpose(0, 1, 3, 2) + elif layout in ODD: + q = _odd_view(draw((B, L, H, D)), layout) + k = _odd_view(draw((B, L, H, D)), layout) + else: + raise ValueError(layout) + mx.eval(q, k) + return q, k + + +def _ref(x, w, cos, sin, eps=EPS): + """float64 rms_norm and interleaved-pair rotation, as [B, H, L, D].""" + x = np.asarray(x.astype(mx.float32), dtype=np.float64) + w = np.asarray(w.astype(mx.float32), dtype=np.float64) + c = np.asarray(cos, dtype=np.float64)[None, :, None, :] + s = np.asarray(sin, dtype=np.float64)[None, :, None, :] + y = x / np.sqrt((x * x).mean(-1, keepdims=True) + eps) * w + o = np.empty_like(y) + o[..., 0::2] = y[..., 0::2] * c - y[..., 1::2] * s + o[..., 1::2] = y[..., 0::2] * s + y[..., 1::2] * c + return o.transpose(0, 2, 1, 3) + + +def _round(ref, dtype): + """ref rounded once, to nearest even, to dtype, as float64.""" + if dtype == mx.float16: + return ref.astype(np.float16).astype(np.float64) + m, e = np.frexp(ref) + return np.ldexp(np.rint(np.ldexp(m, 8)), e - 8) + + +def _bits(a): + b = a.astype(np.int32) + return np.where(b < 0, -(b & 0x7FFF), b) + + +def _ulps(got, ref, dtype): + """Per-element ulp distance of got to ref rounded to dtype.""" + r = mx.array(_round(ref, dtype).astype(np.float32)).astype(dtype) + g = np.asarray(got.view(mx.int16)) + return np.abs(_bits(g) - _bits(np.asarray(r.view(mx.int16)))) + + +def _composed(q, k, q_w, k_w, cos, sin, eps=EPS): + """The unfused composition: mx.fast.rms_norm, then the rotation in f32.""" + + def rope(x): + b, n, h, d = x.shape + x = x.astype(mx.float32).reshape(b, n, h, d // 2, 2) + x0, x1 = x[..., 0], x[..., 1] + c = cos[None, :, None, :] + s = sin[None, :, None, :] + return mx.stack([x0 * c - x1 * s, x0 * s + x1 * c], axis=-1).reshape(b, n, h, d) + + dt = q.dtype + qo = rope(mx.fast.rms_norm(q, q_w, eps)).astype(dt) + ko = rope(mx.fast.rms_norm(k, k_w, eps)).astype(dt) + return qo.transpose(0, 2, 1, 3), ko.transpose(0, 2, 1, 3) + + +def _graph(*outs): + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "graph.dot") + mx.export_to_dot(path, *outs) + with open(path) as f: + return f.read() + + +def _check(q, k, q_w, k_w, cos, sin, against_composed=True, eps=EPS): + """Run the op, check shape, dtype and accuracy, return the outputs. On + the GPU a kernel head dim must build the fused primitive, and any other + head dim the composition.""" + B, L, H, D = q.shape + qo, ko = kq.qk_norm_rope(q, k, q_w, k_w, cos, sin, eps) + fused = _metal() and D in KERNEL_DIMS + assert ("KQuantQkNormRope" in _graph(qo, ko)) == fused + mx.eval(qo, ko) + assert qo.shape == (B, H, L, D) and ko.shape == (B, H, L, D) + assert qo.dtype == q.dtype and ko.dtype == q.dtype + refs = (_ref(q, q_w, cos, sin, eps), _ref(k, k_w, cos, sin, eps)) + comp = _composed(q, k, q_w, k_w, cos, sin, eps) if against_composed else None + for i, (got, ref) in enumerate(zip((qo, ko), refs, strict=True)): + g = np.asarray(got.astype(mx.float32), dtype=np.float64) + assert np.isfinite(g).all() + u = _ulps(got, ref, q.dtype) + assert (u > 1).mean() <= 1e-4, np.bincount(np.minimum(u, 8).ravel()) + if comp is not None: + c = np.asarray(comp[i].astype(mx.float32), dtype=np.float64) + assert np.abs(g - ref).max() <= np.abs(c - ref).max() + assert np.abs(g - ref).mean() <= np.abs(c - ref).mean() + return qo, ko + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize("B", [1, 2]) +def test_layouts(dtype, layout, B): + L, H, D = 77, 3, 128 + q, k = _inputs(B, L, H, D, dtype, layout, seed=B) + q_w, k_w = _weights(D, dtype, 1) + cos, sin = _tables(L, D, 2) + _check(q, k, q_w, k_w, cos, sin) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("D", KERNEL_DIMS + [32, 80]) +def test_head_dims(dtype, D): + # 32 and 80 have no kernel and run the composition on every stream. + B, L, H = 2, 45, 5 + q, k = _inputs(B, L, H, D, dtype, "qkv", seed=D) + q_w, k_w = _weights(D, dtype, 3) + cos, sin = _tables(L, D, 4) + _check(q, k, q_w, k_w, cos, sin) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("B,L", [(1, 33000), (2, 16896)]) +def test_rows_past_int16(dtype, B, L): + # 2 x 16896 is a batch of two 2048-pixel images flattened to rows. + H, D = 2, 128 + q, k = _inputs(B, L, H, D, dtype, "qkv", seed=L) + q_w, k_w = _weights(D, dtype, 5) + cos, sin = _tables(L, D, 6) + _check(q, k, q_w, k_w, cos, sin, against_composed=False) + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_zimage_shape(dtype): + # One Z-Image main-layer call at 1024 pixels: 4096 image rows and 32 + # caption rows, 30 heads of 128, q and k sliced from the fused qkv. + B, L, H, D = 1, 4128, 30, 128 + q, k = _inputs(B, L, H, D, dtype, "qkv", seed=9) + q_w, k_w = _weights(D, dtype, 7) + cos, sin = _tables(L, D, 8) + _check(q, k, q_w, k_w, cos, sin) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("scale", [1e4, 6e4]) +def test_large_activations(dtype, scale): + # Every square is past the float16 range and must not overflow: the mean + # square is accumulated in f32 and the output is back at the weight + # scale. + B, L, H, D = 2, 33, 4, 128 + q, k = _inputs(B, L, H, D, dtype, "qkv", seed=11, scale=scale) + q_w, k_w = _weights(D, dtype, 12) + cos, sin = _tables(L, D, 13) + _check(q, k, q_w, k_w, cos, sin, against_composed=False) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize( + "scale,eps,zero_row", [(1e-3, 1e-5, False), (None, 0.5, False), (None, 1e-5, True)] +) +def test_eps(dtype, scale, eps, zero_row): + # At inputs of 1e-3 the mean square is below the Z-Image eps of 1e-5, and + # eps 0.5 is half the mean square of unit inputs, so in both eps changes + # every output. A head row of zeros comes out as zeros, since eps keeps + # its rsqrt finite. + B, L, H, D = 2, 21, 3, 128 + q, k = _inputs(B, L, H, D, dtype, "packed", seed=19, scale=scale) + if zero_row: + q[1, 4, 2] = 0 + k[0, 7, 0] = 0 + mx.eval(q, k) + q_w, k_w = _weights(D, dtype, 20) + cos, sin = _tables(L, D, 21) + qo, ko = _check(q, k, q_w, k_w, cos, sin, against_composed=False, eps=eps) + if zero_row: + assert (np.asarray(qo[1, 2, 4].astype(mx.float32)) == 0).all() + assert (np.asarray(ko[0, 0, 7].astype(mx.float32)) == 0).all() + + +def test_float32_weights_and_tables(): + # Weights and tables in other float dtypes are used in f32. + B, L, H, D = 1, 20, 2, 64 + q, k = _inputs(B, L, H, D, mx.bfloat16, "qkv", seed=14) + q_w, k_w = _weights(D, mx.float32, 15) + cos, sin = _tables(L, D, 16) + _check(q, k, q_w, k_w, cos, sin, against_composed=False) + qo, ko = kq.qk_norm_rope( + q, k, q_w, k_w, cos.astype(mx.float16), sin.astype(mx.float16), EPS + ) + ref = _ref(q, q_w, cos.astype(mx.float16), sin.astype(mx.float16)) + assert (_ulps(qo, ref, mx.bfloat16) > 1).mean() <= 1e-4 + + +def test_empty_rows(): + q = mx.zeros((1, 0, 2, 64), dtype=mx.bfloat16) + w = mx.ones((64,), dtype=mx.bfloat16) + t = mx.zeros((0, 32)) + qo, ko = kq.qk_norm_rope(q, q, w, w, t, t, EPS) + mx.eval(qo, ko) + assert qo.shape == (1, 2, 0, 64) and ko.shape == (1, 2, 0, 64) + + +def test_rejects_bad_args(): + q = mx.zeros((1, 4, 2, 64), dtype=mx.bfloat16) + w = mx.ones((64,), dtype=mx.bfloat16) + c = mx.zeros((4, 32)) + with pytest.raises(ValueError): + kq.qk_norm_rope(q[..., :63], q[..., :63], w[:63], w[:63], c, c, EPS) + with pytest.raises(ValueError): + kq.qk_norm_rope(q, q[:, :3], w, w, c, c, EPS) + with pytest.raises(ValueError): + kq.qk_norm_rope(q, q.astype(mx.float16), w, w, c, c, EPS) + with pytest.raises(ValueError): + kq.qk_norm_rope(q.astype(mx.float32), q.astype(mx.float32), w, w, c, c, EPS) + with pytest.raises(ValueError): + kq.qk_norm_rope(q[0], q[0], w, w, c, c, EPS) + with pytest.raises(ValueError): + kq.qk_norm_rope(q, q, w[:32], w, c, c, EPS) + with pytest.raises(ValueError): + kq.qk_norm_rope(q, q, w, w, c[:3], c, EPS) + with pytest.raises(ValueError): + kq.qk_norm_rope(q, q, w, w, c, c[:, :16], EPS) + + +def _kernel_names(): + return sorted(n for n in kq.metallib_kernels() if n.startswith("kq_qk_norm_rope_")) + + +@pytest.mark.skipif(not _metal(), reason="Metal kernel names") +def test_kernel_set(): + want = sorted( + f"kq_qk_norm_rope_{t}_{d}" + for t in ("bfloat16_t", "float16_t") + for d in KERNEL_DIMS + ) + assert _kernel_names() == want + + +@pytest.mark.skipif(not _metal(), reason="Metal kernel names") +@pytest.mark.parametrize("name", _kernel_names() if _metal() else []) +def test_every_kernel(name): + # Each kernel runs both load paths: vector loads on the fused qkv slice, + # element loads on a misaligned view and on each ODD view. All read the + # same values and must agree bit for bit. + m = re.fullmatch(r"kq_qk_norm_rope_(bfloat16|float16)_t_(\d+)", name) + assert m, name + dtype = getattr(mx, m.group(1)) + D = int(m.group(2)) + B, L, H = 2, 19, 3 + q_w, k_w = _weights(D, dtype, 17) + cos, sin = _tables(L, D, 18) + qa, ka = _inputs(B, L, H, D, dtype, "qkv", seed=D) + qb = mx.contiguous(qa) + wide = mx.concatenate( + [ + mx.zeros((B, L, 1), dtype=dtype), + mx.concatenate([qa, ka], axis=2).reshape(B, L, 2 * H * D), + mx.zeros((B, L, 1), dtype=dtype), + ], + axis=-1, + ) + qm = wide[..., 1 : 1 + H * D].reshape(B, L, H, D) + km = wide[..., 1 + H * D : 1 + 2 * H * D].reshape(B, L, H, D) + mx.eval(qb, qm, km) + vec = _check(qa, ka, q_w, k_w, cos, sin) + strided = _check(qm, km, q_w, k_w, cos, sin) + packed = _check(qb, ka, q_w, k_w, cos, sin) + for a, b in zip(vec, strided, strict=True): + assert mx.array_equal(a, b).item() + assert mx.array_equal(vec[0], packed[0]).item() + for layout in ODD: + qv, kv = _odd_view(qa, layout), _odd_view(ka, layout) + mx.eval(qv, kv) + for a, b in zip(vec, _check(qv, kv, q_w, k_w, cos, sin), strict=True): + assert mx.array_equal(a, b).item() diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index 8a9fc782..47a79b7f 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -688,11 +688,10 @@ def call(ops): mx.eval(call([ks, kb, vs, vb])) -def test_sdpa_decode_gqa_starts_zero_matches_plain(monkeypatch): - # all-zero starts must match the no-starts call on the same inputs. A - # call with starts keeps the simdgroup pass 1, so the plain call is - # pinned to it too. - monkeypatch.setenv("KQ_GQA_NAX", "0") +def test_sdpa_decode_gqa_starts_zero_matches_plain(gqa_pass1): + # all-zero starts must match the no-starts call on the same inputs, on + # either pass 1. The tensor-op kernel runs the same arithmetic at a + # start of 0, so its output is bit-identical. scale = 1.0 / (512**0.5) q, k, v = _make(4, 24, 4, 1, 2048, 512, mx.bfloat16, seed=3, strided=False) starts = mx.zeros((4,), dtype=mx.int32) @@ -700,6 +699,84 @@ def test_sdpa_decode_gqa_starts_zero_matches_plain(monkeypatch): b = kq.sdpa_decode_gqa(q, k, v, scale) _eval_or_skip(a, b) assert _rel(a, b) < 1e-6 + if gqa_pass1 == "1": + assert mx.array_equal(a, b).item() + + +# Left-padded rows at head_dim 512 on both pass 1 kernels: pad 0, a pad on +# a 256-key tile boundary, a pad mid-tile, the maximum in-contract pad (one +# visible key), pads that leave whole splits empty, the serve pads of three +# rows one prefill tick apart, and the gemma-4 31b global layer at batch 8. +@pytest.mark.parametrize( + "kL,B,Hq,Hkv,pads,splits,dtype", + [ + (3071, 4, 32, 4, [0, 1024, 1543, 3070], 16, mx.bfloat16), + (3071, 4, 32, 4, [0, 1024, 1543, 3070], 0, mx.float16), + (4300, 2, 16, 2, [3000, 4299], 64, mx.bfloat16), + (4300, 3, 32, 4, [0, 1, 2], 0, mx.bfloat16), + (8192, 8, 32, 4, None, 0, mx.bfloat16), + ], +) +def test_sdpa_decode_gqa_hd512_route_starts( + gqa_pass1, kL, B, Hq, Hkv, pads, splits, dtype +): + _check_gqa_starts( + 512, + kL=kL, + dtype=dtype, + B=B, + Hq=Hq, + Hkv=Hkv, + pads=pads, + strided=True, + splits=splits, + ) + + +def test_sdpa_decode_gqa_hd512_starts_takes_nax(monkeypatch): + # With starts the tensor-op pass 1 runs from the same 3072-row switch, + # so KQ_GQA_NAX changes the output there (the kernels round + # differently). + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + scale = 1.0 / (512**0.5) + q, k, v = _make(3, 32, 4, 1, 2000, 512, mx.bfloat16, seed=81, strided=False) + pads = [0, 1, 2] + starts = mx.array(pads, dtype=mx.int32) + outs = [] + for flag in ("0", "1"): + monkeypatch.setenv("KQ_GQA_NAX", flag) + outs.append(kq.sdpa_decode_gqa(q, k, v, scale, starts=starts)) + _eval_or_skip(outs[-1]) + ref = _ref_sdpa_starts(q, k, v, scale, pads, 1) + mx.eval(ref) + assert not mx.array_equal(outs[0], outs[1]).item() + for out in outs: + assert _rel(out, ref) < REL_BOUND[mx.bfloat16] + + +def test_sdpa_decode_gqa_hd512_starts_skips_pad_rows(monkeypatch): + # The tensor-op pass 1 never reads K or V below a row's start: NaN in + # every pad row, including the pad keys of a partly padded tile, leaves + # the output finite and equal to the reference. + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + monkeypatch.setenv("KQ_GQA_NAX", "1") + D, kL, B = 512, 2000, 3 + scale = 1.0 / (D**0.5) + q, k, v = _make(B, 32, 4, 1, kL, D, mx.bfloat16, seed=82, strided=False) + pads = [0, 300, 1031] + pos = mx.arange(kL).reshape(1, 1, kL, 1) + pad = pos < mx.array(pads).reshape(B, 1, 1, 1) + nan = mx.array(float("nan"), dtype=mx.bfloat16) + kn = mx.where(pad, nan, k) + vn = mx.where(pad, nan, v) + starts = mx.array(pads, dtype=mx.int32) + got = kq.sdpa_decode_gqa(q, kn, vn, scale, starts=starts) + ref = _ref_sdpa_starts(q, k, v, scale, pads, 1) + _eval_or_skip(got, ref) + assert mx.isfinite(got.astype(mx.float32)).all().item() + assert _rel(got, ref) < REL_BOUND[mx.bfloat16] @pytest.mark.parametrize("qL", [2, 4]) diff --git a/tests/test_split_kernel_dispatch.py b/tests/test_split_kernel_dispatch.py new file mode 100644 index 00000000..5d6e375d --- /dev/null +++ b/tests/test_split_kernel_dispatch.py @@ -0,0 +1,133 @@ +"""Every split-K kernel name the host can build, dispatched. + +For each codec with split-K kernels, the calls below force each split +route at the row counts that select each of its tiles, on both activation +dtypes, on an N that is and is not a multiple of the column tile, with +partials in the activation dtype and in float (KQ_SPLITK_F32), and with the +half verify_mma kernels (KQ_VERIFY_MMA_HALF). KQ_QMM_LOG names the kernel +each call ran. Together the calls must run every name the host enumerator +lists for the codec, and each result must meet the R1 bound against float64 +with exact kq.dequantize weights, so a kernel that compiles to nothing, or +that no call reaches, fails here. +""" + +from __future__ import annotations + +import re + +import pytest +from mlx_kquant._ext import _split_kernel_names +from qmm_probe import force_cpu, run_probe + +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY + + +def _no_nax(): + if force_cpu(): + return True + try: + return not (kq.metallib_loads() and kq.nax_available()) + except Exception: + return True + + +pytestmark = pytest.mark.skipif(_no_nax(), reason="NAX split-K kernel names") + +R1 = {"bfloat16": 6e-3, "float16": 1e-3} +# The half verify_mma kernels sum each block in half (test_verify_mma_range). +HALF = 1e-2 +K = 4096 +CODEC = re.compile( + r"^kquant_(.+?)_(?:sb_)?(?:qmm_t_splitk|qmm_t_nax_splitk|verify_mma|verify_nax)_" +) +STRICT = {"KQ_QMM_ROUTE_STRICT": "1"} +CLEAR = { + "KQ_QMM_MIDM": None, + "KQ_DISABLE_NAX": None, + "KQ_QMM_ROUTE": None, + "KQ_QMM_ROUTE_STRICT": None, + "KQ_VERIFY_MMA_HALF": None, +} +# (family the arm may be declined for, rows, K, env). The split-K tile is +# BM=8 through M 8, BM=16 through M 16 and the plain tile above; the NAX +# split-K tile is BM=32 below the mid-M band and BM=64 in it. +ARMS = [ + (None, 8, K, {"KQ_QMM_ROUTE": "splitk", **STRICT}), + (None, 16, K, {"KQ_QMM_ROUTE": "splitk", **STRICT}), + (None, 32, K, {"KQ_QMM_ROUTE": "splitk", **STRICT}), + (None, 8, K, {"KQ_QMM_ROUTE": "nax_splitk", **STRICT}), + (None, 64, K, {"KQ_QMM_MIDM": "1"}), + ("_verify_mma_", 8, K, {"KQ_QMM_ROUTE": "verify_mma", **STRICT}), + ( + "_verify_mma_", + 8, + K, + {"KQ_QMM_ROUTE": "verify_mma", "KQ_VERIFY_MMA_HALF": "1", **STRICT}, + ), + ("_verify_nax_", 8, K, {"KQ_QMM_ROUTE": "verify_nax", **STRICT}), +] +# K off the 256 step, for the codecs whose block divides it: q4_0 runs its +# 64-step verify_nax kernel there. +OFF_STEP = ("_verify_nax_", 8, K + 64, {"KQ_QMM_ROUTE": "verify_nax", **STRICT}) + + +def _host_names(): + """Host split-K names per codec, without the shared accumulate kernels.""" + out = {} + for n in _split_kernel_names(): + m = CODEC.match(n) + if m: + out.setdefault(m[1], set()).add(n) + return out + + +HOST = {} if _no_nax() else _host_names() + + +def test_split_codecs_found(): + assert {"q4_k", "q8_0", "q2_0", "pq2_0", "iq4_xs", "iq4_nl"} <= set(HOST) + # Every codec has the three split-K tiles and both NAX split-K tiles. + for codec, names in HOST.items(): + for frag in ( + "_qmm_t_splitk_bm8_", + "_qmm_t_splitk_bm16_", + "_qmm_t_nax_splitk_", + ): + assert any(frag in n for n in names), (codec, frag) + + +@pytest.mark.parametrize("codec", sorted(HOST)) +def test_every_split_kernel_runs(codec): + host = HOST[codec] + arms = list(ARMS) + if (K + 64) % CODEC_GEOMETRY[codec][3] == 0: + arms.append(OFF_STEP) + calls, families = [], [] + for dtype in ("bfloat16", "float16"): + for f32 in ("0", "1"): + for n in (256, 250): + for family, m, k, env in arms: + calls.append( + { + "codec": codec, + "shape": [m, k], + "n": n, + "dtype": dtype, + "seed": 1, + "env": {**CLEAR, **env, "KQ_SPLITK_F32": f32}, + } + ) + families.append(family) + res = run_probe(calls, env={"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"}) + ran = set() + for c, family, r in zip(calls, families, res, strict=True): + if r["error"] is not None: + # A forced verify route declines a codec or a K it has no + # kernel for. + assert family is not None and "does not serve" in r["error"], (c, r) + continue + half = c["env"].get("KQ_VERIFY_MMA_HALF") == "1" + assert r["rel"] < (HALF if half else R1[c["dtype"]]), (c, r["rel"]) + ran |= {name for name, _ in r["kernels"]} + assert sorted(host - ran) == [] diff --git a/tests/test_split_kernel_names.py b/tests/test_split_kernel_names.py new file mode 100644 index 00000000..080ac541 --- /dev/null +++ b/tests/test_split_kernel_names.py @@ -0,0 +1,144 @@ +"""Split-K kernel names, with no dispatch. + +Each split-K kernel has a _pf32 twin that stores float partials (float16 +by default, bfloat16 under KQ_SPLITK_F32=1), and each twin has its base +kernel. The ops cast float32 activations to bfloat16, so no split-K kernel +takes float activations, and a codec outside the split-K codecs has no +split-K kernel. Every name the split-K dispatchers can build must be in the +metallib. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys + +import pytest +from mlx_kquant._ext import _split_kernel_names +from qmm_probe import force_cpu + +import mlx_kquant as kq + + +def _no_metallib(): + """True under KQUANT_FORCE_CPU, on a build without Metal, or when the + bundled metallib does not load.""" + if force_cpu(): + return True + try: + return not kq.metallib_loads() + except Exception: + return True + + +pytestmark = pytest.mark.skipif(_no_metallib(), reason="Metal kernel names") + +FAMILY = re.compile( + r"_(qmm_t_splitk|qmm_t_nax_splitk|verify_mma|verify_nax)_|^kquant_qmm_splitk_accum_" +) +TYPE = re.compile(r"_(bfloat16_t|float16_t|float)(?=_gs_|$)") +CODEC = re.compile( + r"^kquant_(.+?)_(?:sb_)?(?:qmm_t_splitk|qmm_t_nax_splitk|verify_mma|verify_nax)_" +) +BASE = "kquant_q4_k_qmm_t_splitk_bm8_bfloat16_t_gs_256_b_4_alN_true" +FLOAT = "kquant_q4_k_qmm_t_splitk_bm8_float_gs_256_b_4_alN_true" +FP4 = "kquant_mxfp4_qmm_t_splitk_bfloat16_t_gs_32_b_4_alN_true" + + +def _split_codecs(): + """The codecs kq_splitk_codec admits: those the host enumerator names a + qmm_t_splitk kernel for.""" + return {CODEC.match(n)[1] for n in _split_kernel_names() if "_qmm_t_splitk_" in n} + + +def _pairing_errors(names, codecs): + """(reason, name) per break: a split-K kernel on float activations, one + for a codec outside codecs, a kernel without its _pf32 twin, a twin + without its base.""" + names = set(names) + errors = [] + for n in sorted(names): + if not FAMILY.search(n): + continue + base = n.removesuffix("_pf32") + m = TYPE.search(base) + c = CODEC.match(n) + if m and m[1] == "float": + errors.append(("float", n)) + elif c and c[1] not in codecs: + errors.append(("codec", n)) + elif n != base: + if base not in names: + errors.append(("no base", n)) + elif n + "_pf32" not in names: + errors.append(("no twin", n)) + return errors + + +def test_split_kernels_pair(): + names = kq.metallib_kernels() + split = [n for n in names if FAMILY.search(n)] + assert any(n.endswith("_pf32") for n in split), len(split) + assert _pairing_errors(names, _split_codecs()) == [] + + +def test_pairing_check_flags_breaks(): + names = set(kq.metallib_kernels()) + codecs = _split_codecs() + assert "q4_k" in codecs and "mxfp4" not in codecs, sorted(codecs) + twin = BASE + "_pf32" + assert {BASE, twin} <= names and not {FLOAT, FP4} & names + assert _pairing_errors(names - {BASE}, codecs) == [("no base", twin)] + assert _pairing_errors(names - {twin}, codecs) == [("no twin", BASE)] + assert _pairing_errors(names | {FLOAT}, codecs) == [("float", FLOAT)] + bad = FLOAT + "_pf32" + assert _pairing_errors(names | {bad}, codecs) == [("float", bad)] + assert _pairing_errors(names | {FP4}, codecs) == [("codec", FP4)] + + +def test_host_split_names_exist(): + host = set(_split_kernel_names()) + lib = set(kq.metallib_kernels()) + assert sorted(host - lib)[:10] == [] + assert host - (lib - {BASE}) == {BASE} + for frag in ( + "_qmm_t_splitk_bm8_", + "_qmm_t_splitk_bm16_", + "_qmm_t_nax_splitk_", + "_verify_mma_", + "_verify_nax_", + "kquant_q4_0_sb_verify_nax_", + "kquant_qmm_splitk_accum_", + ): + assert any(frag in n for n in host), frag + for end in ("_pf32", "_h", "_h_pf32"): + assert any(n.endswith(end) for n in host), end + assert not any(TYPE.search(n.removesuffix("_pf32"))[1] == "float" for n in host) + + +def test_host_names_follow_the_dispatch_levers(): + # The enumerator takes its tiles and verify_nax steps from the dispatch + # helpers, so the levers that retire a tile or a step retire its names. + env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + env.update(KQ_QMM_SPLITK_BM8="0", KQ_VERIFY_NAX_Q4_0_SB="0") + code = ( + "import json; from mlx_kquant._ext import _split_kernel_names as f; " + "print(json.dumps(f()))" + ) + out = subprocess.run( + [sys.executable, "-c", code], + env=env, + capture_output=True, + text=True, + check=True, + ) + names = set(json.loads(out.stdout.strip().splitlines()[-1])) + host = set(_split_kernel_names()) + assert names < host + retired = { + n for n in host if "_splitk_bm8_" in n or n.startswith("kquant_q4_0_sb_") + } + assert retired and host - names == retired diff --git a/tests/test_splitk_cap.py b/tests/test_splitk_cap.py new file mode 100644 index 00000000..b03cd597 --- /dev/null +++ b/tests/test_splitk_cap.py @@ -0,0 +1,333 @@ +"""KQ_SPLITK_CAP, the grid cap on the small-M split-K slice target. + +The cap holds the split grid of the NAX split-K route and the ALU split +route (M <= 32) to n threadgroups per GPU core. With 40 cores pinned and +n = 1, N 4096 takes one slice at every M 9 to 32, and the call falls +through to the route the split displaced: the BM=32 tile at M 13 to 32, +and at M 9 to 12 the small-M qmm route or mv_ext, by codec and shape, or +per-row qmv under KQ_VERIFY_EXT=0. verify_qmv serves M 8 at most, so it +is no fallthrough here. Each capped call must match the call with the +split route off bit for bit, and float64 within the R1 bound. A partial +cap cuts the slice count, a forced target is not capped, and 0 keeps the +fixed target. The mat-vec routes do not log, so the output digest shows +which of them ran. + +Unset, the cap keeps the fixed target on the ALU route, and the NAX route +takes its default: one slice at a grid of 16 threadgroups per core or more +where the call without the split runs the same BM=32 tile, except q6_k and +iq2_xxs below M 32, and q4_0 and q8_0 on float partials held to 16 +threadgroups per core below a full grid. q8_0 below M 13 keeps the target +on every grid. +""" + +from __future__ import annotations + +import functools +import subprocess + +import pytest +from qmm_probe import force_cpu, run_probe, tile_of + +import mlx_kquant as kq + +pytestmark = pytest.mark.skipif(force_cpu(), reason="GPU routing lever") + +R1 = 6e-3 +N = 4096 +K = 5120 +CORES = {"KQ_GPU_CORES": "40", "KQ_QMM_LOG": "1"} +CODECS = ["q4_k", "q6_k", "q8_0"] +MS = list(range(9, 33)) +# The M 9 to 12 calls whose fallthrough is mv_ext at N 4096, K 5120: q8_0 +# takes no small-M qmm route below N 8192, and q6_k enters it at M 11. +MV_EXT = {("q6_k", 9), ("q6_k", 10)} | {("q8_0", m) for m in range(9, 13)} +# Entry M of the ALU split route on NAX GPUs, at or above the entry on +# other GPUs, so the route runs by default from here on either. +ALU_ENTRY = {"q4_k": 8, "q6_k": 12, "q8_0": 16} + + +def _probe(calls, env): + """run_probe results, or one failure per call if the probe fails, so the + tests that share a probe fail on their own entries.""" + try: + return run_probe(calls, env=env) + except (AssertionError, subprocess.TimeoutExpired) as e: + return [{"error": f"probe failed: {e}", "kernels": []}] * len(calls) + + +def _nax_gpu(): + return kq.nax_available() + + +def _split(res, route): + """(name, grid) of a split-K dispatch of ``route`` with more than one + slice.""" + assert res["error"] is None, res["error"] + assert len(res["kernels"]) == 1, res["kernels"] + name, grid = res["kernels"][0] + assert route in name and grid[2] > 1, (name, grid) + return name, grid + + +def _same(capped, off): + """The capped call ran what the call with the split off ran.""" + for r in (capped, off): + assert r["error"] is None, r["error"] + assert r["rel"] < R1, r["rel"] + assert all("splitk" not in n for n, _ in capped["kernels"]), capped["kernels"] + assert capped["kernels"] == off["kernels"], (capped["kernels"], off["kernels"]) + assert capped["digest"] == off["digest"] + + +def _call(codec, m, env, n=N): + return {"codec": codec, "shape": [m, K], "n": n, "seed": m, "env": env} + + +@functools.cache +def _nax_fallthrough(): + """Default, KQ_SPLITK_CAP=1 and KQ_QMM_SPLITK_NAX=0 results per (codec, + M).""" + arms = [ + {"KQ_SPLITK_CAP": None, "KQ_QMM_SPLITK_NAX": None}, + {"KQ_SPLITK_CAP": "1"}, + {"KQ_SPLITK_CAP": None, "KQ_QMM_SPLITK_NAX": "0"}, + ] + cases = [(c, m) for c in CODECS for m in MS] + calls = [_call(c, m, env) for c, m in cases for env in arms] + res = _probe(calls, CORES) + return {case: res[3 * i : 3 * i + 3] for i, case in enumerate(cases)} + + +@pytest.mark.parametrize("m", MS) +@pytest.mark.parametrize("codec", CODECS) +def test_nax_cap_falls_through(codec, m): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + default, capped, off = _nax_fallthrough()[(codec, m)] + _split(default, "_qmm_t_nax_splitk_") + assert default["rel"] < R1, default["rel"] + _same(capped, off) + if (codec, m) in MV_EXT: + assert capped["kernels"] == [], capped["kernels"] + else: + ((name, grid),) = capped["kernels"] + assert tile_of(name)[0] == 32 and grid[2] == 1, (name, grid) + + +# KQ_VERIFY_EXT is read once, so the per-row qmv fallthrough runs in its own +# process. KQ_QMM_SPLITK_NAX=1 forces the split route with the default +# target, which the cap still holds. +QMV_MS = [9, 10, 11, 12] + + +@functools.cache +def _nax_qmv_fallthrough(): + arms = [ + {"KQ_SPLITK_CAP": None, "KQ_QMM_SPLITK_NAX": "1"}, + {"KQ_SPLITK_CAP": "1", "KQ_QMM_SPLITK_NAX": "1"}, + {"KQ_SPLITK_CAP": None, "KQ_QMM_SPLITK_NAX": "0"}, + ] + calls = [_call("q8_0", m, env) for m in QMV_MS for env in arms] + res = _probe(calls, dict(CORES, KQ_VERIFY_EXT="0")) + return {m: res[3 * i : 3 * i + 3] for i, m in enumerate(QMV_MS)} + + +@pytest.mark.parametrize("m", QMV_MS) +def test_nax_cap_falls_through_to_qmv(m): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + forced, capped, off = _nax_qmv_fallthrough()[m] + _split(forced, "_qmm_t_nax_splitk_") + _same(capped, off) + assert capped["kernels"] == [] + + +# N 1024 at M 16 is 16 BM=32 tiles, and K 5120 is 20 units of 256 weights, +# which the default target of 16 cuts into 10 slices. A cap of 2 per core +# over 40 cores leaves room for 5. A forced target ignores the cap, and 0 +# keeps the default. +def test_nax_cap_partial(): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + arms = [ + {"KQ_SPLITK_CAP": None}, + {"KQ_SPLITK_CAP": "2"}, + {"KQ_SPLITK_CAP": "1", "KQ_QMM_SPLITK_NAX": "10"}, + {"KQ_SPLITK_CAP": "0", "KQ_QMM_SPLITK_NAX": None}, + ] + res = _probe([_call("q4_k", 16, env, n=1024) for env in arms], CORES) + depth = [_split(r, "_qmm_t_nax_splitk_")[1][2] for r in res] + assert depth == [10, 5, 10, 10], depth + for r in res: + assert r["rel"] < R1, r["rel"] + + +def _alu_cases(): + return [(c, m) for c in CODECS for m in MS if m >= max(ALU_ENTRY[c], 9)] + + +@functools.cache +def _alu_fallthrough(): + """Default and KQ_SPLITK_CAP=1 results per (codec, M) with NAX off, and + the same calls with the ALU split route off in a second process + (KQ_QMM_SPLITK is read once).""" + env = dict(CORES, KQ_DISABLE_NAX="1") + cases = _alu_cases() + calls = [_call(c, m, {"KQ_SPLITK_CAP": v}) for c, m in cases for v in (None, "1")] + res = _probe(calls, env) + off = _probe([_call(c, m, {}) for c, m in cases], dict(env, KQ_QMM_SPLITK="0")) + return {case: (res[2 * i], res[2 * i + 1], off[i]) for i, case in enumerate(cases)} + + +@pytest.mark.parametrize("codec,m", _alu_cases()) +def test_alu_cap_falls_through(codec, m): + default, capped, off = _alu_fallthrough()[(codec, m)] + _split(default, "_qmm_t_splitk_") + assert default["rel"] < R1, default["rel"] + _same(capped, off) + + +# N 1024 at M 24 is 32 BM=32 by BN=32 tiles, and K 5120 is 20 blocks, which +# the default target of 16 cuts into 10 slices. A cap of 2 per core over 40 +# cores leaves room for 2. A forced target (KQ_QMM_SPLITK, read once) +# ignores the cap. +def test_alu_cap_partial(): + env = dict(CORES, KQ_DISABLE_NAX="1") + res = _probe( + [_call("q4_k", 24, {"KQ_SPLITK_CAP": v}, n=1024) for v in (None, "2")], env + ) + forced = _probe( + [_call("q4_k", 24, {"KQ_SPLITK_CAP": "1"}, n=1024)], + dict(env, KQ_QMM_SPLITK="4"), + ) + depth = [_split(r, "_qmm_t_splitk_")[1][2] for r in res + forced] + assert depth == [10, 2, 4], depth + for r in res + forced: + assert r["rel"] < R1, r["rel"] + + +# Unset, the NAX route takes a slice target by grid, codec and partials. 4 +# cores pinned make N 4096 a full grid at every M through 32: 64 BM=32 +# tiles, 16 per core. A full grid takes one slice and runs what the call +# with the split route off runs where that call runs the BM=32 tile. +# The split stays where it would run an mv route (iq4_xs at M 9, below its +# small-M qmm route at N 4096) or the BM=64 tile (iq2_xs), for q6_k and +# iq2_xxs below M 32, and for q8_0 below M 13. KQ_SPLITK_CAP=0 keeps the +# fixed target. +HEAD_ENV = {"KQ_GPU_CORES": "4", "KQ_QMM_LOG": "1"} +HEAD_CODECS = ["q4_0", "q4_k", "q6_k", "q8_0", "iq4_xs", "iq2_xs", "iq2_xxs"] +HEAD_MS = [9, 15, 16, 32] +HEAD_DTYPES = ["bfloat16", "float16"] +# iq2_xxs enters the split route at M 13. +HEAD_CASES = [ + (c, m) for c in HEAD_CODECS for m in HEAD_MS if not (c == "iq2_xxs" and m < 13) +] + + +def _head_keeps_split(codec, m): + if codec == "iq2_xs": + return True + if codec in ("q6_k", "iq2_xxs"): + return m < 32 + return codec in ("q8_0", "iq4_xs") and m <= 12 + + +@functools.cache +def _nax_head(): + arms = [ + {"KQ_SPLITK_CAP": None, "KQ_QMM_SPLITK_NAX": None}, + {"KQ_SPLITK_CAP": "0", "KQ_QMM_SPLITK_NAX": None}, + {"KQ_SPLITK_CAP": None, "KQ_QMM_SPLITK_NAX": "0"}, + ] + cases = [(c, m, dt) for c, m in HEAD_CASES for dt in HEAD_DTYPES] + calls = [dict(_call(c, m, env), dtype=dt) for c, m, dt in cases for env in arms] + res = _probe(calls, HEAD_ENV) + return {case: res[3 * i : 3 * i + 3] for i, case in enumerate(cases)} + + +@pytest.mark.parametrize("dtype", HEAD_DTYPES) +@pytest.mark.parametrize(("codec", "m"), HEAD_CASES) +def test_nax_default_full_grid_takes_one_slice(codec, m, dtype): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + default, fixed, off = _nax_head()[(codec, m, dtype)] + _split(fixed, "_qmm_t_nax_splitk_") + assert fixed["rel"] < R1, fixed["rel"] + if _head_keeps_split(codec, m): + assert default["kernels"] == fixed["kernels"], default["kernels"] + assert default["digest"] == fixed["digest"] + else: + _same(default, off) + assert tile_of(off["kernels"][0][0])[0] == 32, off["kernels"] + + +# q8_0 below M 13 keeps the split on a full grid also where the call +# without it runs the BM=32 tile: at N 8192 its small-M qmm route starts +# at M 8. +@pytest.mark.parametrize("dtype", HEAD_DTYPES) +@pytest.mark.parametrize("m", [9, 12]) +def test_nax_default_q8_0_small_m_full_grid_keeps_split(m, dtype): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + arms = [{"KQ_SPLITK_CAP": None}, {"KQ_SPLITK_CAP": "0"}, {"KQ_QMM_SPLITK_NAX": "0"}] + default, fixed, off = _probe( + [dict(_call("q8_0", m, env, n=8192), dtype=dtype) for env in arms], HEAD_ENV + ) + _split(fixed, "_qmm_t_nax_splitk_") + assert tile_of(off["kernels"][0][0])[0] == 32, off["kernels"] + assert default["kernels"] == fixed["kernels"], default["kernels"] + assert default["digest"] == fixed["digest"] + assert default["rel"] < R1, default["rel"] + + +# Below a full grid the target stays 16, except q4_0 and q8_0 on float +# partials, which hold 16 threadgroups per core. With 40 cores at N 4096 and +# M 16 that leaves room for 10 slices. Float partials are the default on +# float16 and KQ_SPLITK_F32=1 sets them on bfloat16. +@pytest.mark.parametrize("codec", ["q4_0", "q8_0", "q4_k"]) +def test_nax_default_float_partials_cap(codec): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + arms = [ + ("float16", {"KQ_SPLITK_CAP": None}), + ("float16", {"KQ_SPLITK_CAP": "0"}), + ("bfloat16", {"KQ_SPLITK_CAP": None}), + ] + calls = [dict(_call(codec, 16, env), dtype=dt) for dt, env in arms] + res = _probe(calls, CORES) + depth = [_split(r, "_qmm_t_nax_splitk_")[1][2] for r in res] + want = {"q4_0": [10, 16, 16], "q8_0": [10, 16, 16], "q4_k": [10, 10, 10]} + assert depth == want[codec], depth + for r in res: + assert r["rel"] < R1, r["rel"] + + +# q8_0 on float partials keeps the target of 16 below M 13. +def test_nax_default_float_partials_q8_0_small_m(): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + res = _probe( + [ + dict(_call("q8_0", m, {"KQ_SPLITK_CAP": None}), dtype="float16") + for m in (9, 12) + ], + CORES, + ) + depth = [_split(r, "_qmm_t_nax_splitk_")[1][2] for r in res] + assert depth == [16, 16], depth + for r in res: + assert r["rel"] < R1, r["rel"] + + +# KQ_SPLITK_F32 is read once, so float partials on bfloat16 run in their own +# process. +@pytest.mark.parametrize("codec", ["q4_0", "q8_0"]) +def test_nax_default_float_partials_cap_bfloat16(codec): + if not _nax_gpu(): + pytest.skip("NAX split-K route") + res = _probe( + [_call(codec, 16, {"KQ_SPLITK_CAP": None})], dict(CORES, KQ_SPLITK_F32="1") + ) + depth = [_split(r, "_qmm_t_nax_splitk_")[1][2] for r in res] + assert depth == [10], depth + assert res[0]["rel"] < R1, res[0]["rel"] diff --git a/tests/test_splitk_f16_range.py b/tests/test_splitk_f16_range.py new file mode 100644 index 00000000..449ab577 --- /dev/null +++ b/tests/test_splitk_f16_range.py @@ -0,0 +1,125 @@ +"""Split-K range on float16 activations. + +The "cancel" design runs float16 activations near 3e4 against weights near ++0.05 on the first half of K and -0.05 on the second, at K 4096 and N 512. +The output peaks near 1.5e4, well inside the float16 range, but the partial +sum of a K slice runs past 65504. Float16 split calls store float partials +by default, so every split route returns a finite result that matches its +KQ_SPLITK_F32=1 run bit for bit. KQ_SPLITK_F32=0 stores the partials in +float16 and returns inf, so the design reaches the overflow. The unsplit +reference is finite: the tile, with the mid-M split off above M 32, and at +M 32 or less the NAX tile, or per-row qmv on GPUs without NAX. +""" + +from __future__ import annotations + +import functools +import math +import subprocess + +import pytest +from qmm_probe import force_cpu, run_probe + +import mlx_kquant as kq + +pytestmark = pytest.mark.skipif(force_cpu(), reason="GPU split-K routes") + +K, N = 4096, 512 +CODECS = ["q4_0", "q8_0", "q4_k"] +STRICT = {"KQ_QMM_ROUTE_STRICT": "1"} +CLEAR = {"KQ_QMM_ROUTE": None, "KQ_QMM_MIDM": None, "KQ_DISABLE_NAX": None} + + +def _cells(codec, nax): + """(name, M, env, kernel name fragment) per split route cell.""" + alu = {"KQ_DISABLE_NAX": "1"} if nax else {} + cells = [ + ("splitk", 16, {"KQ_QMM_ROUTE": "splitk", **STRICT}, "_qmm_t_splitk_"), + ("splitk mid-M", 64, {**alu, "KQ_QMM_MIDM": "1"}, "_qmm_t_splitk_"), + ] + if nax: + cells += [ + ( + "nax_splitk", + 16, + {"KQ_QMM_ROUTE": "nax_splitk", **STRICT}, + "_qmm_t_nax_splitk_", + ), + ("nax mid-M", 64, {"KQ_QMM_MIDM": "1"}, "_qmm_t_nax_splitk_"), + ] + for m in (2, 8): + cells.append( + ( + f"verify_nax M{m}", + m, + {"KQ_QMM_ROUTE": "verify_nax", **STRICT}, + "_verify_nax_", + ) + ) + if codec != "q4_k": + for m in (2, 8): + cells.append( + ( + f"verify_mma M{m}", + m, + {"KQ_QMM_ROUTE": "verify_mma", **STRICT}, + "_verify_mma_", + ) + ) + return cells + + +@functools.cache +def _probe(codec): + """Per cell: the default, KQ_SPLITK_F32=1 and KQ_SPLITK_F32=0 results, + and the unsplit reference at the cell's M.""" + nax = kq.nax_available() + calls, index = [], {} + for name, m, env, _ in _cells(codec, nax): + base = { + "codec": codec, + "shape": [m, K], + "n": N, + "dtype": "float16", + "design": "cancel", + "seed": m, + } + start = len(calls) + for lever in (None, "1", "0"): + calls.append(dict(base, env={**CLEAR, **env, "KQ_SPLITK_F32": lever})) + if m > 32: + ref_env = {"KQ_QMM_MIDM": "0"} + else: + ref_env = {"KQ_QMM_ROUTE": "nax" if nax else "qmv", **STRICT} + calls.append(dict(base, env={**CLEAR, **ref_env, "KQ_SPLITK_F32": None})) + index[name] = start + try: + res = run_probe(calls, env={"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"}) + except (AssertionError, subprocess.TimeoutExpired) as e: + res = [{"error": f"probe failed: {e}", "kernels": []}] * len(calls) + return {name: res[i : i + 4] for name, i in index.items()} + + +def _cell_names(): + nax = kq.nax_available() + return [(c, name) for c in CODECS for name, *_ in _cells(c, nax)] + + +@pytest.mark.parametrize("codec,cell", _cell_names()) +def test_splitk_f16_range(codec, cell): + frag = {name: f for name, _, _, f in _cells(codec, kq.nax_available())}[cell] + default, f32, half, ref = _probe(codec)[cell] + for r in (default, f32, half, ref): + assert r["error"] is None, r["error"] + for r, pf32 in ((default, True), (f32, True), (half, False)): + ((name, grid),) = r["kernels"] + assert frag in name and grid[2] > 1, (name, grid) + assert name.endswith("_pf32") == pf32, name + assert ref["nonfinite"] == 0, ref["nonfinite"] + assert half["nonfinite"] > 0, "the design no longer overflows float16 partials" + assert default["nonfinite"] == 0, default["nonfinite"] + assert default["digest"] == f32["digest"] + # The halves cancel, so the tile's float16 rounding of the weights shows + # as 0.7 to 1.9 percent of the output. A lost or doubled K slice moves + # it by more than the output itself. + assert math.isfinite(default["rel"]) and default["rel"] < 0.05, default["rel"] diff --git a/tests/test_splitk_partials.py b/tests/test_splitk_partials.py new file mode 100644 index 00000000..daacb01a --- /dev/null +++ b/tests/test_splitk_partials.py @@ -0,0 +1,265 @@ +"""Split-K partials: float on float16 and bfloat16 on bfloat16 by default, +float on both under KQ_SPLITK_F32=1, the activation dtype on both under 0. + +The split routes (qmm_splitk, verify_mma, verify_nax, qmm_nax_splitk) fold +their K slices in f32. With float partials the fold rounds the result once, +so each route gives the rounding of the same product run unsplit; with +partials in the activation dtype each slice rounds once more. + +The inputs follow the lead design: nonnegative weights and activations +whose first K column is 1024, so the first slice carries most of every +output and its rounding shows fully. At K 12288 and N 256: + +- qmm_nax_splitk (the default at M 64 on NAX GPUs, and forced at M 8) and + qmm_splitk (forced at M 64 and M 8) keep the tile's R2 floor, R2 matching + the tile run within 2 percent of that floor. +- Every split route stays within 1 ulp and half an output rounding of its + reference: the tile, or the same verify route run as one slice. +- KQ_SPLITK_F32=0 dispatches the T-partial kernels, which sit more than + one output rounding from that reference, with an R2 excess above half + the floor on the qmm routes. + +Dense, lazy transposed, strided and batch-of-two inputs all take the +checks. On the CPU stream the levers do nothing: every pair of calls +matches bit for bit and stays within the rounding floor of its reference. +The CPU calls run with KQ_CPU_NEON=0, since the int8 small-M kernels +quantize the activations by design. +""" + +from __future__ import annotations + +import math + +import pytest +from qmm_probe import force_cpu, run_probe + +import mlx_kquant as kq + +K, N = 12288, 256 +CODECS = ["q4_k", "q5_k", "q6_k", "q8_0"] +DTYPES = ["bfloat16", "float16"] +# (layout, batch): batch 2 splits the M rows over two leading entries. +LAYOUTS = [("dense", 1), ("transposed", 1), ("strided", 1), ("dense", 2)] +CLEAR = { + "KQ_QMM_MIDM": None, + "KQ_DISABLE_NAX": None, + "KQ_QMM_ROUTE": None, + "KQ_QMM_ROUTE_STRICT": None, + "KQ_VERIFY_NAX_SPLITS": None, + "KQ_VERIFY_MMA_SPLITS": None, +} +FORCE = {"KQ_QMM_ROUTE_STRICT": "1"} + + +def _shape(m, batch): + return [batch, m // batch, K] if batch > 1 else [m, K] + + +def _excess(r, floor): + return math.sqrt(max(0.0, r * r - floor * floor)) / floor + + +def _cells(codec, nax): + """(name, M, split env, reference env, family) per route cell. family + is the kernel name fragment of the split kernel.""" + cells = [] + if nax: + cells += [ + ("nax mid-M", 64, {}, {"KQ_QMM_MIDM": "0"}, "_qmm_t_nax_splitk_"), + ( + "nax_splitk", + 8, + {"KQ_QMM_ROUTE": "nax_splitk", **FORCE}, + {"KQ_QMM_ROUTE": "nax", **FORCE}, + "_qmm_t_nax_splitk_", + ), + ( + "verify_nax", + 8, + {"KQ_QMM_ROUTE": "verify_nax", **FORCE}, + {"KQ_QMM_ROUTE": "verify_nax", "KQ_VERIFY_NAX_SPLITS": "1", **FORCE}, + "_verify_nax_", + ), + ] + alu = {"KQ_DISABLE_NAX": "1"} if nax else {} + cells += [ + ( + "splitk mid-M", + 64, + {**alu, "KQ_QMM_MIDM": "1"}, + {**alu, "KQ_QMM_MIDM": "0"}, + "_qmm_t_splitk_", + ), + ( + "splitk", + 8, + {"KQ_QMM_ROUTE": "splitk", **FORCE}, + {"KQ_QMM_ROUTE": "nax", **FORCE} if nax else None, + "_qmm_t_splitk_bm8_", + ), + ] + if codec == "q8_0": + cells.append( + ( + "verify_mma", + 8, + {"KQ_QMM_ROUTE": "verify_mma", **FORCE}, + {"KQ_QMM_ROUTE": "verify_mma", "KQ_VERIFY_MMA_SPLITS": "1", **FORCE}, + "_verify_mma_", + ) + ) + return cells + + +def _run(codec, dtype, cpu): + nax = not cpu and kq.nax_available() + calls, meta = [], [] + + def add(tag, cell, m, layout, batch, env, lever, ulp_vs=None): + calls.append( + { + "codec": codec, + "shape": _shape(m, batch), + "n": N, + "dtype": dtype, + "layout": layout, + "seed": 7 + m, + "design": "lead", + "r2": True, + "env": {**CLEAR, **env, "KQ_SPLITK_F32": lever}, + "ulp_vs": ulp_vs, + } + ) + meta.append((tag, cell, m, layout, batch)) + return len(calls) - 1 + + for name, m, split_env, ref_env, _ in _cells(codec, nax): + for layout, batch in LAYOUTS: + ref = None + if ref_env is not None: + ref = add("ref", name, m, layout, batch, ref_env, "1") + split = add("split", name, m, layout, batch, split_env, "1", ref) + if layout == "dense" and batch == 1: + # Without a reference the old arm compares to the split arm. + base = ref if ref is not None else split + add("old", name, m, layout, batch, split_env, "0", base) + env = {"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"} + if cpu: + env["KQ_CPU_NEON"] = "0" + res = run_probe(calls, env=env) + return calls, meta, res, nax + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("codec", CODECS) +def test_splitk_partials(codec, dtype): + cpu = force_cpu() + calls, meta, res, nax = _run(codec, dtype, cpu) + cells = {c[0]: c for c in _cells(codec, nax)} + refs = {} + for (tag, cell, _, layout, batch), r in zip(meta, res, strict=True): + assert r["error"] is None, (tag, cell, layout, batch, r["error"]) + if tag == "ref": + refs[(cell, layout, batch)] = r + for c, (tag, cell, _, layout, batch), r in zip(calls, meta, res, strict=True): + where = (codec, dtype, tag, cell, layout, batch) + e1 = _excess(r["rel"], r["floor1"]) + e2 = _excess(r["rel2"], r["floor2"]) + if cpu: + # The CPU matmul rounds once, on one of the two models. + assert min(e1, e2) < 0.1, (where, e1, e2) + assert r["kernels"] == [], where + if c["ulp_vs"] is not None: + assert r["ulp"] == 0, (where, r["ulp"]) + continue + (name, grid) = r["kernels"][0] + family = cells[cell][4] + verify = cell.startswith("verify") + if tag == "ref": + if verify: + assert family in name and grid[2] == 1, (where, name, grid) + assert not name.endswith("_pf32"), (where, name) + else: + assert "splitk" not in name, (where, name) + continue + assert family in name and grid[2] > 1, (where, name, grid) + assert name.endswith("_pf32") == (tag == "split"), (where, name) + ref = refs.get((cell, layout, batch)) + # Distance to the reference in output roundings. + dist = r["rdiff"] / r["floor2"] if ref is not None else None + if tag == "split": + if not verify: + # The tile's R2 model: T weights, f32 sum, one rounding. + assert e2 < 0.1, (where, e2) + if ref is not None: + gap = abs(r["rel2"] - ref["rel2"]) / r["floor2"] + assert gap < 0.02, (where, gap) + if ref is not None: + assert r["ulp"] <= 1 and dist < 0.5, (where, r["ulp"], dist) + else: + if not verify: + assert e2 > 0.5, (where, e2) + if ref is not None: + assert dist > 1.0, (where, dist) + + +# With no KQ_SPLITK_F32 set, a float16 split call dispatches the float +# partial kernel and a bfloat16 one the T-partial kernel. KQ_SPLITK_F32=1 +# takes the float twin and 0 the T-partial kernel on both. +@pytest.mark.parametrize("dtype", DTYPES) +def test_splitk_f32_default(dtype): + if force_cpu(): + pytest.skip("GPU split-K kernels") + call = {"codec": "q4_k", "shape": [8, K], "n": N, "dtype": dtype} + env = {"KQ_QMM_LOG": "1", "KQ_QMM_ROUTE": "splitk", **FORCE} + names = [] + for extra in ({}, {"KQ_SPLITK_F32": "1"}, {"KQ_SPLITK_F32": "0"}): + (r,) = run_probe([call], env={**env, **extra}) + assert r["error"] is None, r["error"] + ((name, grid),) = r["kernels"] + assert "_qmm_t_splitk_" in name and grid[2] > 1, (name, grid) + names.append(name) + plain = names[2] + assert not plain.endswith("_pf32"), names + assert names[1] == plain + "_pf32", names + assert names[0] == (plain + "_pf32" if dtype == "float16" else plain), names + + +# Without KQ_TUNE_LIVE, KQ_SPLITK_F32 is read once per process: a change +# after the first split call does not reach the later ones. +@pytest.mark.parametrize("dtype", DTYPES) +def test_splitk_f32_keeps_first_read(dtype): + if force_cpu(): + pytest.skip("GPU split-K kernels") + call = {"codec": "q4_k", "shape": [8, K], "n": N, "dtype": dtype} + env = {"KQ_QMM_LOG": "1", "KQ_QMM_ROUTE": "splitk", **FORCE} + unset_pf32 = dtype == "float16" + arms = [ + ("1", [None, "0"], True), + ("0", [None, "1"], False), + (None, ["0" if unset_pf32 else "1"], unset_pf32), + ] + for first, later, pf32 in arms: + calls = [call] + [dict(call, env={"KQ_SPLITK_F32": v}) for v in later] + first_env = {} if first is None else {"KQ_SPLITK_F32": first} + res = run_probe(calls, env={**env, **first_env}) + names = [] + for r in res: + assert r["error"] is None, r["error"] + ((name, _),) = r["kernels"] + names.append(name) + assert all(n.endswith("_pf32") == pf32 for n in names), (first, names) + + +# The ops cast float32 activations to bfloat16, so a float32 split call +# runs the bfloat16 kernel and returns bfloat16. +def test_splitk_float32_runs_bfloat16_kernel(): + if force_cpu(): + pytest.skip("GPU split-K kernels") + call = {"codec": "q4_k", "shape": [8, K], "n": N, "dtype": "float32"} + env = {"KQ_QMM_LOG": "1", "KQ_QMM_ROUTE": "splitk", **FORCE} + (r,) = run_probe([call], env=env) + assert r["error"] is None, r["error"] + ((name, grid),) = r["kernels"] + assert "_qmm_t_splitk_bm8_bfloat16_t_" in name and grid[2] > 1, (name, grid) + assert r["rel"] < 1e-2, r["rel"] diff --git a/tests/test_tile_override.py b/tests/test_tile_override.py new file mode 100644 index 00000000..794f9ca1 --- /dev/null +++ b/tests/test_tile_override.py @@ -0,0 +1,350 @@ +"""KQ_NAX_TILE and KQ_QMM_TILE: pinning the quantized matmul tile. + +Each case runs in a fresh interpreter (tests/qmm_probe.py) with +``KQ_TUNE_LIVE=1`` unless it tests the cached read, and reads the dispatched +kernel from ``KQ_QMM_LOG``. The numeric cases also run on the CPU stream +under ``KQUANT_FORCE_CPU``, where the overrides must leave the result alone. +The enumerating case dispatches every kernel name in the metallib that an +override can select and checks each against float64, which catches a +silently empty kernel. +""" + +from __future__ import annotations + +import re + +import pytest +from qmm_probe import force_cpu, run_probe, tile_of + +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY + +# R1 bounds: relative Frobenius error against float64 with exact +# kq.dequantize weights. +R1 = {"bfloat16": 6e-3, "float16": 1e-3} +TYPE_DTYPE = {"bfloat16_t": "bfloat16", "float16_t": "float16"} + +K = 256 +LIVE = {"KQ_TUNE_LIVE": "1", "KQ_QMM_LOG": "1"} +TILES = {32: "32,64,2,2", 64: "64,64,2,2", 128: "128,64,2,2"} +MAIN_CODECS = ["q4_k", "q5_k", "q6_k", "q8_0"] + + +def _metal(): + if force_cpu(): + pytest.skip("Metal kernel names only") + + +def _nax(): + _metal() + if not kq.nax_available(): + pytest.skip("NAX tile") + + +def _only_kernel(res): + assert res["error"] is None, res["error"] + assert len(res["kernels"]) == 1, res["kernels"] + return res["kernels"][0] + + +def test_nax_tile_selects_kernel(): + _nax() + # 4128 rows under the parity rule: the policy runs BM=64 (ceil(4128/64) + # is odd). + call = {"codec": "q4_k", "shape": [1, 4128, K], "n": 128} + calls = [ + dict(call), + dict(call, env={"KQ_NAX_TILE": TILES[128]}), + dict(call, env={"KQ_NAX_TILE": TILES[32]}), + dict(call, env={"KQ_NAX_TILE": None}), + ] + res = run_probe(calls, env=dict(LIVE, KQ_NAX_BM128_ODD_M="0"), check=False) + got = [_only_kernel(r) for r in res] + assert [tile_of(n)[0] for n, _ in got] == [64, 128, 32, 64] + assert [g for _, g in got] == [(2, 65, 1), (2, 33, 1), (2, 129, 1), (2, 65, 1)] + + +def test_nax_tile_keeps_first_read(): + _nax() + call = {"codec": "q4_k", "shape": [1, 4128, K], "n": 128} + calls = [dict(call), dict(call, env={"KQ_NAX_TILE": None})] + res = run_probe( + calls, env={"KQ_QMM_LOG": "1", "KQ_NAX_TILE": TILES[128]}, check=False + ) + assert [tile_of(_only_kernel(r)[0])[0] for r in res] == [128, 128] + + +def test_nax_tile_uninstantiated_raises(): + _nax() + call = {"codec": "q4_k", "shape": [300, K], "n": 128} + calls = [ + dict(call, env={"KQ_NAX_TILE": "128,128,2,4"}), + # The process keeps working after the raise. + dict(call, env={"KQ_NAX_TILE": TILES[64]}), + ] + res = run_probe(calls, env=LIVE) + name = ( + "kquant_q4_k_qmm_t_nax_bfloat16_t_gs_256_b_4_bm128_bn128_bk64_wm2_wn4" + "_alN_true_batch_0" + ) + assert res[0]["error"].startswith("ValueError"), res[0]["error"] + assert name in res[0]["error"] and "KQ_NAX_TILE=128,128,2,4" in res[0]["error"] + assert tile_of(_only_kernel(res[1])[0])[0] == 64 + assert res[1]["rel"] < R1["bfloat16"] + + +def test_qmm_tile_selects_and_raises(): + # The path without NAX, reached with KQ_DISABLE_NAX=1 on NAX GPUs. + _metal() + call = {"codec": "q4_k", "shape": [300, K], "n": 128} + calls = [ + dict(call, env={"KQ_QMM_TILE": TILES[64]}), + dict(call, env={"KQ_QMM_TILE": TILES[128]}), + dict(call, env={"KQ_QMM_TILE": None}), + ] + res = run_probe(calls, env={**LIVE, "KQ_DISABLE_NAX": "1"}) + plain = "kquant_q4_k_qmm_t_bfloat16_t_gs_256_b_4_alN_true_batch_0" + assert _only_kernel(res[0])[0] == plain + assert res[0]["rel"] < R1["bfloat16"] + tagged = ( + "kquant_q4_k_qmm_t_bfloat16_t_gs_256_b_4_bm128_bn64_wm2_wn2_alN_true_batch_0" + ) + assert res[1]["error"].startswith("ValueError"), res[1]["error"] + assert tagged in res[1]["error"] and "KQ_QMM_TILE=128,64,2,2" in res[1]["error"] + assert _only_kernel(res[2])[0] == plain + + +@pytest.mark.parametrize( + "lever,value", + [ + ("KQ_NAX_TILE", "128,64"), + ("KQ_NAX_TILE", "0,64,2,2"), + ("KQ_QMM_TILE", "a,b,c,d"), + ], +) +def test_tile_malformed_raises(lever, value): + _metal() + if lever == "KQ_NAX_TILE" and not kq.nax_available(): + pytest.skip("NAX tile") + env = {**LIVE, lever: value} + if lever == "KQ_QMM_TILE": + env["KQ_DISABLE_NAX"] = "1" + res = run_probe([{"codec": "q4_k", "shape": [300, K], "n": 128}], env=env) + err = res[0]["error"] + assert err is not None and err.startswith("ValueError"), err + assert f"{lever}={value}" in err + + +def _numeric_calls(codec, dtype, tiles): + calls = [] + for layout in ("dense", "transposed", "strided"): + for shape in ([1, 260, K], [2, 130, K]): + n = 256 if layout == "dense" else 200 + base = len(calls) + for j, bm in enumerate(tiles): + calls.append( + { + "codec": codec, + "shape": shape, + "n": n, + "dtype": dtype, + "layout": layout, + "seed": base, + "env": {"KQ_NAX_TILE": TILES.get(bm)}, + "ulp_vs": base if j else None, + } + ) + return calls + + +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", MAIN_CODECS) +def test_tile_numerics(codec, dtype): + # Every tile meets the R1 bound on batch 1, batch 2, lazy transposed and + # strided inputs, and agrees with the BM=64 tile within 1 ulp. On the + # CPU stream the override does nothing and the bound still holds. + tiles = [64, None, 32, 128] + calls = _numeric_calls(codec, dtype, tiles) + res = run_probe(calls, env=LIVE) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1[dtype], (c, r["rel"]) + if r["ulp"] is not None: + assert r["ulp"] <= 1, (c, r["ulp"]) + if not force_cpu() and kq.nax_available() and c["env"]["KQ_NAX_TILE"]: + want = int(c["env"]["KQ_NAX_TILE"].split(",")[0]) + assert tile_of(_only_kernel(r)[0])[0] == want + + +def _qmm_codecs(): + return sorted(c for c in CODEC_GEOMETRY if kq.codec_has_matmul(c)) + + +def _nax_tiles(codec): + """BM of every unsplit qmm_t NAX tile of `codec` in the metallib.""" + out = set() + for name in kq.metallib_kernels(): + m = _NAX_NAME.match(name) + if m and m["codec"] == codec and m["db"] is None and m["bn"] == "64": + out.add(int(m["bm"])) + return sorted(out) + + +@pytest.mark.parametrize("codec", _qmm_codecs()) +def test_tile_rows_above_int16(codec): + # 32768 and 32769 rows (an even and an odd ceil(M/64)), and a batch of + # two at 16896 rows flattened to 33792, past the int16 row span the tile + # clamps before narrowing, on every tile the codec carries and on folded + # row-tile orders (KQ_NAX_SWIZZLE). + nax = not force_cpu() and kq.nax_available() + envs = [{}] + tiles = _nax_tiles(codec) if nax else [] + if tiles: + # BM=64 first: every other tile agrees with it within 1 ulp. + tiles.sort(key=lambda bm: bm != 64) + envs = [{"KQ_NAX_TILE": TILES[bm]} for bm in tiles] + envs += [ + {"KQ_NAX_TILE": TILES[64], "KQ_NAX_SWIZZLE": "2"}, + {"KQ_NAX_TILE": TILES[128], "KQ_NAX_SWIZZLE": "1"}, + {"KQ_NAX_TILE": TILES[128], "KQ_NAX_SWIZZLE": "2"}, + ] + calls = [] + for shape in ([1, 32768, K], [1, 32769, K], [2, 16896, K]): + for n in (96, 128): + base = len(calls) + for j, env in enumerate(envs): + full = {"KQ_NAX_TILE": None, "KQ_NAX_SWIZZLE": None, **env} + calls.append( + { + "codec": codec, + "shape": shape, + "n": n, + "seed": base, + "env": full, + "ulp_vs": base if j else None, + } + ) + res = run_probe(calls, env=LIVE) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1["bfloat16"], (c, r["rel"]) + if r["ulp"] is not None: + assert r["ulp"] <= 1, (c, r["ulp"]) + if tiles and c["env"]["KQ_NAX_TILE"]: + want = int(c["env"]["KQ_NAX_TILE"].split(",")[0]) + assert tile_of(_only_kernel(r)[0])[0] == want, c + + +_NAX_NAME = re.compile( + r"^kquant_(?P.+?)_qmm_t_nax_(?Pbfloat16_t|float16_t|float)" + r"_gs_\d+_b_\d+_bm(?P\d+)_bn(?P\d+)_bk(?P64|128)_wm(?P\d+)" + r"_wn(?P\d+)" + r"_alN_(?Ptrue|false)_batch_(?P[01])(?P_db)?$" +) +_ALU_NAME = re.compile( + r"^kquant_(?P.+?)_qmm_t_(?Pbfloat16_t|float16_t|float)" + r"_gs_\d+_b_\d+(?:_bm(?P\d+)_bn(?P\d+)_wm(?P\d+)_wn(?P\d+))?" + r"_alN_(?Ptrue|false)_batch_(?P[01])$" +) + + +def test_tile_names_match_known_patterns(): + # Every qmm_t tile kernel outside the split-K, gather, GLU, low-rank and + # residual families must match a pattern above, so a new tile family cannot + # escape the enumerating test below. test_matmul_glu.py holds the GLU names, + # test_matmul_lora.py the low-rank ones and test_matmul_resid.py the + # residual ones. + _metal() + names = kq.metallib_kernels() + other = ( + "_gather_", + "_qmm_t_nax_splitk_", + "_qmm_t_nax_glu_", + "_qmm_t_nax_lora_", + "_qmm_t_nax_resid_", + ) + nax = [n for n in names if "_qmm_t_nax_" in n and not any(o in n for o in other)] + alu = [ + n + for n in names + if re.search(r"_qmm_t_(bfloat16_t|float16_t|float)_", n) and "_gather_" not in n + ] + assert alu, "no qmm_t kernels in the metallib" + for pat, group in ((_NAX_NAME, nax), (_ALU_NAME, alu)): + bad = [ + n + for n in group + if not (m := pat.match(n)) or m["codec"] not in CODEC_GEOMETRY + ] + assert not bad, bad[:5] + + +def _bk128_lever(codec): + if codec.endswith("_k"): + return "KQ_NAX_" + codec.replace("_", "").upper() + "_BK128" + return "KQ_NAX_" + codec.upper() + "_BK128" + + +def _selectable(codec): + """Calls that dispatch every qmm_t tile kernel of `codec` an override can + select: half and bfloat activations, unbatched (the op flattens a batch + into rows and never builds a float activation for the tile).""" + calls, other = [], [] + for name in kq.metallib_kernels(): + for kind, pat in (("nax", _NAX_NAME), ("alu", _ALU_NAME)): + m = pat.match(name) + if not m or m["codec"] != codec: + continue + if m["type"] == "float" or m["batch"] == "1": + other.append(name) + continue + if kind == "nax" and not kq.nax_available(): + continue + tile = ",".join( + m[g] or d + for g, d in zip("bm bn wm wn".split(), "64 64 2 2".split(), strict=True) + ) + if kind == "nax": + db = m["db"] is not None + rows = 48 if db else max(133, 2 * int(m["bm"]) + 5) + env = { + "KQ_NAX_TILE": tile, + "KQ_NAX_DB64": "1" if db else "0", + _bk128_lever(codec): "1" if m["bk"] == "128" else "0", + "KQ_QMM_MIDM": "0", + "KQ_DISABLE_NAX": None, + "KQ_QMM_TILE": None, + } + else: + rows = 133 + env = { + "KQ_QMM_TILE": tile, + "KQ_DISABLE_NAX": "1", + "KQ_NAX_TILE": None, + } + calls.append( + { + "codec": codec, + "shape": [rows, K], + "n": 128 if m["al"] == "true" else 96, + "dtype": TYPE_DTYPE[m["type"]], + "env": env, + "name": name, + } + ) + return calls, other + + +@pytest.mark.parametrize("codec", sorted(CODEC_GEOMETRY)) +def test_enumerate_tile_kernels(codec): + _metal() + calls, _ = _selectable(codec) + if not kq.codec_has_matmul(codec): + assert not calls + return + assert calls, f"no qmm_t tile kernels found for {codec}" + res = run_probe(calls, env=LIVE) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c["name"], r["error"]) + assert _only_kernel(r)[0] == c["name"] + assert r["rel"] < R1[c["dtype"]], (c["name"], r["rel"]) diff --git a/tests/test_tune_live.py b/tests/test_tune_live.py new file mode 100644 index 00000000..d160c77e --- /dev/null +++ b/tests/test_tune_live.py @@ -0,0 +1,228 @@ +"""KQ_TUNE_LIVE: the NAX tile levers read per dispatch, or once per process. + +``KQ_NAX_SMALL_BM``, ``KQ_NAX_BM128``, ``KQ_NAX_DB64`` and ``KQ_NAX_SWIZZLE`` +keep their first read for the life of the process, and ``KQ_TUNE_LIVE=1`` +makes them read their variables on every dispatch. Each case runs in a fresh +interpreter (tests/qmm_probe.py) and reads the dispatched kernel and grid from +``KQ_QMM_LOG``. The numeric case also runs on the CPU stream under +``KQUANT_FORCE_CPU``, where the levers must leave the result alone. +""" + +from __future__ import annotations + +import pytest +from qmm_probe import force_cpu, run_probe, tile_of + +import mlx_kquant as kq + +# R1 bounds: relative Frobenius error against float64 with exact +# kq.dequantize weights. +R1 = {"bfloat16": 6e-3, "float16": 1e-3} + +K = 256 +N = 128 + +# q6_k takes BM=128 from M 193 and the BM=32 tile at M <= 32; its _db +# band needs N >= 16384, which KQ_NAX_DB64=1 drops. KQ_QMM_MIDM=0 keeps +# M 48 on the qmm tile and KQ_QMM_ROUTE=nax sends M 16 to it; both are +# read per call. +LEVER_CALLS = [ + # 0: the process starts with KQ_NAX_BM128=0. + {"codec": "q6_k", "shape": [256, K], "n": N}, + # 1: BM128 unset. + {"codec": "q6_k", "shape": [256, K], "n": N, "env": {"KQ_NAX_BM128": None}}, + # 2: BM128 off again, a fold of four row tiles. + { + "codec": "q6_k", + "shape": [256, K], + "n": N, + "env": {"KQ_NAX_BM128": "0", "KQ_NAX_SWIZZLE": "2"}, + }, + # 3: the M 33-64 band with the N floor dropped. + { + "codec": "q6_k", + "shape": [48, K], + "n": N, + "env": {"KQ_NAX_SWIZZLE": None, "KQ_NAX_DB64": "1", "KQ_QMM_MIDM": "0"}, + }, + # 4: M 16 on the qmm tile with the BM=32 pick off. + { + "codec": "q6_k", + "shape": [16, K], + "n": N, + "env": {"KQ_NAX_SMALL_BM": "0", "KQ_QMM_ROUTE": "nax"}, + }, +] + + +def _gpu_nax(): + if force_cpu(): + pytest.skip("Metal kernel names only") + if not kq.nax_available(): + pytest.skip("NAX tile levers") + + +def _only_kernel(res): + assert res["error"] is None, res["error"] + assert len(res["kernels"]) == 1, res["kernels"] + return res["kernels"][0] + + +def test_levers_keep_first_read(): + _gpu_nax() + res = run_probe(LEVER_CALLS, env={"KQ_QMM_LOG": "1", "KQ_NAX_BM128": "0"}) + first = _only_kernel(res[0]) + assert tile_of(first[0])[0] == 64 + # Every later change is ignored: the first read holds. + for r in res[1:3]: + assert _only_kernel(r) == first + db = _only_kernel(res[3])[0] + assert tile_of(db)[0] == 64 and not db.endswith("_db") + assert tile_of(_only_kernel(res[4])[0])[0] == 32 + + +def test_levers_live_under_tune_live(): + _gpu_nax() + res = run_probe( + LEVER_CALLS, + env={"KQ_QMM_LOG": "1", "KQ_NAX_BM128": "0", "KQ_TUNE_LIVE": "1"}, + ) + name, grid = _only_kernel(res[0]) + assert tile_of(name)[0] == 64 and grid == (2, 4, 1) + name, grid = _only_kernel(res[1]) + assert tile_of(name)[0] == 128 and grid == (2, 2, 1) + # Four row tiles fold into grid.x with log 2. + name, grid = _only_kernel(res[2]) + assert tile_of(name)[0] == 64 and grid == (8, 1, 1) + assert _only_kernel(res[3])[0].endswith("_db") + assert tile_of(_only_kernel(res[4])[0])[0] == 64 + + +def test_default_levers_unchanged_by_tune_live(): + # KQ_TUNE_LIVE with no lever set dispatches what the cached read does. + _gpu_nax() + calls = [dict(c, env={}) for c in LEVER_CALLS] + calls[3]["env"] = {"KQ_QMM_MIDM": "0"} + calls[4]["env"] = {"KQ_QMM_ROUTE": "nax"} + base = run_probe(calls, env={"KQ_QMM_LOG": "1"}) + live = run_probe(calls, env={"KQ_QMM_LOG": "1", "KQ_TUNE_LIVE": "1"}) + for b, v in zip(base, live, strict=True): + assert _only_kernel(b) == _only_kernel(v) + + +# The routing at DiT row counts: per image, and with a batch of two +# flattened into M. From 4096 rows the default runs BM=128 at either +# parity; under KQ_NAX_BM128_ODD_M=0 BM=128 runs where ceil(M/64) is even. +DIT_TILES = { + 4096: 128, + 4128: 64, + 4192: 128, + 4256: 64, + 4608: 128, + 4640: 64, + 16384: 128, + 16896: 128, +} +DIT_TILES_B2 = { + 4096: 128, + 4128: 64, + 4192: 64, + 4256: 64, + 4608: 128, + 4640: 64, + 16384: 128, + 16896: 128, +} + + +@pytest.mark.parametrize("parity", [False, True]) +@pytest.mark.parametrize("codec", ["q4_k", "q5_k", "q6_k", "q8_0"]) +def test_default_tile_at_dit_rows(codec, parity): + _gpu_nax() + calls = [] + expect = [] + for rows, bm in DIT_TILES.items(): + calls.append({"codec": codec, "shape": [1, rows, K], "n": 64}) + expect.append((rows, bm if parity else 128)) + calls.append({"codec": codec, "shape": [2, rows, K], "n": 64}) + expect.append((2 * rows, DIT_TILES_B2[rows] if parity else 128)) + env = {"KQ_QMM_LOG": "1"} + if parity: + env["KQ_NAX_BM128_ODD_M"] = "0" + res = run_probe(calls, env=env, check=False) + for r, (m, bm) in zip(res, expect, strict=True): + name, grid = _only_kernel(r) + assert tile_of(name)[0] == bm, (m, name) + assert grid == (1, -(-m // bm), 1), (m, grid) + + +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +def test_live_levers_numerics(dtype): + # Every lever setting gives the same numbers within the R1 bound, on + # batch 1, batch 2, lazy transposed and strided inputs. On the CPU + # stream the levers do nothing and the bound still holds. + settings = [ + {}, + {"KQ_NAX_BM128": "0"}, + {"KQ_NAX_BM128": "2"}, + {"KQ_NAX_BM128": "0", "KQ_NAX_SWIZZLE": "1"}, + ] + calls = [] + for layout in ("dense", "transposed", "strided"): + for shape in ([1, 260, K], [2, 130, K], [320, K]): + for env in settings: + full = { + "KQ_NAX_BM128": None, + "KQ_NAX_SWIZZLE": None, + **env, + } + calls.append( + { + "codec": "q6_k", + "shape": shape, + "n": 200, + "dtype": dtype, + "layout": layout, + "seed": len(calls) // len(settings), + "env": full, + } + ) + res = run_probe(calls, env={"KQ_TUNE_LIVE": "1"}) + for c, r in zip(calls, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1[dtype], (c, r["rel"]) + + +# q2_0 takes the tile levers of every NAX codec: BM=128 from M 449 at an +# even ceil(M/64), the BM=32 tile at M 16. +Q2_0_CALLS = [ + {"codec": "q2_0", "shape": [512, K], "n": N}, + {"codec": "q2_0", "shape": [512, K], "n": N, "env": {"KQ_NAX_BM128": "0"}}, + {"codec": "q2_0", "shape": [512, K], "n": N, "env": {"KQ_NAX_SWIZZLE": "1"}}, + { + "codec": "q2_0", + "shape": [16, K], + "n": N, + "env": {"KQ_NAX_SWIZZLE": None, "KQ_QMM_ROUTE": "nax"}, + }, + {"codec": "q2_0", "shape": [16, K], "n": N, "env": {"KQ_NAX_SMALL_BM": "0"}}, + { + "codec": "q2_0", + "shape": [512, K], + "n": N, + "env": {"KQ_QMM_ROUTE": None, "KQ_NAX_TILE": "32,64,2,2"}, + }, +] + + +def test_q2_0_levers_live(): + res = run_probe(Q2_0_CALLS, env={"KQ_QMM_LOG": "1", "KQ_TUNE_LIVE": "1"}) + for c, r in zip(Q2_0_CALLS, res, strict=True): + assert r["error"] is None, (c, r["error"]) + assert r["rel"] < R1["bfloat16"], (c, r["rel"]) + if force_cpu() or not kq.nax_available(): + return + got = [_only_kernel(r) for r in res] + assert [tile_of(n)[0] for n, _ in got] == [128, 64, 64, 32, 64, 32] + assert [g for _, g in got[:3]] == [(2, 4, 1), (2, 8, 1), (4, 4, 1)] + assert got[5][1] == (2, 16, 1) diff --git a/tests/test_verify_mma_range.py b/tests/test_verify_mma_range.py new file mode 100644 index 00000000..f058955b --- /dev/null +++ b/tests/test_verify_mma_range.py @@ -0,0 +1,130 @@ +"""verify_mma past the half range. + +The "big" design draws activations from 2e4 up, with one bfloat16 column +at 1e5, against small weights, so each block's sum runs past the half +range while the output stays inside it. The "normal" design draws +ordinary activations. verify_mma stages the activations in their own +dtype and sums each block in float. With float partials +(KQ_SPLITK_F32=1) it keeps the R1 model on both designs: the result sits +within a tenth of an output rounding of float64 over the exact weights. +The tile rounds the weights to the activation dtype (R2), so the two sit +within two output roundings of each other. With partials in the +activation dtype (KQ_SPLITK_F32=0) each K slice rounds once more, which +adds about one output rounding: the result stays finite and within one and +a half output roundings of the R1 model. KQ_VERIFY_MMA_HALF=1 selects the +half kernels, whose range stops at 65504, on every codec but ptq1_0, which +has none. With the same partials on ordinary float16 activations, their +error beyond the R1 model is at least three times that of the float +kernels. +""" + +from __future__ import annotations + +import functools +import math +import re +import subprocess + +import pytest +from qmm_probe import force_cpu, run_probe + +import mlx_kquant as kq + +pytestmark = pytest.mark.skipif(force_cpu(), reason="GPU verify route") + +CODECS = ["pq2_0", "ptq1_0", "q2_0", "q4_0", "q8_0"] +DTYPES = ["bfloat16", "float16"] +# design: (K, N) +DESIGNS = {"normal": (4096, 512), "big": (1024, 256)} +MS = [8, 3] +FORCE = {"KQ_QMM_ROUTE_STRICT": "1", "KQ_QMM_LOG": "1"} + + +def _half(name): + """Whether a verify_mma kernel name is a KQ_VERIFY_MMA_HALF kernel.""" + return re.search(r"_b_\d+_h(_pf32)?$", name) is not None + + +def _excess(r, floor): + return math.sqrt(max(0.0, r * r - floor * floor)) / floor + + +@functools.cache +def _probe(): + """Tile, verify_mma with float and with T partials, and + KQ_VERIFY_MMA_HALF=1 results per case.""" + tile = "nax" if kq.nax_available() else "splitk" + arms = [ + {"KQ_QMM_ROUTE": tile, "KQ_VERIFY_MMA_HALF": None, "KQ_SPLITK_F32": "1"}, + {"KQ_QMM_ROUTE": "verify_mma"}, + {"KQ_SPLITK_F32": "0"}, + {"KQ_VERIFY_MMA_HALF": "1"}, + ] + cases = [(d, c, t, m) for d in DESIGNS for c in CODECS for t in DTYPES for m in MS] + calls = [] + for design, codec, dtype, m in cases: + k, n = DESIGNS[design] + base = len(calls) + for i, env in enumerate(arms): + calls.append( + { + "codec": codec, + "shape": [m, k], + "n": n, + "dtype": dtype, + "seed": 5 + m, + "design": design, + "r2": True, + "env": env, + "ulp_vs": base if i else None, + } + ) + try: + res = run_probe(calls, env={"KQ_TUNE_LIVE": "1", **FORCE}) + except (AssertionError, subprocess.TimeoutExpired) as e: + res = [{"error": f"probe failed: {e}", "kernels": []}] * len(calls) + w = len(arms) + return {case: res[w * i : w * i + w] for i, case in enumerate(cases)} + + +@pytest.mark.parametrize("m", MS) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("codec", CODECS) +@pytest.mark.parametrize("design", list(DESIGNS)) +def test_verify_mma_range(design, codec, dtype, m): + tile, verify, verify_t, _ = _probe()[(design, codec, dtype, m)] + for r in (tile, verify, verify_t): + assert r["error"] is None, r["error"] + assert len(r["kernels"]) == 1, r["kernels"] + for r, pf32 in ((verify, True), (verify_t, False)): + name = r["kernels"][0][0] + assert "_verify_mma_" in name and not _half(name), name + assert name.endswith("_pf32") == pf32, name + e1 = _excess(verify["rel"], verify["floor1"]) + assert e1 < 0.1, e1 + # Distance to the tile in output roundings. + dist = verify["rdiff"] / verify["floor2"] + assert dist < 2.0, dist + assert math.isfinite(verify_t["rel"]), verify_t["rel"] + assert _excess(verify_t["rel"], verify_t["floor1"]) < 1.5 + + +# The half kernels behind KQ_VERIFY_MMA_HALF=1, on ordinary activations. +# ptq1_0 has none, so the lever leaves it on its float kernels. +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("codec", CODECS) +def test_verify_mma_half_lever(codec, dtype): + for m in MS: + _, _, verify_t, half = _probe()[("normal", codec, dtype, m)] + assert half["error"] is None, half["error"] + ((name, _),) = half["kernels"] + if codec == "ptq1_0": + assert "_verify_mma_" in name and not _half(name), name + assert half["digest"] == verify_t["digest"] + continue + assert "_verify_mma_" in name and _half(name), name + assert math.isfinite(half["rel"]) and half["rel"] < 1e-2, half["rel"] + if dtype == "float16": + e_half = _excess(half["rel"], half["floor1"]) + e_float = _excess(verify_t["rel"], verify_t["floor1"]) + assert e_half > 3 * e_float, (e_half, e_float) diff --git a/tests/test_weight_base.py b/tests/test_weight_base.py index fd967dec..af7a2009 100644 --- a/tests/test_weight_base.py +++ b/tests/test_weight_base.py @@ -164,6 +164,38 @@ def test_gpu_weight_base(codec): assert _equal(y, kq.quantized_matmul(x, w, s, codec, transpose=True)) +@GPU +@pytest.mark.parametrize("codec", ["q4_k", "q5_k", "q6_k", "q8_0"]) +def test_gpu_weight_base_glu(codec): + """quantized_matmul_glu checks the start of both weights: in its fused + kernel, which runs from 1024 rows on a NAX GPU, and through + quantized_matmul elsewhere.""" + a = _align(codec) + wg, _ = _wire(codec, (N, K)) + wu, _ = _wire(codec, (N, K), seed=4) + for m in (64, 1100): + x = (mx.random.normal((m, K), key=mx.random.key(m)) * 0.5).astype(mx.bfloat16) + fused = m >= 1024 and kq.nax_available() + msg = _msg( + "quantized_matmul_glu" if fused else "quantized_matmul", codec, a // 2 + ) + for j in range(2): + + def call(o, j=j, x=x): + g = wg if o is None or j != 0 else _shift(wg, o) + u = wu if o is None or j != 1 else _shift(wu, o) + return kq.quantized_matmul_glu(x, g, u, codec) + + _check(call, sorted({a, 16}), a // 2, msg) + # A weight that is not yet evaluated raises from the eval instead. + y = kq.quantized_matmul_glu(x, wg, _shift(wu, a // 2, evaluate=False), codec) + with pytest.raises(ValueError, match=msg): + mx.eval(y) + # The stream keeps working. + ref = kq.quantized_matmul_glu(x, wg, wu, codec) + assert _equal(kq.quantized_matmul_glu(x, _shift(wg, 16), wu, codec), ref) + + @GPU def test_gpu_weight_base_qmv_bias(): codec = "q8_0" diff --git a/tests/test_weight_views.py b/tests/test_weight_views.py new file mode 100644 index 00000000..50cdd37f --- /dev/null +++ b/tests/test_weight_views.py @@ -0,0 +1,342 @@ +"""Views of a packed weight and of the other operands a kernel reads dense. + +quantized_matmul and the ops built on it walk the last two dims of a +packed weight as dense rows. A view with other strides has to give the +values of its contiguous copy, whether the view has memory when the op is +built or not. An unevaluated view reports dense strides whatever it +becomes, so the op decides for it at eval, in a KQuantMatrixDense node. A +weight that has memory and dense rows adds no node. The same holds for the +bytes dequantize decodes, the floats quantize encodes and the bias of the +fused mat-vec. + +The cases run on the default device, so the suite covers the GPU and, under +KQUANT_FORCE_CPU, the CPU stream. +""" + +from __future__ import annotations + +import re + +import mlx.core as mx +import numpy as np +import pytest + +import mlx_kquant as kq + +N, K = 96, 512 +CODECS = ["q4_k", "q6_k", "q8_0", "q4_0"] +# 1 row takes qmv, 4 the verify kernel, 64 and 1100 the matmul tiles. +ROWS = [1, 4, 64, 1100] +SCALES = mx.zeros((1,), dtype=mx.uint8) +_W = {} + + +def _wire(codec, n, k, seed): + key = (codec, n, k, seed) + if key not in _W: + wf = mx.random.normal((n, k), key=mx.random.key(seed)) * 0.1 + w, _ = kq.quantize(wf, codec) + mx.eval(w) + _W[key] = w + return _W[key] + + +def _x(m, k, dtype=mx.bfloat16): + x = (mx.random.normal((m, k), key=mx.random.key(m)) * 0.5).astype(dtype) + mx.eval(x) + return x + + +def _dense(view): + """A copy of a view in its own row-major buffer.""" + return mx.array(np.ascontiguousarray(np.array(view))) + + +def _views(codec): + """name -> a function that builds a fresh [N, K bytes] view.""" + w = _wire(codec, 2 * N, K, 1) + w2 = _wire(codec, 2 * N, K, 2) + kb = w.shape[1] + wide = mx.concatenate([w, w2], axis=1) + stack = mx.stack([w, w2]) + turned = _dense(mx.swapaxes(w[:N], 0, 1)) + flipped = _dense(w[:N, ::-1]) + mx.eval(wide, stack, turned, flipped) + return { + # Rows of the right length that run backwards in memory. + "reversed_cols": lambda: flipped[:, ::-1], + # Rows that are not dense: every other row, from row 0 and row 1. + "even_rows": lambda: w[0::2], + "odd_rows": lambda: w[1::2], + # Dense rows that start inside a larger buffer. + "row_band": lambda: w[N // 2 : N // 2 + N], + "picked": lambda: stack[1][:N], + # A column band of a wider buffer: the row stride is the wider row. + "left_cols": lambda: wide[:N, :kb], + "right_cols": lambda: wide[:N, kb:], + # Rows along the stride-1 axis of the buffer. + "transposed": lambda: mx.swapaxes(turned, 0, 1), + "whole": lambda: w[:N], + } + + +def _nodes(*outputs): + """Names of the primitives in the graph of the outputs.""" + import tempfile + + with tempfile.NamedTemporaryFile("r", suffix=".dot") as f: + mx.export_to_dot(f.name, *outputs) + text = f.read() + return sorted(re.findall(r'label ="([A-Za-z0-9_]+)"', text)) + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("codec", CODECS) +def test_quantized_matmul_weight_views(codec, state): + views = _views(codec) + for m in ROWS: + x = _x(m, K) + for name, make in views.items(): + ref = kq.quantized_matmul(x, _dense(make()), SCALES, codec, True) + view = make() + if state == "evaluated": + mx.eval(view) + y = kq.quantized_matmul(x, view, SCALES, codec, True) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (codec, m, name, state) + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) +@pytest.mark.parametrize("m", [64, 1100]) +@pytest.mark.parametrize("codec", ["q4_k", "q8_0"]) +def test_glu_weight_views(codec, m, dtype, state): + """64 rows build the composition on every device, 1100 rows run the + fused kernel where there is one.""" + both = _wire(codec, 2 * N, 768, 5) + x = _x(m, 768, dtype) + pairs = { + "interleaved": (lambda: both[0::2], lambda: both[1::2]), + "bands": (lambda: both[:N], lambda: both[N:]), + } + for activation in ("silu", "gelu_tanh"): + for name, (gate, up) in pairs.items(): + ref = kq.quantized_matmul_glu( + x, _dense(gate()), _dense(up()), codec, activation + ) + wg, wu = gate(), up() + if state == "evaluated": + mx.eval(wg, wu) + y = kq.quantized_matmul_glu(x, wg, wu, codec, activation) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (codec, m, name, activation) + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) +def test_qmv_bias_weight_views(dtype, state): + """The weight of the bias-fused mat-vec, which runs q8_0.""" + x = _x(1, K, dtype) + b = (mx.random.normal((N,), key=mx.random.key(9)) * 4).astype(dtype) + mx.eval(b) + for name, make in _views("q8_0").items(): + ref = kq.quantized_matmul_qmv_bias(x, _dense(make()), SCALES, b, "q8_0") + view = make() + if state == "evaluated": + mx.eval(view) + y = kq.quantized_matmul_qmv_bias(x, view, SCALES, b, "q8_0") + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (name, state) + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("sorted_indices", [False, True], ids=["unsorted", "sorted"]) +@pytest.mark.parametrize("rows", [6, 300]) +def test_gather_qmm_weight_views(rows, sorted_indices, state): + """An expert stack that is every other row of each expert. The rows + of a matrix are dense and the matrices are not, with and without the + sorted_indices promise.""" + experts = 4 + stack = _dense(_wire("q8_0", experts * 2 * N, K, 12).reshape(experts, 2 * N, -1)) + mx.eval(stack) + rng = np.random.default_rng(rows) + rhs = mx.array(np.sort(rng.integers(0, experts, size=rows)).astype(np.uint32)) + x = _x(rows, K, mx.float16).reshape(rows, 1, K) + + def call(w): + return kq.gather_qmm( + x, + w, + SCALES, + "q8_0", + rhs_indices=rhs, + transpose=True, + sorted_indices=sorted_indices, + ) + + for make in (lambda: stack[:, 0::2], lambda: stack[:, 1::2]): + ref = call(_dense(make())) + view = make() + if state == "evaluated": + mx.eval(view) + y = call(view) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item() + + +def test_weight_with_memory_adds_no_node(): + """The call gmlx makes in decode: an evaluated weight with dense rows. + Its graph is the cast-free x layout node and the matmul, and nothing + for the weight.""" + w = _wire("q4_k", 2 * N, K, 1) + x = _x(4, K) + assert _nodes(kq.quantized_matmul(x, w, SCALES, "q4_k", True)) == [ + "Contiguous", + "KQuantMatmul", + ] + # A row band with memory has dense rows too. + band = w[N // 2 : N // 2 + N] + mx.eval(band) + assert _nodes(kq.quantized_matmul(x, band, SCALES, "q4_k", True)) == [ + "Contiguous", + "KQuantMatmul", + ] + + +def test_lazy_weight_gets_the_eval_time_node(): + w = _wire("q4_k", 2 * N, K, 1) + x = _x(4, K) + names = _nodes(kq.quantized_matmul(x, w[0::2], SCALES, "q4_k", True)) + assert names.count("KQuantMatrixDense") == 1 + assert "KQuantMatmul" in names + + +def test_dense_lazy_view_keeps_its_buffer(): + """A lazy row band is passed through, not copied: the call allocates the + output and nothing the size of the weight.""" + n, k = 2048, 2048 + w = _wire("q8_0", 2 * n, k, 7) + x = _x(4, k) + y = kq.quantized_matmul(x, w[n:], SCALES, "q8_0", True) + mx.eval(y) + del y + mx.synchronize() + mx.clear_cache() + base = mx.get_active_memory() + mx.reset_peak_memory() + y = kq.quantized_matmul(x, w[n:], SCALES, "q8_0", True) + mx.eval(y) + grown = mx.get_peak_memory() - base + # The band is 4.4 MB. The output and x layout node are a few KB. + assert grown < w.nbytes // 8, (grown, w.nbytes) + + +# dequantize decodes its input as one row-major run of bytes, quantize +# encodes its float input the same way, and the bias of the fused mat-vec +# is read as dense elements. A view passes through when it is row-major in +# all of its dims and is copied otherwise, lazy or not. + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("codec", CODECS) +def test_dequantize_weight_views(codec, state): + views = _views(codec) + w = _wire(codec, 2 * N, K, 1) + views["reversed"] = lambda: w[::-1][:N] + for name, make in views.items(): + ref = kq.dequantize(_dense(make()), SCALES, codec, dtype=mx.float32) + view = make() + if state == "evaluated": + mx.eval(view) + y = kq.dequantize(view, SCALES, codec, dtype=mx.float32) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (codec, name, state) + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +def test_dequantize_and_quantize_3d_views(state): + """Every other matrix of a 3-D input. Each matrix has dense rows, and + the run of bytes or floats as a whole is not row-major.""" + wf = mx.random.normal((8, 16, K), key=mx.random.key(3)) * 0.1 + w, _ = kq.quantize(wf, "q8_0") + mx.eval(wf, w) + for step in (slice(0, None, 2), slice(1, None, 2), slice(None, None, -1)): + ref = kq.dequantize(_dense(w[step]), SCALES, "q8_0", dtype=mx.float32) + view = w[step] + if state == "evaluated": + mx.eval(view) + y = kq.dequantize(view, SCALES, "q8_0", dtype=mx.float32) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (step, state) + ref, _ = kq.quantize(_dense(wf[step]), "q8_0") + view = wf[step] + if state == "evaluated": + mx.eval(view) + y, _ = kq.quantize(view, "q8_0") + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (step, state) + + +def test_dequantize_nodes(): + """A weight with memory adds no node. A lazy view gets the node that + decides at eval.""" + w = _wire("q4_k", 2 * N, K, 1) + assert _nodes(kq.dequantize(w, SCALES, "q4_k")) == ["KQuantDequantize"] + names = _nodes(kq.dequantize(w[0::2], SCALES, "q4_k")) + assert names.count("KQuantMatrixDense") == 1 + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("dtype", [mx.float32, mx.float16], ids=["f32", "f16"]) +@pytest.mark.parametrize("codec", ["q4_k", "q8_0"]) +def test_quantize_float_views(codec, dtype, state): + wf = (mx.random.normal((2 * N, 2 * K), key=mx.random.key(3)) * 0.1).astype(dtype) + turned = _dense(mx.swapaxes(wf[:N, :K], 0, 1)) + im2 = mx.abs(mx.random.normal((2 * K,), key=mx.random.key(4))) + 0.5 + mx.eval(wf, turned, im2) + views = { + "even_rows": lambda: wf[0::2, :K], + "row_band": lambda: wf[N // 2 : N // 2 + N, :K], + "even_cols": lambda: wf[:N, 0::2], + "reversed": lambda: wf[::-1][:N, K:], + "transposed": lambda: mx.swapaxes(turned, 0, 1), + } + for name, make in views.items(): + # q4_k weighs its error by an importance vector. + for imatrix in (None, lambda: im2[0::2]) if codec == "q4_k" else (None,): + ref_im = None if imatrix is None else _dense(imatrix()) + ref, _ = kq.quantize(_dense(make()), codec, ref_im) + view = make() + im = None if imatrix is None else imatrix() + if state == "evaluated": + mx.eval(view) + if im is not None: + mx.eval(im) + y, _ = kq.quantize(view, codec, im) + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (codec, name, state) + + +@pytest.mark.parametrize("state", ["lazy", "evaluated"]) +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) +def test_qmv_bias_views(dtype, state): + # The bias-fused mat-vec runs q8_0 weights. + w = _wire("q8_0", 2 * N, K, 1)[:N] + mx.eval(w) + x = _x(1, K, dtype) + b2 = (mx.random.normal((2 * N,), key=mx.random.key(9)) * 4).astype(dtype) + mx.eval(b2) + for name, make in { + "every_other": lambda: b2[0::2], + "reversed": lambda: b2[::-1][:N], + "band": lambda: b2[N // 2 : N // 2 + N], + }.items(): + # numpy has no bfloat16, and float32 holds both dtypes exactly. + dense = _dense(make().astype(mx.float32)).astype(dtype) + ref = kq.quantized_matmul_qmv_bias(x, w, SCALES, dense, "q8_0") + bias = make() + if state == "evaluated": + mx.eval(bias) + y = kq.quantized_matmul_qmv_bias(x, w, SCALES, bias, "q8_0") + mx.eval(y, ref) + assert mx.array_equal(y, ref).item(), (name, state)