diff --git a/CHANGELOG.md b/CHANGELOG.md index 0259f21..17aa076 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,20 @@ 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] + +### Added +- `kq.quantized_matmul_multi` computes the products of several packed weights + with one input, such as the q, k and v projections, in one dispatch on NAX + GPUs (M5) at speculative-verify widths of 2 to 8 rows. + +### Changed +- `kq.sdpa_fa_verify` runs its key pass on the NAX tensor-op units (M5) at head + dims 128 and 256, so verify attention over a deep cache takes less time. +- On NAX GPUs (M5), the `verify_nax` matmul kernels of `q4_0` at a K that is a + multiple of 256, `q8_0`, the k-quants and `iq4_*` read two steps of activations + per load, so speculative verify at 2 to 8 rows takes less time. + ## [0.5.0] ### Added @@ -30,6 +44,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). 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_QMM_LOG` also logs the mat-vec, gather and MoE dispatches with `xa`, the + activation's byte offset modulo 16, and each operand an op copies to align 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. @@ -99,6 +115,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - 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. +- float16 and bfloat16 inputs that start off a 16-byte boundary, such as a row + slice, no longer give wrong output from matmul, norm, attention and DSA ops. ## [0.4.18] diff --git a/benchmarks/dit_common.py b/benchmarks/dit_common.py index 0d0db26..cc77395 100644 --- a/benchmarks/dit_common.py +++ b/benchmarks/dit_common.py @@ -134,7 +134,7 @@ class KernelLog: 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+)$" + r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+) xa=\d+$" ) def __enter__(self): diff --git a/bindings.cpp b/bindings.cpp index 260a009..f15632d 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -666,6 +666,42 @@ NB_MODULE(_ext, m) { the input dtype. )"); + m.def( + "quantized_matmul_multi", + &mlx_kquant::quantized_matmul_multi, + "x"_a, + "weights"_a, + "kquant_type"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + Quantized matmuls of several weights that share x: + [x @ W_i^T for W_i in weights]. + + On a NAX GPU, at 2 to 8 rows of x where quantized_matmul would run + its verify_nax route on the widest weight, one dispatch per four + weights computes the products. Each weight is read from its own + buffer and each product is written to its own output, so the + weights need no joining copy. One split-K count covers the joint + grid. The narrower weights take the widest weight's route, so + their outputs can differ from quantized_matmul in the last bits. + Every other case runs quantized_matmul per weight. + KQ_MATMUL_MULTI=0 runs quantized_matmul per weight on every call + (read once per process unless KQ_TUNE_LIVE=1). + mx.grad gives the gradient with respect to x; the weights are + frozen, as in quantized_matmul. + + Args: + x (array): [..., K], float16/bfloat16. float32 runs in bfloat16. + weights (list[array]): 2-D uint8 wire bytes [N_i, K bytes] of + one codec. + kquant_type (str): the codec of every weight. + + Returns: + list[array]: [..., N_i] per weight, in the dtype + quantized_matmul gives for x. + )"); + m.def( "quantized_matmul_glu", &mlx_kquant::quantized_matmul_glu, @@ -1830,9 +1866,13 @@ NB_MODULE(_ext, m) { GQA fold as sdpa_fa_verify (q [1, Hkv, G*q_len, D], kv-major heads, original q_len; G*q_len <= 64, or 32 at head_dim 512) in the rotated domain, and the sealed records dequantize at tile stage - exactly as in sdpa_decode_gqa_kvarn, so the output matches - sdpa_fa_verify over the materialized cache bit for bit. Operands, - n, n_attend and full_visibility follow sdpa_decode_gqa_kvarn; the + exactly as in sdpa_decode_gqa_kvarn. The key pass runs on the + simdgroup-matrix kernel at every head dim, so the output matches + sdpa_fa_verify over the materialized cache bit for bit where that + call runs the same kernel (KQ_FA_VERIFY_NAX=0). Where + sdpa_fa_verify takes its NAX key pass (head_dim 128 and 256 on NAX + GPUs), the two agree to rounding. Operands, n, n_attend and + full_visibility follow sdpa_decode_gqa_kvarn; the result is rotated-domain like the query (un-rotate it). This is the verify-width route: the vector kernel's cost climbs steeply past two queries, while the matrix tile prices extra rows at nearly diff --git a/docs/integration.md b/docs/integration.md index 77476d4..e79d772 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -38,7 +38,8 @@ boundary. It raises from the call when the weight is evaluated and row-contiguou `mx.eval` otherwise. Arrays that MLX allocates and tensors from `kq.load_gguf` always meet the boundary, as do slices of them that start on a block boundary. A view that starts inside a block can miss it, and a copy in a new buffer, such as `mx.array(np.array(w))`, runs. The CPU kernels accept -any start. +any start. Activations, KV caches and the other float operands can start anywhere: a GPU op copies +one that starts off the boundary of its kernel's vector loads to a new buffer first. ## Codecs diff --git a/docs/kernels.md b/docs/kernels.md index af3452c..881341c 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -217,18 +217,19 @@ 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 + `KQ_NAX_BM128_ODD_M`, `KQ_NAX_DB64`, the `_BK128` levers, `KQ_NAX_SWIZZLE`), the `verify_nax` load + lever (`KQ_VNAX_X16`), 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_multi` lever (`KQ_MATMUL_MULTI`), 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_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 @@ -248,14 +249,19 @@ Tuning levers (defaults are right for normal use): 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_QMM_LOG` - set to any value, logs every quantized matmul dispatch to stderr, one + `[kq qmm] M= N= K= grid=,, xa=` line each: the NAX and ALU qmm + tiles, the mat-vec routes (`qmv`, `qvm`, `verify_qmv`, `mv_ext`), split-K, `verify_mma` and + `verify_nax`, split or not, the fused `quantized_matmul_glu`, `quantized_matmul_lora` and + `quantized_matmul_residual` kernels, `skinny_matmul`, the `gather_qmm` and `gather_qmm_seg` + routes, the fused MoE gathers, `gather_mix` and `conv2d_3x3`. The grid depth is the K slice + count. `xa` is the byte offset of the activation's start modulo 16, as the kernel reads it. A + gather or MoE line gives the rows and last dim of the output as M and N. Any GPU op that copies + an operand to an aligned buffer, because the operand starts or steps off the boundary of its + kernel's vector loads, logs one `[kq align] xin=` line, where `n` is the + operand's byte offset modulo that boundary. A call that takes the dequantize route logs one + `[kq dequant] M= N= K=` line when the op builds its graph. 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 @@ -397,6 +403,10 @@ Tuning levers (defaults are right for normal use): 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_VNAX_X16` - probe lever: `0` keeps the `verify_nax` kernels that load their activations per + NAX op on one 8-byte load per op where the activations start on a 16-byte boundary, in place of + one 16-byte load per two ops. The output is the same. In `KQ_QMM_LOG` the kernel name + carries the variant, `_x16` or `_x8`. Read once per process unless `KQ_TUNE_LIVE=1`. - `KQ_MV_EXT_SB` / `KQ_MV_EXT_NX` / `KQ_MV_EXT_HD` - `mv_ext` activation-traffic experiments: shuffle-broadcast (`1`), wide nxpsg (`16`/`32`), half-precision chunk dots (`1`). q6_k M 4-12 only. `HD` measured +4-5% at M 8; the rest flat to negative on M5 Max. Kept as probes. Default @@ -554,7 +564,11 @@ eight-block kernel at a K that is a multiple of 256 and the two-block kernel els kernels fold each 16- or 32-weight sub-block's scale, and its min where the codec has one, into the half weights. Loading a whole K step of activations ahead measured faster for `pq2_0` and the two-block `q4_0` kernel, while the kernels with larger decoded words load them one NAX op at a time, -since loading ahead ran `q4_k` 1.6-2x slower. +since loading ahead ran `q4_k` 1.6-2x slower. Two consecutive NAX ops read 8 contiguous activations +of a row, so where the activations start on a 16-byte boundary these kernels read them in one +16-byte load per two NAX ops, and where they start on an 8-byte boundary in one 8-byte load per op. +Activations that start elsewhere are read from a copy on a 16-byte boundary. The fewer loads take +less time in a chain of dependent calls, and both forms give the same bits. Split counts for `pq2_0`, `q8_0` and the two-block `q4_0` kernel are the largest divisor of the K steps at or under their target, at most 16. The two-block kernel's K walk of 640 per simdgroup @@ -642,6 +656,45 @@ 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. +## Quantized matmuls that share an input + +- `quantized_matmul_multi` - `[x @ W^T for W in weights]` for packed weights of one codec and one + inner dim that read the same x, such as the q, k and v projections or the gate and up projections + of a transformer block. It returns one output per weight, as `quantized_matmul` would. + + On a NAX GPU (M5), at 2 to 8 activation rows, where `quantized_matmul` would take its verify_nax + route on the widest weight, one dispatch computes the products of up to four weights. The row + tiles of the weights lie end to end in one grid, each threadgroup owns 128 rows of one weight, + and the split-K partials of all of them share one temporary and one fold dispatch. More than four + weights run one dispatch per four. The split count is chosen for the joint grid. The call saves + the dispatch, the wait and the tail of every weight but the first, which counts most for narrow + weights next to a wide one. + + The narrower weights take the route of the widest. A weight that `quantized_matmul` alone would + run on per-row qmv or on another kernel runs on verify_nax here, so its output can differ from + `quantized_matmul` in the last bits, inside the same bound against a float64 reference. This + holds for every codec with a verify_nax kernel, including those whose verify_nax entry starts at + a width that the narrow weight does not reach. Where every weight would take verify_nax at one + split count, the outputs equal `quantized_matmul` bit for bit. + + Every other case runs `quantized_matmul` per weight: one row or more than 8, a single weight, + codecs without a verify_nax kernel, an inner dim that kernel does not step over, a weight whose + start it cannot read, CPU streams and GPUs without NAX. x is float16 or bfloat16, float32 x runs + in bfloat16 as in `quantized_matmul`, and leading batch dims fold into rows. + + The route levers of `quantized_matmul` decide the route on the widest weight the same way: + `KQ_NAX_QMV`, `KQ_VERIFY_NAX`, `KQ_VERIFY_MMA`, `KQ_QMM_SPLITK_NAX`, `KQ_QMM_SPLITK`, + `KQ_VERIFY_EXT` and `KQ_DISABLE_NAX`. A set `KQ_QMM_ROUTE` or `KQ_FORCE_QMM_MIN_M` runs + `quantized_matmul` per weight, so a forced route times the route it names. + + Lever, read when the op builds its graph, once per process unless `KQ_TUNE_LIVE=1`: + + - `KQ_MATMUL_MULTI` - `0` builds `quantized_matmul` per weight on every call. Unset (on) by + default. + + `mx.grad` gives the gradient with respect to x through `quantized_matmul` per weight. The weights + are frozen, as in `quantized_matmul`. + ## 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 @@ -1019,7 +1072,16 @@ into a packed layout first. blocks instead of the full key axis and no `[L, S]` mask is materialized. Causality inside the diagonal page is enforced in-kernel. - **`sdpa_fa_verify`** - speculative-verify attention on the matrix units for a GQA-folded query tile. - Head dims 64 through 512; `return_lse` as above. + Head dims 64 through 512; `return_lse` as above. On NAX GPUs, at head dims 128 and 256 with up to + 64 folded rows on plain K and V, the key pass runs on the tensor-op units: each 16-row strip of + the fold is one row of simdgroups, and the head dim splits over 2 simdgroups (4 at head dim 256 + for folds of up to 32 rows), which add their partial scores in threadgroup memory. It writes the + partials of the simdgroup-matrix kernel to the same merge pass. Its automatic split count aims at + about 64 threadgroups over the KV heads, between 4 and 64, and halves while a split would hold + fewer than 64 keys. The two kernels round differently, so their outputs are not bit-identical. + `KQ_FA_VERIFY_NAX=0` keeps the simdgroup-matrix kernel, and `KQ_FA_VERIFY_NAX_WN` sets the head + dim split (2 or 4 at 256, 1 or 2 at 128; another value warns once on stderr and keeps the + default); both are read at each call. - **`sdpa_fa_indexed`** - index-gathered attention over one shared K/V latent at head dim 512, the absorbed-MLA sparse decode step: query j of `q [1, Hq, Q, 512]` attends the rows of `kv [1, 1, N, 512]` that `idx[j]` (int32, -1 pads) lists, so the selected rows are read once from the latent @@ -1049,8 +1111,11 @@ their cache layout. row's region map follows its own end and `tail_rows` is the per-row form of `n_attend`, so a ragged batch merges per-row tails the same way. Head dims 128, 256 and 512, q_len 1 to 4. - **`sdpa_fa_verify_kvarn`** - the same matrix-unit verify pass over a KVarN cache: sealed records - dequantize at tile stage through the loaders `sdpa_decode_gqa_kvarn` uses, so the result matches - `sdpa_fa_verify` over the materialized cache bit for bit. The verify-width route (q_len 2 to 8): + dequantize at tile stage through the loaders `sdpa_decode_gqa_kvarn` uses, and its key pass runs + on the simdgroup-matrix kernel at every head dim. The result matches `sdpa_fa_verify` over the + materialized cache bit for bit where that call runs the same kernel (`KQ_FA_VERIFY_NAX=0`). Where + `sdpa_fa_verify` takes its NAX key pass (head dims 128 and 256 on NAX GPUs), the two agree to + rounding. The verify-width route (q_len 2 to 8): the vector kernel's cost climbs steeply past two queries, the matrix tile prices extra rows at nearly zero. Head dims 128, 256 and 512; `n_attend` / `full_visibility` / `return_lse` as the decode op. diff --git a/metal/kq_quantized_nax.metal b/metal/kq_quantized_nax.metal index 1378f89..9994f5d 100644 --- a/metal/kq_quantized_nax.metal +++ b/metal/kq_quantized_nax.metal @@ -443,4 +443,41 @@ instantiate_kquant_verify_nax(bfloat16_t, 256, 3, q3_k) instantiate_kquant_verify_nax(float16_t, 256, 3, q3_k) instantiate_kquant_verify_nax(bfloat16_t, 256, 2, q2_k) instantiate_kquant_verify_nax(float16_t, 256, 2, q2_k) + +// Several weights that share x in one dispatch (kq_verify_nax.h). +#define instantiate_kquant_verify_nax_multi(type, gs, bits, codec) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_nax_multi_" #type "_gs_" #gs "_b_" \ + #bits, \ + kq_ ## codec ## _verify_nax_multi, \ + type, gs, bits) \ + instantiate_kernel( \ + "kquant_" #codec "_verify_nax_multi_" #type "_gs_" #gs "_b_" \ + #bits "_pf32", \ + kq_ ## codec ## _verify_nax_multi, \ + type, gs, bits, float) +instantiate_kquant_verify_nax_multi(bfloat16_t, 128, 2, pq2_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 64, 2, q2_0) +instantiate_kquant_verify_nax_multi(float16_t, 128, 2, pq2_0) +instantiate_kquant_verify_nax_multi(float16_t, 64, 2, q2_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 4, q4_0) +instantiate_kquant_verify_nax_multi(float16_t, 32, 4, q4_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 4, q4_0_sb) +instantiate_kquant_verify_nax_multi(float16_t, 32, 4, q4_0_sb) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 8, q8_0) +instantiate_kquant_verify_nax_multi(float16_t, 32, 8, q8_0) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 4, q4_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 4, q4_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 4, iq4_xs) +instantiate_kquant_verify_nax_multi(float16_t, 256, 4, iq4_xs) +instantiate_kquant_verify_nax_multi(bfloat16_t, 32, 4, iq4_nl) +instantiate_kquant_verify_nax_multi(float16_t, 32, 4, iq4_nl) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 5, q5_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 5, q5_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 6, q6_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 6, q6_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 3, q3_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 3, q3_k) +instantiate_kquant_verify_nax_multi(bfloat16_t, 256, 2, q2_k) +instantiate_kquant_verify_nax_multi(float16_t, 256, 2, q2_k) // clang-format on diff --git a/metal/kq_sdpa_nax.metal b/metal/kq_sdpa_nax.metal index 81c354e..04c2ef6 100644 --- a/metal/kq_sdpa_nax.metal +++ b/metal/kq_sdpa_nax.metal @@ -20,3 +20,24 @@ instantiate_kq_sdpa_fa_indexed_nax(float16_t) instantiate_kq_sdpa_gqa_nax(bfloat16_t) instantiate_kq_sdpa_gqa_nax(float16_t) + +#define instantiate_kq_sdpa_fa_verify_nax(type, D, WM, WN) \ + instantiate_kernel( \ + "kq_sdpa_fa_verify_2pass_1_nax_" #type "_" #D "_wm" #WM "_wn" #WN, \ + kq_sdpa_fa_verify_2pass_1_nax, \ + type, D, WM, WN) + +#define instantiate_kq_sdpa_fa_verify_nax_rows(type, D, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 1, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 2, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 3, WN) \ + instantiate_kq_sdpa_fa_verify_nax(type, D, 4, WN) + +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 256, 2) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 256, 2) +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 256, 4) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 256, 4) +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 128, 1) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 128, 1) +instantiate_kq_sdpa_fa_verify_nax_rows(bfloat16_t, 128, 2) +instantiate_kq_sdpa_fa_verify_nax_rows(float16_t, 128, 2) diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h b/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h index b7a5dc6..82ca870 100644 --- a/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_sdpa_nax.h @@ -607,3 +607,261 @@ template } } } + +// Speculative-verify attention on the tensor-op units, pass 1 of the +// sdpa_fa_verify split-K form, after MLX's head-dim split NAX attention +// (attention_nax_dsplit). It writes the partials of +// kq_sdpa_fa_verify_2pass_1 (kq_sdpa.h), so kq_sdpa_gqa_2pass_2 merges +// either. The caller folds the GQA group into the query rows: q is +// [B, Hkv, n_rows, D], and folded row r attends keys up to +// N - q_len + (r % q_len). +// +// One threadgroup per (kv head, batch row, key split), WM x WN simdgroups. +// Simdgroup (row group, d group) owns query rows 16 * row group onward and +// the d group's D / WN columns: its Q slice stays in registers, it reads +// its columns of each 32-key K and V tile from device memory into +// fragments, and keeps O for its columns only. The WN partial scores of a +// row group are summed in threadgroup memory in a fixed order, and every +// d group runs the same row softmax on the sum. A tile past a row's limit +// leaves the row's running max at finite_min and its P row at zero; an +// empty split writes O = 0, sum = 0 and max = finite_min. +template +[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +kq_sdpa_fa_verify_2pass_1_nax( + const device T* queries [[buffer(0)]], + const device T* keys [[buffer(1)]], + const device T* values [[buffer(2)]], + device float* out [[buffer(3)]], + device float* sums [[buffer(4)]], + device float* maxs [[buffer(5)]], + const constant int& N [[buffer(6)]], + const constant size_t& k_head_stride [[buffer(7)]], + const constant size_t& k_seq_stride [[buffer(8)]], + const constant size_t& v_head_stride [[buffer(9)]], + const constant size_t& v_seq_stride [[buffer(10)]], + const constant float& scale [[buffer(11)]], + const constant int& q_len [[buffer(12)]], + const constant int& n_rows [[buffer(13)]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tpg [[threadgroups_per_grid]]) { + using namespace kq_sdpa_nax; + constexpr short kU = 16; + constexpr int BK = 32; // keys per tile + constexpr short TK = BK / kU; + constexpr short TDh = D / kU / WN; // column fragments per simdgroup + constexpr int BDh = D / WN; + static_assert(TDh % 2 == 0, "P @ V accumulates column fragments in pairs"); + + using otile_t = NAXTile; + using stile_t = NAXTile; + constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; + constexpr short kRowsPT = otile_t::kRowsPerThread; + + threadgroup float s_xchg[WM][WN][TK * kEPF * 32]; + + const short row_group = simd_group_id / WN; + const short d_group = simd_group_id % WN; + const int n_kv_heads = tpg.x; + const int split_idx = tid.z; + const size_t qbh = (size_t)tid.y * n_kv_heads + tid.x; + const int tm = kU * row_group; + + 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); + + const int k_ld = int(k_seq_stride); + const int v_ld = int(v_seq_stride); + const device T* K = + keys + qbh * k_head_stride + (size_t)k0 * k_seq_stride + d_group * BDh; + const device T* V = + values + qbh * v_head_stride + (size_t)k0 * v_seq_stride + d_group * BDh; + + const short2 sc = otile_t::NAXFrag_t::get_coord(); + const short sm = sc.y; + const short sn = sc.x; + + // Q slice resident in registers; rows past n_rows load as zero. + const short q_rows = short(metal::clamp(n_rows - tm, 0, int(kU))); + NAXTile Qtiles[TDh]; + { + const device T* qs = queries + (qbh * n_rows + tm) * D + d_group * BDh; + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + Qtiles[id].load_rows(qs + id * kU, D, q_rows); + } + } + + otile_t Otile; + Otile.clear(); + metal::vec max_score; + metal::vec sum_score{0}; + // Highest key of this lane's rows (tm + sm and tm + sm + 8). + int lim[kRowsPT]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = Limits::finite_min; + lim[i] = N - q_len + ((tm + sm + i * 8) % q_len); + } + const int lim_min = N - q_len; // every row reaches it + const float scale2 = scale * M_LOG2E_F; + + for (int kt = k0; kt < k1; kt += BK) { + const int krem = min(k1 - kt, BK); + const bool full = krem == BK; + + stile_t Stile; + Stile.clear(); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik += 2) { + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + NAXTile Ktile; + const device T* kp = K + ik * kU * k_ld + id * kU; + if (full) { + Ktile.load(kp, k_ld); + } else { + Ktile.load_rows(kp, k_ld, short(krem - ik * kU)); + } + stile_t::NAXFrag_t::mma( + Stile.frag_at(0, ik), + Stile.frag_at(0, ik + 1), + Qtiles[id].frag_at(0, 0), + metal::false_type{}, + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), + metal::true_type{}); + } + } + + if (WN > 1) { + // Sum the WN partial scores of the row group in a fixed order. + threadgroup float* slot = s_xchg[row_group][d_group]; + const short base = short(simd_lane_id) * (TK * kEPF); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + slot[base + ik * kEPF + i] = Stile.frag_at(0, ik)[i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] = s_xchg[row_group][0][base + ik * kEPF + i]; + } + STEEL_PRAGMA_UNROLL + for (short p = 1; p < WN; p++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] += s_xchg[row_group][p][base + ik * kEPF + i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { + Stile.elems()[ii] *= scale2; + } + // Only the split tail or a tile reaching past N - q_len masks. + if (!full || kt + BK - 1 > lim_min) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& fg = Stile.frag_at(0, ik); + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const int c = kt + ik * kU + sn + jj; + const short loc = ii * stile_t::kFragThrCols + jj; + if (c >= k1 || c > lim[ii]) { + fg[loc] = Limits::finite_min; + } + } + } + } + } + + // Online softmax (exp2 space). A row with no valid key yet keeps its + // max at finite_min and contributes a zero P row. + metal::vec new_max = max_score; + Stile.template row_reduce(new_max); + Stile.template row_bin_op(new_max); + metal::vec factor; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + if (new_max[i] > Limits::finite_min) { + factor[i] = fast::exp2(max_score[i] - new_max[i]); + max_score[i] = new_max[i]; + } else { + factor[i] = 1.0f; + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + Stile.frag_at(0, ik)[i * stile_t::kFragThrCols + jj] = 0.0f; + } + } + } + } + sum_score = sum_score * factor; + Stile.template row_reduce(sum_score); + Otile.template row_bin_op(factor); + + simdgroup_barrier(mem_flags::mem_none); + + // O += P @ V over this simdgroup's columns. + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id += 2) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + NAXTile Vtile; + const device T* vp = V + ik * kU * v_ld + id * kU; + if (full) { + Vtile.load(vp, v_ld); + } else { + Vtile.load_rows(vp, v_ld, short(krem - ik * kU)); + } + otile_t::NAXFrag_t::mma( + Otile.frag_at(0, id), + Otile.frag_at(0, id + 1), + Stile.frag_at(0, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); + } + } + + K += BK * k_ld; + V += BK * v_ld; + } + + // Unnormalized partials [B, Hkv, n_rows, splits, D]; d group 0 writes the + // row stats, with the max in natural log for the shared merge. + if (q_rows > 0) { + device float* dst = out + + ((qbh * n_rows + tm) * gqa_splits + split_idx) * D + d_group * BDh; + Otile.store_rows(dst, gqa_splits * D, q_rows); + if (d_group == 0 && sn == 0) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const int row = tm + sm + i * 8; + if (row < n_rows) { + const size_t po = (qbh * n_rows + row) * gqa_splits + split_idx; + sums[po] = sum_score[i]; + maxs[po] = max_score[i] == Limits::finite_min + ? Limits::finite_min + : max_score[i] * M_LN2_F; + } + } + } + } +} diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index dc32757..4efae33 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -42,9 +42,34 @@ // pair(words, f): fragment f's element pair, block scale applied // xstep(xb, q, s): left elements 0..3 of step s, xb the unit's // first activation of the lane's row +// xoff(q, s): where xstep reads, in elements from xb. A codec that loads +// one step at a time reads steps s and s + 1 (s even) from 8 +// contiguous elements at a 16-byte boundary of the row, so where x +// starts on one (kq_vnax_x16) it loads both in one 16-byte load. + +// Function constant: x starts on a 16-byte boundary. Unset keeps the +// 8-byte load per step. +constant bool kq_vnax_x16_fc [[function_constant(40)]]; +constant bool kq_vnax_x16 = + is_function_constant_defined(kq_vnax_x16_fc) && kq_vnax_x16_fc; MLX_MTL_CONST int KQ_VNAX_NSG = 4; +// The 16-byte load's geometry: for every lane-quad q and even step s, +// steps s and s + 1 sit in 8 contiguous elements at an 8-element boundary. +template +constexpr bool kq_vnax_x16_pairs() { + for (short q = 0; q < 8; ++q) { + for (short s = 0; s < SPB; s += 2) { + if (Codec::xoff(q, s) % 8 != 0 || + Codec::xoff(q, s + 1) != Codec::xoff(q, s) + 4) { + return false; + } + } + } + return true; +} + // PQ2_0: lane-quad q owns the two 16-code words at qs bytes 8q..8q+7. // Pair f is code f % 8 of word f / 8 in each half (KqPq2_0Mma). struct KqPq2_0Nax { @@ -53,6 +78,10 @@ struct KqPq2_0Nax { static constant constexpr int block_bytes = KQ_PQ2_0_BLOCK_BYTES; static constant constexpr int ub = 1; static constant constexpr bool x_ahead = true; + // Unused: x_ahead codecs load their steps through xstep. + static constexpr METAL_FUNC int xoff(short, short) { + return 0; + } struct Words { uint w[2]; half d; @@ -98,6 +127,10 @@ struct KqQ2_0Nax { static constant constexpr int block_bytes = KQ_Q2_0_BLOCK_BYTES; static constant constexpr int ub = 2; static constant constexpr bool x_ahead = true; + // Unused: x_ahead codecs load their steps through xstep. + static constexpr METAL_FUNC int xoff(short, short) { + return 0; + } struct Words { uint w; half d; @@ -163,9 +196,12 @@ struct KqQ4_0Nax { half2(1032.0h)) * o.d; } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 16 * s + 4 * q; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 16 * s + 4 * q); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -216,9 +252,12 @@ struct KqQ4_0SbNax { half2(1032.0h)) * o.d[b]; } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 64 * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 64 * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -266,10 +305,12 @@ struct KqQ8_0Nax { half2(1152.0h)) * o.d; } + static constexpr METAL_FUNC int xoff(short q, short s) { + return KQ_Q8_0_GROUP * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + KQ_Q8_0_GROUP * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -346,10 +387,12 @@ struct KqQ4_KNax { const uint wv = (hi ? (o.w[s & 7] >> 4) : o.w[s & 7]) & 0x0F0F0F0Fu; return kq_vnax_kpair(wv >> (8 * (f & 1)), o.ds, o.dm, hi); } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 64 * q + 32 * (s >> 3) + 4 * (s & 7); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 64 * q + 32 * (s >> 3) + 4 * (s & 7)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -409,11 +452,12 @@ struct KqQ5_KNax { lo4 | (((o.hb >> (2 * i + (hi ? 1 : 0))) & 0x01010101u) << 4); return kq_vnax_kpair(q5 >> (8 * (f & 1)), o.ds, o.dm, hi); } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 64 * (q >> 1) + 16 * (q & 1) + 32 * (s >> 2) + 4 * (s & 3); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 64 * (q >> 1) + 16 * (q & 1) + - 32 * (s >> 2) + 4 * (s & 3)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -485,11 +529,12 @@ struct KqQ6_KNax { half2(1056.0h); return v * (hi ? o.sc.y : o.sc.x); } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 32 * (q & 1) + 16 * (q >> 1) + 64 * (s >> 2) + 4 * (s & 3); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 32 * (q & 1) + 16 * (q >> 1) + - 64 * (s >> 2) + 4 * (s & 3)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -561,10 +606,12 @@ struct KqQ3_KNax { const half2 sc = o.sc[j >> 1]; return v * ((j & 1) ? sc.y : sc.x); } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 32 * (s >> 1) + 8 * q + 4 * (s & 1); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 32 * (s >> 1) + 8 * q + 4 * (s & 1)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -624,10 +671,12 @@ struct KqQ2_KNax { const half m = (j & 1) ? o.dm[k].y : o.dm[k].x; return fma(v, half2(d), -half2(m)); } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 32 * (s >> 1) + 8 * q + 4 * (s & 1); + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = - *(const device vec*)(xb + 32 * (s >> 1) + 8 * q + 4 * (s & 1)); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -742,9 +791,12 @@ struct KqIq4XsNax { const uint src = (((s >> 2) & 1) ? (w >> 4) : w) >> (8 * (f & 1)); return kq_vnax_iq4pair(src, o.dl[b]); } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 64 * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 64 * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -788,9 +840,12 @@ struct KqIq4NlNax { const uint src = (((s >> 2) & 1) ? (w >> 4) : w) >> (8 * (f & 1)); return kq_vnax_iq4pair(src, o.d[b]); } + static constexpr METAL_FUNC int xoff(short q, short s) { + return 64 * q + 4 * s; + } template static METAL_FUNC vec xstep(const device T* xb, short q, short s) { - const vec p = *(const device vec*)(xb + 64 * q + 4 * s); + const vec p = *(const device vec*)(xb + xoff(q, s)); return vec(p[0], p[2], p[1], p[3]); } }; @@ -800,10 +855,10 @@ METAL_FUNC void kq_verify_nax_impl( const device uint8_t* w, const device T* x, device U* y, - const constant int& K, - const constant int& N, - const constant int& M, - const constant int& k_partition_size, + const int K, + const int N, + const int M, + const int k_partition_size, uint3 tid, uint simd_gid, uint simd_lid) { @@ -811,6 +866,9 @@ METAL_FUNC void kq_verify_nax_impl( constexpr int BK = Codec::block_k; constexpr int UB = Codec::ub; constexpr int SPB = BK / 16; + static_assert( + Codec::x_ahead || (SPB % 2 == 0 && kq_vnax_x16_pairs()), + "x16 loads step pairs"); const int n0 = (tid.x * KQ_VNAX_NSG + simd_gid) * 32; if (n0 >= N) { return; @@ -874,13 +932,28 @@ METAL_FUNC void kq_verify_nax_impl( } } } + uint4 xu = uint4(0); #pragma unroll for (short u = 0; u < UB; ++u) { #pragma unroll for (short s = 0; s < SPB; ++s) { - const vec xv = Codec::x_ahead ? xa[u][s] - : xrow ? Codec::template xstep(xr + kb + u * BK, q, s) - : vec(0); + vec xv; + if (Codec::x_ahead) { + xv = xa[u][s]; + } else if (kq_vnax_x16) { + // Steps s and s + 1 from one 16-byte load, in xstep's order. + if ((s & 1) == 0) { + xu = xrow + ? *(const device uint4*)(xr + kb + u * BK + Codec::xoff(q, s)) + : uint4(0); + } + const vec a = as_type>((s & 1) ? xu.z : xu.x); + const vec c = as_type>((s & 1) ? xu.w : xu.y); + xv = vec(a[0], c[0], a[1], c[1]); + } else { + xv = xrow ? Codec::template xstep(xr + kb + u * BK, q, s) + : vec(0); + } ca[0] = xv[0]; ca[1] = xv[1]; ca[2] = xv[2]; @@ -944,6 +1017,54 @@ METAL_FUNC void kq_verify_nax_impl( simd_lid); \ } +// Several weights that share x in one dispatch: weight i owns the +// threadgroups [tg_end[i - 1], tg_end[i]) of grid x, each weight from its +// own buffer and into its own output, with one split count over grid z. +// part_stride[i] is output i's partial stride. Unused weight slots bind +// slot 0's buffers, and their tg_end is the grid end, so they own no +// threadgroup. +#define KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(CODEC, TRAITS) \ + template \ + [[kernel]] void kq_##CODEC##_verify_nax_multi( \ + const device uint8_t* w0, \ + const device uint8_t* w1, \ + const device uint8_t* w2, \ + const device uint8_t* w3, \ + const device T* x, \ + device U* y0, \ + device U* y1, \ + device U* y2, \ + device U* y3, \ + const constant int& K, \ + const constant int& M, \ + const constant int& k_partition_size, \ + const constant int4& n_rows, \ + const constant int4& tg_end, \ + const constant int4& part_stride, \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]) { \ + static_assert(group_size == TRAITS::group, #CODEC " block width"); \ + const int t = int(tid.x); \ + /* tg_end is cumulative, so the weight is the ends passed. */ \ + const short i = \ + short(t >= tg_end[0]) + short(t >= tg_end[1]) + short(t >= tg_end[2]); \ + const int t0 = i == 0 ? 0 : tg_end[i - 1]; \ + const device uint8_t* w = i == 0 ? w0 : i == 1 ? w1 : i == 2 ? w2 : w3; \ + device U* y = i == 0 ? y0 : i == 1 ? y1 : i == 2 ? y2 : y3; \ + kq_verify_nax_impl( \ + w, \ + x, \ + y + tid.z * static_cast(part_stride[i]), \ + K, \ + n_rows[i], \ + M, \ + k_partition_size, \ + uint3(uint(t - t0), tid.y, tid.z), \ + simd_gid, \ + simd_lid); \ + } + KQ_DEFINE_VERIFY_NAX_KERNEL(pq2_0, KqPq2_0Nax) KQ_DEFINE_VERIFY_NAX_KERNEL(q2_0, KqQ2_0Nax) KQ_DEFINE_VERIFY_NAX_KERNEL(q4_0, KqQ4_0Nax) @@ -956,3 +1077,16 @@ KQ_DEFINE_VERIFY_NAX_KERNEL(q5_k, KqQ5_KNax) KQ_DEFINE_VERIFY_NAX_KERNEL(q6_k, KqQ6_KNax) KQ_DEFINE_VERIFY_NAX_KERNEL(q3_k, KqQ3_KNax) KQ_DEFINE_VERIFY_NAX_KERNEL(q2_k, KqQ2_KNax) + +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(pq2_0, KqPq2_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q2_0, KqQ2_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q4_0, KqQ4_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q4_0_sb, KqQ4_0SbNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q8_0, KqQ8_0Nax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q4_k, KqQ4_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(iq4_xs, KqIq4XsNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(iq4_nl, KqIq4NlNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q5_k, KqQ5_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q6_k, KqQ6_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q3_k, KqQ3_KNax) +KQ_DEFINE_VERIFY_NAX_MULTI_KERNEL(q2_k, KqQ2_KNax) diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index c757ce5..2436ce5 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -87,6 +87,7 @@ quantized_matmul, quantized_matmul_glu, quantized_matmul_lora, + quantized_matmul_multi, quantized_matmul_qmv_bias, quantized_matmul_residual, residency_commit, @@ -196,6 +197,7 @@ "quantized_matmul", "quantized_matmul_glu", "quantized_matmul_lora", + "quantized_matmul_multi", "quantized_matmul_qmv_bias", "quantized_matmul_residual", "rmsnorm2_add", diff --git a/src/kquant.h b/src/kquant.h index 27d53f0..7379f7b 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -1706,6 +1706,67 @@ class KQuantMatmulFused : public mx::Primitive { int act_; // 0 = silu, 1 = gelu (tanh form) }; +// Quantized matmuls of several weights that share x: y_i = x @ W_i^T. +// x [..., K] (float16/bfloat16; float32 runs in bfloat16), weights 2-D +// uint8 wire bytes of one codec, [N_i, K bytes] each, in separate buffers. +// Returns one [..., N_i] array per weight, in the dtype quantized_matmul +// gives for x. On a NAX GPU, at 2 to 8 rows of x where quantized_matmul +// would run its verify_nax route on the widest weight, one dispatch per +// four weights computes the products: each weight is read in place, each +// output is written in its own buffer, and one split-K count covers the +// joint grid. Every other case runs quantized_matmul per weight. +std::vector quantized_matmul_multi( + mx::array x, + std::vector weights, + const std::string& kquant_type, + mx::StreamOrDevice s = {}); + +// Primitive of quantized_matmul_multi: inputs {x, scales, w_0, ...}, +// one output per weight. vjp and jvp are those of the per-weight +// matmuls, wrt x only. +class KQuantMatmulMulti : public mx::Primitive { + public: + KQuantMatmulMulti( + mx::Stream stream, + std::string kquant_type, + int group_size, + int bits) + : mx::Primitive(stream), + kquant_type_(std::move(kquant_type)), + group_size_(group_size), + bits_(bits) {} + + 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 "KQuantMatmulMulti"; + } + 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 group_size_; + int bits_; +}; + // 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 diff --git a/src/kquant_conv_nax.cpp b/src/kquant_conv_nax.cpp index 3d42e73..7d4350a 100644 --- a/src/kquant_conv_nax.cpp +++ b/src/kquant_conv_nax.cpp @@ -138,14 +138,15 @@ void KQuantConv3x3::eval_gpu( if (conv_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, O, 9 * C, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname); diff --git a/src/kquant_dense.cpp b/src/kquant_dense.cpp index fdeaac6..cebe267 100644 --- a/src/kquant_dense.cpp +++ b/src/kquant_dense.cpp @@ -17,6 +17,7 @@ // 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 @@ -152,6 +153,56 @@ mx::array kq_row_major_gpu(const mx::array& a, const mx::Stream& s) { return kq_dense_copy_gpu(a, s); } +namespace { + +// The copy of an operand that does not start on its kernels' alignment, +// logged under KQ_QMM_LOG with the byte offset it came at. +mx::array aligned_copy( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align) { + if (kq_qmm_log_on()) { + std::fprintf( + stderr, + "[kq align] %s xin=%d\n", + what, + static_cast(a.offset() % align)); + } + return kq_dense_copy_gpu(a, s); +} + +} // namespace + +mx::array kq_x_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align) { + if (a.size() == 0 || (a.flags().row_contiguous && a.offset() % align == 0)) { + return a; + } + return aligned_copy(a, s, what, align); +} + +mx::array kq_strides_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align) { + if (a.size() == 0) { + return a; + } + const int nd = static_cast(a.ndim()); + bool ok = a.offset() % align == 0 && + (nd == 0 || a.shape(nd - 1) == 1 || a.strides()[nd - 1] == 1); + for (int i = 0; ok && i + 1 < nd; i++) { + ok = a.shape(i) == 1 || + (a.strides()[i] * static_cast(a.itemsize())) % align == 0; + } + return ok ? a : aligned_copy(a, s, what, align); +} + 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; @@ -186,6 +237,16 @@ 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_x_aligned_gpu(const mx::array&, const mx::Stream&, const char*, int) { + throw std::runtime_error("[mlx_kquant] GPU eval requires a Metal build."); +} + +mx::array +kq_strides_aligned_gpu(const mx::array&, const mx::Stream&, const char*, int) { + 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."); } diff --git a/src/kquant_dsa_indexer.cpp b/src/kquant_dsa_indexer.cpp index 875fa91..940b08b 100644 --- a/src/kquant_dsa_indexer.cpp +++ b/src/kquant_dsa_indexer.cpp @@ -230,11 +230,16 @@ void KQDsaIndexerScoreDecode::eval_gpu( auto& d = mx::metal::device(s.device); auto& ce = mx::metal::get_command_encoder(s); - const auto& q = inputs[0]; + // q and k are read as 8-wide vectors. + const int align = 8 * static_cast(inputs[0].itemsize()); + const mx::array q = + kq_x_aligned_gpu(inputs[0], s, "dsa_indexer_score_decode q", align); // [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); + // with an element stride other than 1, which is copied here, as is a + // view whose start or batch stride is off the vector alignment. + const mx::array k = + kq_strides_aligned_gpu(inputs[1], s, "dsa_indexer_score_decode k", align); const auto& weights = inputs[2]; auto& out = outputs[0]; diff --git a/src/kquant_dsa_sparse_attn.cpp b/src/kquant_dsa_sparse_attn.cpp index 415ac71..383d6b5 100644 --- a/src/kquant_dsa_sparse_attn.cpp +++ b/src/kquant_dsa_sparse_attn.cpp @@ -36,9 +36,13 @@ void KQDsaSparseAttention::eval_gpu( auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& q = inputs[0]; - const auto& local_kv = inputs[1]; - const auto& pooled = inputs[2]; + // q and the key rows are read as 16-byte vectors through their strides. + const mx::array q = + kq_strides_aligned_gpu(inputs[0], s, "dsa_sparse_attention q"); + const mx::array local_kv = + kq_strides_aligned_gpu(inputs[1], s, "dsa_sparse_attention local_kv"); + const mx::array pooled = + kq_strides_aligned_gpu(inputs[2], s, "dsa_sparse_attention pooled"); const auto& topk = inputs[3]; const auto& sinks = inputs[4]; auto& o = outputs[0]; diff --git a/src/kquant_gather.cpp b/src/kquant_gather.cpp index 8bf2b40..219e1a5 100644 --- a/src/kquant_gather.cpp +++ b/src/kquant_gather.cpp @@ -121,6 +121,7 @@ void gather_qmm_nax( c = add_strides_and_shapes(ce, false, x, w, scales, c); add_gather_strides_and_shapes(ce, lhs_indices, rhs_indices, c); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -199,6 +200,7 @@ void gather_qmm( c = add_strides_and_shapes(ce, false, x, w, scales, c); add_gather_strides_and_shapes(ce, lhs_indices, rhs_indices, c); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -255,6 +257,7 @@ void gather_qmv( c = add_strides_and_shapes(ce, false, x, w, scales, c); add_gather_strides_and_shapes(ce, lhs_indices, rhs_indices, c); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -358,6 +361,7 @@ void gather_qmm_rhs_nax( ce.set_bytes(N, c++); ce.set_bytes(K, c++); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -438,6 +442,7 @@ void gather_qmm_rhs( ce.set_bytes(N, c++); ce.set_bytes(K, c++); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -825,7 +830,8 @@ void KQuantGatherQMM::eval_gpu( // inputs are row-contiguous (ensured by the op): // x (float), w (uint8), scales, lhs_indices (uint32), rhs_indices (uint32). - const auto& x = inputs[0]; + // x is bound from a 16-byte boundary (kq_x_aligned_gpu). + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "gather_qmm x"); const auto& w = inputs[1]; const auto& scales = inputs[2]; const auto& lhs_indices = inputs[3]; @@ -1050,8 +1056,9 @@ void KQuantGatherQMMSeg::eval_gpu( auto& out = outputs[0]; // inputs are row-contiguous (ensured by the op): x [R, K] (float), - // w [E, N, bpr] (uint8), scales, map [T, 3] (uint32), counts [1]. - const auto& x = inputs[0]; + // w [E, N, bpr] (uint8), scales, map [T, 3] (uint32), counts [1]. x is + // bound from a 16-byte boundary (kq_x_aligned_gpu). + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "gather_qmm_seg x"); const auto& w = inputs[1]; const auto& scales = inputs[2]; const auto& map = inputs[3]; @@ -1107,6 +1114,7 @@ void KQuantGatherQMMSeg::eval_gpu( ce.set_bytes(K, c++); MTL::Size group_dims(32, 2, 2); MTL::Size grid_dims((N + bn - 1) / bn, map.shape(0), 1); + kq_log_dispatch(nname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); return; } @@ -1135,6 +1143,7 @@ void KQuantGatherQMMSeg::eval_gpu( ce.set_bytes(N, c++); MTL::Size group_dims(32, 2, 2); MTL::Size grid_dims((N + bn - 1) / bn, map.shape(0), 1); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/src/kquant_hc_glue.cpp b/src/kquant_hc_glue.cpp index b4458ff..cdacc40 100644 --- a/src/kquant_hc_glue.cpp +++ b/src/kquant_hc_glue.cpp @@ -89,6 +89,12 @@ void check_fn(const mx::array& fn, int D, const char* op) { namespace { +// The kernels read the stream, residual and collapsed rows, fn, the mixes +// and the base biases as vectors of four elements. +mx::array hc_vec(const mx::array& a, const mx::Stream& s, const char* what) { + return kq_x_aligned_gpu(a, s, what, 4 * static_cast(a.itemsize())); +} + // Threads per threadgroup for one launch: the smaller of the count the // kernel prefers and what this pipeline supports, rounded down to a whole // simdgroup. Register pressure puts the pipeline limit below the preferred @@ -175,7 +181,10 @@ void KQuantHcFrontReduce::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& x = inputs[0]; + // The operands the kernels read as vectors are bound from boundaries of + // four of their elements (hc_vec), here and in the evals below. + const mx::array x = hc_vec(inputs[0], s, "hc_front_reduce x"); + const mx::array fn = hc_vec(inputs[1], s, "hc_front_reduce fn"); int D = x.shape(-1); int rows = int(x.size() / (HC * D)); @@ -184,7 +193,7 @@ void KQuantHcFrontReduce::eval_gpu( auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); ce.set_input_array(x, 0); - ce.set_input_array(inputs[1], 1); + ce.set_input_array(fn, 1); ce.set_output_array(outputs[0], 2); ce.set_output_array(outputs[1], 3); ce.set_bytes(D, 4); @@ -201,7 +210,9 @@ void KQuantHcFrontExpandReduce::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& resid = inputs[1]; + const mx::array x_sub = hc_vec(inputs[0], s, "hc_front_expand_reduce x_sub"); + const mx::array resid = hc_vec(inputs[1], s, "hc_front_expand_reduce resid"); + const mx::array fn = hc_vec(inputs[4], s, "hc_front_expand_reduce fn"); int D = resid.shape(-1); int rows = int(resid.size() / (HC * D)); @@ -210,11 +221,11 @@ void KQuantHcFrontExpandReduce::eval_gpu( auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); - ce.set_input_array(inputs[0], 0); + ce.set_input_array(x_sub, 0); ce.set_input_array(resid, 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(fn, 4); ce.set_output_array(outputs[0], 5); ce.set_output_array(outputs[1], 6); ce.set_output_array(outputs[2], 7); @@ -232,7 +243,9 @@ void KQuantHcSinkhornCollapse::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& x = inputs[0]; + const mx::array x = hc_vec(inputs[0], s, "hc_sinkhorn_collapse x"); + const mx::array mixes = hc_vec(inputs[1], s, "hc_sinkhorn_collapse mixes"); + const mx::array base = hc_vec(inputs[4], s, "hc_sinkhorn_collapse base"); int D = x.shape(-1); int rows = int(x.size() / (HC * D)); @@ -241,10 +254,10 @@ void KQuantHcSinkhornCollapse::eval_gpu( auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); ce.set_input_array(x, 0); - ce.set_input_array(inputs[1], 1); + ce.set_input_array(mixes, 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(base, 4); ce.set_input_array(inputs[5], 5); ce.set_output_array(outputs[0], 6); ce.set_output_array(outputs[1], 7); @@ -265,7 +278,12 @@ void KQuantHcFrontExpandCollapse::eval_gpu( for (auto& out : outputs) { out.set_data(mx::allocator::malloc(out.nbytes())); } - const auto& resid = inputs[1]; + const mx::array x_sub = + hc_vec(inputs[0], s, "hc_front_expand_collapse x_sub"); + const mx::array resid = + hc_vec(inputs[1], s, "hc_front_expand_collapse resid"); + const mx::array fn = hc_vec(inputs[4], s, "hc_front_expand_collapse fn"); + const mx::array base = hc_vec(inputs[6], s, "hc_front_expand_collapse base"); int D = resid.shape(-1); int rows = int(resid.size() / (HC * D)); @@ -293,13 +311,13 @@ void KQuantHcFrontExpandCollapse::eval_gpu( "kq_hc_front_expand_collapse_" + kq_type_string(resid.dtype()); auto kernel = kq_get_kernel(d, kname); ce.set_compute_pipeline_state(kernel); - ce.set_input_array(inputs[0], 0); + ce.set_input_array(x_sub, 0); ce.set_input_array(resid, 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(fn, 4); ce.set_input_array(inputs[5], 5); - ce.set_input_array(inputs[6], 6); + ce.set_input_array(base, 6); ce.set_input_array(inputs[7], 7); ce.set_output_array(outputs[0], 8); ce.set_output_array(mixes_raw, 9); @@ -324,7 +342,8 @@ void KQuantHcExpand::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& resid = inputs[1]; + const mx::array x = hc_vec(inputs[0], s, "hc_expand x"); + const mx::array resid = hc_vec(inputs[1], s, "hc_expand resid"); int D = resid.shape(-1); int rows = int(resid.size() / (HC * D)); @@ -332,7 +351,7 @@ void KQuantHcExpand::eval_gpu( auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); - ce.set_input_array(inputs[0], 0); + ce.set_input_array(x, 0); ce.set_input_array(resid, 1); ce.set_input_array(inputs[2], 2); ce.set_input_array(inputs[3], 3); diff --git a/src/kquant_internal.h b/src/kquant_internal.h index 2c5e85e..e0ee98d 100644 --- a/src/kquant_internal.h +++ b/src/kquant_internal.h @@ -61,6 +61,12 @@ int kq_dequant_route_mode(); // only. bool kq_qmm_log_on(); +// Byte offset of x's start modulo 16, for KQ_QMM_LOG. A buffer starts on a +// 16-byte boundary, so this is the alignment the kernels see. +inline int kq_x_align(const mx::array& x) { + return static_cast(x.offset() % 16); +} + // 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 || @@ -179,6 +185,32 @@ mx::array kq_dense_copy_gpu(const mx::array& a, const mx::Stream& s); // 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); +// An operand as kernels that read it as vectors need it: a itself when it +// is row-major and starts on an align-byte boundary, else the row-major +// copy, whose buffer starts on a 256-byte boundary. Metal defines a vector +// load only at an address aligned to the vector's size, and a row-major +// view can start at any element of its buffer. The ops bind activations +// and float operands they read as vectors through this, with align the +// widest load (16 bytes unless noted at the call). what names the op and +// operand for KQ_QMM_LOG, which prints "[kq align] xin=" for +// each copy, n being the byte offset of a's start modulo align. +mx::array kq_x_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align = 16); + +// 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, such as a KV +// cache: a itself when its last stride is 1 and its start and the stride +// of each dim longer than 1 fall on align-byte boundaries, else the +// row-major copy. +mx::array kq_strides_aligned_gpu( + const mx::array& a, + const mx::Stream& s, + const char* what, + int align = 16); + // 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. diff --git a/src/kquant_latent_fp4.cpp b/src/kquant_latent_fp4.cpp index d4c24ed..eb46114 100644 --- a/src/kquant_latent_fp4.cpp +++ b/src/kquant_latent_fp4.cpp @@ -29,7 +29,12 @@ void KQLatentFp4Pack::eval_gpu( auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& x = inputs[0]; + // The kernel reads each group of 16 as two 8-wide vectors. + const mx::array x = kq_x_aligned_gpu( + inputs[0], + s, + "latent_fp4_pack x", + 8 * static_cast(inputs[0].itemsize())); auto& codes = outputs[0]; auto& scales = outputs[1]; codes.set_data(mx::allocator::malloc(codes.nbytes())); diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index f6ab404..4007153 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -10,6 +10,7 @@ // 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 @@ -748,8 +749,8 @@ static KqSplitkPartials kq_splitk_partials(mx::Dtype dt) { 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. +// KQ_QMM_LOG set (any value): the quantized matmul and mat-vec kernels 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; @@ -813,24 +814,27 @@ static void kq_require_tile_kernel( 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. +// One KQ_QMM_LOG line: kernel name, shape and grid of a dispatch, and xa, +// the byte offset of x's start modulo 16. static void kq_log_qmm( const std::string& kname, int M, int N, int K, - const MTL::Size& grid) { + const MTL::Size& grid, + const array& x) { if (kq_qmm_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, K, static_cast(grid.width), static_cast(grid.height), - static_cast(grid.depth)); + static_cast(grid.depth), + kq_x_align(x)); } } @@ -1012,7 +1016,7 @@ void qmm_nax( if (tile) { kq_require_tile_kernel(d, "KQ_NAX_TILE", *tile, kname); } - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -1110,7 +1114,7 @@ void qmm( if (tile) { kq_require_tile_kernel(d, "KQ_QMM_TILE", *tile, kname); } - kq_log_qmm(kname, M, N, K, grid_dims); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -1331,7 +1335,7 @@ 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); + kq_log_qmm(kname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); @@ -1411,7 +1415,7 @@ 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); + kq_log_qmm(kname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; @@ -1432,6 +1436,29 @@ void verify_mma( ce.dispatch_threads(agrid, agroup); } +// The several-weight form of kq_verify_nax_kname. +static std::string kq_verify_nax_multi_kname( + const std::string& t, + const std::string& ts, + int gs, + int bits, + int kstep, + const std::string& suffix) { + std::string kname; + kname.reserve(72); + mx::concatenate( + kname, + (t == "q4_0" && kstep == 256 ? "kquant_q4_0_sb_" : kq_kname_prefix(t)) + + "verify_nax_multi_", + ts, + "_gs_", + gs, + "_b_", + bits, + suffix); + return kname; +} + // K per verify_nax iteration for this call. q4_0 runs its eight-block // kernel (q4_0_sb, 256 weights per step) where K is a multiple of 256, // and its two-block kernel (64) elsewhere. KQ_VERIFY_NAX_Q4_0_SB=0 keeps @@ -1472,6 +1499,32 @@ kq_verify_nax_fits(const array& w, int K, int kstep, const std::string& t) { // device memory per block and nothing staged. Same partial/fold shape as // verify_mma. Non-batched transpose shapes, bf16 or f16 activations, NAX // GPUs. +// KQ_VNAX_X16=0 keeps verify_nax on one 8-byte activation load per step. +// Read once per process unless KQ_TUNE_LIVE=1. +static bool kq_vnax_x16_enabled() { + static const bool cached = kq_env_flag("KQ_VNAX_X16", 1) != 0; + return kq_tune_live() ? kq_env_flag("KQ_VNAX_X16", 1) != 0 : cached; +} + +// verify_nax pipeline: the kernels that load one activation step at a +// time load two in one 16-byte load where x starts on a 16-byte boundary +// (function constant 40, kq_vnax_x16). vname, kname with the variant +// (_x16 or _x8) ahead of its partials suffix, names the pipeline and the +// KQ_QMM_LOG line. +static MTL::ComputePipelineState* kq_verify_nax_pipeline( + Device& d, + const std::string& kname, + const std::string& suffix, + const array& x, + std::string& vname) { + bool x16 = reinterpret_cast(x.data()) % 16 == 0 && + kq_vnax_x16_enabled(); + mx::metal::MTLFCList fc = {{&x16, MTL::DataType::DataTypeBool, 40}}; + vname = kname.substr(0, kname.size() - suffix.size()) + + (x16 ? "_x16" : "_x8") + suffix; + return kq_get_kernel(d, kname, vname, fc); +} + void verify_nax( const array& x, const array& w, @@ -1503,7 +1556,8 @@ void verify_nax( 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); + std::string vname; + auto kernel = kq_verify_nax_pipeline(d, kname, pp.suffix, x, vname); ce.set_compute_pipeline_state(kernel); int c = 0; @@ -1518,7 +1572,7 @@ 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); + kq_log_qmm(vname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); if (splits == 1) { return; @@ -1539,6 +1593,107 @@ void verify_nax( ce.dispatch_threads(agrid, agroup); } +static int kq_verify_nax_splits(int N, int K, int kstep, const std::string& t); + +// verify_nax over up to four weights that share x, in one dispatch +// (kq_verify_nax_multi): weight i takes ceil(N_i / 128) threadgroups of +// grid x after those of the weights before it, and reads its own buffer. +// The split count is the one verify_nax takes for the joint grid. Split +// partials of weight i sit in one temporary at splits * M * (rows of the +// weights before it), and one fold per weight writes its output. +void verify_nax_multi( + const array& x, + const std::vector& ws, + std::vector& outs, + int group_size, + int bits, + int M, + int K, + int kstep, + Device& d, + const Stream& s, + const std::string& kquant_type) { + const int n = static_cast(ws.size()); + std::array n_rows{0, 0, 0, 0}; + std::array tg_end{0, 0, 0, 0}; + std::array part_stride{0, 0, 0, 0}; + std::array row_off{0, 0, 0, 0}; + int tgs = 0; + int64_t rows = 0; + for (int i = 0; i < 4; ++i) { + const int N = i < n ? ws[i].shape(0) : 0; + n_rows[i] = N; + row_off[i] = rows; + rows += N; + tgs += (N + 127) / 128; + tg_end[i] = tgs; + part_stride[i] = M * N; + } + const int splits = kq_verify_nax_splits(tgs * 128, K, kstep, kquant_type); + const int k_partition = (K / kstep / splits) * kstep; + + auto& ce = mx::metal::get_command_encoder(s); + const KqSplitkPartials pp = splits > 1 ? kq_splitk_partials(x.dtype()) + : KqSplitkPartials{x.dtype(), ""}; + array partials({splits, M, static_cast(rows)}, pp.dtype, nullptr, {}); + if (splits > 1) { + partials.set_data(mx::allocator::malloc(partials.nbytes())); + ce.add_temporary(partials); + } + + const std::string type_string = kq_type_string(x.dtype()); + const std::string kname = kq_verify_nax_multi_kname( + kquant_type, type_string, group_size, bits, kstep, pp.suffix); + std::string vname; + auto kernel = kq_verify_nax_pipeline(d, kname, pp.suffix, x, vname); + ce.set_compute_pipeline_state(kernel); + + int c = 0; + for (int i = 0; i < 4; ++i) { + ce.set_input_array(ws[i < n ? i : 0], c++); + } + ce.set_input_array(x, c++); + for (int i = 0; i < 4; ++i) { + const int j = i < n ? i : 0; + if (splits > 1) { + ce.set_output_array( + partials, c++, splits * M * row_off[j] * partials.itemsize()); + } else { + ce.set_output_array(outs[j], c++); + } + } + ce.set_bytes(K, c++); + ce.set_bytes(M, c++); + ce.set_bytes(k_partition, c++); + ce.set_bytes(n_rows.data(), 4, c++); + ce.set_bytes(tg_end.data(), 4, c++); + ce.set_bytes(part_stride.data(), 4, c++); + MTL::Size group_dims(128, 1, 1); + MTL::Size grid_dims(tgs, 1, splits); + kq_log_qmm(vname, M, static_cast(rows), K, grid_dims, x); + ce.dispatch_threadgroups(grid_dims, group_dims); + if (splits == 1) { + return; + } + + const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); + auto accum = kq_get_kernel(d, aname); + for (int i = 0; i < n; ++i) { + ce.set_compute_pipeline_state(accum); + const int n_elems = M * n_rows[i]; + c = 0; + ce.set_input_array( + partials, c++, splits * M * row_off[i] * partials.itemsize()); + ce.set_output_array(outs[i], c++); + ce.set_bytes(n_elems, c++); + ce.set_bytes(splits, c++); + ce.set_bytes(part_stride[i], c++); + MTL::Size agrid(static_cast(n_elems), 1, 1); + MTL::Size agroup(256, 1, 1); + ce.dispatch_threads(agrid, agroup); + } +} + // Split-K qmm_t on the NAX BM=32 or BM=64 tile (KQ_QMM_SPLITK_NAX). Same // partial/fold shape as qmm_splitk, but slices run the tensor-core tile: // the steel splitk probe measured per-TG pipeline bound (~140-160 GB/s flat @@ -1593,7 +1748,7 @@ 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); + kq_log_qmm(kname, M, N, K, grid_dims, x); ce.dispatch_threadgroups(grid_dims, group_dims); const std::string aname = kq_splitk_accum_kname(type_string, pp.suffix); @@ -1645,6 +1800,7 @@ void qvm( bits, B > 1 ? "_batch_1" : "_batch_0"); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -1702,6 +1858,7 @@ static bool qmv_splitk( if (int(kernel->maxTotalThreadsPerThreadgroup()) < 32 * nsg) { return false; } + kq_log_qmm(kname, 1, N, K, MTL::Size(1, N / rps, 1), x); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); int c = 0; @@ -1778,6 +1935,7 @@ void qmv( bits, B > 1 ? "_batch_1" : "_batch_0"); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -1842,6 +2000,7 @@ void qmv_bias( bits, "_batch_0"); + kq_log_qmm(kname, 1, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -1921,6 +2080,7 @@ void verify_qmv( bits, "_batch_0"); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -2019,6 +2179,7 @@ void verify_mv_ext( : (use_nx ? (mv_ext_nx == 16 ? "_x16" : "_x32") : (use_hd ? "_hd" : "")))); + kq_log_qmm(kname, M, N, K, grid_dims, x); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); @@ -2339,7 +2500,7 @@ static int kq_verify_nax_splits(int N, int K, int kstep, const std::string& t) { // non-f32 shapes with M <= 32 only (the caller checks). bool kq_forced_route( const char* route, - const array& x, + const array& x_in, const array& w, const array& scales, array& out, @@ -2352,23 +2513,25 @@ bool kq_forced_route( const Stream& s, const std::string& t) { const std::string r(route); + // x is bound once the route serves the call, as in eval_gpu_base. + auto x16 = [&] { return kq_x_aligned_gpu(x_in, s, "quantized_matmul x"); }; const KQuantCodec* codec = codec_by_name(t); const int wpb = codec ? codec->weights_per_block : 256; if (r == "qmv") { - dispatch_qmv(x, w, scales, out, group_size, bits, M, N, K, d, s, t); + dispatch_qmv(x16(), w, scales, out, group_size, bits, M, N, K, d, s, t); return true; } if (r == "verify_qmv") { if (M >= 2 && M <= verify_qmv_max_rows() && N % kquant_qmv_bn(t) == 0 && K % qmv_fast_k_align() == 0 && codec_has_verify_qmv(t)) { - verify_qmv(x, w, scales, out, group_size, bits, M, N, K, d, s, t); + verify_qmv(x16(), w, scales, out, group_size, bits, M, N, K, d, s, t); return true; } return false; } if (r == "mv_ext") { if (kq_codec_has_mv_ext(t) && M >= 2 && M <= 12 && K % wpb == 0) { - verify_mv_ext(x, w, scales, out, group_size, bits, M, N, K, d, s, t); + verify_mv_ext(x16(), w, scales, out, group_size, bits, M, N, K, d, s, t); return true; } return false; @@ -2377,7 +2540,7 @@ bool kq_forced_route( const int rows = codec_verify_mma_rows(t); if (rows > 0 && M <= 8) { const int sp = kq_verify_mma_splits(N, K, group_size, rows); - verify_mma(x, w, scales, out, group_size, bits, M, N, K, sp, d, s, t); + verify_mma(x16(), w, scales, out, group_size, bits, M, N, K, sp, d, s, t); return true; } return false; @@ -2388,7 +2551,20 @@ bool kq_forced_route( kq_verify_nax_fits(w, K, kstep, t)) { const int sp = kq_verify_nax_splits(N, K, kstep, t); verify_nax( - x, w, scales, out, group_size, bits, M, N, K, kstep, sp, d, s, t); + kq_x_aligned_gpu(x_in, s, "quantized_matmul x", 8), + w, + scales, + out, + group_size, + bits, + M, + N, + K, + kstep, + sp, + d, + s, + t); return true; } return false; @@ -2398,7 +2574,7 @@ bool kq_forced_route( const int target = (e != nullptr && std::atoi(e) > 1) ? std::atoi(e) : 16; const int sp = kq_split_count(target, K / group_size); if (kq_splitk_codec(t) && sp > 1) { - qmm_splitk(x, w, scales, out, group_size, bits, M, N, K, sp, d, s, t); + qmm_splitk(x16(), w, scales, out, group_size, bits, M, N, K, sp, d, s, t); return true; } return false; @@ -2407,7 +2583,7 @@ bool kq_forced_route( kq_is_nax_available() && codec_has_nax(t) && (K % 64 == 0); if (r == "nax") { if (nax_ok) { - qmm(x, w, scales, out, true, group_size, bits, M, N, K, d, s, t); + qmm(x16(), w, scales, out, true, group_size, bits, M, N, K, d, s, t); return true; } return false; @@ -2422,7 +2598,20 @@ bool kq_forced_route( const int sp = kq_nax_split(target, K / q, per); if (nax_ok && sp > 1) { qmm_nax_splitk( - x, w, scales, out, group_size, bits, M, N, K, sp, d, s, t, per * q); + x16(), + w, + scales, + out, + group_size, + bits, + M, + N, + K, + sp, + d, + s, + t, + per * q); return true; } return false; @@ -2453,6 +2642,18 @@ bool kq_qmm_log_on() { return kq_qmm_log(); } +void kq_log_dispatch( + const std::string& kname, + const array& x, + const array& out, + const MTL::Size& grid) { + if (kq_qmm_log()) { + const int N = out.ndim() > 0 ? out.shape(-1) : 1; + const int M = N > 0 ? static_cast(out.size() / N) : 0; + kq_log_qmm(kname, M, N, x.ndim() > 0 ? x.shape(-1) : 1, grid, x); + } +} + #endif // _METAL_ std::vector KQuantMatmul::output_shapes( @@ -2596,6 +2797,160 @@ std::vector KQuantMatmul::vjp( #ifdef _METAL_ +// Probe lever: KQ_FORCE_QMM_MIN_M= routes transpose shapes with M >= m +// straight to qmm (-> NAX), bypassing the mv_ext/verify_qmv claims, for +// qmm-vs-mv_ext crossover measurement below M=13. Read once per process. +static int kq_force_qmm_min_m() { + static const int m = []() { + const char* e = std::getenv("KQ_FORCE_QMM_MIN_M"); + return e != nullptr ? std::atoi(e) : 0; + }(); + return m; +} + +// KQ_QMM_ROUTE= is set (see eval_gpu_base). +static bool kq_qmm_route_forced() { + const char* e = std::getenv("KQ_QMM_ROUTE"); + return e != nullptr && e[0] != '\0'; +} + +enum class KqSmallMRoute { other, qmv, verify_nax }; + +// The NAX small-M route quantized_matmul takes for a transposed, +// non-batched call on a float16 or bfloat16 x of M rows, once the forced +// KQ_FORCE_QMM_MIN_M and KQ_QMM_ROUTE routes are out of the way: per-row +// qmv, verify_nax (still subject to kq_verify_nax_fits), or a later route. +// quantized_matmul_multi joins weights into one verify_nax dispatch only +// where this returns verify_nax for the widest weight. +// +// Per-row qmv on NAX GPUs, M 2 through kq_nax_qmv_max_m for the call's +// N. KQ_NAX_QMV= sets the limit to m at every N, and 0 disables the +// route. The default yields to a forced KQ_VERIFY_NAX or KQ_VERIFY_MMA +// on a codec they serve, a forced KQ_QMM_SPLITK_NAX or KQ_QMM_SPLITK, +// and a set KQ_VERIFY_EXT, so their arms time the route they name. +// KQ_DISABLE_NAX turns it off, since the limits are measured against +// the NAX routes. KQ_NAX_QMV is read live so an A/B can flip arms on +// one generator. +// +// Register-fed NAX verify, from kq_verify_nax_min_m through M8; see +// verify_nax. KQ_VERIFY_NAX= forces the entry at M >= m, and 0 +// disables it. Read live so an A/B can flip arms on one generator. +// KQ_DISABLE_NAX turns it off with the other NAX routes. Decided ahead +// of verify_mma and NAX split-K, but the default entry yields to a +// forced KQ_VERIFY_MMA on a codec verify_mma serves, or a forced +// KQ_QMM_SPLITK_NAX, so their arms time the route they name. +static KqSmallMRoute +kq_nax_small_m_route(const std::string& t, int M, int N, int K, bool f16) { + if (M < 1 || M > 32 || !kq_is_nax_available() || !codec_has_nax(t)) { + return KqSmallMRoute::other; + } + 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); + if (M >= 2) { + int qmv_max_m = kq_env_int("KQ_NAX_QMV", -1); + if (qmv_max_m == -1) { + const int ext = kq_verify_ext_env(); + const bool forced = + (vnax_env >= 1 && M >= vnax_env && codec_verify_nax_kstep(t) > 0) || + (vmma_force >= 2 && M >= vmma_force && + codec_verify_mma_rows(t) > 0) || + sk_force >= 1 || kq_qmm_splitk_env() > 1 || ext == 0 || ext == 1; + qmv_max_m = forced ? 0 : kq_nax_qmv_max_m(t, N); + } + if (M <= qmv_max_m) { + return KqSmallMRoute::qmv; + } + } + if (M <= 8) { + const bool other_forced = + (vmma_force >= 2 && M >= vmma_force && codec_verify_mma_rows(t) > 0) || + sk_force >= 1; + int vnax_min_m = vnax_env; + if (vnax_env == -1) { + vnax_min_m = other_forced ? 0 : kq_verify_nax_min_m(t, N, K, f16); + } + if (vnax_min_m > 0 && M >= vnax_min_m) { + return KqSmallMRoute::verify_nax; + } + } + return KqSmallMRoute::other; +} + +// One verify_nax dispatch per four weights where quantized_matmul would +// take the verify_nax route on the widest weight at this M, and +// quantized_matmul per weight otherwise. The narrower weights ride along on +// that route. Every route lever of quantized_matmul applies, through +// kq_nax_small_m_route, and a forced KQ_QMM_ROUTE or KQ_FORCE_QMM_MIN_M +// runs quantized_matmul per weight. +void KQuantMatmulMulti::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + const auto& x = inputs[0]; + const auto& scales = inputs[1]; + const int n = static_cast(outputs.size()); + std::vector ws; + ws.reserve(n); + for (int i = 0; i < n; ++i) { + ws.push_back(kq_matrix_dense_gpu(inputs[2 + i], s)); + kq_check_weight_base(ws.back(), kquant_type_, "quantized_matmul_multi"); + } + const int K = x.shape(-1); + const int M = K > 0 ? static_cast(x.size() / K) : 0; + int n_max = 0; + for (const auto& w : ws) { + n_max = std::max(n_max, w.shape(0)); + } + const int kstep = kq_verify_nax_call_kstep(K, kquant_type_); + const int force_min_m = kq_force_qmm_min_m(); + bool multi = n_max > 0 && kstep > 0 && x.flags().row_contiguous && + x.dtype() != mx::float32 && !kq_qmm_route_forced() && + !(force_min_m > 0 && M >= force_min_m) && + kq_nax_small_m_route( + kquant_type_, M, n_max, K, x.dtype() == mx::float16) == + KqSmallMRoute::verify_nax; + if (multi) { + for (const auto& w : ws) { + multi = multi && kq_verify_nax_fits(w, K, kstep, kquant_type_); + } + } + if (!multi) { + // x is bound once, so that each per-weight call reads it in place. + const array x16 = kq_x_aligned_gpu(x, s, "quantized_matmul_multi x"); + KQuantMatmul one(s, kquant_type_, group_size_, bits_, true); + for (int i = 0; i < n; ++i) { + std::vector o{outputs[i]}; + one.eval_gpu({x16, ws[i], scales}, o); + } + return; + } + // verify_nax reads x where it lands on an 8-byte boundary and from the + // aligned copy otherwise, as in KQuantMatmul::eval_gpu_base. + const array xv = kq_x_aligned_gpu(x, s, "quantized_matmul_multi x", 8); + std::vector gw; + std::vector go; + for (int i = 0; i < n; ++i) { + outputs[i].set_data(mx::allocator::malloc(outputs[i].nbytes())); + if (outputs[i].size() == 0) { + continue; + } + gw.push_back(ws[i]); + go.push_back(outputs[i]); + if (gw.size() == 4 || i == n - 1) { + verify_nax_multi( + xv, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); + gw.clear(); + go.clear(); + } + } + if (!gw.empty()) { + verify_nax_multi( + xv, gw, go, group_size_, bits_, M, K, kstep, d, s, kquant_type_); + } +} + void KQuantMatmul::eval_gpu( const std::vector& inputs, std::vector& outputs) { @@ -2621,7 +2976,14 @@ void KQuantMatmul::eval_gpu_base( auto& out = outputs[0]; // inputs are row-contiguous (ensured by the op): x, w (uint8), scales. - const auto& x = inputs[0]; + // x is bound once the route is chosen, from a 16-byte boundary + // (kq_x_aligned_gpu), except by verify_nax: it reads x where it lands on + // an 8-byte boundary, picking its x16 or x8 load pipeline by the start of + // x (kq_verify_nax_pipeline), and from the aligned copy otherwise. + const auto& x_in = inputs[0]; + auto bind_x = [&](int align) { + return kq_x_aligned_gpu(x_in, s, "quantized_matmul x", align); + }; const auto& w = inputs[1]; const auto& scales = inputs[2]; kq_check_weight_base(w, kquant_type_, "quantized_matmul"); @@ -2630,16 +2992,16 @@ void KQuantMatmul::eval_gpu_base( return; } - bool non_batched = w.ndim() == 2 && x.flags().row_contiguous; - int K = x.shape(-1); - int M = non_batched ? static_cast(x.size()) / K : x.shape(-2); + bool non_batched = w.ndim() == 2 && x_in.flags().row_contiguous; + int K = x_in.shape(-1); + int M = non_batched ? static_cast(x_in.size()) / K : x_in.shape(-2); int N = out.shape(-1); int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4; // KQuant special cases. if (!transpose_ && M < vector_limit) { - qmm(x, + qmm(bind_x(16), w, scales, out, @@ -2661,15 +3023,10 @@ void KQuantMatmul::eval_gpu_base( // through to dispatch_qmv below rather than throw; a qmv_quad kernel would be // a perf-only path for an essentially-dead case. - // Probe lever: KQ_FORCE_QMM_MIN_M= routes transpose shapes with M >= m - // straight to qmm (-> NAX), bypassing the mv_ext/verify_qmv claims, for - // qmm-vs-mv_ext crossover measurement below M=13. - static const int force_qmm_min_m = []() { - const char* e = std::getenv("KQ_FORCE_QMM_MIN_M"); - return e != nullptr ? std::atoi(e) : 0; - }(); + // Probe lever KQ_FORCE_QMM_MIN_M (kq_force_qmm_min_m). + const int force_qmm_min_m = kq_force_qmm_min_m(); if (force_qmm_min_m > 0 && transpose_ && M >= force_qmm_min_m) { - qmm(x, + qmm(bind_x(16), w, scales, out, @@ -2695,10 +3052,10 @@ void KQuantMatmul::eval_gpu_base( if (transpose_) { const char* route_e = std::getenv("KQ_QMM_ROUTE"); if (route_e != nullptr && route_e[0] != '\0') { - if (non_batched && M <= 32 && x.dtype() != mx::float32 && + if (non_batched && M <= 32 && x_in.dtype() != mx::float32 && kq_forced_route( route_e, - x, + x_in, w, scales, out, @@ -2722,64 +3079,32 @@ void KQuantMatmul::eval_gpu_base( } } - // Per-row qmv on NAX GPUs, M 2 through kq_nax_qmv_max_m for the call's - // N. KQ_NAX_QMV= sets the limit to m at every N, and 0 disables the - // route. The default yields to a forced KQ_VERIFY_NAX or KQ_VERIFY_MMA - // on a codec they serve, a forced KQ_QMM_SPLITK_NAX or KQ_QMM_SPLITK, - // and a set KQ_VERIFY_EXT, so their arms time the route they name. - // KQ_DISABLE_NAX turns it off, since the limits are measured against - // the NAX routes. KQ_NAX_QMV is read live so an A/B can flip arms on - // one generator. - if (transpose_ && non_batched && M >= 2 && M <= 32 && - x.dtype() != mx::float32 && kq_is_nax_available() && - codec_has_nax(kquant_type_)) { - int qmv_max_m = kq_env_int("KQ_NAX_QMV", -1); - if (qmv_max_m == -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) || - (vmma_force >= 2 && M >= vmma_force && - codec_verify_mma_rows(kquant_type_) > 0) || - sk_force >= 1 || kq_qmm_splitk_env() > 1 || ext == 0 || ext == 1; - qmv_max_m = forced ? 0 : kq_nax_qmv_max_m(kquant_type_, N); - } - if (M <= qmv_max_m) { + // Per-row qmv and verify_nax on NAX GPUs: kq_nax_small_m_route. + if (transpose_ && non_batched && x_in.dtype() != mx::float32) { + const KqSmallMRoute route = kq_nax_small_m_route( + kquant_type_, M, N, K, x_in.dtype() == mx::float16); + if (route == KqSmallMRoute::qmv) { dispatch_qmv( - x, w, scales, out, group_size_, bits_, M, N, K, d, s, kquant_type_); + bind_x(16), + w, + scales, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + kquant_type_); return; } - } - - // Register-fed NAX verify, from kq_verify_nax_min_m through M8; see - // verify_nax. KQ_VERIFY_NAX= forces the entry at M >= m, and 0 - // disables it. Read live so an A/B can flip arms on one generator. - // KQ_DISABLE_NAX turns it off with the other NAX routes. Decided ahead - // of verify_mma and NAX split-K, but the default entry yields to a - // forced KQ_VERIFY_MMA on a codec verify_mma serves, or a forced - // KQ_QMM_SPLITK_NAX, so their arms time the route they name. - 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 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; - int vnax_min_m = vnax_env; - if (vnax_env == -1) { - vnax_min_m = other_forced - ? 0 - : kq_verify_nax_min_m(kquant_type_, N, K, x.dtype() == mx::float16); - } - if (vnax_min_m > 0 && M >= vnax_min_m && + if (route == KqSmallMRoute::verify_nax && kq_verify_nax_fits(w, K, kstep, kquant_type_)) { const int sp = kq_verify_nax_splits(N, K, kstep, kquant_type_); verify_nax( - x, + bind_x(8), w, scales, out, @@ -2797,6 +3122,9 @@ void KQuantMatmul::eval_gpu_base( } } + // Every other route reads x from a 16-byte boundary. + const mx::array x = bind_x(16); + // Register-resident MMA verify, from the per-device entry // (kq_verify_mma_min_m, kq_verify_mma_min_m_nax) through M8. // KQ_VERIFY_MMA= forces the entry at M >= m, 0 or 1 disables; read @@ -3126,6 +3454,13 @@ void KQuantMatmul::eval_gpu( "[mlx_kquant] quantized_matmul has no GPU implementation."); } +void KQuantMatmulMulti::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant] quantized_matmul_multi has no GPU implementation."); +} + #endif // _METAL_ std::vector KQuantQmvBias::output_shapes( @@ -3217,7 +3552,8 @@ void KQuantQmvBias::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; - const auto& x = inputs[0]; + const mx::array x = + kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_qmv_bias x"); const auto& w = inputs[1]; const auto& scales = inputs[2]; const auto& bias = inputs[3]; @@ -3328,6 +3664,7 @@ std::vector split_kernel_names() { for (int k : ksteps) { each_suffix(dt, [&](const std::string& sfx) { out.push_back(kq_verify_nax_kname(t, ts, gs, bits, k, sfx)); + out.push_back(kq_verify_nax_multi_kname(t, ts, gs, bits, k, sfx)); }); } } diff --git a/src/kquant_matmul_fused.cpp b/src/kquant_matmul_fused.cpp index 565aac6..6b56583 100644 --- a/src/kquant_matmul_fused.cpp +++ b/src/kquant_matmul_fused.cpp @@ -216,7 +216,7 @@ void KQuantMatmulFused::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& x = inputs[0]; + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_glu x"); // 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. @@ -279,14 +279,15 @@ void KQuantMatmulFused::eval_gpu( if (glu_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, K, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname, hash_name, func_consts); @@ -461,4 +462,182 @@ mx::array quantized_matmul_glu( {std::move(x_c), std::move(w_gate), std::move(w_up)}); } +// ----------------------------- several weights, one x ----------------- + +namespace { + +const char* kMultiOp = "[mlx_kquant.quantized_matmul_multi]"; + +// KQ_MATMUL_MULTI=0 builds quantized_matmul per weight on every call. Read +// once per process, or on every call under KQ_TUNE_LIVE=1, as KQ_GLU. +bool multi_on() { + static const bool cached = kq_env_flag("KQ_MATMUL_MULTI", 1) != 0; + return glu_tune_live() ? kq_env_flag("KQ_MATMUL_MULTI", 1) != 0 : cached; +} + +std::vector multi_composed( + const mx::array& x, + const std::vector& weights, + const std::string& kquant_type, + mx::Stream s) { + auto scales = mx::zeros({1}, mx::uint8, s); + std::vector out; + out.reserve(weights.size()); + for (const auto& w : weights) { + out.push_back(quantized_matmul( + x, + w, + scales, + kquant_type, + true, + std::nullopt, + std::nullopt, + std::nullopt, + s)); + } + return out; +} + +} // namespace + +void KQuantMatmulMulti::eval_cpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + std::string(kMultiOp) + " CPU streams build the per-weight ops."); +} + +std::vector KQuantMatmulMulti::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_multi vjp: no gradient wrt the " + "quantized weights (the kquant base is frozen)."); + } + } + std::vector ws(primals.begin() + 2, primals.end()); + auto fun = [&](const std::vector& in) { + return multi_composed(in[0], ws, kquant_type_, stream()); + }; + auto vjps = + mx::vjp(fun, std::vector{primals[0]}, cotangents).second; + return std::vector(argnums.size(), vjps[0]); +} + +std::vector KQuantMatmulMulti::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_multi jvp: no derivative wrt the " + "quantized weights (the kquant base is frozen)."); + } + } + std::vector ws(primals.begin() + 2, primals.end()); + auto fun = [&](const std::vector& in) { + return multi_composed(in[0], ws, kquant_type_, stream()); + }; + return mx::jvp( + fun, + std::vector{primals[0]}, + std::vector{tangents[0]}) + .second; +} + +bool KQuantMatmulMulti::is_equivalent(const mx::Primitive& other) const { + const auto& o = static_cast(other); + return kquant_type_ == o.kquant_type_; +} + +std::vector KQuantMatmulMulti::output_shapes( + const std::vector& inputs) { + std::vector out; + for (size_t i = 2; i < inputs.size(); ++i) { + auto shape = inputs[0].shape(); + shape.back() = inputs[i].shape(0); + out.push_back(std::move(shape)); + } + return out; +} + +std::vector quantized_matmul_multi( + mx::array x, + std::vector weights, + const std::string& kquant_type, + mx::StreamOrDevice s_) { + const std::string op = kMultiOp; + const KQuantCodec* codec = codec_by_name(kquant_type); + if (codec == nullptr) { + throw std::invalid_argument( + op + " Unknown kquant_type: '" + kquant_type + "'."); + } + if (weights.empty()) { + throw std::invalid_argument(op + " weights must hold at least one weight."); + } + for (const auto& w : weights) { + if (w.dtype() != mx::uint8 || w.ndim() != 2) { + throw std::invalid_argument( + op + " weights must be 2-D uint8 (raw GGUF wire bytes)."); + } + if (w.shape(-1) != weights[0].shape(-1)) { + throw std::invalid_argument( + op + " weights must have one row length in bytes."); + } + } + if (weights[0].shape(-1) % codec->bytes_per_block != 0) { + throw std::invalid_argument( + op + " weight last dim (" + std::to_string(weights[0].shape(-1)) + + " bytes) is not a whole number of " + + std::to_string(codec->bytes_per_block) + "-byte " + codec->name + + " blocks."); + } + const int K = (weights[0].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_); + const mx::Dtype out_type = + x.dtype() == mx::float32 ? mx::bfloat16 : x.dtype(); + const int64_t rows = K > 0 ? static_cast(x.size()) / K : 0; + const bool fused = s.device == mx::Device::gpu && nax_available() && + codec->has_matmul_kernel && multi_on() && weights.size() >= 2 && + rows >= 2 && rows <= 8; + if (!fused) { + return multi_composed(x, weights, kquant_type, s); + } + std::vector inputs; + inputs.reserve(weights.size() + 2); + inputs.push_back(mx::contiguous(mx::astype(x, out_type, s), false, s)); + inputs.push_back(mx::zeros({1}, mx::uint8, s)); + std::vector shapes; + std::vector dtypes; + for (auto& w : weights) { + kq_check_weight_base_at_build(w, kquant_type, s, "quantized_matmul_multi"); + auto shape = x.shape(); + shape.back() = w.shape(0); + shapes.push_back(std::move(shape)); + dtypes.push_back(out_type); + inputs.push_back(kq_matrix_dense(w, s)); + } + return mx::array::make_arrays( + std::move(shapes), + dtypes, + std::make_shared( + s, kquant_type, codec->weights_per_block, codec->bits), + std::move(inputs)); +} + } // namespace mlx_kquant diff --git a/src/kquant_matmul_lora.cpp b/src/kquant_matmul_lora.cpp index eb74db4..9982537 100644 --- a/src/kquant_matmul_lora.cpp +++ b/src/kquant_matmul_lora.cpp @@ -271,7 +271,7 @@ void KQuantMatmulLora::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& x = inputs[0]; + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_lora x"); // 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. @@ -312,7 +312,7 @@ void KQuantMatmulLora::eval_gpu( if (lora_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d R=%d bands=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d R=%d bands=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, @@ -321,7 +321,8 @@ void KQuantMatmulLora::eval_gpu( n_bands, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname); diff --git a/src/kquant_matmul_resid.cpp b/src/kquant_matmul_resid.cpp index 0d91fae..d4240f2 100644 --- a/src/kquant_matmul_resid.cpp +++ b/src/kquant_matmul_resid.cpp @@ -118,7 +118,8 @@ void KQuantMatmulResid::eval_gpu( 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 x = + kq_x_aligned_gpu(inputs[0], s, "quantized_matmul_residual x"); // 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. @@ -160,7 +161,7 @@ void KQuantMatmulResid::eval_gpu( if (resid_log()) { std::fprintf( stderr, - "[kq qmm] %s M=%d N=%d K=%d B=%d grid=%lu,%lu,%lu\n", + "[kq qmm] %s M=%d N=%d K=%d B=%d grid=%lu,%lu,%lu xa=%d\n", kname.c_str(), M, N, @@ -168,7 +169,8 @@ void KQuantMatmulResid::eval_gpu( B, static_cast(grid_dims.width), static_cast(grid_dims.height), - static_cast(grid_dims.depth)); + static_cast(grid_dims.depth), + kq_x_align(x)); } auto kernel = kq_get_kernel(d, kname); diff --git a/src/kquant_metal_internal.h b/src/kquant_metal_internal.h index b238b3f..26c031c 100644 --- a/src/kquant_metal_internal.h +++ b/src/kquant_metal_internal.h @@ -388,6 +388,14 @@ inline int add_gather_strides_and_shapes( return offset; } +// One KQ_QMM_LOG line for a gather or MoE dispatch: the kernel name, M and +// N of out's rows, K of x's last dim, the grid and xa. +void kq_log_dispatch( + const std::string& kname, + const array& x, + const array& out, + const MTL::Size& grid); + // Fetch a kq kernel from the bundled metallib. // // The library handle is resolved exactly once and cached in a function-local diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index e1e34e5..bc5e3dd 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -265,6 +265,7 @@ inline void kq_gather_qmv_pair_gpu( const int fsg = kq_moe_fit_sg(kernel, sg); MTL::Size group_dims(32, fsg, 1); MTL::Size grid_dims(N / (fsg * 32 / nx), 1, (T + 1) / 2); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } #endif // _METAL_ @@ -392,8 +393,9 @@ void KQuantMoEGLU::eval_gpu( 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); + // x is read as dense rows from a 16-byte boundary: any other view is + // copied first (kq_x_aligned_gpu). + const mx::array x = kq_x_aligned_gpu(inputs[6], s, "moe_glu_gather x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -423,6 +425,7 @@ void KQuantMoEGLU::eval_gpu( ce.set_bytes(limit, 12); MTL::Size group_dims(32, 2, 1); MTL::Size grid_dims(N / 8, R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -438,7 +441,7 @@ void KQuantGatherQMVBias::eval_gpu( const auto& b = inputs[2]; const auto& indices = inputs[4]; moe_check_stack("gather_qmv_bias", w, sc); - const mx::array x = kq_row_major_gpu(inputs[3], s); + const mx::array x = kq_x_aligned_gpu(inputs[3], s, "gather_qmv_bias x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -463,6 +466,7 @@ void KQuantGatherQMVBias::eval_gpu( ce.set_bytes(N, 7); MTL::Size group_dims(32, 2, 1); MTL::Size grid_dims(N / (fine ? 2 : 8), R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -479,7 +483,7 @@ void KQuantGatherQMVMixBias::eval_gpu( 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); + const mx::array x = kq_x_aligned_gpu(inputs[3], s, "gather_qmv_mix_bias x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -506,6 +510,7 @@ void KQuantGatherQMVMixBias::eval_gpu( ce.set_bytes(S, 9); MTL::Size group_dims(32, 2, 1); MTL::Size grid_dims(N / (fine ? 2 : 8), 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -525,7 +530,8 @@ void KQuantMoEGLUKQ::eval_gpu( 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); + const mx::array x = + kq_x_aligned_gpu(inputs[biased ? 4 : 2], s, "moe_glu_gather_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -580,6 +586,7 @@ void KQuantMoEGLUKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sg * 32 / nx), R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -597,7 +604,8 @@ void KQuantGatherQMVKQ::eval_gpu( 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); + const mx::array x = + kq_x_aligned_gpu(inputs[biased ? 2 : 1], s, "gather_qmv_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -645,6 +653,7 @@ void KQuantGatherQMVKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (fine ? 2 : (sg * 32 / nx)), R, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); if (lora_flags_ & KQ_LORA_PRESENT) { kq_lora_epilogue_rows_gpu( @@ -679,7 +688,8 @@ void KQuantMoEGLUShexpKQ::eval_gpu( 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); + const mx::array x = + kq_x_aligned_gpu(inputs[4], s, "moe_glu_gather_shexp_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = indices.shape(0); @@ -714,6 +724,7 @@ void KQuantMoEGLUShexpKQ::eval_gpu( const int sg = kq_moe_fit_sg(kernel, kq_moe_pick_sg(kquant_type_, N, nx)); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sg * 32 / nx), R + 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); return; } @@ -762,6 +773,7 @@ void KQuantMoEGLUShexpKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sg * 32 / nx), R + 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -780,7 +792,7 @@ void KQuantGatherQMVMixKQ::eval_gpu( 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); + const mx::array x = kq_x_aligned_gpu(inputs[2], s, "gather_qmv_mix_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = x.shape(0); @@ -826,6 +838,7 @@ void KQuantGatherQMVMixKQ::eval_gpu( ce.set_bytes(SC, 9); MTL::Size group_dims(32, 2 * S, 1); MTL::Size grid_dims(N / 8, 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); return; } @@ -898,6 +911,7 @@ void KQuantGatherQMVMixKQ::eval_gpu( : (tuned ? 2 : kq_moe_pick_sg(kquant_type_, N, nx))); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (fine ? 2 : (sg * 32 / nx)), 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } @@ -913,7 +927,7 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( 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); + const mx::array x = kq_x_aligned_gpu(inputs[1], s, "gather_qmv_mix_ns_kq x"); out.set_data(mx::allocator::malloc(out.nbytes())); int T = x.shape(0); @@ -992,6 +1006,7 @@ void KQuantGatherQMVMixNSKQ::eval_gpu( sp ? 2 * S : kq_moe_fit_sg(kernel, kq_moe_pick_sg(kquant_type_, N, nx)); MTL::Size group_dims(32, sg, 1); MTL::Size grid_dims(N / (sp ? 8 : (sg * 32 / nx)), 1, T); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } if (lora_z) { diff --git a/src/kquant_moe_mix.cpp b/src/kquant_moe_mix.cpp index 5bdabd6..85d4761 100644 --- a/src/kquant_moe_mix.cpp +++ b/src/kquant_moe_mix.cpp @@ -92,7 +92,8 @@ void KQuantGatherMix::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& y = inputs[0]; + // y is the activation, bound from a 16-byte boundary. + const mx::array y = kq_x_aligned_gpu(inputs[0], s, "gather_mix y"); const int N = y.shape(1); const int Tn = inputs[2].shape(0); const int K = inputs[2].shape(1); @@ -100,7 +101,7 @@ void KQuantGatherMix::eval_gpu( auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); - ce.set_input_array(inputs[0], 0); + ce.set_input_array(y, 0); ce.set_input_array(inputs[1], 1); ce.set_input_array(inputs[2], 2); ce.set_output_array(out, 3); @@ -113,6 +114,7 @@ void KQuantGatherMix::eval_gpu( } MTL::Size group_dims(threads, 1, 1); MTL::Size grid_dims(Tn, 1, 1); + kq_log_dispatch(kname, y, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/src/kquant_norm_fused.cpp b/src/kquant_norm_fused.cpp index 7ff2b29..81fe4a8 100644 --- a/src/kquant_norm_fused.cpp +++ b/src/kquant_norm_fused.cpp @@ -756,7 +756,12 @@ void KQuantRMSNormGate::eval_gpu( auto& d = mx::metal::device(s.device); auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& x = inputs[0]; + // x, w and gate are read as vectors of up to four elements. + const int align = 4 * static_cast(inputs[0].itemsize()); + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "rmsnorm_gate x", align); + const mx::array w = kq_x_aligned_gpu(inputs[1], s, "rmsnorm_gate w", align); + const mx::array gate = + kq_x_aligned_gpu(inputs[2], s, "rmsnorm_gate gate", align); const int D = x.shape(-1); const int nrows = int(x.size() / D); std::string kname = "kq_rmsnorm_gate_" + kq_type_string(x.dtype()) + "_" + @@ -764,9 +769,9 @@ void KQuantRMSNormGate::eval_gpu( auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); ce.set_compute_pipeline_state(kernel); - ce.set_input_array(inputs[0], 0); - ce.set_input_array(inputs[1], 1); - ce.set_input_array(inputs[2], 2); + ce.set_input_array(x, 0); + ce.set_input_array(w, 1); + ce.set_input_array(gate, 2); ce.set_output_array(out, 3); ce.set_bytes(eps_, 4); ce.set_bytes(nrows, 5); diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp index 1e07d03..671bc4e 100644 --- a/src/kquant_sdpa.cpp +++ b/src/kquant_sdpa.cpp @@ -56,6 +56,18 @@ void kq_sdpa_check_layout( } } +// The kernels read q, k, v and the KVarN stages as vectors of four +// elements (four uint32 words of a q8 cache), so each is bound from a +// boundary of four of its elements, k, v and the stages with their strides +// on such boundaries too. +array sdpa_q(const array& a, const Stream& s, const char* what) { + return kq_x_aligned_gpu(a, s, what, 4 * static_cast(a.itemsize())); +} + +array sdpa_rows(const array& a, const Stream& s, const char* what) { + return kq_strides_aligned_gpu(a, s, what, 4 * static_cast(a.itemsize())); +} + // Mirror of the Metal-side KQKvarnMeta (kq_sdpa.h, buffer 28); the layouts // must match field for field. struct KQKvarnMeta { @@ -191,6 +203,23 @@ int kq_fa_verify_splits( return std::min({splits, cap, 128}); } +// Split count for sdpa_fa_verify on the tensor-op pass 1 when the caller +// passes 0. That kernel runs a split in a fraction of the time of the +// simdgroup-matrix kernel, so the float32 partials and their merge weigh +// more: the count aims at about 64 threadgroups over the kv heads, at least +// 4 and at most 64, and drops while a split would hold under 64 keys. Tuned +// on M5 Max at 1 to 32 kv heads. +int kq_fa_verify_nax_splits(int n_heads, int kL) { + int splits = 64; + for (int h = 1; h < n_heads && splits > 4; h *= 2) { + splits /= 2; + } + while (splits > 4 && kL < 64 * splits) { + splits /= 2; + } + return splits; +} + // Split count for sdpa_fa_indexed when the caller passes 0: the decode // bucket over the M listed keys. The NAX kernel on the s, c and d classes // writes 32 heads by 512 float32 partials per split, so once about 512 @@ -254,10 +283,12 @@ void KQuantSDPA::eval_gpu( out.set_data(mx::allocator::malloc(out.nbytes())); // q row-contiguous [B, Hq, qL, D]; k/v [B, Hkv, kL, D], D contiguous. - const auto& q = inputs[0]; - const auto& k = inputs[1]; - const auto& v = inputs[2]; - kq_sdpa_check_layout("sdpa_vector", q, k, v); + kq_sdpa_check_layout("sdpa_vector", inputs[0], inputs[1], inputs[2]); + // q, k and v are bound from vector boundaries (sdpa_q, sdpa_rows), here + // and in the evals below. + const mx::array q = sdpa_q(inputs[0], s, "sdpa_vector q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_vector k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_vector v"); int B = q.shape(0); int n_q_heads = q.shape(1); @@ -368,9 +399,10 @@ void KQuantSDPAGQA::eval_gpu( outputs[1].set_data(mx::allocator::malloc(outputs[1].nbytes())); } - const auto& q = inputs[0]; - const auto& k = inputs[1]; - const auto& v = inputs[2]; + kq_sdpa_check_layout("sdpa_decode_gqa", inputs[0], inputs[1], inputs[2]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_decode_gqa q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_decode_gqa k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_decode_gqa v"); const bool sinks = has_sinks_; const bool starts = has_starts_; const bool kv_q8 = has_kv_q8_; @@ -380,7 +412,12 @@ void KQuantSDPAGQA::eval_gpu( const size_t starts_idx = 3 + (sinks ? 1 : 0); const size_t ends_idx = starts_idx + (starts ? 1 : 0); const size_t qkv_idx = ends_idx + (ends ? 1 : 0); - kq_sdpa_check_layout("sdpa_decode_gqa", q, k, v); + // The KVarN stages are key and value rows too. A copy dispatches on the + // encoder, so each operand is bound before the first kernel is set up. + const mx::array stk = + kvarn ? sdpa_rows(inputs[qkv_idx + 2], s, "sdpa_decode_gqa stage_k") : q; + const mx::array stv = + kvarn ? sdpa_rows(inputs[qkv_idx + 3], s, "sdpa_decode_gqa stage_v") : q; if (kv_q8 || kvarn) { for (int i = 0; i < 4; i++) { kq_sdpa_check_batch_heads("sdpa_decode_gqa", inputs[qkv_idx + i]); @@ -595,8 +632,6 @@ void KQuantSDPAGQA::eval_gpu( // KVarN axes/stage operands and region map (dummies when compiled out). KQKvarnMeta kvm = {}; if (kvarn) { - const auto& stk = inputs[qkv_idx + 2]; - const auto& stv = inputs[qkv_idx + 3]; const int sink_cap = stk.shape(2) - KQ_KVARN_GROUP; kvm.sink_rows = std::min(kvarn_n_, sink_cap); kvm.n_rec_keys = @@ -674,24 +709,38 @@ void KQuantSDPAFAVerify::eval_gpu( // q is the GQA-folded query tile [B, Hkv, n_rows, D], row-contiguous; // k/v [B, Hkv, kL, D] are read in place via their head/seq strides. - const auto& q = inputs[0]; - const auto& k = inputs[1]; - const auto& v = inputs[2]; - kq_sdpa_check_layout("sdpa_fa_verify", q, k, v); + kq_sdpa_check_layout("sdpa_fa_verify", inputs[0], inputs[1], inputs[2]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_fa_verify q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_fa_verify k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_fa_verify v"); // With kvarn, k/v are the code slabs and the key count comes from the // op; a tail-merge body walks only the first n_attend keys while the // region map stays on the full layout (as sdpa_decode_gqa). const bool kvarn = has_kvarn_; + const mx::array stk = + kvarn ? sdpa_rows(inputs[5], s, "sdpa_fa_verify stage_k") : q; + const mx::array stv = + kvarn ? sdpa_rows(inputs[6], s, "sdpa_fa_verify stage_v") : q; int B = q.shape(0); int n_kv_heads = k.shape(1); int n_rows = q.shape(2); int D = q.shape(3); int kL = kvarn ? (kvarn_n_attend_ ? kvarn_n_attend_ : kvarn_n_) : k.shape(2); int q_len = q_len_; + // Pass 1 on the tensor-op units (kq_sdpa_fa_verify_2pass_1_nax) at + // head_dim 128 and 256 on plain K/V: a 16-row strip of the fold per + // simdgroup row group, the head dim split over WN simdgroups. It writes + // the partials of the simdgroup-matrix kernel below. WN is 2 at 128, and + // at 256 it is 4 for folds of up to 32 rows and 2 above. + // KQ_FA_VERIFY_NAX=0 keeps that kernel, and KQ_FA_VERIFY_NAX_WN sets WN + // (2 or 4 at 256, 1 or 2 at 128). Both are read per call. + const bool nax_pass1 = !kvarn && (D == 256 || D == 128) && n_rows <= 64 && + kq_env_flag("KQ_FA_VERIFY_NAX", 1) != 0 && kq_is_nax_available(); int splits = splits_; if (splits == 0) { - splits = kq_fa_verify_splits(B, n_kv_heads, n_rows, D, kL, d); + splits = nax_pass1 ? kq_fa_verify_nax_splits(B * n_kv_heads, kL) + : kq_fa_verify_splits(B, n_kv_heads, n_rows, D, kL, d); } size_t k_head_stride = @@ -743,7 +792,43 @@ void KQuantSDPAFAVerify::eval_gpu( // per 8 tile rows ((BQ/8)*32 threads, BQ the smallest of 32, 48, 64 that // holds the fold); 512 runs the 256-thread d-split variant (BQ 32 only, // capped at the op). - { + // Pass 1 on the tensor-op units, set up above. + if (nax_pass1) { + const int wm = (n_rows + 15) / 16; + int wn = D == 128 ? 2 : wm <= 2 ? 4 : 2; + // Another value warns once and keeps the default. + static const std::vector wn_128 = {"1", "2"}; + static const std::vector wn_256 = {"2", "4"}; + const auto& wn_set = D == 128 ? wn_128 : wn_256; + const int wn_env = kq_env_choice("KQ_FA_VERIFY_NAX_WN", wn_set, -1); + if (wn_env >= 0) { + wn = std::stoi(wn_set[wn_env]); + } + std::string kname = "kq_sdpa_fa_verify_2pass_1_nax_" + ts + "_" + + std::to_string(D) + "_wm" + std::to_string(wm) + "_wn" + + std::to_string(wn); + std::string hash = kname + "_s" + std::to_string(splits); + auto kernel = kq_get_kernel(d, kname, hash, fc); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(q, 0); + ce.set_input_array(k, 1); + ce.set_input_array(v, 2); + ce.set_output_array(partials, 3); + ce.set_output_array(sums, 4); + ce.set_output_array(maxs, 5); + ce.set_bytes(kL, 6); + ce.set_bytes(k_head_stride, 7); + ce.set_bytes(k_seq_stride, 8); + ce.set_bytes(v_head_stride, 9); + ce.set_bytes(v_seq_stride, 10); + ce.set_bytes(scale, 11); + ce.set_bytes(q_len, 12); + ce.set_bytes(n_rows, 13); + MTL::Size group_dims(32, wm * wn, 1); + MTL::Size grid_dims(n_kv_heads, B, splits); + ce.dispatch_threadgroups(grid_dims, group_dims); + } + if (!nax_pass1) { const int bq = n_rows <= 32 ? 32 : n_rows <= 48 ? 48 : 64; std::string kname = "kq_sdpa_fa_verify_2pass_1_" + ts + "_" + std::to_string(D) + "_bq" + std::to_string(bq); @@ -800,8 +885,6 @@ void KQuantSDPAFAVerify::eval_gpu( // KVarN axes/stage operands and region map (dummies when compiled out). KQKvarnMeta kvm = {}; if (kvarn) { - const auto& stk = inputs[5]; - const auto& stv = inputs[6]; const int sink_cap = stk.shape(2) - KQ_KVARN_GROUP; kvm.sink_rows = std::min(kvarn_n_, sink_cap); kvm.n_rec_keys = @@ -870,10 +953,10 @@ void KQuantSDPAFAIndexed::eval_gpu( // q [1, Hq, Q, D] row-contiguous; kv [1, 1, N, D] read in place through // its seq stride; idx [Q, M] int32 row-contiguous. - const auto& q = inputs[0]; - const auto& kv = inputs[1]; + kq_sdpa_check_layout("sdpa_fa_indexed", inputs[0], inputs[1], inputs[1]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_fa_indexed q"); + const mx::array kv = sdpa_rows(inputs[1], s, "sdpa_fa_indexed kv"); const auto& idx = inputs[2]; - kq_sdpa_check_layout("sdpa_fa_indexed", q, kv, kv); int n_heads = q.shape(1); int n_queries = q.shape(2); @@ -1005,19 +1088,23 @@ void KQuantSDPACascade::eval_gpu( outputs[1].set_data(mx::allocator::malloc(outputs[1].nbytes())); } - const auto& q = inputs[0]; - const auto& qf = inputs[1]; - const auto& k_sh = inputs[2]; - const auto& v_sh = inputs[3]; - const auto& k_pr = inputs[4]; - const auto& v_pr = inputs[5]; + kq_sdpa_check_layout( + "sdpa_decode_gqa_cascade", inputs[0], inputs[4], inputs[5]); + kq_sdpa_check_layout( + "sdpa_decode_gqa_cascade", inputs[1], inputs[2], inputs[3]); + const mx::array q = sdpa_q(inputs[0], s, "sdpa_decode_gqa_cascade q"); + const mx::array qf = sdpa_q(inputs[1], s, "sdpa_decode_gqa_cascade q_folded"); + const mx::array k_sh = + sdpa_rows(inputs[2], s, "sdpa_decode_gqa_cascade k_shared"); + const mx::array v_sh = + sdpa_rows(inputs[3], s, "sdpa_decode_gqa_cascade v_shared"); + const mx::array k_pr = sdpa_rows(inputs[4], s, "sdpa_decode_gqa_cascade k"); + const mx::array v_pr = sdpa_rows(inputs[5], s, "sdpa_decode_gqa_cascade v"); const bool starts = has_starts_; const bool kv_q8 = has_kv_q8_; // Optional trailing inputs: [starts], then the eight q8 scale/bias // arrays (shared k/v pairs first, then private). const int q8_base = 6 + int(starts); - kq_sdpa_check_layout("sdpa_decode_gqa_cascade", q, k_pr, v_pr); - kq_sdpa_check_layout("sdpa_decode_gqa_cascade", qf, k_sh, v_sh); if (kv_q8) { for (int i = 4; i < 8; i++) { kq_sdpa_check_batch_heads("sdpa_decode_gqa_cascade", inputs[q8_base + i]); @@ -1283,14 +1370,16 @@ void KQuantSDPABSPrefill::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = mx::metal::device(s.device); - const auto& q = inputs[0]; // folded [1, Hkv, n_qt, 48, D] - const auto& k = inputs[1]; - const auto& v = inputs[2]; + kq_sdpa_check_layout( + "sdpa_prefill_block_sparse", inputs[0], inputs[1], inputs[2]); + // q folded [1, Hkv, n_qt, 48, D]. + const mx::array q = sdpa_q(inputs[0], s, "sdpa_prefill_block_sparse q"); + const mx::array k = sdpa_rows(inputs[1], s, "sdpa_prefill_block_sparse k"); + const mx::array v = sdpa_rows(inputs[2], s, "sdpa_prefill_block_sparse v"); const auto& pages = inputs[3]; 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 9ab3c2c..eaff51c 100644 --- a/src/kquant_sdpa_sparse_decode.cpp +++ b/src/kquant_sdpa_sparse_decode.cpp @@ -52,7 +52,8 @@ mx::array rows_flat(const mx::array& a, const char* op, mx::StreamOrDevice s) { // 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 +// than 1, which is copied here, as is a window or pool that starts or +// steps off a 16-byte boundary. 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 { @@ -65,10 +66,19 @@ SparseRows sparse_rows( const mx::array& win, const mx::array& pool, const mx::array* scales, - const mx::Stream& s) { + const mx::Stream& s, + bool prefill) { + // The kernels read window and pool rows as 16-byte vectors (4-byte words + // for packed codes), so the rows start on 16-byte boundaries too. SparseRows r{ - kq_unit_stride_gpu(win, s), - kq_unit_stride_gpu(pool, s), + kq_strides_aligned_gpu( + win, + s, + prefill ? "sdpa_sparse_prefill win" : "sdpa_sparse_decode win"), + kq_strides_aligned_gpu( + pool, + s, + prefill ? "sdpa_sparse_prefill pool" : "sdpa_sparse_decode pool"), 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)) { @@ -89,11 +99,13 @@ void KQSdpaSparseDecode::eval_gpu( auto& d = mx::metal::device(s.device); auto& ce = mx::metal::get_command_encoder(s); - const auto& q = inputs[0]; + // q is read as 16-byte vectors through its strides. + const mx::array q = + kq_strides_aligned_gpu(inputs[0], s, "sdpa_sparse_decode q"); const auto& idx = inputs[3]; size_t next = 4; - const SparseRows rows = - sparse_rows(inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s); + const SparseRows rows = sparse_rows( + inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s, false); 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; @@ -481,11 +493,12 @@ void KQSdpaSparsePrefill::eval_gpu( auto& d = mx::metal::device(s.device); auto& ce = mx::metal::get_command_encoder(s); - const auto& q = inputs[0]; + const mx::array q = + kq_strides_aligned_gpu(inputs[0], s, "sdpa_sparse_prefill q"); const auto& idx = inputs[3]; size_t next = 4; - const SparseRows rows = - sparse_rows(inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s); + const SparseRows rows = sparse_rows( + inputs[1], inputs[2], packed_ ? &inputs[next++] : nullptr, s, true); 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; diff --git a/src/kquant_skinny_mv.cpp b/src/kquant_skinny_mv.cpp index ec6f523..aa89e12 100644 --- a/src/kquant_skinny_mv.cpp +++ b/src/kquant_skinny_mv.cpp @@ -35,8 +35,8 @@ void KQuantSkinnyMV::eval_gpu( auto& out = outputs[0]; out.set_data(mx::allocator::malloc(out.nbytes())); - const auto& x = inputs[0]; - const auto& w = inputs[1]; + const mx::array x = kq_x_aligned_gpu(inputs[0], s, "skinny_matmul x"); + const mx::array w = kq_x_aligned_gpu(inputs[1], s, "skinny_matmul w"); int K = x.shape(-1); int M = x.shape(-2); int N = w.shape(0); @@ -55,6 +55,7 @@ void KQuantSkinnyMV::eval_gpu( ce.set_bytes(K, 5); MTL::Size group_dims(32 * 8, 1, 1); MTL::Size grid_dims((N + 7) / 8, T, 1); + kq_log_dispatch(kname, x, out, grid_dims); ce.dispatch_threadgroups(grid_dims, group_dims); } diff --git a/tests/qmm_probe.py b/tests/qmm_probe.py index d4e108b..17a5afe 100644 --- a/tests/qmm_probe.py +++ b/tests/qmm_probe.py @@ -6,6 +6,8 @@ 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 +``"multi": [n, ...]`` runs ``quantized_matmul_multi`` over its weight and +one more of each listed N, and gets the largest error over the outputs. 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``, @@ -16,8 +18,11 @@ 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. +``digest``, a hash of its output bits (of every output on a ``"multi"`` +call), to compare calls across probe processes. ``"x_offset": E`` passes +x as a row-contiguous view E elements into its buffer, and ``copies`` +lists the ``[kq align]`` lines of the call, (operand, xin) for each +operand an op copied to an aligned buffer. ``"design": "lead"`` draws nonnegative weights and activations whose first K column is 1024. ``"design": "big"`` draws small weights and dense @@ -242,8 +247,54 @@ def ref_of(lo): x = activations( list(c["shape"]), dtype, c.get("layout", "dense"), c.get("seed", i), design ) + if c.get("x_offset"): + # The same values as a row-contiguous view x_offset elements into + # its buffer. The buffer is allocated with the cache off: a reused + # one can carry the slack that makes the op copy the view to a + # fresh one, and buffers freed late can refill an emptied cache. + off = c["x_offset"] + limit = mx.set_cache_limit(0) + try: + mx.clear_cache() + flat = mx.concatenate([mx.zeros((off,), dtype), x.reshape(-1)]) + mx.eval(flat) + finally: + mx.set_cache_limit(limit) + x = flat[off:].reshape(x.shape) os.write(2, f"[probe] {i}\n".encode()) res = {"i": i, "error": None, "rel": None, "ulp": None} + if c.get("multi"): + # quantized_matmul_multi over weights of N n and each of "multi"; + # rel is the largest over the outputs. + ns = [c["n"], *c["multi"]] + wl = [weights(c["codec"], n, c["shape"][-1], design)[0] for n in ns] + try: + ys = kq.quantized_matmul_multi(x, wl, c["codec"]) + mx.eval(ys) + except Exception as e: + res["error"] = f"{type(e).__name__}: {e}" + outs.append(None) + results.append(res) + continue + bits = [y.view(mx.uint16) if y.dtype in DT.values() else y for y in ys] + h = hashlib.sha1() + for b in bits: + h.update(memoryview(np.array(b))) + res["digest"] = h.hexdigest() + if not spec.get("check", True): + outs.append(None) + results.append(res) + mx.clear_cache() + continue + xn = as64(x) + res["rel"] = max( + rel(as64(y), xn @ dequantized(c["codec"], n, c["shape"][-1], design).T) + for y, n in zip(ys, ns, strict=True) + ) + outs.append(None) + results.append(res) + mx.clear_cache() + continue try: y = kq.quantized_matmul(x, w, s, c["codec"], transpose=True) mx.eval(y) @@ -314,14 +365,19 @@ def ref_of(lo): print(json.dumps(results)) """ -_LOG = re.compile(r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+)$") +_LOG = re.compile( + r"^\[kq qmm\] (\S+) M=(\d+) N=(\d+) K=(\d+) grid=(\d+),(\d+),(\d+) xa=(\d+)$" +) +_ALIGN = re.compile(r"^\[kq align\] (.+) xin=(\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``.""" + ``kernels`` (list of (name, grid) from KQ_QMM_LOG), ``xa`` (the logged + byte offset of x modulo 16 per kernel), ``copies`` (the operands copied + to an aligned buffer, with the byte offset they came at), ``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} @@ -335,12 +391,16 @@ def run_probe(calls, env=None, check=True, timeout=900): assert r.returncode == 0, r.stderr[-6000:] results = json.loads(r.stdout.strip().splitlines()[-1]) logs = {} + aligns = {} + copies = {} other = {} cur = None for line in r.stderr.splitlines(): if line.startswith("[probe] "): cur = int(line.split()[1]) logs[cur] = [] + aligns[cur] = [] + copies[cur] = [] other[cur] = [] continue if cur is None: @@ -349,10 +409,17 @@ def run_probe(calls, env=None, check=True, timeout=900): if m: grid = tuple(int(g) for g in m.group(5, 6, 7)) logs[cur].append((m.group(1), grid)) + aligns[cur].append(int(m.group(8))) + continue + m = _ALIGN.match(line) + if m: + copies[cur].append((m.group(1), int(m.group(2)))) else: other[cur].append(line) for res in results: res["kernels"] = logs.get(res["i"], []) + res["xa"] = aligns.get(res["i"], []) + res["copies"] = copies.get(res["i"], []) res["stderr"] = other.get(res["i"], []) return results diff --git a/tests/test_env_levers.py b/tests/test_env_levers.py index 8dee9aa..4d1e77c 100644 --- a/tests/test_env_levers.py +++ b/tests/test_env_levers.py @@ -9,6 +9,8 @@ from __future__ import annotations +import re + import pytest from qmm_probe import force_cpu, run_probe @@ -86,3 +88,29 @@ def test_verify_nax_q4_0_sb_values(): ((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"] + + +# KQ_VNAX_X16 is on unless 0, read live under KQ_TUNE_LIVE: x at a 16-byte +# boundary takes the _x16 pipeline unless it is off. Values outside 0 and 1 +# warn and keep it on. +@pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") +def test_vnax_x16_values(): + values = [ + (None, True, False), + ("", True, False), + ("on", True, False), + ("1", True, False), + ("off", False, False), + ("0", False, False), + ("2", True, True), + ("x", True, True), + ] + call = {"codec": "q4_k", "shape": [8, 3840], "n": 1000, "dtype": "bfloat16"} + calls = [dict(call, env={"KQ_VNAX_X16": v}) for v, _, _ in values] + res = run_probe(calls, env={**LOG, "KQ_QMM_ROUTE": "verify_nax"}) + for (v, x16, warns), r in zip(values, res, strict=True): + assert r["error"] is None, (v, r["error"]) + ((name, _),) = r["kernels"] + assert "_verify_nax_" in name, (v, name) + assert re.search("_x16(_pf32)?$" if x16 else "_x8(_pf32)?$", name), (v, name) + assert bool(_warnings(r, "KQ_VNAX_X16", v)) == warns, (v, r["stderr"]) diff --git a/tests/test_kvarn_sdpa.py b/tests/test_kvarn_sdpa.py index ac68fad..649d011 100644 --- a/tests/test_kvarn_sdpa.py +++ b/tests/test_kvarn_sdpa.py @@ -6,8 +6,8 @@ therefore materializes the cache (sink + dequantized records + live) into a plain fp16 KV and runs sdpa_decode_gqa with identical splits/tile_c: outputs must match bit for bit, including lse. The reference keeps the simdgroup -pass 1 (KQ_GQA_NAX=0), since at head dim 512 the plain op may otherwise take -the tensor-op pass, whose rounding differs. +pass 1 (KQ_GQA_NAX=0, KQ_FA_VERIFY_NAX=0), since on NAX GPUs the plain ops +may otherwise take a tensor-op pass, whose rounding differs. """ from __future__ import annotations @@ -32,6 +32,7 @@ @pytest.fixture(autouse=True) def _simdgroup_reference(monkeypatch): monkeypatch.setenv("KQ_GQA_NAX", "0") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "0") def quantize_head(x, bits, kind): @@ -772,7 +773,8 @@ def test_dwide_rejects_wrong_tile_c_and_stale_shapes(d): # sdpa_fa_verify_kvarn stages the same values as the decode kernel through # the shared kvarn loaders, and everything downstream is sdpa_fa_verify # unchanged, so it must match sdpa_fa_verify over the materialized twin bit -# for bit, lse included. +# for bit, lse included, where that call runs the simdgroup pass 1 (the +# autouse fixture above). def fold(q, hkv, ql): @@ -856,6 +858,29 @@ def test_fa_verify_explicit_splits(): assert_bit_equal(out, ref) +@pytest.mark.parametrize("d", [128, 256]) +def test_fa_verify_default_lever_keeps_kvarn_kernel(d, monkeypatch): + # At the default KQ_FA_VERIFY_NAX, plain sdpa_fa_verify takes the NAX + # pass 1 at these head dims on NAX GPUs, which reads plain K/V. The + # kvarn op keeps the simdgroup kernel and its split rule, so it gives + # the bits it gives with the lever off, and those of the simdgroup + # reference. The lever is read when the kernel is dispatched, so each + # arm is evaluated before the lever changes. + st = build_state(1, 2, 1000, 6, 6, seed=123 + d, d=d) + q = make_q(1, 8, 4, d=d) + monkeypatch.delenv("KQ_FA_VERIFY_NAX") + out_default, ref_nax = run_fa_pair(st, q, 1000, 6, 6, 4, return_lse=True) + mx.eval(out_default, ref_nax) + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "0") + out_off, ref = run_fa_pair(st, q, 1000, 6, 6, 4, return_lse=True) + mx.eval(out_off, ref) + # The plain op takes the NAX pass at the default on NAX GPUs only. + took_nax = not np.array_equal(np.array(ref_nax[0]), np.array(ref[0])) + assert took_nax == kq.nax_available() + assert_bit_equal(out_default, out_off) + assert_bit_equal(out_default, ref) + + def test_fa_verify_n_attend_full_visibility(): # Body walk of a tail merge at verify width: the first 800 of 1000 keys # with the clamp lifted equals a q_len-1 fold (every row sees every diff --git a/tests/test_matmul_empty.py b/tests/test_matmul_empty.py index 6ddc4ef..1e6b185 100644 --- a/tests/test_matmul_empty.py +++ b/tests/test_matmul_empty.py @@ -1,4 +1,5 @@ -"""quantized_matmul and quantized_matmul_glu on empty inputs. +"""quantized_matmul, quantized_matmul_glu and quantized_matmul_multi 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 @@ -27,11 +28,19 @@ 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), + "multi": lambda x, w: _multi(x, w, codec), } +def _multi(x, w, codec): + # Two weights through one call; both outputs are the same product. + a, b = kq.quantized_matmul_multi(x, [w, w], codec) + assert a.shape == b.shape and a.dtype == b.dtype + return mx.concatenate([a[None], b[None]]).max(axis=0) + + @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) -@pytest.mark.parametrize("op", ["matmul", "glu"]) +@pytest.mark.parametrize("op", ["matmul", "glu", "multi"]) @pytest.mark.parametrize("codec", ["q4_k", "q8_0", "q3_k"]) def test_empty_inputs(codec, op, dtype): f = _ops(codec)[op] diff --git a/tests/test_matmul_multi.py b/tests/test_matmul_multi.py new file mode 100644 index 0000000..bba316e --- /dev/null +++ b/tests/test_matmul_multi.py @@ -0,0 +1,495 @@ +"""quantized_matmul_multi against quantized_matmul per weight. + +On a NAX GPU the op runs one verify_nax dispatch per four weights at 2 to +8 rows, where quantized_matmul would run verify_nax on the widest weight. +With the split count forced to one value and the verify_nax entry forced +at M 2, quantized_matmul per weight runs the same kernel body and fold, so +the outputs match bit for bit. At the default +levers the joint grid can take another split count than each weight +alone, so the outputs are checked +against a float64 reference built from kq.dequantize weights. Every other +case runs quantized_matmul per weight, which the CPU stream checks. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys + +import mlx.core as mx +import numpy as np +import pytest +from kqref import synth_wire +from qmm_probe import force_cpu, run_probe + +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY + +DTYPES = [mx.bfloat16, mx.float16] +R1 = {mx.bfloat16: 6e-3, mx.float16: 1e-3} +# Gate of quantized_matmul on the CPU stream, as in test_cpu_neon. +CPU_REL = 5e-2 +# Codecs with a verify_nax kernel, and an inner dim each serves. +CODECS = { + "q4_0": 512, + "q8_0": 512, + "iq4_nl": 512, + "q4_k": 512, + "q5_k": 512, + "q6_k": 512, + "q3_k": 512, + "q2_k": 512, + "iq4_xs": 512, +} +# Row counts of the joint grid's weights: tiny, partial row tiles, a +# 32-row remainder and a weight wider than the rest. +NS = [48, 160, 1000, 300] + +_W = {} + + +def _weight(codec, n, k, seed): + key = (codec, n, k, seed) + if key not in _W: + wf = mx.random.normal((n, k), key=mx.random.key(97 * seed + n + k)) * 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) + _W[key] = (w, wd) + return _W[key] + + +def _weights(codec, ns, k): + return [_weight(codec, n, k, i) for i, n in enumerate(ns)] + + +def _x(m, k, dtype, seed=0): + return mx.random.normal((m, k), key=mx.random.key(seed + 7 * m)).astype(dtype) + + +def _per_weight(x, ws, codec): + sc = mx.zeros((1,), mx.uint8) + return [kq.quantized_matmul(x, w, sc, codec, transpose=True) for w in ws] + + +@pytest.fixture +def exact_levers(monkeypatch): + """One split count and verify_nax forced from M 2, which also turns the + default per-row qmv route off, so quantized_matmul runs verify_nax on + every weight the op joins.""" + + def set_splits(n): + monkeypatch.setenv("KQ_VERIFY_NAX_SPLITS", str(n)) + + monkeypatch.setenv("KQ_VERIFY_NAX", "2") + return set_splits + + +def _synth(codec, n, k, seed): + # Codecs kq.quantize does not encode: random valid wire. + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + rng = np.random.default_rng(seed) + wire = synth_wire(rng, codec, bpb, n * (k // wpb)) + return mx.array(wire.reshape(n, (k // wpb) * bpb)) + + +def _equal(a, b): + return np.array_equal( + np.array(a.astype(mx.float32)), np.array(b.astype(mx.float32)) + ) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("codec", sorted(CODECS)) +@pytest.mark.parametrize("splits", [1, 2]) +def test_matches_per_weight_bits(codec, dtype, splits, exact_levers): + exact_levers(splits) + k = CODECS[codec] + ws = [w for w, _ in _weights(codec, NS, k)] + for m in range(1, 9): + x = _x(m, k, dtype) + got = kq.quantized_matmul_multi(x, ws, codec) + want = _per_weight(x, ws, codec) + mx.eval(got, want) + assert len(got) == len(ws) + for g, w in zip(got, want, strict=True): + assert g.dtype == w.dtype and g.shape == w.shape + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ), f"{codec} M={m} N={w.shape[-1]}" + + +@pytest.mark.parametrize("codec", ["pq2_0", "q2_0"]) +@pytest.mark.parametrize("splits", [1, 2, 4]) +def test_synth_codecs_match_per_weight_bits(codec, splits, exact_levers): + exact_levers(splits) + ws = [_synth(codec, n, 1024, i) for i, n in enumerate(NS)] + for dtype in DTYPES: + for m in range(1, 9): + x = _x(m, 1024, dtype) + got = kq.quantized_matmul_multi(x, ws, codec) + want = _per_weight(x, ws, codec) + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert _equal(g, w), f"{codec} M={m} N={w.shape[-1]}" + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("codec", ["q4_0", "q4_k", "q6_k"]) +def test_reference_default_levers(codec, dtype): + k = CODECS[codec] + pairs = _weights(codec, NS, k) + ws = [w for w, _ in pairs] + for m in (1, 2, 3, 4, 6, 8, 9, 16): + x = _x(m, k, dtype, seed=3) + got = kq.quantized_matmul_multi(x, ws, codec) + mx.eval(got) + xd = np.array(x.astype(mx.float32)).astype(np.float64) + for g, (_, wd) in zip(got, pairs, strict=True): + ref = xd @ wd.T + err = np.linalg.norm(np.array(g.astype(mx.float32)) - ref) + rel = CPU_REL if force_cpu() else R1[dtype] + assert err <= rel * np.linalg.norm(ref), f"{codec} M={m}" + + +def test_q4_0_two_block_inner_dim(exact_levers): + # K not a multiple of 256 runs the two-block q4_0 kernel. + exact_levers(1) + ws = [w for w, _ in _weights("q4_0", [64, 200], 320)] + for m in (3, 5, 8): + x = _x(m, 320, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ) + + +@pytest.mark.parametrize("count", [1, 4, 5, 9]) +def test_weight_counts(count, exact_levers): + # More than four weights run one dispatch per four. + exact_levers(1) + ns = [64 + 32 * i for i in range(count)] + ws = [w for w, _ in _weights("q4_0", ns, 512)] + x = _x(6, 512, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ) + + +def test_batch_dims_views_and_float32(exact_levers): + exact_levers(1) + big, _ = _weight("q4_0", 512, 512, 11) + ws = [big[100:228], big[300:348], _weight("q4_0", 96, 512, 12)[0]] + x = _x(6, 512, mx.float32).reshape(2, 3, 512) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert g.shape == w.shape and g.dtype == mx.bfloat16 + assert np.array_equal( + np.array(g.astype(mx.float32)), np.array(w.astype(mx.float32)) + ) + + +def test_empty_weight(exact_levers): + exact_levers(1) + ws = [_weight("q4_0", 128, 512, 0)[0], mx.zeros((0, 288), mx.uint8)] + x = _x(4, 512, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + mx.eval(got) + assert got[1].shape == (4, 0) + want = _per_weight(x, ws[:1], "q4_0")[0] + assert np.array_equal( + np.array(got[0].astype(mx.float32)), np.array(want.astype(mx.float32)) + ) + + +def test_misaligned_base_falls_back(exact_levers): + # verify_nax reads q4_0 in 4-byte words: a weight 2 bytes past a 4-byte + # boundary sends the whole group per weight, off verify_nax. + exact_levers(1) + w = _weight("q4_0", 64, 512, 3)[0] + flat = mx.concatenate([mx.zeros((2,), mx.uint8), w.reshape(-1)]) + mx.eval(flat) + wv = flat[2:].reshape(w.shape) + mx.eval(wv) + ws = [_weight("q4_0", 128, 512, 4)[0], wv] + x = _x(6, 512, mx.bfloat16) + got = kq.quantized_matmul_multi(x, ws, "q4_0") + want = _per_weight(x, ws, "q4_0") + mx.eval(got, want) + for g, w in zip(got, want, strict=True): + assert np.isfinite(np.array(g.astype(mx.float32))).all() + assert _equal(g, w) + + +def test_grad_matches_per_weight(): + ws = [w for w, _ in _weights("q4_0", NS, 512)] + x = _x(6, 512, mx.float32) + cot = [mx.random.normal((6, n), key=mx.random.key(n)) for n in NS] + + def f_multi(x): + return sum( + (y.astype(mx.float32) * c).sum() + for y, c in zip(kq.quantized_matmul_multi(x, ws, "q4_0"), cot, strict=True) + ) + + def f_each(x): + return sum( + (y.astype(mx.float32) * c).sum() + for y, c in zip(_per_weight(x, ws, "q4_0"), cot, strict=True) + ) + + g1 = mx.grad(f_multi)(x) + g2 = mx.grad(f_each)(x) + mx.eval(g1, g2) + assert np.allclose(np.array(g1), np.array(g2), rtol=2e-2, atol=2e-2) + + +def test_argument_errors(): + w, _ = _weight("q4_0", 64, 512, 0) + x = _x(4, 512, mx.bfloat16) + with pytest.raises(ValueError): + kq.quantized_matmul_multi(x, [], "q4_0") + with pytest.raises(ValueError): + kq.quantized_matmul_multi(x, [w, _weight("q4_0", 64, 256, 0)[0]], "q4_0") + with pytest.raises(ValueError): + kq.quantized_matmul_multi(_x(4, 256, mx.bfloat16), [w], "q4_0") + with pytest.raises(ValueError): + kq.quantized_matmul_multi(x, [w], "no_such_codec") + + +_LOG_CHILD = r""" +import json, os, sys +import mlx.core as mx +import mlx_kquant as kq +codec, m, k, ns, steps = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), \ + json.loads(sys.argv[4]), json.loads(sys.argv[5]) +w = [kq.quantize(mx.random.normal((n, k), key=mx.random.key(i)) * 0.1, codec)[0] + for i, n in enumerate(ns)] +x = mx.random.normal((m, k), key=mx.random.key(9)).astype(mx.bfloat16) +sc = mx.zeros((1,), mx.uint8) +mx.eval(w, x) +for i, (op, env) in enumerate(steps): + for name, value in env.items(): + os.environ[name] = value + sys.stderr.write(f"=== {i}\n") + sys.stderr.flush() + if op == "multi": + out = kq.quantized_matmul_multi(x, w, codec) + else: + out = [kq.quantized_matmul(x, wi, sc, codec, transpose=True) for wi in w] + mx.eval(out) +""" + + +def _census(steps, codec="q4_0", m=6, k=512, ns=(4096, 640, 48, 48), env=None): + """KQ_QMM_LOG kernel names per step of a child process. Each step is + (op, env): the env is set before the step's call, which is + quantized_matmul_multi ("multi") or quantized_matmul per weight.""" + child = {kk: v for kk, v in os.environ.items() if not kk.startswith("KQ_")} + child["KQ_QMM_LOG"] = "1" + child.update(env or {}) + r = subprocess.run( + [ + sys.executable, + "-c", + _LOG_CHILD, + codec, + str(m), + str(k), + json.dumps(list(ns)), + json.dumps(steps), + ], + env=child, + capture_output=True, + text=True, + timeout=300, + ) + assert r.returncode == 0, r.stderr + out = [] + for part in r.stderr.split("=== ")[1:]: + lines = part.splitlines()[1:] + out.append([ln.split()[2] for ln in lines if ln.startswith("[kq qmm]")]) + return out + + +NAX_ONLY = pytest.mark.skipif( + force_cpu() or not kq.nax_available(), reason="NAX GPU route" +) + + +@NAX_ONLY +def test_one_dispatch_on_nax(): + (names,) = _census([["multi", {}]], env={"KQ_VERIFY_NAX_SPLITS": "1"}) + assert len(names) == 1 and "verify_nax_multi" in names[0], names + + +@NAX_ONLY +@pytest.mark.parametrize( + "lever", + [ + {"KQ_VERIFY_MMA": "2"}, + {"KQ_QMM_SPLITK_NAX": "1"}, + {"KQ_QMM_ROUTE": "qmv"}, + {"KQ_QMM_ROUTE": "verify_mma"}, + {"KQ_FORCE_QMM_MIN_M": "2"}, + {"KQ_VERIFY_NAX": "0"}, + {"KQ_NAX_QMV": "8"}, + ], + ids=lambda e: ",".join(f"{k}={v}" for k, v in e.items()), +) +def test_route_levers_are_followed(lever): + # A lever that sends quantized_matmul off verify_nax sends the op to + # quantized_matmul per weight, which then runs the same kernels. + ns = (4096, 4096) + per, multi = _census([["per", {}], ["multi", {}]], ns=ns, env=lever) + assert not any("verify_nax" in n for n in per), per + assert multi == per + + +@NAX_ONLY +def test_forced_verify_nax_fuses_narrow_weights(): + # KQ_VERIFY_NAX=2 alone takes quantized_matmul to verify_nax at widths + # whose default is per-row qmv; the op follows and joins the weights. + ns = (512, 256) + per, multi = _census( + [["per", {}], ["multi", {}]], m=4, ns=ns, env={"KQ_VERIFY_NAX": "2"} + ) + assert len(per) == 2 and all("_verify_nax_" in n for n in per), per + assert len(multi) == 1 and "verify_nax_multi" in multi[0], multi + + +@NAX_ONLY +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("codec", [*CODECS, "pq2_0", "q2_0"]) +def test_x16_load_bits(codec, dtype): + # x at offset 0 (paired 16-byte activation loads), at offset 0 with + # KQ_VNAX_X16=0, at every 2-byte offset from 2 to 14 bytes (8 bytes: one + # load per step; the others: an aligned copy, paired) and at 16 bytes: + # the same bits from one verify_nax_multi dispatch. + arms = [(0, None), (0, "0"), *[(e, None) for e in range(1, 9)]] + cases = [(m, sp) for m in range(2, 9) for sp in ("1", "2")] + calls = [ + { + "codec": codec, + "shape": [m, 1024], + "n": 1000, + "multi": [48, 300], + "dtype": dtype, + "seed": m, + "x_offset": off, + "env": {"KQ_VNAX_X16": lever, "KQ_VERIFY_NAX_SPLITS": sp}, + } + for m, sp in cases + for off, lever in arms + ] + env = {"KQ_QMM_LOG": "1", "KQ_TUNE_LIVE": "1", "KQ_VERIFY_NAX": "2"} + res = run_probe(calls, env=env, check=False) + for i, (m, sp) in enumerate(cases): + group = res[len(arms) * i : len(arms) * (i + 1)] + for (off, lever), r in zip(arms, group, strict=True): + tag = (codec, m, sp, off, lever) + assert r["error"] is None, (tag, r["error"]) + b = 2 * off + x8 = lever is not None or b % 16 == 8 + end = "_x8(_pf32)?$" if x8 else "_x16(_pf32)?$" + names = [n for n, _ in r["kernels"]] + assert names and all( + "verify_nax_multi" in n and re.search(end, n) for n in names + ), (tag, names) + copied = [("quantized_matmul_multi x", b % 8)] if b % 8 else [] + assert r["copies"] == copied, (tag, r["copies"]) + assert set(r["xa"]) == {0 if b % 8 else b % 16}, (tag, r["xa"]) + assert r["digest"] == group[0]["digest"], tag + + +@NAX_ONLY +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize( + "lever", + [{"KQ_VERIFY_NAX": "0"}, {"KQ_QMM_ROUTE": "qmv"}], + ids=lambda e: ",".join(f"{k}={v}" for k, v in e.items()), +) +def test_per_weight_copies_x_once(lever, dtype): + # Off verify_nax on a NAX GPU, the op runs quantized_matmul per weight + # on one aligned copy of x: one copy for the three weights, every kernel + # reads x on a 16-byte boundary, and the bits match x at offset 0. + calls = [ + { + "codec": "q4_k", + "shape": [4, 1024], + "n": 1024, + "multi": [256, 512], + "dtype": dtype, + "seed": 4, + "x_offset": off, + } + for off in range(9) + ] + res = run_probe(calls, env={"KQ_QMM_LOG": "1", **lever}, check=False) + for off, r in enumerate(res): + b = 2 * off + assert r["error"] is None, (off, r["error"]) + names = [n for n, _ in r["kernels"]] + assert names and not any("multi" in n for n in names), (off, names) + copied = [("quantized_matmul_multi x", b % 16)] if b % 16 else [] + assert r["copies"] == copied, (off, r["copies"]) + assert set(r["xa"]) == {0}, (off, r["xa"]) + assert r["digest"] == res[0]["digest"], off + + +@NAX_ONLY +def test_lever_read_once_per_process(): + # KQ_MATMUL_MULTI keeps its first read, as KQ_GLU, unless KQ_TUNE_LIVE=1. + steps = [["multi", {}], ["multi", {"KQ_MATMUL_MULTI": "0"}]] + first, second = _census(steps) + assert first == second and "verify_nax_multi" in first[0], (first, second) + first, second = _census(steps, env={"KQ_TUNE_LIVE": "1"}) + assert "verify_nax_multi" in first[0], first + assert not any("verify_nax_multi" in n for n in second), second + + +def test_lever_off(): + # KQ_MATMUL_MULTI=0 builds quantized_matmul per weight: the same bits. + r = subprocess.run( + [sys.executable, "-c", _LEVER_OFF_CHILD], + env={**os.environ, "KQ_MATMUL_MULTI": "0"}, + capture_output=True, + text=True, + timeout=300, + ) + assert r.returncode == 0, r.stderr + assert r.stdout.split()[-1] == "EQUAL", r.stdout + + +_LEVER_OFF_CHILD = r""" +import os +import re +import mlx.core as mx +import numpy as np +if os.environ.get("KQUANT_FORCE_CPU"): + mx.set_default_device(mx.cpu) +import mlx_kquant as kq +ns = (48, 160, 1000, 300) +w = [kq.quantize(mx.random.normal((n, 512), key=mx.random.key(n)) * 0.1, "q4_0")[0] + for n in ns] +x = mx.random.normal((6, 512), key=mx.random.key(1)).astype(mx.bfloat16) +sc = mx.zeros((1,), mx.uint8) +got = kq.quantized_matmul_multi(x, w, "q4_0") +want = [kq.quantized_matmul(x, wi, sc, "q4_0", transpose=True) for wi in w] +mx.eval(got, want) +f32 = [np.array(a.astype(mx.float32)) for a in got + want] +same = all(np.array_equal(a, b) for a, b in zip(f32[: len(w)], f32[len(w) :])) +print("EQUAL" if same else "DIFFER") +""" diff --git a/tests/test_nax_smallm.py b/tests/test_nax_smallm.py index 943e653..0687e5f 100644 --- a/tests/test_nax_smallm.py +++ b/tests/test_nax_smallm.py @@ -24,6 +24,7 @@ import functools import json import os +import re import subprocess import sys @@ -681,8 +682,9 @@ 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. +# Kernel name fragment of each route but the mat-vec routes (qmv, mv_ext, +# verify_qmv), whose kernels hold one of MATVEC_KERNEL. +MATVEC_KERNEL = ("_qmv_", "_mv_ext_") ROUTE_KERNEL = { "verify_nax": "_verify_nax_", "verify_mma": "_verify_mma_", @@ -692,14 +694,14 @@ def _depth(entry): 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.""" + """The route an entry of _logged ran: "matvec" for a mat-vec kernel, a + ROUTE_KERNEL key, or "tile" for another kernel.""" error, kernels = entry assert error is None, error - if not kernels: - return "matvec" assert len(kernels) == 1, kernels name = kernels[0][0] + if any(frag in name for frag in MATVEC_KERNEL): + return "matvec" return next((r for r, frag in ROUTE_KERNEL.items() if frag in name), "tile") @@ -1321,6 +1323,83 @@ def test_verify_nax_disable_nax(codec, monkeypatch): _sweep(codec, w, s, ref_w, 1000, ms=[3, 5, 8]) +# The verify_nax codecs that read one activation step at a time read two +# steps in one 16-byte load where x starts on a 16-byte boundary, one step +# per 8-byte load where it starts 8 bytes past one, and x from an aligned +# copy elsewhere. Each arm multiplies the same values in the same order: x +# at offset 0 (the paired load), at offset 0 with KQ_VNAX_X16=0, at every +# 2-byte offset from 2 to 14 bytes (8 bytes: one load per step; the others: +# the copy, paired) and at 16 bytes (the paired load off the buffer start). +X16_ARMS = [(0, None), (0, "0"), *[(e, None) for e in range(1, 9)]] + + +def x16_variant(off, lever): + # The KQ_QMM_LOG name end of the pipeline an arm takes. + x8 = lever is not None or (2 * off) % 16 == 8 + return "_x8(_pf32)?$" if x8 else "_x16(_pf32)?$" + + +def x16_reads(off, what): + # (copies, xa): the copy an arm makes of x and the x offset its kernel + # logs. x off an 8-byte boundary is copied; x8 reads x 8 bytes in. + b = 2 * off + return ([(what, b % 8)], [0]) if b % 8 else ([], [b % 16]) + + +X16_SPLITS = [None, "3"] +X16_DTYPES = ["bfloat16", "float16"] + + +@functools.cache +def _x16_digests(): + cases = [ + (codec, m, dtype, splits) + for codec in VNAX_CODECS + VNAX_IQ4 + for m in range(2, 9) + for dtype in X16_DTYPES + for splits in X16_SPLITS + ] + calls = [ + { + "codec": codec, + "shape": [m, K], + "n": 1000, + "dtype": dtype, + "seed": m, + "x_offset": off, + "env": {"KQ_VNAX_X16": lever, "KQ_VERIFY_NAX_SPLITS": splits}, + } + for codec, m, dtype, splits in cases + for off, lever in X16_ARMS + ] + env = { + "KQ_QMM_LOG": "1", + "KQ_TUNE_LIVE": "1", + "KQ_QMM_ROUTE": "verify_nax", + "KQ_QMM_ROUTE_STRICT": "1", + } + res = run_probe(calls, env=env, check=False) + n = len(X16_ARMS) + return {c: res[n * i : n * i + n] for i, c in enumerate(cases)} + + +@pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") +@pytest.mark.parametrize("codec", VNAX_CODECS + VNAX_IQ4) +def test_verify_nax_x16_bits(codec): + for (c, m, dtype, splits), arms in _x16_digests().items(): + if c != codec: + continue + tag = f"{codec} M{m} {dtype} splits {splits}" + for (off, lever), r in zip(X16_ARMS, arms, strict=True): + assert r["error"] is None, (tag, off, lever, r["error"]) + ((name, _),) = r["kernels"] + assert "_verify_nax_" in name, (tag, name) + assert re.search(x16_variant(off, lever), name), (tag, off, lever, name) + copies, xa = x16_reads(off, "quantized_matmul x") + assert (r["copies"], r["xa"]) == (copies, xa), (tag, off, lever) + assert r["digest"] == arms[0]["digest"], (tag, off, lever) + + # verify_nax folds the block scale into half weights and accumulates in # float32, so one activation channel near the top of the half range stays # finite and within the contract. diff --git a/tests/test_prism_mv_range.py b/tests/test_prism_mv_range.py index 82d1c12..8144799 100644 --- a/tests/test_prism_mv_range.py +++ b/tests/test_prism_mv_range.py @@ -75,7 +75,8 @@ def test_q2_0_mv_range(design, dtype, m, route): def _check(r, dtype): assert r["error"] is None, r["error"] - # The mat-vec routes log no tile kernel. - assert r["kernels"] == [], r["kernels"] + # One mat-vec kernel, no tile. + ((name, _),) = r["kernels"] + assert "_qmv_" in name, name assert r["nonfinite"] == 0, r["nonfinite"] assert _excess(r["rel"], r["floor1"]) < EXCESS[dtype], (r["rel"], r["floor1"]) diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py index 47a79b7..0110361 100644 --- a/tests/test_sdpa.py +++ b/tests/test_sdpa.py @@ -997,10 +997,22 @@ def _check_fa(D, qL, kL, dtype, Hkv=4, G=6, strided=False, splits=0): assert got.shape == q.shape +# At head_dim 128 and 256 with up to 64 folded rows, sdpa_fa_verify runs +# pass 1 on the tensor-op units on NAX GPUs, and KQ_FA_VERIFY_NAX=0 keeps +# the simdgroup-matrix kernel. The op reads the variable at each call, so +# these cases run both kernels against the reference. +@pytest.fixture(params=["0", "1"], ids=["simdgroup", "nax"]) +def fa_pass1(request, monkeypatch): + if request.param == "1" and not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", request.param) + return request.param + + @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -@pytest.mark.parametrize("qL", [2, 3, 4, 5, 6]) -def test_sdpa_fa_verify(D, qL, dtype): +@pytest.mark.parametrize("qL", [2, 3, 4, 5, 6, 7, 8]) +def test_sdpa_fa_verify(D, qL, dtype, fa_pass1): _check_fa(D, qL, kL=4096, dtype=dtype, G=4) @@ -1011,7 +1023,7 @@ def test_sdpa_fa_verify_qwen_geometry(): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_decode(D, dtype): +def test_sdpa_fa_verify_decode(D, dtype, fa_pass1): # q_len == 1: plain GQA decode on the matrix units (every folded row # attends the full KV). 122b shape: G16 x qL1 at hd256, 2 kv heads; # same fold on the hd512 d-split kernel. @@ -1024,39 +1036,39 @@ def test_sdpa_fa_verify_gemma_geometry(): @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_full_tile(D): +def test_sdpa_fa_verify_full_tile(D, fa_pass1): # n_rows == 32 fills the tile exactly (no padding rows), qL at the cap _check_fa(D, 8, kL=4096, dtype=mx.bfloat16, Hkv=2, G=4) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_partial_warp(D): +def test_sdpa_fa_verify_partial_warp(D, fa_pass1): # n_rows == 18: the third row strip covers rows 16..17 plus padding _check_fa(D, 6, kL=2048, dtype=mx.bfloat16, G=3) @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_strided_unaligned(D, dtype): +def test_sdpa_fa_verify_strided_unaligned(D, dtype, fa_pass1): # strided KV-cache prefix + a key length off every tile/split boundary _check_fa(D, 4, kL=3071, dtype=dtype, strided=True, splits=16) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_short_kv(D): +def test_sdpa_fa_verify_short_kv(D, fa_pass1): # kL small enough that most splits stage zero keys: their empty partials # (max = finite_min, sum = 0) must merge as weight zero _check_fa(D, 4, kL=17, dtype=mx.bfloat16, splits=16, G=6) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_min_kv(D): +def test_sdpa_fa_verify_min_kv(D, fa_pass1): # kL == qL floor: row 0 attends exactly one key, every row masked hard _check_fa(D, 4, kL=4, dtype=mx.bfloat16, G=6) @pytest.mark.parametrize("D", [64, 128, 256, 512]) -def test_sdpa_fa_verify_causal_split_straddle(D): +def test_sdpa_fa_verify_causal_split_straddle(D, fa_pass1): # the last qL keys straddle a split boundary (splits=128, kL=4098 puts # keys 4096..4097 alone in the final split): that split is entirely past # the low rows' causal limits, exercising the dead-row guard in a @@ -1065,38 +1077,38 @@ def test_sdpa_fa_verify_causal_split_straddle(D): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) -def test_sdpa_fa_verify_bq64_full_tile(dtype): +def test_sdpa_fa_verify_bq64_full_tile(dtype, fa_pass1): # gqa16 x qL4 = 64 rows (qwen3.5-122b-a10b verify fold): fills the BQ=64 # tile exactly _check_fa(256, 4, kL=4096, dtype=dtype, Hkv=2, G=16) @pytest.mark.parametrize("G,qL", [(10, 6), (7, 7), (14, 4)]) -def test_sdpa_fa_verify_bq64_padded(G, qL): +def test_sdpa_fa_verify_bq64_padded(G, qL, fa_pass1): # 49..63 rows: BQ=64 with padding rows in the upper simdgroups _check_fa(256, qL, kL=2048, dtype=mx.bfloat16, Hkv=2, G=G) @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("D", [64, 128, 256]) -def test_sdpa_fa_verify_bq48_full_tile(D, dtype): +def test_sdpa_fa_verify_bq48_full_tile(D, dtype, fa_pass1): # gqa6 x qL8 = 48 rows (qwen3.x full attention at a DFlash 2 block of # 8): fills the BQ=48 tile exactly _check_fa(D, 8, kL=4096, dtype=dtype, Hkv=4, G=6) @pytest.mark.parametrize("G,qL", [(8, 5), (6, 6), (6, 7), (12, 4)]) -def test_sdpa_fa_verify_bq48_padded(G, qL): +def test_sdpa_fa_verify_bq48_padded(G, qL, fa_pass1): # 33..48 rows: BQ=48, padding rows in the top simdgroup below 48 _check_fa(256, qL, kL=2048, dtype=mx.bfloat16, Hkv=2, G=G) -def test_sdpa_fa_verify_bq48_strided_split_straddle(): +def test_sdpa_fa_verify_bq48_strided_split_straddle(fa_pass1): # the last qL keys alone in the final split, over a strided cache view _check_fa(256, 8, kL=4098, dtype=mx.bfloat16, Hkv=4, G=6, strided=True, splits=128) -def test_sdpa_fa_verify_bq64_strided_kv(): +def test_sdpa_fa_verify_bq64_strided_kv(fa_pass1): _check_fa(256, 4, kL=3071, dtype=mx.bfloat16, Hkv=2, G=16, strided=True, splits=16) @@ -1112,11 +1124,17 @@ def test_sdpa_fa_verify_bq64_strided_kv(): (64, 8, 4, 8, 8192, 32), # 1024 simdgroups at head_dim <= 128 ], ) -def test_sdpa_fa_verify_default_splits(D, Hkv, G, qL, kL, want): +def test_sdpa_fa_verify_default_splits(D, Hkv, G, qL, kL, want, monkeypatch): # On the large GPU classes splits=0 raises the count past the decode - # bucket to fill the GPU; elsewhere it keeps the bucket. + # bucket to fill the GPU; elsewhere it keeps the bucket. Counts of the + # simdgroup-matrix pass 1. + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "0") if mx.device_info()["architecture"][-1] not in "scd": want = 16 if kL <= 8192 else 32 + _check_default_splits(D, Hkv, G, qL, kL, want) + + +def _check_default_splits(D, Hkv, G, qL, kL, want): q, k, v = _make(1, Hkv, Hkv, G * qL, kL, D, mx.bfloat16, seed=kL + D, strided=False) scale = 1.0 / (D**0.5) got = kq.sdpa_fa_verify(q, k, v, scale, q_len=qL) @@ -1127,9 +1145,100 @@ def test_sdpa_fa_verify_default_splits(D, Hkv, G, qL, kL, want): assert _rel(got, ref) < REL_BOUND[mx.bfloat16] +@pytest.mark.parametrize( + "D,Hkv,G,qL,kL,want", + [ + (256, 1, 8, 4, 8192, 64), # 64 threadgroups on one kv head + (256, 2, 12, 2, 4096, 32), + (256, 4, 6, 8, 16385, 16), # no rise past 16k keys + (128, 8, 4, 4, 4096, 8), + (128, 16, 2, 4, 4096, 4), # never under 4 + (128, 32, 1, 8, 4096, 4), + (256, 4, 6, 4, 1023, 8), # at least 64 keys per split + (256, 2, 12, 2, 700, 8), + (256, 4, 6, 6, 17, 4), + (128, 3, 8, 2, 4096, 16), # 3 kv heads count as 4 + ], +) +def test_sdpa_fa_verify_nax_default_splits(D, Hkv, G, qL, kL, want): + # splits=0 on the tensor-op pass 1 aims at about 64 threadgroups. + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + _check_default_splits(D, Hkv, G, qL, kL, want) + + +# KQ_FA_VERIFY_NAX_WN sets the simdgroups across the head dim of the +# tensor-op pass 1: every value it takes, over row strips 1 to 4, partial +# strips, split counts that leave splits without keys, and a strided cache +# view with a key count off every tile. +@pytest.mark.parametrize( + "D,wn", [(128, "1"), (128, "2"), (256, "2"), (256, "4")], ids=str +) +@pytest.mark.parametrize( + "G,qL,kL,splits,strided", + [ + (6, 1, 3071, 0, True), + (6, 4, 3071, 16, True), + (6, 6, 4098, 128, False), + (6, 8, 300, 1, False), + (16, 4, 2048, 2, True), + (3, 7, 17, 16, False), + (1, 5, 5, 0, False), + (64, 1, 4300, 0, True), + ], +) +def test_sdpa_fa_verify_nax_wn(D, wn, G, qL, kL, splits, strided, monkeypatch): + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "1") + monkeypatch.setenv("KQ_FA_VERIFY_NAX_WN", wn) + for dtype in (mx.bfloat16, mx.float16): + _check_fa(D, qL, kL=kL, dtype=dtype, Hkv=2, G=G, strided=strided, splits=splits) + + +# A KQ_FA_VERIFY_NAX_WN value outside the head dim's set warns once on +# stderr and keeps the default split: the same bits as the lever unset. +@pytest.mark.parametrize("D,bad", [(128, "4"), (256, "1"), (256, "3"), (128, "x")]) +def test_sdpa_fa_verify_nax_wn_out_of_set(D, bad, monkeypatch, capfd): + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + monkeypatch.setenv("KQ_FA_VERIFY_NAX", "1") + q, k, v = _make(1, 2, 2, 6 * 4, 1000, D, mx.bfloat16, seed=D, strided=False) + scale = 1.0 / (D**0.5) + out = {} + for value in (None, bad): + if value is None: + monkeypatch.delenv("KQ_FA_VERIFY_NAX_WN", raising=False) + else: + monkeypatch.setenv("KQ_FA_VERIFY_NAX_WN", value) + out[value] = kq.sdpa_fa_verify(q, k, v, scale, q_len=4, splits=8) + _eval_or_skip(out[value]) + assert mx.array_equal(out[None], out[bad]).item() + err = capfd.readouterr().err + assert f"KQ_FA_VERIFY_NAX_WN={bad} is not" in err, err + + +@pytest.mark.parametrize("D", [128, 256]) +@pytest.mark.parametrize("qL", [1, 2, 3, 4, 5, 6, 7, 8]) +def test_sdpa_fa_verify_nax_matches_simdgroup(D, qL, monkeypatch): + # Both pass 1 kernels write the same partials layout to the same merge: + # at one split count they agree to bfloat16 rounding of the output. + if not kq.nax_available(): + pytest.skip("the tensor-op pass 1 needs NAX") + G = 6 + q, k, v = _make(1, 4, 4, G * qL, 4300, D, mx.bfloat16, seed=qL + D, strided=True) + scale = 1.0 / (D**0.5) + out = {} + for lever in ("0", "1"): + monkeypatch.setenv("KQ_FA_VERIFY_NAX", lever) + out[lever] = kq.sdpa_fa_verify(q, k, v, scale, q_len=qL, splits=16) + _eval_or_skip(out[lever]) + assert _rel(out["1"], out["0"]) < REL_BOUND[mx.bfloat16] + + @pytest.mark.parametrize("D", [64, 128]) @pytest.mark.parametrize("G", [32, 64]) -def test_sdpa_fa_verify_cascade_fold(D, G): +def test_sdpa_fa_verify_cascade_fold(D, G, fa_pass1): # shared-prefix cascade: B batch rows folded kv-head-major into the row # axis at q_len 1 (every row attends the full prefix); G=64 fills the # BQ=64 tile at the register-resident head dims diff --git a/tests/test_split_kernel_dispatch.py b/tests/test_split_kernel_dispatch.py index 5d6e375..20bbcd3 100644 --- a/tests/test_split_kernel_dispatch.py +++ b/tests/test_split_kernel_dispatch.py @@ -4,11 +4,15 @@ 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 +half verify_mma kernels (KQ_VERIFY_MMA_HALF), and run quantized_matmul_multi +for its joined verify_nax kernels. Every verify_nax call also runs with x +8 bytes into its buffer, so each verify_nax kernel runs both its 16-byte +(_x16) and 8-byte (_x8) x load pipelines. KQ_QMM_LOG names the pipeline 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. +lists for the codec, each verify_nax name in both pipelines, 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 @@ -48,6 +52,10 @@ def _no_nax(): "KQ_QMM_ROUTE": None, "KQ_QMM_ROUTE_STRICT": None, "KQ_VERIFY_MMA_HALF": None, + "KQ_NAX_QMV": None, + "KQ_VERIFY_NAX": None, + "KQ_VERIFY_NAX_SPLITS": None, + "KQ_VNAX_X16": 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 @@ -70,6 +78,11 @@ def _no_nax(): # 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}) +# quantized_matmul_multi runs the verify_nax_multi kernels: verify_nax from +# M 2 and more than one split, so the partials dtype shows in the name. The +# split count divides the K steps: 16 of 256 at K, 65 of 64 at K + 64. +MULTI = {"KQ_NAX_QMV": "0", "KQ_VERIFY_NAX": "2"} +MULTI_ARMS = [(8, K, "2"), (8, K + 64, "5")] def _host_names(): @@ -82,6 +95,16 @@ def _host_names(): return out +def _pipelines(name): + """The pipeline names KQ_QMM_LOG gives a host kernel name: a verify_nax + kernel has an _x16 and an _x8 pipeline, named ahead of its partials + suffix.""" + if "_verify_nax_" not in name: + return {name} + base, suffix = (name[:-5], "_pf32") if name.endswith("_pf32") else (name, "") + return {base + "_x16" + suffix, base + "_x8" + suffix} + + HOST = {} if _no_nax() else _host_names() @@ -108,17 +131,41 @@ def test_every_split_kernel_runs(codec): 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) + for off in (0, 4) if family == "_verify_nax_" else (0,): + calls.append( + { + "codec": codec, + "shape": [m, k], + "n": n, + "dtype": dtype, + "seed": 1, + "x_offset": off, + "env": {**CLEAR, **env, "KQ_SPLITK_F32": f32}, + } + ) + families.append(family) + for m, k, splits in MULTI_ARMS: + if k % CODEC_GEOMETRY[codec][3]: + continue + for off in (0, 4): + calls.append( + { + "codec": codec, + "shape": [m, k], + "n": n, + "multi": [n + 64], + "dtype": dtype, + "seed": 1, + "x_offset": off, + "env": { + **CLEAR, + **MULTI, + "KQ_VERIFY_NAX_SPLITS": splits, + "KQ_SPLITK_F32": f32, + }, + } + ) + families.append(None) 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): @@ -130,4 +177,5 @@ def test_every_split_kernel_runs(codec): 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) == [] + want = set().union(*(_pipelines(n) for n in host)) + assert sorted(want - ran) == [] diff --git a/tests/test_splitk_cap.py b/tests/test_splitk_cap.py index b03cd59..0f04f83 100644 --- a/tests/test_splitk_cap.py +++ b/tests/test_splitk_cap.py @@ -9,8 +9,7 @@ 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. +fixed target. The logged kernel shows which route 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 @@ -107,10 +106,10 @@ def test_nax_cap_falls_through(codec, m): _split(default, "_qmm_t_nax_splitk_") assert default["rel"] < R1, default["rel"] _same(capped, off) + ((name, grid),) = capped["kernels"] if (codec, m) in MV_EXT: - assert capped["kernels"] == [], capped["kernels"] + assert "_mv_ext_" in name, name else: - ((name, grid),) = capped["kernels"] assert tile_of(name)[0] == 32 and grid[2] == 1, (name, grid) @@ -139,7 +138,8 @@ def test_nax_cap_falls_through_to_qmv(m): forced, capped, off = _nax_qmv_fallthrough()[m] _split(forced, "_qmm_t_nax_splitk_") _same(capped, off) - assert capped["kernels"] == [] + ((name, _),) = capped["kernels"] + assert "_qmv_" in name, name # N 1024 at M 16 is 16 BM=32 tiles, and K 5120 is 20 units of 256 weights, diff --git a/tests/test_weight_base.py b/tests/test_weight_base.py index af7a200..415853f 100644 --- a/tests/test_weight_base.py +++ b/tests/test_weight_base.py @@ -196,6 +196,47 @@ def call(o, j=j, x=x): assert _equal(kq.quantized_matmul_glu(x, _shift(wg, 16), wu, codec), ref) +@GPU +@pytest.mark.parametrize("codec", ["q4_k", "q6_k", "q8_0", "q4_0"]) +def test_gpu_weight_base_multi(codec): + """quantized_matmul_multi checks the start of every weight, at the rows + that join them in one verify kernel and at those that do not.""" + a = _align(codec) + ws = [_wire(codec, (N, K), seed=j)[0] for j in range(3)] + # verify_nax declines these starts, so the group runs another route. + vnax = VNAX_BASE.get(codec, 1) + for m in (1, 6): + x = (mx.random.normal((m, K), key=mx.random.key(m)) * 0.5).astype(mx.bfloat16) + # One row builds quantized_matmul per weight, which names itself. + op = ( + "quantized_matmul_multi" + if m > 1 and kq.nax_available() + else ("quantized_matmul") + ) + ref = kq.quantized_matmul_multi(x, ws, codec) + mx.eval(ref) + for j in range(3): + for off in sorted({a, 16}): + moved = [_shift(w, off) if i == j else w for i, w in enumerate(ws)] + got = kq.quantized_matmul_multi(x, moved, codec) + same = _close if m > 1 and off % vnax else _equal + for g, r in zip(got, ref, strict=True): + assert same(g, r), (m, j, off) + moved = [_shift(w, a // 2) if i == j else w for i, w in enumerate(ws)] + msg = _msg(op, codec, a // 2) + with pytest.raises(ValueError, match=msg): + kq.quantized_matmul_multi(x, moved, codec) + # A weight that is not yet evaluated raises from the eval instead. + lazy = [ws[0], _shift(ws[1], a // 2, evaluate=False), ws[2]] + ys = kq.quantized_matmul_multi(x, lazy, codec) + with pytest.raises(ValueError, match=_msg(op, codec, a // 2)): + mx.eval(ys) + # The stream keeps working. + got = kq.quantized_matmul_multi(x, ws, codec) + for g, r in zip(got, ref, strict=True): + assert _equal(g, r) + + @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 index 50cdd37..2226f2a 100644 --- a/tests/test_weight_views.py +++ b/tests/test_weight_views.py @@ -106,6 +106,27 @@ def test_quantized_matmul_weight_views(codec, state): 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("codec", ["q4_k", "q8_0", "q4_0"]) +def test_multi_weight_views(codec, dtype, state): + """1 row runs quantized_matmul per weight, 4 and 6 rows the joined + verify kernel where there is one.""" + views = _views(codec) + for m in (1, 4, 6): + x = _x(m, K, dtype) + for pair in (("even_rows", "odd_rows"), ("row_band", "left_cols")): + makes = [views[name] for name in pair] + [views["whole"]] + ref = kq.quantized_matmul_multi(x, [_dense(f()) for f in makes], codec) + ws = [f() for f in makes] + if state == "evaluated": + mx.eval(ws) + ys = kq.quantized_matmul_multi(x, ws, codec) + mx.eval(ys, ref) + for y, r in zip(ys, ref, strict=True): + assert mx.array_equal(y, r).item(), (codec, m, pair, state) + + @pytest.mark.parametrize("state", ["lazy", "evaluated"]) @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16], ids=["bf16", "f16"]) @pytest.mark.parametrize("m", [64, 1100]) diff --git a/tests/test_x_offset_views.py b/tests/test_x_offset_views.py new file mode 100644 index 0000000..6dbffd4 --- /dev/null +++ b/tests/test_x_offset_views.py @@ -0,0 +1,1030 @@ +"""Operands at an element offset into their buffer. + +A row-contiguous view can start at any element of its buffer, and an op's +contiguous node passes it on in place. Kernels load activations, caches and +some float operands as vectors, and Metal defines a vector load only at an +address aligned to its size. So each op binds such an operand that starts +off the boundary of its kernel's widest load from a copy, and a strided +operand whose rows step off it too; the matmul ops bind x from a 16-byte +boundary. Each case runs an op on views at offsets of 1 to 8 elements and +requires the output bits of the same values in their own buffer. +KQ_QMM_LOG prints a "[kq align] xin=" line for each copy, +and every matmul dispatch logs x on a 16-byte boundary (xa=0), but +verify_nax's: it reads x where it lands on an 8-byte boundary (xa=8) and +from the copy otherwise. The cases require all three: each view that +starts off the boundary reaches its op there and is copied or read on its +kernel's boundary, each view on the boundary is read in place, and no +matmul kernel reads x off it. + +- quantized_matmul for every codec with a matmul kernel, at 1 to 64 rows, + on the default routes, with the NAX mat-vec off (KQ_NAX_QMV=0), with NAX + off (KQ_DISABLE_NAX=1, the routes of a GPU without NAX), and on each + forced route (KQ_QMM_ROUTE); +- the transpose=False, GLU, LoRA (also with the fused LoRA kernel, + KQ_QMM_LORA=2), residual and bias matmuls, skinny_matmul, gather_qmm, + gather_qmm_seg, the fused MoE gathers and gather_mix, in a child per route + environment; +- rmsnorm_gate, latent_fp4_pack, dsa_indexer_score_decode, the hc glue ops, + the attention ops and the sparse attention ops, in a child with NAX on and + off, each operand read as vectors at the offsets and, for the keys and + values the attention ops read through their strides, as rows that step + one element past their length. + +Float32 activations of the matmul ops are cast to bfloat16 into a new +buffer, so their views reach the kernels aligned. Under KQUANT_FORCE_CPU +the quantized_matmul cases run on the CPU stream, which needs no alignment; +the other ops are Metal-only kernels with no CPU eval, and are skipped. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys + +import mlx.core as mx +import numpy as np +import pytest +from kqref import is_synth, synth_wire +from qmm_probe import TESTS_DIR, force_cpu, run_probe + +import mlx_kquant as kq +from mlx_kquant.codec_geometry import CODEC_GEOMETRY + +OFFSETS = tuple(range(1, 9)) +DTYPES = ("float16", "bfloat16", "float32") +ROWS = (1, 2, 4, 8, 12, 16, 32, 64) +STRICT = {"KQ_QMM_ROUTE_STRICT": "1"} +# (name, env, rows). The forced routes run at the rows they serve; a route +# that declines a codec or shape raises, on the offset views too. +ROUTES = [ + ("default", {}, ROWS), + ("nax_qmv_off", {"KQ_NAX_QMV": "0"}, ROWS), + ("nax_off", {"KQ_DISABLE_NAX": "1"}, ROWS), + ("qmv", {"KQ_QMM_ROUTE": "qmv", **STRICT}, (1, 2, 8)), + ("verify_qmv", {"KQ_QMM_ROUTE": "verify_qmv", **STRICT}, (2, 8)), + ("mv_ext", {"KQ_QMM_ROUTE": "mv_ext", **STRICT}, (2, 8, 12)), + ("verify_mma", {"KQ_QMM_ROUTE": "verify_mma", **STRICT}, (2, 8)), + ("verify_nax", {"KQ_QMM_ROUTE": "verify_nax", **STRICT}, (2, 8)), + ("splitk", {"KQ_QMM_ROUTE": "splitk", **STRICT}, (8, 32)), + ("nax", {"KQ_QMM_ROUTE": "nax", **STRICT}, (8, 32)), + ("nax_splitk", {"KQ_QMM_ROUTE": "nax_splitk", **STRICT}, (8, 32)), +] +CPU_ROUTES = [("default", {}, ROWS)] +MATMUL_CODECS = [c for c in kq.codecs() if kq.codec_has_matmul(c)] +ITEMSIZE = {"float16": 2, "bfloat16": 2, "float32": 4} + + +def _unaligned_x(name, xa): + """Whether a kernel read x off its boundary: 16 bytes, or 8 for + verify_nax, which reads x where it lands on an 8-byte boundary.""" + return xa % (8 if "_verify_nax_" in name else 16) != 0 + + +metal_only = pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")), + reason="Metal-only kernel; no CPU path.", +) + + +def quantized_matmul_offsets(codec): + """(bad, unaligned, unreached, copied) over the routes: the offset calls + whose output or error differs from the call on x in its own buffer, the + kernels that read x off a 16-byte boundary with the byte offset they + logged, the calls on a view off the boundary that log neither a copy of + x nor a kernel reading x there, and the calls on a view on the boundary + that log a copy.""" + bad, unaligned, unreached, copied = [], set(), [], [] + for route, env, rows in CPU_ROUTES if force_cpu() else ROUTES: + calls, cells = [], [] + for dtype in DTYPES: + # float32 activations are cast to bfloat16 into a new buffer: + # the default route covers them. + if dtype == "float32" and route != "default": + continue + for m in rows: + for off in (0, *OFFSETS): + calls.append( + { + "codec": codec, + "shape": [m, 4096], + "n": 256, + "dtype": dtype, + "seed": m, + "x_offset": off, + } + ) + cells.append((dtype, m, off)) + res = run_probe(calls, env={"KQ_QMM_LOG": "1", **env}, check=False) + base = {} + for (dtype, m, off), r in zip(cells, res, strict=True): + got = r["error"] or r["digest"] + names = [k for k, _ in r["kernels"]] + if off == 0: + base[(dtype, m)] = got + elif got != base[(dtype, m)]: + bad.append((route, dtype, m, off, names, r["xa"])) + for name, xa in zip(names, r["xa"], strict=True): + if _unaligned_x(name, xa): + unaligned.add((name, xa)) + off_boundary = (off * ITEMSIZE[dtype]) % 16 != 0 + if not off_boundary and r["copies"]: + copied.append((route, dtype, m, off, r["copies"])) + if ( + off_boundary + and dtype != "float32" + and not force_cpu() + and not r["error"] + and not any(xin for _, xin in r["copies"]) + and not any(r["xa"]) + ): + unreached.append((route, dtype, m, off, names)) + return bad, sorted(unaligned), unreached, copied + + +@pytest.mark.parametrize("codec", MATMUL_CODECS) +def test_quantized_matmul_x_offsets(codec): + """Bit-exact output at every offset, every kernel reads x from a + 16-byte boundary, every view off it reaches the op there, and every + view on it is read in place.""" + bad, unaligned, unreached, copied = quantized_matmul_offsets(codec) + assert bad == [], bad + assert unaligned == [], unaligned + assert unreached == [], unreached + assert copied == [], copied + + +# -------------------------------- other ops -------------------------------- +# Each group runs in a child interpreter per environment, since the route +# levers are read once per process. The child logs its dispatches and +# copies, and marks each offset call on stderr with the copies it needs. + +ENVS = { + "default": {}, + "nax_qmv_off": {"KQ_NAX_QMV": "0"}, + "nax_off": {"KQ_DISABLE_NAX": "1"}, + "lora2": {"KQ_QMM_LORA": "2"}, +} +K, N = 768, 256 +OP_CODECS = ("q4_0", "q4_k", "q6_k", "q8_0", "iq4_nl", "iq2_xxs") +MOE_CODECS = ("q4_k", "q8_0", "iq4_nl", "iq2_xxs") +MX_DT = {"float16": mx.float16, "bfloat16": mx.bfloat16, "float32": mx.float32} +HALF = ("float16", "bfloat16") +_W = {} + + +def _own_buffer(make): + """make() evaluated into a buffer of its own size. A buffer reused from + the cache can run more than 16 KiB past the view, and the op's + contiguous node then copies the view before the eval sees it; buffers + freed late by earlier work could refill an emptied cache, so the cache + stays off while the buffer is allocated.""" + limit = mx.set_cache_limit(0) + mx.clear_cache() + try: + a = make() + mx.eval(a) + finally: + mx.set_cache_limit(limit) + return a + + +def _at_offset(x, off): + """The values of x as a row-contiguous view off elements into a buffer + of their own.""" + flat = _own_buffer( + lambda: mx.concatenate([mx.zeros((off,), x.dtype), x.reshape(-1)]) + ) + return flat[off:].reshape(x.shape) + + +def _rows_view(x): + """The values of x in rows that step one element past their length, so + that every row after the first starts off a 16-byte boundary.""" + pad = mx.zeros((*x.shape[:-1], 1), x.dtype) + return _own_buffer(lambda: mx.concatenate([x, pad], axis=-1))[..., :-1] + + +def _mark(label, need, names=()): + """need: "off" when the call must copy each operand in names (or, with + no names, log a copy or a kernel reading x) at a nonzero offset, "copy" + when it must copy each of them at any offset, "none" when it must copy + no operand, "any" for no requirement.""" + print(f"[arm] {label} | {need} | {','.join(names)}", file=sys.stderr, flush=True) + + +def _listify(out): + return list(out) if isinstance(out, (list, tuple)) else [out] + + +def _same(got, want): + return all( + g.dtype == w.dtype and mx.array_equal(g, w).item() + for g, w in zip(got, want, strict=True) + ) + + +def _x(shape, dtype, seed): + x = (mx.random.normal(shape, key=mx.random.key(seed)) * 0.5).astype(MX_DT[dtype]) + mx.eval(x) + return x + + +def _weight(codec, n, k, seed=0, experts=None): + """Weights [n, k], or [experts, n, k], and scales. The codecs that + quantize only with an importance matrix get random valid wire.""" + key = (codec, n, k, seed, experts) + if key not in _W: + e = experts or 1 + if is_synth(codec): + _, _, bpb, wpb = CODEC_GEOMETRY[codec] + rng = np.random.default_rng(seed * 131 + n + k) + wire = synth_wire(rng, codec, bpb, e * n * (k // wpb)) + w = mx.array(wire.reshape(e, n, (k // wpb) * bpb)) + w = w if experts else w[0] + s = mx.zeros((1,), dtype=mx.uint8) + else: + shape = (n, k) if experts is None else (experts, n, k) + wf = mx.random.normal(shape, key=mx.random.key(seed)) * 0.1 + w, s = kq.quantize(wf, codec) + w = mx.contiguous(w) + mx.eval(w, s) + _W[key] = (w, s) + return _W[key] + + +def _offset_failures(call, shape, dtype, seed=0, label="", names=()): + """The offsets at which call(view) differs from call(x) in its own + buffer. A call the op refuses is skipped. A float32 view is cast into a + new buffer, so only the 16-bit views must reach the op off the + boundary; every view on it must reach the op uncopied.""" + x = _x(shape, dtype, seed) + try: + want = _listify(call(x)) + mx.eval(*want) + except (ValueError, RuntimeError): + return [] + bad = [] + for off in OFFSETS: + view = _at_offset(x, off) + need = "none" + if (off * ITEMSIZE[dtype]) % 16: + need = "off" if dtype in HALF else "any" + _mark(f"{label} {dtype} {list(shape)} off{off}", need, names) + got = _listify(call(view)) + mx.eval(*got) + if not _same(got, want): + bad.append(off) + return bad + + +def _operand_failures(label, call, ops, cases): + """call(**ops) with one operand at a time replaced by its views. cases: + (operand, the copies it needs, vec, strided), vec the elements of the + kernel's widest vector load of the operand. Each view at an offset off + a boundary of vec elements must be copied, and each view on one must + not; a strided operand also runs as rows that step one element past + that boundary, which must be copied too.""" + want = _listify(call(**ops)) + mx.eval(*want) + bad = [] + for name, names, vec, strided in cases: + a = ops[name] + arms = [(f"off{off}", off, "off" if off % vec else "none") for off in OFFSETS] + if strided: + arms.append(("rows", None, "copy")) + for arm, off, need in arms: + view = _rows_view(a) if off is None else _at_offset(a, off) + _mark(f"{label} {a.dtype} {name} {arm}", need, names) + got = _listify(call(**{**ops, name: view})) + mx.eval(*got) + if not _same(got, want): + bad.append([label, str(a.dtype), name, arm]) + return bad + + +def _matmul_ops(): + """transpose=False, GLU, LoRA and residual matmuls.""" + bad = [] + a = (mx.random.normal((K, 8), key=mx.random.key(5)) / np.sqrt(K)).astype( + mx.bfloat16 + ) + b = (mx.random.normal((8, N), key=mx.random.key(6)) * 0.3).astype(mx.bfloat16) + mx.eval(a, b) + for codec in OP_CODECS: + w, _ = _weight(codec, N, K) + wg, _ = _weight(codec, N, K, 1) + wt, st = _weight(codec, K, N, 2) + + def resid(x, w=w, codec=codec): + out = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + r = (mx.random.normal((x.shape[0], N), key=mx.random.key(7)) * 0.5).astype( + out + ) + return kq.quantized_matmul_residual(x, w, codec, r) + + def lora(x, w=w, codec=codec): + dt = mx.bfloat16 if x.dtype == mx.float32 else x.dtype + return kq.quantized_matmul_lora(x, w, codec, a.astype(dt), b.astype(dt)) + + def glu(x, w=w, wg=wg, codec=codec): + return kq.quantized_matmul_glu(x, wg, w, codec) + + def matmul_n(x, wt=wt, st=st, codec=codec): + return kq.quantized_matmul(x, wt, st, codec, False) + + ops = [ + ("matmul_n", (1, 8, 64), matmul_n), + ("glu", (1, 8, 64, 1024), glu), + ("lora", (1, 8, 64, 1024), lora), + ("residual", (1, 8, 64, 1024), resid), + ] + for name, rows, call in ops: + for dtype in DTYPES: + for m in rows: + offs = _offset_failures(call, (m, K), dtype, m, f"{name} {codec}") + if offs: + bad.append([name, codec, dtype, m, offs]) + return bad + + +def _bias_and_skinny(): + """The q8_0 bias mat-vec and skinny_matmul, whose dense weight is read + as vectors too.""" + bad = [] + w, s = _weight("q8_0", N, K) + bias = mx.random.normal((N,), key=mx.random.key(3)).astype(mx.float32) + for dtype in DTYPES: + offs = _offset_failures( + lambda x: kq.quantized_matmul_qmv_bias(x, w, s, bias, "q8_0"), + (1, K), + dtype, + label="qmv_bias", + ) + if offs: + bad.append(["qmv_bias", dtype, 1, offs]) + wd = (mx.random.normal((64, K), key=mx.random.key(4)) * 0.5).astype( + MX_DT[dtype] + ) + mx.eval(wd) + for m in (1, 4, 16): + offs = _offset_failures( + lambda x, wd=wd: kq.skinny_matmul(x, wd), (m, K), dtype, label="skinny" + ) + if offs: + bad.append(["skinny", dtype, m, offs]) + if dtype in HALF: + x = _x((m, K), dtype, m) + bad += _operand_failures( + f"skinny w m{m}", + lambda x, w: kq.skinny_matmul(x, w), + {"x": x, "w": wd}, + [("w", ("skinny_matmul w",), 8, False)], + ) + return bad + + +def _gathers(): + """gather_qmm: the mat-vec, the tiled gather, the sorted per-expert + GEMM and transpose=False; and gather_qmm_seg.""" + experts, bad = 4, [] + rng = np.random.default_rng(4) + for codec in MOE_CODECS: + wt, st = _weight(codec, N, K, 3, experts) + wn, sn = _weight(codec, K, N, 4, experts) + for rows, m, transpose, sort in [ + (8, 1, True, False), + (4, 32, True, False), + (256, 1, True, True), + (8, 1, False, False), + ]: + idx = rng.integers(0, experts, rows).astype(np.uint32) + rhs = mx.array(np.sort(idx) if sort else idx) + w, s = (wt, st) if transpose else (wn, sn) + + def call(x, w=w, s=s, rhs=rhs, transpose=transpose, sort=sort, c=codec): + return kq.gather_qmm( + x, + w, + s, + c, + rhs_indices=rhs, + transpose=transpose, + sorted_indices=sort, + ) + + for dtype in DTYPES: + offs = _offset_failures( + call, (rows, m, K), dtype, label=f"gather_qmm {codec}" + ) + if offs: + bad.append(["gather_qmm", codec, dtype, rows, m, transpose, offs]) + idx = np.sort(rng.integers(0, experts, 200).astype(np.uint32)) + maps = kq.expert_tile_map(mx.array(idx), experts) + for dtype in DTYPES: + offs = _offset_failures( + lambda x, wt=wt, st=st, c=codec, mp=maps: kq.gather_qmm_seg( + x, wt, st, c, *mp + ), + (len(idx), K), + dtype, + label=f"gather_qmm_seg {codec}", + ) + if offs: + bad.append(["gather_qmm_seg", codec, dtype, len(idx), offs]) + return bad + + +T, R = 3, 2 +S = R + 1 +ME, MN, MK = 4, 64, 512 + + +def _routing(seed, slots): + rng = np.random.default_rng(seed) + inds = mx.array(rng.integers(0, ME, 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_moe_ops(codec): + """name -> (shape of x, call(x)) for the K-quant fused MoE ops.""" + gw, uw = _weight(codec, MN, MK, 1, ME)[0], _weight(codec, MN, MK, 2, ME)[0] + shexp = _weight(codec, MN, MK, 3, ME)[0] + sgw, suw = mx.contiguous(shexp[0]), mx.contiguous(shexp[1]) + inds, _ = _routing(3, R) + inds_s, sc_s = _routing(4, S) + la = (mx.random.normal((ME, MK, 4), key=mx.random.key(8)) * 0.05).astype(mx.float16) + lb = (mx.random.normal((ME, 4, MN), key=mx.random.key(9)) * 0.05).astype(mx.float16) + mx.eval(sgw, suw, la, lb) + return { + "moe_glu_gather_kq": ( + (T, MK), + lambda x: kq.moe_glu_gather_kq(x, gw, uw, codec, inds), + ), + "gather_qmv_kq": ((T, R, MK), lambda x: kq.gather_qmv_kq(x, gw, codec, inds)), + "gather_qmv_kq_lora": ( + (T, R, MK), + lambda x: kq.gather_qmv_kq(x, gw, codec, inds, lora_a=la, lora_b=lb), + ), + "moe_glu_gather_shexp_kq": ( + (T, MK), + lambda x: kq.moe_glu_gather_shexp_kq(x, gw, uw, sgw, suw, codec, inds), + ), + "gather_qmv_mix_kq": ( + (T, S, MK), + lambda x: kq.gather_qmv_mix_kq(x, gw, sgw, codec, inds, sc_s), + ), + "gather_qmv_mix_ns_kq": ( + (T, S, MK), + lambda x: kq.gather_qmv_mix_ns_kq(x, gw, codec, inds_s, sc_s), + ), + "gather_qmv_mix_ns_kq_lora": ( + (T, S, MK), + lambda x: kq.gather_qmv_mix_ns_kq( + x, gw, codec, inds_s, sc_s, lora_a=la, lora_b=lb + ), + ), + } + + +def _mxfp4_moe_ops(): + """name -> (shape of x, call(x)) for the mxfp4 fused MoE ops.""" + rng = np.random.default_rng(1) + + def stack(): + w = mx.array(rng.integers(0, 2**32, size=(ME, MN, MK // 8), dtype=np.uint32)) + s = mx.array(rng.integers(118, 124, size=(ME, MN, MK // 32), dtype=np.uint8)) + b = mx.array(rng.uniform(-0.2, 0.2, (ME, MN)).astype(np.float32)) + mx.eval(w, s, b) + return w, s, b + + gw, gs, gb = stack() + uw, us, ub = stack() + inds, sc = _routing(3, R) + return { + "moe_glu_gather": ( + (T, MK), + lambda x: kq.moe_glu_gather(x, gw, gs, gb, uw, us, ub, inds), + ), + "gather_qmv_bias": ( + (T, R, MK), + lambda x: kq.gather_qmv_bias(x, gw, gs, gb, inds), + ), + "gather_qmv_mix_bias": ( + (T, R, MK), + lambda x: kq.gather_qmv_mix_bias(x, gw, gs, gb, inds, sc), + ), + } + + +def _moe(): + """The fused MoE gathers, K-quant and mxfp4.""" + bad = [] + for codec in (*MOE_CODECS, "mxfp4"): + ops = _mxfp4_moe_ops() if codec == "mxfp4" else _kq_moe_ops(codec) + for name, (shape, call) in ops.items(): + for dtype in DTYPES: + offs = _offset_failures(call, shape, dtype, label=f"{name} {codec}") + if offs: + bad.append([name, codec, dtype, offs]) + return bad + + +def _gather_mix(): + """gather_mix reads the expert outputs y, its activation, as vectors.""" + tokens, k, n = 5, 4, 256 + order = np.random.default_rng(6).permutation(tokens * k).astype(np.uint32) + inv = mx.array(order) + sc = mx.array( + np.random.default_rng(7).uniform(0.1, 0.9, (tokens, k)).astype(np.float32) + ) + mx.eval(inv, sc) + bad = [] + for dtype in HALF: + offs = _offset_failures( + lambda y: kq.gather_mix(y, inv, sc), + (tokens * k, n), + dtype, + label="gather_mix", + ) + if offs: + bad.append(["gather_mix", dtype, offs]) + return bad + + +def _rand(shape, dtype, seed, scale=0.5): + a = (mx.random.normal(shape, key=mx.random.key(seed)) * scale).astype(dtype) + mx.eval(a) + return a + + +def _norm_latent_indexer(): + """rmsnorm_gate, latent_fp4_pack and dsa_indexer_score_decode.""" + bad = [] + for dt in (mx.float16, mx.bfloat16): + for d in (64, 256): + ops = { + "x": _rand((6, d), dt, 1), + "w": _rand((d,), dt, 2) + 1.0, + "gate": _rand((6, d), dt, 3), + } + bad += _operand_failures( + f"rmsnorm_gate d{d}", + lambda x, w, gate: kq.rmsnorm_gate(x, w, gate, 1e-6), + ops, + [(n, (f"rmsnorm_gate {n}",), 4, False) for n in ("x", "w", "gate")], + ) + rng = np.random.default_rng(11) + for heads, ql, rows, qo, ratio in ((4, 1, 64, 63, 1), (64, 2, 1027, 4106, 4)): + ops = { + "q": mx.array(rng.standard_normal((1, heads, ql, 128)) * 0.4).astype( + dt + ), + "k": mx.array(rng.standard_normal((1, rows, 128)) * 0.4).astype(dt), + "w": mx.array(rng.standard_normal((1, ql, heads)) * 0.2).astype(dt), + } + bad += _operand_failures( + f"dsa_indexer_score_decode h{heads}", + lambda q, k, w, qo=qo, ratio=ratio: kq.dsa_indexer_score_decode( + q, k, w, qo, ratio + ), + ops, + [(n, (f"dsa_indexer_score_decode {n}",), 8, False) for n in ("q", "k")], + ) + for dt in (mx.float16, mx.bfloat16, mx.float32): + bad += _operand_failures( + "latent_fp4_pack", + lambda x: kq.latent_fp4_pack(x), + {"x": _rand((8, 64), dt, 4)}, + [("x", ("latent_fp4_pack x",), 8, False)], + ) + return bad + + +def _hc(): + """The hc glue ops: the stream, residual and collapsed rows, the + projection fn, the mixes and the base biases.""" + return [b for dt in (mx.float16, mx.bfloat16) for b in _hc_dtype(dt)] + + +def _hc_dtype(dt): + from test_hc_glue import HC, HC_EPS, ITERS, MIX, NORM_EPS + + d = 1024 + sk = {"iters": ITERS, "hc_eps": HC_EPS, "norm_eps": NORM_EPS} + bad = [] + x = _rand((1, 1, HC, d), dt, 1, 0.05) + fn = _rand((MIX, HC * d), mx.float32, 2, 0.02) + scale = mx.array([1.1, 0.9, 1.3], dtype=mx.float32) + base = _rand((MIX,), mx.float32, 3, 0.1) + w = _rand((d,), dt, 4, 0.1) + 1.0 + mr, ssq = kq.hc_front_reduce(x, fn) + col, post, comb = kq.hc_sinkhorn_collapse(x, mr, ssq, scale, base, w, **sk) + mx.eval(scale, w, mr, ssq, col, post, comb) + + def case(op, *names): + return [(n, (f"{op} {n}",), 4, False) for n in names] + + bad += _operand_failures( + "hc_front_reduce", + lambda x, fn: kq.hc_front_reduce(x, fn), + {"x": x, "fn": fn}, + case("hc_front_reduce", "x", "fn"), + ) + bad += _operand_failures( + "hc_sinkhorn_collapse", + lambda x, mixes, base: kq.hc_sinkhorn_collapse( + x, mixes, ssq, scale, base, w, **sk + ), + {"x": x, "mixes": mr, "base": base}, + case("hc_sinkhorn_collapse", "x", "mixes", "base"), + ) + bad += _operand_failures( + "hc_expand", + lambda x, resid: kq.hc_expand(x, resid, post, comb), + {"x": col, "resid": x}, + case("hc_expand", "x", "resid"), + ) + bad += _operand_failures( + "hc_front_expand_reduce", + lambda x_sub, resid, fn: kq.hc_front_expand_reduce( + x_sub, resid, post, comb, fn + ), + {"x_sub": col, "resid": x, "fn": fn}, + case("hc_front_expand_reduce", "x_sub", "resid", "fn"), + ) + bad += _operand_failures( + "hc_front_expand_collapse", + lambda x_sub, resid, fn, base: kq.hc_front_expand_collapse( + x_sub, resid, post, comb, fn, scale, base, w, **sk + ), + {"x_sub": col, "resid": x, "fn": fn, "base": base}, + case("hc_front_expand_collapse", "x_sub", "resid", "fn", "base"), + ) + return bad + + +def _qkv(op, strided=("k", "v"), q="q"): + """Cases for an attention op's query and its keys and values.""" + return [(q, (f"{op} {q}",), 4, False)] + [ + (n, (f"{op} {n}",), 4, True) for n in strided + ] + + +def _sdpa(): + """The attention ops: q, k and v, the q8 cache words and the KVarN + stages.""" + bad = [b for dt in (mx.float16, mx.bfloat16) for b in _sdpa_dtype(dt)] + return bad + _sdpa_kvarn() + + +def _sdpa_dtype(dt): + from test_sdpa import _bs_prefill_lists, _make_indexed + + bad = [] + q = _rand((1, 8, 1, 128), dt, 1) + k = _rand((1, 2, 512, 128), dt, 2) + v = _rand((1, 2, 512, 128), dt, 3) + qkv = {"q": q, "k": k, "v": v} + bad += _operand_failures( + "sdpa_vector", + lambda q, k, v: kq.sdpa_vector(q, k, v, 0.0625), + { + "q": _rand((1, 8, 2, 256), dt, 15), + "k": _rand((1, 2, 300, 256), dt, 16), + "v": _rand((1, 2, 300, 256), dt, 17), + }, + _qkv("sdpa_vector"), + ) + bad += _operand_failures( + "sdpa_decode_gqa", + lambda q, k, v: kq.sdpa_decode_gqa(q, k, v, 0.0884), + qkv, + _qkv("sdpa_decode_gqa"), + ) + pages = mx.array(np.broadcast_to(np.arange(16, dtype=np.int32), (1, 2, 16))) + bad += _operand_failures( + "sdpa_decode_gqa_paged", + lambda q, k, v: kq.sdpa_decode_gqa_paged(q, k, v, 0.0884, pages), + qkv, + _qkv("sdpa_decode_gqa"), + ) + # Head dim 512 at the depth that runs pass 1 on NAX. + wide = { + "q": _rand((1, 8, 1, 512), dt, 4), + "k": _rand((1, 1, 4096, 512), dt, 5), + "v": _rand((1, 1, 4096, 512), dt, 6), + } + bad += _operand_failures( + "sdpa_decode_gqa d512", + lambda q, k, v: kq.sdpa_decode_gqa(q, k, v, 0.0442), + wide, + _qkv("sdpa_decode_gqa"), + ) + # A q8 cache: k and v are the packed words. + kw, ks, kb = mx.quantize(k, group_size=64, bits=8) + vw, vs, vb = mx.quantize(v, group_size=64, bits=8) + mx.eval(kw, ks, kb, vw, vs, vb) + bad += _operand_failures( + "sdpa_decode_gqa q8", + lambda k, v: kq.sdpa_decode_gqa( + q, k, v, 0.0884, k_scales=ks, k_biases=kb, v_scales=vs, v_biases=vb + ), + {"k": kw, "v": vw}, + _qkv("sdpa_decode_gqa")[1:], + ) + qf = _rand((1, 2, 12, 128), dt, 7) + bad += _operand_failures( + "sdpa_fa_verify", + lambda q, k, v: kq.sdpa_fa_verify(q, k, v, 0.0884, q_len=3), + {"q": qf, "k": k, "v": v}, + _qkv("sdpa_fa_verify"), + ) + qi, kv, idx = _make_indexed(8, 1, 100, 64, 0, dt, seed=9) + bad += _operand_failures( + "sdpa_fa_indexed", + lambda q, kv: kq.sdpa_fa_indexed(q, kv, idx, 0.0442), + {"q": qi, "kv": kv}, + _qkv("sdpa_fa_indexed", strided=("kv",)), + ) + # One sequence: the shared pass's folded query is q itself. + bad += _operand_failures( + "sdpa_decode_gqa_cascade", + lambda q, k_shared, v_shared, k, v: kq.sdpa_decode_gqa_cascade( + q, k_shared, v_shared, k, v, 0.0884 + ), + { + "q": q, + "k_shared": k, + "v_shared": v, + "k": _rand((1, 2, 64, 128), dt, 10), + "v": _rand((1, 2, 64, 128), dt, 11), + }, + [ + ( + "q", + ("sdpa_decode_gqa_cascade q", "sdpa_decode_gqa_cascade q_folded"), + 4, + False, + ) + ] + + _qkv("sdpa_decode_gqa_cascade", ("k_shared", "v_shared", "k", "v"))[1:], + ) + # A q8 cache: the shared pass stages the words in the fa_verify kernel. + kp, kps, kpb = mx.quantize(_rand((1, 2, 64, 128), dt, 12), group_size=64, bits=8) + vp, vps, vpb = mx.quantize(_rand((1, 2, 64, 128), dt, 13), group_size=64, bits=8) + mx.eval(kp, kps, kpb, vp, vps, vpb) + bad += _operand_failures( + "sdpa_decode_gqa_cascade q8", + lambda k_shared, v_shared, k, v: kq.sdpa_decode_gqa_cascade( + q, + k_shared, + v_shared, + k, + v, + 0.0884, + k_shared_scales=ks, + k_shared_biases=kb, + v_shared_scales=vs, + v_shared_biases=vb, + k_priv_scales=kps, + k_priv_biases=kpb, + v_priv_scales=vps, + v_priv_biases=vpb, + ), + {"k_shared": kw, "v_shared": vw, "k": kp, "v": vp}, + _qkv("sdpa_decode_gqa_cascade", ("k_shared", "v_shared", "k", "v"))[1:], + ) + # Four queries: the folded query is a view of q. + length, s, topk = 4, 1024, 16 + offset = s - length + sel = np.zeros((length, topk), dtype=np.int64) + rng = np.random.default_rng(3) + for t in range(length): + sel[t] = rng.choice((offset + t + 1) // 4, size=topk, replace=False) + lists = [mx.array(a) for a in _bs_prefill_lists(sel, offset, s)] + bad += _operand_failures( + "sdpa_prefill_block_sparse", + lambda q, k, v: kq.sdpa_prefill_block_sparse(q, k, v, 0.0625, *lists, offset), + { + "q": _rand((1, 24, length, 256), dt, 12), + "k": _rand((1, 2, s, 256), dt, 13), + "v": _rand((1, 2, s, 256), dt, 14), + }, + _qkv("sdpa_prefill_block_sparse"), + ) + return bad + + +def _sdpa_kvarn(): + """The KVarN stages, float16.""" + from test_kvarn_sdpa import build_state, make_q + + bad = [] + st = build_state(1, 2, 300, 6, 6, seed=21) + codes = {n: st[n] for n in ("codes_k", "axes_k", "codes_v", "axes_v")} + qk = mx.array(make_q(1, 8, 1)).astype(mx.float16) + qv = mx.array(make_q(1, 8, 4)).astype(mx.float16).reshape(1, 2, 16, 128) + stages = {"stage_k": st["stage_k"], "stage_v": st["stage_v"]} + bad += _operand_failures( + "sdpa_decode_gqa_kvarn", + lambda stage_k, stage_v: kq.sdpa_decode_gqa_kvarn( + qk, + codes["codes_k"], + codes["axes_k"], + codes["codes_v"], + codes["axes_v"], + stage_k, + stage_v, + 300, + 0.0884, + 6, + 6, + ), + stages, + [(n, (f"sdpa_decode_gqa {n}",), 4, False) for n in stages], + ) + bad += _operand_failures( + "sdpa_fa_verify_kvarn", + lambda stage_k, stage_v: kq.sdpa_fa_verify_kvarn( + qv, + codes["codes_k"], + codes["axes_k"], + codes["codes_v"], + codes["axes_v"], + stage_k, + stage_v, + 300, + 0.0884, + 6, + 6, + 4, + ), + stages, + [(n, (f"sdpa_fa_verify {n}",), 4, False) for n in stages], + ) + return bad + + +def _sparse(): + """The sparse attention ops: q and the window, pool and key rows.""" + return [b for dt in (mx.float16, mx.bfloat16) for b in _sparse_dtype(dt)] + + +def _sparse_dtype(dt): + from test_dsa_sparse_attn import _build as dsa_build + from test_sdpa_sparse_decode import _case as decode_case + from test_sdpa_sparse_decode import _on_grid_pool + from test_sdpa_sparse_prefill import _case as prefill_case + + bad = [] + (q, win, pool, idx, scale, kw), _ = decode_case( + 1, 16, 2, 128, 24, 300, 40, dt, np.int32, True, True, True + ) + mx.eval(q, win, pool, idx) + names = [(n, (f"sdpa_sparse_decode {n}",), 8, False) for n in ("q", "win", "pool")] + bad += _operand_failures( + "sdpa_sparse_decode", + lambda q, win, pool: kq.sdpa_sparse_decode(q, win, pool, idx, scale, **kw), + {"q": q, "win": win, "pool": pool}, + names, + ) + if dt == mx.float16: + codes, scales = kq.latent_fp4_pack(_on_grid_pool(pool)) + mx.eval(codes, scales) + bad += _operand_failures( + "sdpa_sparse_decode packed", + lambda pool: kq.sdpa_sparse_decode( + q, win, pool, idx, scale, pool_scales=scales, **kw + ), + {"pool": codes}, + [("pool", ("sdpa_sparse_decode pool",), 16, False)], + ) + (q, win, pool, idx, scale, band, kw), _ = prefill_case( + 1, 12, 9, 128, 30, 96, 24, 8, dt, np.int32, True, True, True + ) + mx.eval(q, win, pool, idx) + bad += _operand_failures( + "sdpa_sparse_prefill", + lambda q, win, pool: kq.sdpa_sparse_prefill( + q, win, pool, idx, scale, band, **kw + ), + {"q": q, "win": win, "pool": pool}, + [(n, (f"sdpa_sparse_prefill {n}",), 8, False) for n in ("q", "win", "pool")], + ) + q, local_kv, pooled, topk, sinks, scale = dsa_build( + 1, 128, 640, 512, 2560, 4, 128, dt, 5, "random" + ) + mx.eval(q, local_kv, pooled, topk, sinks) + bad += _operand_failures( + "dsa_sparse_attention", + lambda q, local_kv, pooled: kq.dsa_sparse_attention( + q, local_kv, pooled, topk, sinks, scale, 2560, 4, 128 + ), + {"q": q, "local_kv": local_kv, "pooled": pooled}, + [ + (n, (f"dsa_sparse_attention {n}",), 8, False) + for n in ("q", "local_kv", "pooled") + ], + ) + return bad + + +GROUPS = { + "matmul_ops": _matmul_ops, + "bias_skinny": _bias_and_skinny, + "gathers": _gathers, + "moe": _moe, + "gather_mix": _gather_mix, + "norm_latent_indexer": _norm_latent_indexer, + "hc": _hc, + "sdpa": _sdpa, + "sparse": _sparse, +} +# The matmul groups run under each route environment, the fused LoRA kernel +# (KQ_QMM_LORA=2) with the LoRA matmul; the other ops with NAX on and off. +MATMUL_ENVS = ("default", "nax_qmv_off", "nax_off") +GROUP_ENVS = [ + *((g, e) for g in list(GROUPS)[:5] for e in MATMUL_ENVS), + ("matmul_ops", "lora2"), + *((g, e) for g in list(GROUPS)[5:] for e in ("default", "nax_off")), +] + +_CHILD = r""" +import json, sys +sys.path.insert(0, sys.argv[1]) +import test_x_offset_views as t +print("RESULT " + json.dumps(t.GROUPS[sys.argv[2]]())) +""" +_XA = re.compile(r"^\[kq qmm\] (\S+) .* xa=(\d+)$") +_ALIGN = re.compile(r"^\[kq align\] (.+) xin=(\d+)$") +_ARM = re.compile(r"^\[arm\] (.*) \| (\S+) \| (.*)$") + + +def _unreached(arm, need, names, copies, xa): + """True when an arm that needs copies has not logged them.""" + if need in ("none", "any"): + return False + want = [n for n in names.split(",") if n] + if need == "copy": + return any(n not in copies for n in want) if want else not copies + off = {n for n, xin in copies.items() if xin} + if want: + return any(n not in off for n in want) + return not off and not any(xa) + + +def ops_offsets(group, env): + """(bad, unaligned, unreached, copied) of an op group in a child under + one environment, as quantized_matmul_offsets gives them.""" + child_env = {k: v for k, v in os.environ.items() if not k.startswith("KQ_")} + child_env.update({"KQ_QMM_LOG": "1", **ENVS[env]}) + r = subprocess.run( + [sys.executable, "-c", _CHILD, TESTS_DIR, group], + env=child_env, + capture_output=True, + text=True, + timeout=900, + ) + assert r.returncode == 0, r.stderr[-6000:] + line = [ln for ln in r.stdout.splitlines() if ln.startswith("RESULT ")][-1] + bad = json.loads(line[len("RESULT ") :]) + unaligned, unreached, copied = set(), [], [] + arm = None + for ln in [*r.stderr.splitlines(), "[arm] end | any | "]: + m = _ARM.match(ln) + if m: + if arm and _unreached(*arm): + unreached.append(arm[0]) + if arm and arm[1] == "none" and arm[3]: + copied.append((arm[0], sorted(arm[3]))) + arm = (m.group(1), m.group(2), m.group(3), {}, []) + continue + m = _XA.match(ln) + if m: + if _unaligned_x(m.group(1), int(m.group(2))): + unaligned.add((m.group(1), int(m.group(2)))) + if arm: + arm[4].append(int(m.group(2))) + continue + m = _ALIGN.match(ln) + if m and arm: + arm[3][m.group(1)] = max(arm[3].get(m.group(1), 0), int(m.group(2))) + return bad, sorted(unaligned), unreached, copied + + +@metal_only +@pytest.mark.parametrize("group,env", GROUP_ENVS) +def test_ops_x_offsets(group, env): + """Bit-exact output at every offset, every matmul kernel reads x from a + 16-byte boundary, every view off the boundary reaches its op there and + is copied, and every view on it is read in place.""" + bad, unaligned, unreached, copied = ops_offsets(group, env) + assert bad == [], bad + assert unaligned == [], unaligned + assert unreached == [], unreached[:40] + assert copied == [], copied[:40]