diff --git a/CHANGELOG.md b/CHANGELOG.md index a802678..4badb5f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +## [0.5.2] + ### Changed - `q8_0` `quantized_matmul` at one activation row takes less time on every GPU. `KQ_QMV_SPLITK=0` restores the previous kernels. @@ -19,6 +21,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - On M5 GPUs `quantized_matmul` at one activation row takes up to 26% less time for the iq1, iq2 and iq3 codecs. `KQ_QMV_C32=0` restores the previous kernels. +- `nvfp4` `quantized_matmul` at one activation row takes about a third of the + time, and from 4 to 32 rows on GPUs without NAX it runs on split-K tiles, + about half the time at 8 rows. An `nvfp4` weight must now start on a 4-byte + boundary, as GGUF tensors do. +- `nvfp4` runs on the NAX tensor units of M5 GPUs, for matmuls, MoE gathers + and speculative verify, with small-batch routing tuned for M5. ## [0.5.1] diff --git a/benchmarks/bench_qmm_bm128_ab.py b/benchmarks/bench_qmm_bm128_ab.py index fa1949a..695100a 100644 --- a/benchmarks/bench_qmm_bm128_ab.py +++ b/benchmarks/bench_qmm_bm128_ab.py @@ -17,6 +17,11 @@ Cells: codec x (N,K) shape x M, all M chosen inside the even-ceil(M/64) dispatch window. GB/s = weight-bytes / call. Markdown + JSON out. + +--tile-ab is the source of the large-M K ceilings (`kq_largem_max_k`): the +arms force KQ_NAX_TILE to 128,64,2,2 and 64,64,2,2, which bypasses the +current ceiling and the parity rule, so M may sit at either parity. Run it +at N 4096 and M 4128, 8256 and 16384 over K 4096, 10240, 12288 and 16384. """ import argparse @@ -102,23 +107,26 @@ def call(xx, w=w, s=s, codec=codec): def probe_layout(codec, N, K): - """Synthesize (w, scales) in the codec's quantized layout: IQ codecs - via the tests' synthetic-wire helper (their encoder demands an - imatrix), encodable codecs by random-filling the shape of a probe - quantize.""" + """Synthesize (w, scales) in the codec's quantized layout: the IQ + codecs, stq1_0 and the decode-only codecs via the tests' synthetic-wire + helper (the IQ encoders demand an imatrix, and nvfp4 has no encoder; + its group scales are 1/64 to 1/32), encodable codecs by random-filling + the shape of a probe quantize.""" import numpy as np - if codec.startswith("iq") or codec == "stq1_0": + from mlx_kquant.codec_geometry import DECODE_ONLY_CODECS, geometry + + if codec.startswith("iq") or codec == "stq1_0" or codec in DECODE_ONLY_CODECS: tests_dir = os.path.join( os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tests", ) sys.path.insert(0, tests_dir) - from test_codecs import CODECS, _synth_iq_wire + from kqref import synth_wire - _, wpb, bpb, _, _ = CODECS[codec] + _, _, bpb, wpb = geometry(codec) rng = np.random.default_rng(N + K) - wire = _synth_iq_wire(rng, bpb, N * (K // wpb)) + wire = synth_wire(rng, codec, bpb, N * (K // wpb), nvfp4_scales=(0x08, 0x11)) return ( wire.reshape(N, (K // wpb) * bpb), np.zeros((1,), dtype=np.uint8), @@ -159,18 +167,28 @@ def main(): help="on-arm runs with KQ_NAX_BM128 unset (per-codec " "policy) instead of 1 (force floor 193)", ) + ap.add_argument( + "--tile-ab", + action="store_true", + help="arms force KQ_NAX_TILE=128,64,2,2 and 64,64,2,2 (K ceiling " + "sweep; M at either parity)", + ) ap.add_argument("--json-out") args = ap.parse_args() shapes = [tuple(int(v) for v in s.split("x")) for s in args.shapes.split(",")] for m in args.ms: - assert ((m + 63) // 64) % 2 == 0, f"M={m} outside even-ceil window" + assert args.tile_ab or ((m + 63) // 64) % 2 == 0, ( + f"M={m} outside even-ceil window" + ) tmp = tempfile.mkdtemp(prefix="bm128ab-") env_base = dict(os.environ) def spawn(bm128): env = dict(env_base) - if bm128 is None: # policy mode: leave the tri-state env unset + if args.tile_ab: + env["KQ_NAX_TILE"] = "128,64,2,2" if bm128 != "0" else "64,64,2,2" + elif bm128 is None: # policy mode: leave the tri-state env unset env.pop("KQ_NAX_BM128", None) else: env["KQ_NAX_BM128"] = bm128 diff --git a/benchmarks/bench_verify_routes.py b/benchmarks/bench_verify_routes.py index cefc157..ee56393 100644 --- a/benchmarks/bench_verify_routes.py +++ b/benchmarks/bench_verify_routes.py @@ -20,7 +20,10 @@ odd rounds. Weights are real codec output: --seed-rows float rows are quantized with -kq.quantize and tiled to N. The chain cycles through enough distinct +kq.quantize and tiled to N. The decode-only codecs (mxfp4, nvfp4), which +kq.quantize cannot encode, take --seed-rows rows of synthetic wire from +tests/kqref.py synth_wire instead, with nvfp4 group scales of 1/64 to +1/32, the range of real weights. The chain cycles through enough distinct copies of the weights to exceed --stream-mb, so each call streams its weights from DRAM the way a forward does instead of reading them from the system cache. Before timing, every arm's output is checked against an f32 @@ -84,8 +87,9 @@ "q6_k": 128, "q3_k": 128, "q2_k": 128, + "nvfp4": 128, } -NO_NAX = {"mxfp4", "nvfp4"} +NO_NAX = {"mxfp4"} def applies(route, codec, M, K, nax): @@ -107,13 +111,43 @@ def applies(route, codec, M, K, nax): return False +def synth_weights(codec, rows, K): + """Synthetic wire for a codec kq.quantize cannot encode.""" + import sys + + import mlx.core as mx + import numpy as np + + from mlx_kquant.codec_geometry import geometry + + tests = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tests") + sys.path.insert(0, os.path.abspath(tests)) + from kqref import synth_wire + + _, _, bpb, wpb = geometry(codec) + wire = synth_wire( + np.random.default_rng(K), + codec, + bpb, + rows * (K // wpb), + nvfp4_scales=(0x08, 0x11), + ) + return mx.array(wire.reshape(rows, (K // wpb) * bpb)), mx.zeros((1,), mx.uint8) + + def make_weights(codec, N, K, seed_rows): import mlx.core as mx import mlx_kquant as kq + from mlx_kquant.codec_geometry import DECODE_ONLY_CODECS - wf = mx.random.normal((seed_rows, K), key=mx.random.key(K)).astype(mx.float32) - w, s = kq.quantize(wf, codec) + if codec in DECODE_ONLY_CODECS: + wf = None + w, s = synth_weights(codec, seed_rows, K) + else: + wf = mx.random.normal((seed_rows, K), key=mx.random.key(K)) + wf = wf.astype(mx.float32) + w, s = kq.quantize(wf, codec) mx.eval(w, s) reps = (N + seed_rows - 1) // seed_rows wt = mx.tile(w, (reps, 1))[:N] if w.ndim == 2 else w diff --git a/docs/kernels.md b/docs/kernels.md index 16bfa20..e4e9958 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -67,7 +67,7 @@ ceil(M/64) is even. The k-quants `q3_k` to `q6_k`, the flat block codecs `q4_0`, and `iq4_nl`, `iq3_xxs`, `iq3_s`, `iq2_xs` and `iq2_s` run the BM=64 and BM=128 tiles on a K tile of 128 (the `_BK128` levers). Above the qmv limit, the verify kernels take the widths through M 8 on the codecs that have them. The -NAX verify kernel serves `pq2_0`, `q4_0`, `q8_0` and `q2_k` to `q6_k` from a per-codec entry, +NAX verify kernel serves `pq2_0`, `q4_0`, `q8_0`, `q2_k` to `q6_k` and `nvfp4` from a per-codec entry, and the register-resident MMA verify kernel serves `ptq1_0`, as described at the end of this section. Every limit and floor is a measured per-codec policy (`kq_nax_small_m` and `kq_smallbm_policy` in `src/kquant_matmul.cpp`). @@ -443,7 +443,8 @@ it is the largest K before the first K (of 10240, 12288 and 16384) where BM=128 K 3840 and 4096. On the K tile of 128, `q3_k`, `q5_k`, `q6_k`, `iq3_xxs`, `iq3_s`, `iq2_xs` and `iq2_s` win at every K measured (`iq3_s` 0.91-0.94x, the other IQ codecs 0.93-1.00x), while the flat block codecs lose at K 12288 and 16384: `q4_0` and `q4_1` 1.02-1.06x on the K tile of 128, -`q5_1` and `q8_0` 1.04-1.11x on 64. A ceiling moves only where the new tile takes at most the old +`q5_1` and `q8_0` 1.04-1.11x on 64. `nvfp4` on the K tile of 64 takes 0.89-0.96x over the three M at +every K, the median of three runs. A ceiling moves only where the new tile takes at most the old time at every M measured: on the K tile of 128 at K 12288, `q4_k` and `q5_0` take 1.00-1.02x over the three M and 1.02-1.03x at M 4128, so they keep K 10240. @@ -459,7 +460,7 @@ the three M and 1.02-1.03x at M 4128, so they keep K 10240. | `q2_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | 4096 | inherited, grid class; the K ceiling is that of `pq2_0`, not measured | | `stq1_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | 12288 | inherited, grid class | | `iq2_s`, `iq2_xs`, `iq3_s`, `iq3_xxs` | 128, 64, 2, 2 | 0 | 4096 | 0 | none | inherited, grid class | -| `iq1_s`, `iq1_m`, `iq2_xxs`, `iq4_xs`, `ptq1_0` | 128, 64, 2, 2 | 0 | 4096 | 0 | none | inherited, grid class | +| `iq1_s`, `iq1_m`, `iq2_xxs`, `iq4_xs`, `ptq1_0`, `nvfp4` | 128, 64, 2, 2 | 0 | 4096 | 0 | none | inherited, grid class | `stq1_0` (structured-sparse ternary, llama.cpp PR #22836) ships the full ALU and NAX kernel set. Its pre-NAX floors are measured (M3 Max: plain split-K entry M 5), and so are its NAX small-M routes in @@ -476,8 +477,8 @@ with the float kernels by default and the half kernels under `KQ_VERIFY_MMA_HALF `KQ_QMM_TILE` and the tile levers under `KQ_TUNE_LIVE` apply to it as to every NAX codec. Left out: `quantized_matmul_glu` and `quantized_matmul_lora` have no `q2_0` or `pq2_0` kernel and build the product from `quantized_matmul` and stock ops, and like every codec `q2_0` has no kernel on float -activations (float32 runs the bfloat16 kernels). Their NAX split-K entry is M 9, where the verify kernels hand off. The -`bm128_min_m` and db64 floors are still inherited from `iq1_s`. The `ptq1_0` M=1 kernel gives each +activations (float32 runs the bfloat16 kernels). Their NAX split-K entry is M 9, where the verify kernels hand off. +The `bm128_min_m` and db64 floors are still inherited from `iq1_s`. The `ptq1_0` M=1 kernel gives each 28-byte block four lanes. Each lane reads its six trit bytes in two loads plus the word that holds the high-trit bytes and the scale, and decodes the trits with the exact-float base-3 coefficient collapse, in which floors of the scaled byte value multiply precomputed activation coefficients @@ -493,6 +494,14 @@ Both codecs have a `verify_qmv` sibling that decodes each row's block once and d activation row. It serves M 2 below the `verify_mma` entry on GPUs without NAX, where the `mv_ext` re-decode per row costs `ptq1_0` 3x. On NAX GPUs per-row qmv serves M 2. +`nvfp4` (64-wide blocks of four 16-weight groups) runs its NAX tiles through the same loader as the +Prism codecs and takes the grid class large-M row. It has a `verify_nax` kernel and no `verify_mma` +kernel. Its NAX routes are measured on M5 Max. Per-row qmv keeps only M 2 at N 1025 to 2048, NAX +split-K enters at M 2 for N 1024 and below and at M 6 above, and `verify_nax` enters by shape, as +described with the other `verify_nax` entries. Its `bm128_min_m` floor is 449: BM=128 takes +1.13-1.16x the BM=64 time at M 224, 0.95-0.98x at M 256 and 0.86-0.90x at M 1024. Its large-M row +has no K ceiling, and it has no db64 band. + The register-resident MMA verify kernels (`verify_mma`) serve `pq2_0`, `ptq1_0`, `q2_0`, `q4_0` and `q8_0` through M 8. Each simdgroup decodes a block of its weight rows straight into 8x8 simdgroup-matrix fragments, with the k order inside the block permuted so every lane extracts the code pairs it holds @@ -544,7 +553,7 @@ The `q8_0` kernel has no default entry on either class of GPU, so it runs only w `KQ_VERIFY_MMA` or `KQ_QMM_ROUTE=verify_mma` forces it. On NAX GPUs, the register-fed NAX verify kernels (`verify_nax`) serve `pq2_0`, `q4_0`, `q8_0`, -`q2_k` to `q6_k`, `iq4_xs` and `iq4_nl` from a per-codec entry through M 8, above the per-row qmv limit where that limit is +`q2_k` to `q6_k`, `iq4_xs`, `iq4_nl` and `nvfp4` from a per-codec entry through M 8, above the per-row qmv limit where that limit is wider. Each simdgroup owns 32 weight rows and decodes every block of them straight into the right operand of a 16x32x16 NAX matmul, with the block scale folded into the half weights. The left operand holds the activation rows padded to 16, read from device memory in the k order of the @@ -569,6 +578,7 @@ split rule, as the table lists. | `q2_k` | 128 | 8 bytes of codes and 8 bytes of scales and mins | per NAX op | waves, partial charge 2560 | 4-byte | | `iq4_xs` | 256 | the 8-byte header and a 32-byte quarter of one row's codes, as 8-byte words | per NAX op | waves, partial charge 640 | 8-byte | | `iq4_nl` | 256 | two whole blocks of one row, as nine 4-byte words | per NAX op | waves, partial charge 640 | 4-byte | +| `nvfp4` | 128 | the scale byte and 8 code bytes of one 16-weight group of one row | a K step ahead | waves, partial charge 1280 | 4-byte | The `pq2_0` and two-block `q4_0` kernels decode each block in the permuted k order of `verify_mma`. The other kernels keep each lane's loads on one row, which measured faster. Spreading a `q8_0` @@ -576,10 +586,11 @@ block across 8 rows ran 1.04-1.24x slower at every Qwen3.8-27B shape except N 10 eight-block kernel at a K that is a multiple of 256 and the two-block kernel elsewhere. The K-quant 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. 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. +two-block `q4_0` kernel, and `nvfp4` loads them the same way, 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. 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. @@ -593,13 +604,16 @@ and the `q6_k` head ran 1.06-1.07x slower on one split than on two. The other kernels take their count from a model of GPU waves of 640 simdgroups. A count costs its waves times the K steps each simdgroup walks, plus the simdgroups of all its splits divided by the codec's partial charge. The charge is 640 for `q4_k`, the eight-block `q4_0` kernel, `iq4_xs` and -`iq4_nl`, and 2560 for `q2_k`, `q3_k` and `q5_k`. The cheapest power of two that divides the K steps wins, unless another divisor costs +`iq4_nl`, 1280 for `nvfp4`, and 2560 for `q2_k`, `q3_k` and `q5_k`. The cheapest power of two that divides the K steps wins, unless another divisor costs at most 3/4 as much, and the count has no cap. N 1024, K 5120 therefore takes 20 splits, which ran 1.13-1.23x faster than 8 on `q3_k` and `q2_k`. The `q5_k` vocab head takes 4 splits and runs 1.14-1.32x faster than on 1. Across the Qwen3.8-27B projection and head shapes and N 1024 to 14336 at K 2048 to 14336, timed on M5 Max with bfloat16 activations, the count ran within 1.03x of the best measured count on 164 of 200 cells and 1.13x slower at worst. The sweep ran at M 3 for `q4_0`, -`q4_k` and `q5_k`, M 4 for `q2_k` and `q3_k`, and M 8. +`q4_k` and `q5_k`, M 4 for `q2_k` and `q3_k`, and M 8. On `nvfp4`, at M 3, 5 and 8 over 14 shapes of +N 512 to 248320 and K 2048 to 17408, the count runs at 1.01x the best measured count on average and +1.07x at worst. A target of about 600 simdgroups ran its vocab head on one split, 1.14-1.17x slower +than on 5 to 10. A K that is not a whole number of K steps declines the route, and the K-quant kernels also need K in whole 256-weight superblocks. A two-block `q4_0` K with an odd count of blocks therefore runs on the @@ -669,6 +683,20 @@ 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. +The `nvfp4` entry also moves with the shape, fitted on both activation types over 24 shapes of N 256 +to 248320 and K 2048 to 17408, 8 of them held out: + +| shape | entry, bfloat16 | entry, float16 | time against the default route at the entry | +|---|---|---|---| +| N >= 100000 | M 3 | M 3 | 0.89x on bfloat16, 1.01x on float16 | +| N 512 and under | none | none | NAX split-K ran 1.07-1.42x faster through M 8 | +| K >= 6144, or N 513 to 1024 from K 4096 | M 2 | M 2 | 0.83-1.04x | +| N 1025 to 2048 from K 4096 | M 3 | M 3 | 0.85x on bfloat16, 0.90x on float16 | +| N 513 to 2048 under K 4096 | M 5 | M 5 | 0.99-1.02x | +| N over 2048 under K 6144 | M 4 | M 5 | 0.84-1.05x | + +At M 8 the kernel runs 1.04-1.29x faster than NAX split-K from N 2048, and 1.5x on a vocab head. + ## Quantized matmuls that share an input - `quantized_matmul_multi` - `[x @ W^T for W in weights]` for packed weights of one codec and one diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index 5e992b6..081046c 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -708,6 +708,12 @@ instantiate_kquant_mxfp4_for_type(float16_t) instantiate_kquant_qmm_t(type, 64, 4, true, 1, nvfp4) \ instantiate_kquant_qmm_t(type, 64, 4, false, 0, nvfp4) \ instantiate_kquant_qmm_t(type, 64, 4, false, 1, nvfp4) \ + instantiate_kquant_qmm_t_splitk(type, 64, 4, true, nvfp4) \ + instantiate_kquant_qmm_t_splitk(type, 64, 4, false, nvfp4) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 64, 4, true, nvfp4) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 64, 4, false, nvfp4) \ + instantiate_kquant_qmm_t_splitk_bm8(type, 64, 4, true, nvfp4) \ + instantiate_kquant_qmm_t_splitk_bm8(type, 64, 4, false, nvfp4) \ instantiate_kquant_qmm_n(type, 64, 4, 0, nvfp4) \ instantiate_kquant_qmm_n(type, 64, 4, 1, nvfp4) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 64, 4, nvfp4) \ diff --git a/metal/kq_quantized_nax.metal b/metal/kq_quantized_nax.metal index 9994f5d..acaf1dc 100644 --- a/metal/kq_quantized_nax.metal +++ b/metal/kq_quantized_nax.metal @@ -164,6 +164,7 @@ instantiate_kquant_nax_splitk(iq1_m, 256, 1) instantiate_kquant_nax_splitk(stq1_0, 256, 1) instantiate_kquant_nax_splitk(pq2_0, 128, 2) instantiate_kquant_nax_splitk(q2_0, 64, 2) +instantiate_kquant_nax_splitk(nvfp4, 64, 4) instantiate_kquant_nax_splitk(ptq1_0, 128, 1) // Double-buffered BM=64 qmm_t, name-suffixed _db: dispatched by the host @@ -243,6 +244,7 @@ instantiate_kquant_nax_bm128(iq1_m, 256, 1) instantiate_kquant_nax_bm128(stq1_0, 256, 1) instantiate_kquant_nax_bm128(pq2_0, 128, 2) instantiate_kquant_nax_bm128(q2_0, 64, 2) +instantiate_kquant_nax_bm128(nvfp4, 64, 4) instantiate_kquant_nax_bm128(ptq1_0, 128, 1) // The BM=64 and BM=128 tiles on a K tile of 128 (half super-block or pair // loaders), per codec. KQ_NAX__BK128=0 keeps the BK=64 kernels. @@ -321,6 +323,7 @@ instantiate_kquant_nax_gather_rhs_bm32(iq1_m, 256, 1) instantiate_kquant_nax_gather_rhs_bm32(stq1_0, 256, 1) instantiate_kquant_nax_gather_rhs_bm32(pq2_0, 128, 2) instantiate_kquant_nax_gather_rhs_bm32(q2_0, 64, 2) +instantiate_kquant_nax_gather_rhs_bm32(nvfp4, 64, 4) instantiate_kquant_nax_gather_rhs_bm32(ptq1_0, 128, 1) instantiate_kquant_nax_smallbm(q6_k, 256, 6) @@ -345,6 +348,7 @@ instantiate_kquant_nax_smallbm(iq1_m, 256, 1) instantiate_kquant_nax_smallbm(stq1_0, 256, 1) instantiate_kquant_nax_smallbm(pq2_0, 128, 2) instantiate_kquant_nax_smallbm(q2_0, 64, 2) +instantiate_kquant_nax_smallbm(nvfp4, 64, 4) instantiate_kquant_nax_smallbm(ptq1_0, 128, 1) instantiate_kquant_nax_codec(q8_0, 32, 8) @@ -369,6 +373,7 @@ instantiate_kquant_nax_codec(iq1_m, 256, 1) instantiate_kquant_nax_codec(stq1_0, 256, 1) instantiate_kquant_nax_codec(pq2_0, 128, 2) instantiate_kquant_nax_codec(q2_0, 64, 2) +instantiate_kquant_nax_codec(nvfp4, 64, 4) instantiate_kquant_nax_codec(ptq1_0, 128, 1) // Expert-major sorted gather GEMM (gather_qmm_seg on NAX): BM=64 matches @@ -405,6 +410,7 @@ instantiate_kquant_nax_gather_seg_codec(iq1_m, 256, 1) instantiate_kquant_nax_gather_seg_codec(stq1_0, 256, 1) instantiate_kquant_nax_gather_seg_codec(pq2_0, 128, 2) instantiate_kquant_nax_gather_seg_codec(q2_0, 64, 2) +instantiate_kquant_nax_gather_seg_codec(nvfp4, 64, 4) instantiate_kquant_nax_gather_seg_codec(ptq1_0, 128, 1) // Register-fed NAX verify kernels (kq_verify_nax.h), M <= 8, half and @@ -423,6 +429,8 @@ instantiate_kquant_verify_nax(bfloat16_t, 128, 2, pq2_0) instantiate_kquant_verify_nax(bfloat16_t, 64, 2, q2_0) instantiate_kquant_verify_nax(float16_t, 128, 2, pq2_0) instantiate_kquant_verify_nax(float16_t, 64, 2, q2_0) +instantiate_kquant_verify_nax(bfloat16_t, 64, 4, nvfp4) +instantiate_kquant_verify_nax(float16_t, 64, 4, nvfp4) instantiate_kquant_verify_nax(bfloat16_t, 32, 4, q4_0) instantiate_kquant_verify_nax(float16_t, 32, 4, q4_0) instantiate_kquant_verify_nax(bfloat16_t, 32, 4, q4_0_sb) @@ -460,6 +468,8 @@ 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, 64, 4, nvfp4) +instantiate_kquant_verify_nax_multi(float16_t, 64, 4, nvfp4) 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) diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_fp.h b/metal/mlx/backend/metal/kernels/kq_quantized_fp.h index 1c1e3cd..8aacdfc 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_fp.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_fp.h @@ -565,19 +565,34 @@ template kq_nvfp4_dequantize_impl(w, out, num_weights, gid); } +// A group's 8 qs bytes as two words: the 36-byte superblock keeps them +// 4-byte aligned. +METAL_FUNC uint2 kq_nvfp4_group_qs(const device uint8_t* block, int g) { + const device uint32_t* q = + (const device uint32_t*)(block + KQ_NVFP4_QS_OFFSET + 8 * g); + return uint2(q[0], q[1]); +} + +// E2M1 values of nibble `shift` (0 low, 4 high) of each byte of `q`. +METAL_FUNC float4 kq_nvfp4_nibbles4(uint32_t q, uint shift) { + return float4( + kq_fp_e2m1_lut[(q >> shift) & 0x0F], + kq_fp_e2m1_lut[(q >> (8 + shift)) & 0x0F], + kq_fp_e2m1_lut[(q >> (16 + shift)) & 0x0F], + kq_fp_e2m1_lut[(q >> (24 + shift)) & 0x0F]); +} + // Natural-order chunk [il*16, il*16+16) == sub-scale group il; il in [0,4). METAL_FUNC void kq_nvfp4_deq_chunk16( const device uint8_t* block, short il, thread float4x4& reg) { const float d = kq_fp_ue4m3_scale(block[il]); - const device uint8_t* qs = block + KQ_NVFP4_QS_OFFSET + 8 * il; -#pragma unroll - for (int i = 0; i < 16; ++i) { - const uint8_t b = qs[i % 8]; - const int nib = (i < 8) ? (b & 0x0F) : (b >> 4); - reg[i / 4][i % 4] = d * kq_fp_e2m1_lut[nib]; - } + const uint2 q = kq_nvfp4_group_qs(block, il); + reg[0] = d * kq_nvfp4_nibbles4(q.x, 0); + reg[1] = d * kq_nvfp4_nibbles4(q.y, 0); + reg[2] = d * kq_nvfp4_nibbles4(q.x, 4); + reg[3] = d * kq_nvfp4_nibbles4(q.y, 4); } struct KqNvfp4Ext { @@ -605,8 +620,9 @@ template w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); } -// NVFP4 mat-vec: lane `simd_lid` owns weights `lane` and `lane+32` of every -// 64-superblock (groups 0-1 and 2-3 respectively); simd_sum reduces. +// NVFP4 mat-vec on the MXFP4 lane layout: lane `simd_lid` owns sub-scale +// group simd_lid % 4 of superblock simd_lid / 4, 8 superblocks per pass; x is +// read once into four float4 and reused across the unrolled rows. template METAL_FUNC void kq_nvfp4_qmv_impl( const device uint8_t* w, @@ -620,6 +636,7 @@ METAL_FUNC void kq_nvfp4_qmv_impl( static_assert(group_size == KQ_NVFP4_SUPERBLOCK, "NVFP4 requires gs=64"); static_assert(bits == 4, "NVFP4 requires bits=4"); constexpr int num_simdgroups = 2; + constexpr int blocks_per_pass = 8; typedef float U; const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; @@ -632,14 +649,15 @@ METAL_FUNC void kq_nvfp4_qmv_impl( const int nb = in_vec_size / KQ_NVFP4_SUPERBLOCK; x += tid.x * in_vec_size; y += tid.x * out_vec_size; - const int g0 = int(simd_lid) / 16; // 0 or 1 - const int j = int(simd_lid) % 16; - const bool is_high = j >= 8; - const int byteidx = 8 * g0 + (j % 8); + const int ix = simd_lid / 4; + const int g = simd_lid % 4; U result[results_per_simdgroup] = {0}; - for (int ib = 0; ib < nb; ib++) { - const U xv0 = U(x[ib * KQ_NVFP4_SUPERBLOCK + simd_lid]); - const U xv1 = U(x[ib * KQ_NVFP4_SUPERBLOCK + 32 + simd_lid]); + for (int ib = ix; ib < nb; ib += blocks_per_pass) { + const device T* xb = x + ib * KQ_NVFP4_SUPERBLOCK + 16 * g; + const float4 yl0(xb[0], xb[1], xb[2], xb[3]); + const float4 yl1(xb[4], xb[5], xb[6], xb[7]); + const float4 yl2(xb[8], xb[9], xb[10], xb[11]); + const float4 yl3(xb[12], xb[13], xb[14], xb[15]); // static trip count; the clamped tail row is dropped at the store #pragma unroll for (int row = 0; row < results_per_simdgroup; row++) { @@ -647,14 +665,13 @@ METAL_FUNC void kq_nvfp4_qmv_impl( static_cast(min(out_row + row, out_vec_size - 1)) * row_bytes + ib * KQ_NVFP4_BLOCK_BYTES; - const U d0 = U(kq_fp_ue4m3_scale(blk[g0])); - const U d1 = U(kq_fp_ue4m3_scale(blk[g0 + 2])); - const uint8_t b0 = blk[KQ_NVFP4_QS_OFFSET + byteidx]; - const uint8_t b1 = blk[KQ_NVFP4_QS_OFFSET + 16 + byteidx]; - const int nib0 = is_high ? (b0 >> 4) : (b0 & 0x0F); - const int nib1 = is_high ? (b1 >> 4) : (b1 & 0x0F); - result[row] += d0 * U(kq_fp_e2m1_lut[nib0]) * xv0 + - d1 * U(kq_fp_e2m1_lut[nib1]) * xv1; + const U d = U(kq_fp_ue4m3_scale(blk[g])); + const uint2 q = kq_nvfp4_group_qs(blk, g); + float4 acc = yl0 * kq_nvfp4_nibbles4(q.x, 0); + acc += yl1 * kq_nvfp4_nibbles4(q.y, 0); + acc += yl2 * kq_nvfp4_nibbles4(q.x, 4); + acc += yl3 * kq_nvfp4_nibbles4(q.y, 4); + result[row] += d * ((acc.x + acc.y) + (acc.z + acc.w)); } } for (int row = 0; row < results_per_simdgroup; row++) { @@ -739,15 +756,36 @@ struct KqNvfp4BlockLoader { void load_unsafe() const thread { const short sb = (reduction_dim == 0) ? fixed_sub_block_idx : sub_block_idx; -#pragma unroll - for (short i = 0; i < n_reads; i++) { - const short t = sb * sub_block_size + bj + i; // weight in superblock - const short g = t / 16; - const short j = t % 16; + if constexpr (reduction_dim == 1 && (n_reads == 16 || n_reads == 8)) { + // The thread's run is one group (16) or one nibble half of it (8): + // one scale and the group's two qs words. + const short t0 = sb * sub_block_size + bj; + const short g = t0 / 16; const float d = kq_fp_ue4m3_scale(src[g]); - const uint8_t b = src[KQ_NVFP4_QS_OFFSET + 8 * g + (j % 8)]; - const int nib = (j < 8) ? (b & 0x0F) : (b >> 4); - dst[i] = T(d * kq_fp_e2m1_lut[nib]); + const uint2 q = kq_nvfp4_group_qs(src, g); + const uint shift = (n_reads == 8 && (t0 % 16) != 0) ? 4 : 0; + float4 v[n_reads / 4]; + v[0] = d * kq_nvfp4_nibbles4(q.x, shift); + v[1] = d * kq_nvfp4_nibbles4(q.y, shift); + if constexpr (n_reads == 16) { + v[2] = d * kq_nvfp4_nibbles4(q.x, 4); + v[3] = d * kq_nvfp4_nibbles4(q.y, 4); + } +#pragma unroll + for (short i = 0; i < n_reads; i++) { + dst[i] = T(v[i / 4][i % 4]); + } + } else { +#pragma unroll + for (short i = 0; i < n_reads; i++) { + const short t = sb * sub_block_size + bj + i; // weight in superblock + const short g = t / 16; + const short j = t % 16; + const float d = kq_fp_ue4m3_scale(src[g]); + const uint8_t b = src[KQ_NVFP4_QS_OFFSET + 8 * g + (j % 8)]; + const int nib = (j < 8) ? (b & 0x0F) : (b >> 4); + dst[i] = T(d * kq_fp_e2m1_lut[nib]); + } } } @@ -830,6 +868,58 @@ template < w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } +template < + typename T, + int group_size, + int bits, + bool aligned_N, + int small_bm = 0, + typename U = T> +[[kernel]] void kq_nvfp4_qmm_t_splitk( + const device uint8_t* w, + const device uint8_t* /* scales */, + 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 constant int& split_k_partition_stride, + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(group_size == KQ_NVFP4_SUPERBLOCK, "NVFP4 requires gs=64"); + static_assert(bits == 4, "NVFP4 requires bits=4"); + constexpr int BM = small_bm ? small_bm : 32; + constexpr int BK = 32, BN = small_bm ? 64 : 32; + constexpr int WM = BM == 8 ? 1 : 2, WN = 4 / WM; + constexpr int BK_padded = (BK + 16 / sizeof(T)); + threadgroup T Xs[BM * BK_padded]; + threadgroup T Ws[BN * BK_padded]; + using LoaderW = + KqNvfp4BlockLoader; + const int k_start = tid.z * k_partition_size; + x += k_start; + auto wl = w; + wl += (k_start / LoaderW::weights_per_block) * LoaderW::bytes_per_block; + y += tid.z * static_cast(split_k_partition_stride); + kq_qmm_t_impl( + wl, + x, + y, + Xs, + Ws, + K, + N, + M, + k_partition_size, + tid, + lid, + simd_gid, + simd_lid); +} + template [[kernel]] void kq_nvfp4_qmm_n( const device uint8_t* w, diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h index b9c15b0..9a30419 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h @@ -3893,7 +3893,7 @@ struct KqNaxStq1_0BlockLoader { } }; -// Prism codecs (128-weight blocks) and q2_0 (64): a 32-weight k-tile is +// Prism codecs (128-weight blocks), q2_0 and nvfp4 (64): a 32-weight k-tile is // chunks 2s and 2s+1 of deq_chunk16; four or two k-tiles per block, two per // 64-wide BCOLS step. template < @@ -4042,6 +4042,22 @@ using KqNaxPtq1_0BlockLoader = KqNaxPrismBlockLoader< reduction_dim, tgp_size>; +template < + typename T, + short BROWS, + short BCOLS, + short dst_ld, + short reduction_dim, + short tgp_size> +using KqNaxNvfp4BlockLoader = KqNaxPrismBlockLoader< + KqNvfp4Ext, + T, + BROWS, + BCOLS, + dst_ld, + reduction_dim, + tgp_size>; + KQ_NAX_DEFINE_KERNELS(iq4_nl, 32, 4, KqNaxIq4_nlBlockLoader) KQ_NAX_DEFINE_KERNELS(iq4_xs, 256, 4, KqNaxIq4_xsBlockLoader) KQ_NAX_DEFINE_KERNELS(iq3_xxs, 256, 3, KqNaxIq3_xxsBlockLoader) @@ -4054,6 +4070,7 @@ KQ_NAX_DEFINE_KERNELS(iq1_m, 256, 1, KqNaxIq1_mBlockLoader) KQ_NAX_DEFINE_KERNELS(stq1_0, 256, 1, KqNaxStq1_0BlockLoader) KQ_NAX_DEFINE_KERNELS(pq2_0, 128, 2, KqNaxPq2_0BlockLoader) KQ_NAX_DEFINE_KERNELS(q2_0, 64, 2, KqNaxQ2_0BlockLoader) +KQ_NAX_DEFINE_KERNELS(nvfp4, 64, 4, KqNaxNvfp4BlockLoader) KQ_NAX_DEFINE_KERNELS(ptq1_0, 128, 1, KqNaxPtq1_0BlockLoader) KQ_NAX_DEFINE_KERNELS(q4_0, 32, 4, KqNaxQ4_0BlockLoader) KQ_NAX_DEFINE_KERNELS(q4_1, 32, 4, KqNaxQ4_1BlockLoader) @@ -4314,6 +4331,7 @@ KQ_NAX_DEFINE_SPLITK_KERNEL(iq1_m, 256, 1, KqNaxIq1_mBlockLoader) KQ_NAX_DEFINE_SPLITK_KERNEL(stq1_0, 256, 1, KqNaxStq1_0BlockLoader) KQ_NAX_DEFINE_SPLITK_KERNEL(pq2_0, 128, 2, KqNaxPq2_0BlockLoader) KQ_NAX_DEFINE_SPLITK_KERNEL(q2_0, 64, 2, KqNaxQ2_0BlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(nvfp4, 64, 4, KqNaxNvfp4BlockLoader) KQ_NAX_DEFINE_SPLITK_KERNEL(ptq1_0, 128, 1, KqNaxPtq1_0BlockLoader) template < @@ -4608,6 +4626,7 @@ KQ_NAX_DEFINE_GATHER_RHS(iq1_m, 256, 1, KqNaxIq1_mBlockLoader) KQ_NAX_DEFINE_GATHER_RHS(stq1_0, 256, 1, KqNaxStq1_0BlockLoader) KQ_NAX_DEFINE_GATHER_RHS(pq2_0, 128, 2, KqNaxPq2_0BlockLoader) KQ_NAX_DEFINE_GATHER_RHS(q2_0, 64, 2, KqNaxQ2_0BlockLoader) +KQ_NAX_DEFINE_GATHER_RHS(nvfp4, 64, 4, KqNaxNvfp4BlockLoader) KQ_NAX_DEFINE_GATHER_RHS(ptq1_0, 128, 1, KqNaxPtq1_0BlockLoader) KQ_NAX_DEFINE_GATHER_RHS(q4_0, 32, 4, KqNaxQ4_0BlockLoader) KQ_NAX_DEFINE_GATHER_RHS(q4_1, 32, 4, KqNaxQ4_1BlockLoader) @@ -4861,6 +4880,7 @@ KQ_NAX_DEFINE_GATHER_SEG(iq1_m, 256, 1, KqNaxIq1_mBlockLoader) KQ_NAX_DEFINE_GATHER_SEG(stq1_0, 256, 1, KqNaxStq1_0BlockLoader) KQ_NAX_DEFINE_GATHER_SEG(pq2_0, 128, 2, KqNaxPq2_0BlockLoader) KQ_NAX_DEFINE_GATHER_SEG(q2_0, 64, 2, KqNaxQ2_0BlockLoader) +KQ_NAX_DEFINE_GATHER_SEG(nvfp4, 64, 4, KqNaxNvfp4BlockLoader) KQ_NAX_DEFINE_GATHER_SEG(ptq1_0, 128, 1, KqNaxPtq1_0BlockLoader) KQ_NAX_DEFINE_GATHER_SEG(q4_0, 32, 4, KqNaxQ4_0BlockLoader) KQ_NAX_DEFINE_GATHER_SEG(q4_1, 32, 4, KqNaxQ4_1BlockLoader) diff --git a/metal/mlx/backend/metal/kernels/kq_verify_nax.h b/metal/mlx/backend/metal/kernels/kq_verify_nax.h index 4efae33..4288f3d 100644 --- a/metal/mlx/backend/metal/kernels/kq_verify_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_verify_nax.h @@ -164,6 +164,50 @@ struct KqQ2_0Nax { } }; +// NVFP4: lane-quad q owns 16-weight group q, its UE4M3 scale byte q and qs +// bytes 8q..8q+7. Pair f is byte f's low and high nibbles, weights f and +// f + 8 of the group, so step s reads Q2_0's activation order. An E2M1 code +// e1 e0 m in half bits 11..9 is its value times 2^-14; every product with +// the scale is exact in half. Two 64-weight blocks per iteration. +struct KqNvfp4Nax { + static constant constexpr int group = KQ_NVFP4_SUPERBLOCK; + static constant constexpr int block_k = KQ_NVFP4_SUPERBLOCK; + static constant constexpr int block_bytes = KQ_NVFP4_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[2]; + half d; + }; + static METAL_FUNC Words load(const device uint8_t* row, int u, short q) { + const device uint8_t* bp = row + u * block_bytes; + Words o; + o.d = half(kq_fp_ue4m3_scale(bp[q])); + const device uint* qs = + (const device uint*)(bp + KQ_NVFP4_QS_OFFSET + 8 * q); + o.w[0] = qs[0]; + o.w[1] = qs[1]; + return o; + } + static METAL_FUNC half2 pair(thread const Words& o, short f) { + const uint b = (o.w[f >> 2] >> (8 * (f & 3))) & 0xFFu; + const uint n = (b & 0x0Fu) | ((b & 0xF0u) << 12); + const uint bits = ((n & 0x00070007u) << 9) | ((n & 0x00080008u) << 12); + return as_type(bits) * half2(16384.0h) * o.d; + } + template + static METAL_FUNC vec xstep(const device T* xb, short q, short s) { + const int o = 16 * q + 2 * s; + const vec p0 = *(const device vec*)(xb + o); + const vec p1 = *(const device vec*)(xb + o + 8); + return vec(p0.x, p1.x, p0.y, p1.y); + } +}; + // Q4_0: lane-quad q owns qs bytes 4q..4q+3. Pairs (n0, n2), (n1, n3), // (h0, h2), (h1, h3) of its low and high nibbles (KqQ4_0Mma), as // 1024 + code in the half mantissa, minus 1032, times d. Two 32-weight @@ -1067,6 +1111,7 @@ METAL_FUNC void kq_verify_nax_impl( KQ_DEFINE_VERIFY_NAX_KERNEL(pq2_0, KqPq2_0Nax) KQ_DEFINE_VERIFY_NAX_KERNEL(q2_0, KqQ2_0Nax) +KQ_DEFINE_VERIFY_NAX_KERNEL(nvfp4, KqNvfp4Nax) KQ_DEFINE_VERIFY_NAX_KERNEL(q4_0, KqQ4_0Nax) KQ_DEFINE_VERIFY_NAX_KERNEL(q4_0_sb, KqQ4_0SbNax) KQ_DEFINE_VERIFY_NAX_KERNEL(q8_0, KqQ8_0Nax) @@ -1080,6 +1125,7 @@ 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(nvfp4, KqNvfp4Nax) 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) diff --git a/mlx_kquant/_version.py b/mlx_kquant/_version.py index dd9b22c..7225152 100644 --- a/mlx_kquant/_version.py +++ b/mlx_kquant/_version.py @@ -1 +1 @@ -__version__ = "0.5.1" +__version__ = "0.5.2" diff --git a/src/kquant_codec.cpp b/src/kquant_codec.cpp index f40d251..bec556a 100644 --- a/src/kquant_codec.cpp +++ b/src/kquant_codec.cpp @@ -35,7 +35,7 @@ static const std::vector& registry() { {"ptq1_0", 128, 28, 1, true, true, true, false}, {"q2_0", 64, 18, 2, true, true, true, false}, {"mxfp4", 32, 17, 4, true, false, false, false}, - {"nvfp4", 64, 36, 4, true, false, false, false}, + {"nvfp4", 64, 36, 4, true, false, true, false}, }; return codecs; } diff --git a/src/kquant_internal.h b/src/kquant_internal.h index e0ee98d..d315428 100644 --- a/src/kquant_internal.h +++ b/src/kquant_internal.h @@ -81,7 +81,7 @@ inline std::string kq_kname_prefix(const std::string& kquant_type) { // Byte alignment the GPU kernels need at the start of a packed weight: the // widest load any kernel of the codec makes through the weight pointer. // q4_k and q5_k read 16-byte words, iq4_xs and iq1_m 8-byte words, and the -// codecs at 4 read 4-byte words. mxfp4 and nvfp4 read bytes. A start off +// codecs at 4 read 4-byte words. mxfp4 reads bytes. A start off // the value is undefined for Metal loads, and on M5 Max the kernels return // wrong values or NaN from every odd start, and from starts 2 mod 4 on the // codecs at 4 or more. The CPU kernels accept any start. Each value divides @@ -95,10 +95,11 @@ inline int kq_weight_base_align(const std::string& kquant_type) { if (t == "iq4_xs" || t == "iq1_m") { return 8; } - if (t == "q2_k" || t == "q4_1" || t == "q5_1" || t == "ptq1_0") { + if (t == "q2_k" || t == "q4_1" || t == "q5_1" || t == "ptq1_0" || + t == "nvfp4") { return 4; } - if (t == "mxfp4" || t == "nvfp4") { + if (t == "mxfp4") { return 1; } return 2; diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 3933df8..2f771cf 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -168,6 +168,14 @@ static KqSmallBmPolicy kq_smallbm_base(const std::string& t) { // bench_qmm_bm128_ab. return {true, 0, 0, 449}; } + if (t == "nvfp4") { + // Measured (median of three ABBA runs): BM=128 takes 1.13-1.16x the + // BM=64 time at M224, 0.95-0.98x at M256, 0.90-1.01x at M512 and + // 0.86-0.90x at M1024. bm32 keeps the vocab head at M 13-32 on the + // un-split BM=32 tile, 1.12-1.19x faster than NAX split-K there. The + // other routes beat the un-split tile through M 12, so no route_min. + return {true, 0, 0, 449}; + } return {false, 0, 0, 193}; // iq2_xs, iq2_s, iq1_m } @@ -191,8 +199,10 @@ static KqLoaderClass kq_loader_class(const std::string& t) { // 4128, 8256 and 16384 at N 4096, K 10240, 12288 and 16384; K 3840 and // 4096 win for every codec). 0: no loss up to K 16384. On the BK=128 tiles // q3_k, q5_k, q6_k, iq2_s, iq2_xs, iq3_s and iq3_xxs win at every K, and so -// do iq1_s, iq1_m, iq2_xxs, iq4_xs and ptq1_0 on BK=64. A ceiling holds for -// the codec's default K tile. q2_0 takes the pq2_0 ceiling; it is not +// do iq1_s, iq1_m, iq2_xxs, iq4_xs, ptq1_0 and nvfp4 on BK=64 (nvfp4 +// 0.89-0.96x at every K from 4096 to 16384, median of three runs, with +// benchmarks/bench_qmm_bm128_ab.py --tile-ab). A ceiling holds for the +// codec's default K tile. q2_0 takes the pq2_0 ceiling; it is not // measured. static int kq_largem_max_k(const std::string& t) { if (t == "q2_k" || t == "pq2_0" || t == "q2_0") { @@ -294,7 +304,7 @@ static bool kq_splitk_codec(const std::string& t) { t == "iq3_xxs" || t == "iq3_s" || t == "iq2_xxs" || t == "iq2_xs" || t == "iq2_s" || t == "iq1_s" || t == "iq1_m" || t == "stq1_0" || t == "pq2_0" || t == "q2_0" || t == "ptq1_0" || t == "q4_0" || - t == "q4_1" || t == "q5_0" || t == "q5_1"; + t == "q4_1" || t == "q5_0" || t == "q5_1" || t == "nvfp4"; } // Non-NAX default split-K entry M per codec (0 = env lever only). @@ -440,6 +450,19 @@ static int kq_verify_mma_min_m_nax(const std::string& t) { // up one row: at the bf16 entries the kernel ran up to 1.10x slower on K // 4096 and 1.21x at N 12288, K 2560. A vocab head keeps M 3 on both types // (0.97x). +// +// nvfp4 was fitted the same way on bf16 and f16 over 24 shapes of N 256 to +// 248320 and K 2048 to 17408, against the routes of its kq_nax_small_m row. +// Times below are of the kernel against the default route it displaces at +// the entry. A vocab head enters at M 3 (0.89x on bf16, 1.01x on f16). N +// 512 and below has no entry: NAX split-K runs 1.07-1.42x faster there +// through M 8. K 6144 and wider, and N 1024 from K 4096, enter at M 2 +// (0.83-1.04x; N 4096, K 8192 runs 1.04x slower at M 2 on f16). N 1025 to +// 2048 enters at M 3 from K 4096 (0.85-0.90x), and N 513 to 2048 at M 5 +// below K 4096 (0.99-1.02x). Wider N under K 6144 enters at M 4 on bf16 +// and M 5 on f16 (0.84-1.05x; N 6144, K 5120 runs 1.05x slower at M 4 on +// bf16). At M 8 the kernel takes 0.78-0.96x the time of NAX split-K from N +// 2048, and 0.67x on a vocab head. static int kq_verify_nax_min_m(const std::string& t, int N, int K, bool f16) { if (t == "iq4_xs" || t == "iq4_nl") { if (N >= 100000) { @@ -455,6 +478,21 @@ static int kq_verify_nax_min_m(const std::string& t, int N, int K, bool f16) { t == "q5_k") { return 3; } + if (t == "nvfp4") { + if (N >= 100000) { + return 3; + } + if (N <= 512) { + return 0; + } + if (K >= 6144 || (N <= 1024 && K >= 4096)) { + return 2; + } + if (N <= 2048) { + return K >= 4096 ? 3 : 5; + } + return f16 ? 5 : 4; + } if (t == "q8_0") { return 6; } @@ -475,7 +513,7 @@ static int kq_splitk_min_m(const std::string& t) { return kq_splitk_min_m_nax_alu(t); } if (t == "q4_k" || t == "q6_k" || t == "q2_k" || t == "iq2_xxs" || - t == "q4_1" || t == "q5_1") { + t == "q4_1" || t == "q5_1" || t == "nvfp4") { return 4; } if (t == "q5_k" || t == "q8_0" || t == "iq4_nl" || t == "iq3_xxs" || @@ -508,7 +546,7 @@ static int kq_splitk_min_m(const std::string& t) { // env lever only). The un-split BM=32 grid is ceil(N/64) threadgroups // with a serial in-tile K walk, and split-K multiplies the threadgroup // count. Between qmv and split-K, the verify routes (verify_nax for -// pq2_0, q4_0, q8_0 and q2_k to q6_k, verify_mma for ptq1_0) keep +// pq2_0, q4_0, q8_0, q2_k to q6_k and nvfp4, verify_mma for ptq1_0) keep // their own entries through M 8, so those codecs' split-K entries take // only the widths above M 8 and the calls verify_nax declines. Other // codecs run the mat-vec kernels (mv_ext, verify_qmv) there, or the BM=32 @@ -528,7 +566,11 @@ static int kq_splitk_min_m(const std::string& t) { // (1.1-3x). Above N 1024, iq2_xxs keeps the mat-vec kernels through // M 12, where its split-K tile runs 1.2-3.6x slower. These figures cover // the table's own routes. verify_nax takes its band from them where it -// measured faster. +// measured faster. The nvfp4 row was fitted the same way on bf16 and f16 +// over 16 shapes (N 256 to 248320, K 2048 to 17408) and held out on 8 (N +// 512 to 14336, K 2048 to 14336), with its verify_nax entries. It +// runs at 1.006x the fastest route on the fitted cells and 1.002x on the +// held-out ones, against 1.065x and 1.109x for the q2_0 row it replaces. struct KqNaxSmallM { int qmv_m[4]; int splitk_small; @@ -596,6 +638,13 @@ static KqNaxSmallM kq_nax_small_m(const std::string& t) { if (t == "pq2_0" || t == "q2_0" || t == "ptq1_0") { return {{8, 5, 3, 2}, 9, 9}; } + if (t == "nvfp4") { + // Per-row qmv keeps M 2 only for N 1025 to 2048. At N 1024 it runs + // 1.14-1.26x faster than the table's route at K 2048, M 2 to 4, but + // 1.1-1.4x slower than verify_nax from K 4096, so that bucket keeps it + // at M 1. + return {{1, 1, 2, 1}, 2, 6}; + } return {{1, 1, 1, 1}, 0, 0}; } @@ -2503,12 +2552,21 @@ static int kq_verify_nax_wave_splits(int sgs, int steps, double partial_c) { // and 1.65x. It cut the verify_nax matmul time of one Qwen3.8-27B forward // by 2.4-2.8% on q5_k and 0.05-0.4% on q4_0 and q4_k. q6_k and q8_0 keep // their simdgroup targets, which beat the wave count on that matmul time -// by 0.3-0.9%, the largest share at N 17408, K 5120. +// by 0.3-0.9%, the largest share at N 17408, K 5120. nvfp4 fits best at +// 1280 (M5 Max, bfloat16, M 3, 5 and 8 over 14 shapes of N 512 to 248320 +// and K 2048 to 17408): 1.01x the best measured count on average and 1.07x +// at worst, against 1.04x and 1.20x for the simdgroup target, which left +// the vocab head unsplit at 1.14-1.17x the time of 5 to 10 splits. 640 +// takes N 17408, K 5120 unsplit (1.08x) and 2560 takes N 4096, K 8192 at +// 32 splits (1.08-1.10x). static double kq_verify_nax_partial_c(const std::string& t, int kstep) { if (t == "q4_k" || (t == "q4_0" && kstep == 256) || t == "iq4_xs" || t == "iq4_nl") { return 640.0; } + if (t == "nvfp4") { + return 1280.0; + } if (t == "q2_k" || t == "q3_k" || t == "q5_k") { return 2560.0; } diff --git a/src/kquant_metal_internal.h b/src/kquant_metal_internal.h index e516875..ef66076 100644 --- a/src/kquant_metal_internal.h +++ b/src/kquant_metal_internal.h @@ -225,7 +225,7 @@ inline int codec_verify_mma_rows(const std::string& kquant_type) { // where K is a multiple of 256 (kq_verify_nax_call_kstep). inline int codec_verify_nax_kstep(const std::string& kquant_type) { if (kquant_type == "pq2_0" || kquant_type == "q2_0" || - kquant_type == "q8_0") { + kquant_type == "q8_0" || kquant_type == "nvfp4") { return 128; } if (kquant_type == "q4_0") { diff --git a/tests/test_nax.py b/tests/test_nax.py index ce7192e..fa8bc19 100644 --- a/tests/test_nax.py +++ b/tests/test_nax.py @@ -61,6 +61,7 @@ "stq1_0": 1.0e-2, "pq2_0": 1.0e-2, "q2_0": 1.0e-2, + "nvfp4": 1.0e-2, "ptq1_0": 1.0e-2, } diff --git a/tests/test_nax_bm128_rows.py b/tests/test_nax_bm128_rows.py index 69ccba6..4c4004b 100644 --- a/tests/test_nax_bm128_rows.py +++ b/tests/test_nax_bm128_rows.py @@ -40,7 +40,7 @@ "q2_k", "q3_k", "q4_k", "q5_k", "q6_k", "q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq1_s", "iq1_m", "iq2_xxs", "iq2_xs", "iq2_s", "iq3_xxs", "iq3_s", - "iq4_nl", "iq4_xs", "stq1_0", "pq2_0", "ptq1_0", "q2_0", + "iq4_nl", "iq4_xs", "stq1_0", "pq2_0", "ptq1_0", "q2_0", "nvfp4", ] ) # fmt: skip @@ -163,7 +163,7 @@ def test_large_rows_k_ceiling(codec): # Codecs whose large-M row has no K ceiling, on BK=64 and on BK=128 tiles. -NO_MAX_K = ["q6_k", "iq2_s", "iq2_xs", "iq3_s", "iq3_xxs", "iq4_xs"] +NO_MAX_K = ["q6_k", "iq2_s", "iq2_xs", "iq3_s", "iq3_xxs", "iq4_xs", "nvfp4"] @pytest.mark.parametrize("codec", NO_MAX_K) diff --git a/tests/test_nax_smallm.py b/tests/test_nax_smallm.py index 0687e5f..3fe41d2 100644 --- a/tests/test_nax_smallm.py +++ b/tests/test_nax_smallm.py @@ -105,12 +105,12 @@ def _encodable_setup(codec, n_out): return w, s, ref_w -def _iq_setup(codec, n_out): +def _iq_setup(codec, n_out, **synth): from kqref import quants gtype, wpb, bpb, _, _ = CODECS[codec] rng = np.random.default_rng(7) - wire = synth_wire(rng, codec, bpb, n_out * (K // wpb)) + wire = synth_wire(rng, codec, bpb, n_out * (K // wpb), **synth) wire = wire.reshape(n_out, (K // wpb) * bpb) ref = quants.dequantize(np.ascontiguousarray(wire), gtype) ref_w = mx.array(ref.astype(np.float32)).T @@ -528,10 +528,12 @@ def test_qmm_route_probe(codec, route, monkeypatch): VMMA_CODECS = ["pq2_0", "ptq1_0", "q2_0", "q4_0", "q8_0"] # verify_nax entries on NAX GPUs below N 100000 (kq_verify_nax_min_m). -# q2_k enters one row earlier above N 4096. +# q2_k enters one row earlier above N 4096. The nvfp4 entry moves with the +# shape (_nvfp4_vnax_entry); listed here at N 1000, K 1024. VNAX_ENTRY = { "pq2_0": 3, "q2_0": 3, + "nvfp4": 5, "q4_0": 3, "q8_0": 6, "q4_k": 3, @@ -546,15 +548,30 @@ def test_qmm_route_probe(codec, route, monkeypatch): VNAX_IQ4 = ["iq4_xs", "iq4_nl"] -def _vnax_entry(codec, n): +def _nvfp4_vnax_entry(n, k, f16): + """kq_verify_nax_min_m for nvfp4; 9 = no entry.""" + if n >= 100000: + return 3 + if n <= 512: + return 9 + if k >= 6144 or (n <= 1024 and k >= 4096): + return 2 + if n <= 2048: + return 3 if k >= 4096 else 5 + return 5 if f16 else 4 + + +def _vnax_entry(codec, n, k=K, f16=False): + if codec == "nvfp4": + return _nvfp4_vnax_entry(n, k, f16) e = VNAX_ENTRY.get(codec, 9) return e - 1 if n > VNAX_WIDE_N.get(codec, n) else e -def _setup(codec, n_out): +def _setup(codec, n_out, **synth): if codec in ENCODABLE: return _encodable_setup(codec, n_out) - return _iq_setup(codec, n_out) + return _iq_setup(codec, n_out, **synth) def _route_serves(route, codec, m): @@ -1307,6 +1324,56 @@ def test_verify_nax_iq4_entry(codec, n_out, k, entry, dtype, monkeypatch): assert _bits_equal(y_def, y_nax), tag +# The nvfp4 entry by shape (kq_verify_nax_min_m): M 3 on a vocab head, +# none at N 512 and below, M 2 at K 6144 and wider and at N 1024 and below +# from K 4096, M 3 up to N 2048 from K 4096 and M 5 below it, and M 4 on +# bfloat16 and M 5 on float16 at wider N. Checked as the iq4 entry is, with +# per-row qmv off, which takes M 2 at N 1025 to 2048 by default. +VNAX_NVFP4_SHAPES = [ + (100000, 256), + (100000, 6144), + (99999, 256), + (512, 8192), + (513, 8192), + (2560, 6144), + (2560, 6016), + (1024, 4096), + (1025, 4096), + (1024, 3968), + (2048, 4096), + (2049, 4096), + (2048, 3968), +] + + +@pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) +@pytest.mark.parametrize("n_out, k", VNAX_NVFP4_SHAPES) +def test_verify_nax_nvfp4_entry(n_out, k, dtype, monkeypatch): + _clear_levers(monkeypatch) + monkeypatch.setenv("KQ_NAX_QMV", "0") + f16 = dtype == mx.float16 + entry = _nvfp4_vnax_entry(n_out, k, f16) + rng = np.random.default_rng(3) + seed = synth_wire(rng, "nvfp4", 36, 64 * (k // 64), nvfp4_scales=(0x08, 0x11)) + seed = seed.reshape(64, -1) + w = mx.contiguous(mx.tile(mx.array(seed), (n_out // 64 + 1, 1))[:n_out]) + s = mx.zeros((1,), dtype=mx.uint8) + mx.eval(w) + ms = (2, 3, 4, 5, 8) + name = "float16" if f16 else "bfloat16" + calls = [{"codec": "nvfp4", "shape": [m, k], "n": n_out, "dtype": name} for m in ms] + routes = [_route_of(e) for e in _logged(calls, {"KQ_NAX_QMV": "0"})] + for m, route in zip(ms, routes, strict=True): + tag = f"nvfp4 N{n_out} K{k} M{m} entry {entry}" + assert (route == "verify_nax") == (entry <= m), (tag, route) + if entry <= m: + x = (mx.random.normal((m, k)) * 0.5).astype(dtype) + y_def = _run_route(monkeypatch, x, w, s, "nvfp4") + y_nax = _run_route(monkeypatch, x, w, s, "nvfp4", "verify_nax") + assert _bits_equal(y_def, y_nax), tag + + # KQ_DISABLE_NAX turns the route off with the other NAX routes. @pytest.mark.skipif(not kq.nax_available(), reason="NAX verify only") @pytest.mark.parametrize("codec", VNAX_CODECS + VNAX_IQ4) @@ -1407,7 +1474,9 @@ def test_verify_nax_x16_bits(codec): @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16]) @pytest.mark.parametrize("codec", VNAX_CODECS) def test_verify_nax_outlier_channel(codec, dtype, monkeypatch): - w, s, ref_w = _setup(codec, 1000) + # nvfp4's default synthetic scales (0.5 to 2) put the exact outputs past + # the float16 range; 1/64 to 1/32 matches the other codecs' weights. + w, s, ref_w = _setup(codec, 1000, nvfp4_scales=(0x08, 0x11)) monkeypatch.setenv("KQ_QMM_ROUTE", "verify_nax") monkeypatch.setenv("KQ_QMM_ROUTE_STRICT", "1") for m in (1, 5, 8): @@ -1582,6 +1651,7 @@ def test_verify_nax_head_entry(codec, monkeypatch): "q2_k": ((8, 5, 3, 2), 2, 5), "pq2_0": ((8, 5, 3, 2), 9, 9), "q2_0": ((8, 5, 3, 2), 9, 9), + "nvfp4": ((1, 1, 2, 1), 2, 6), "ptq1_0": ((8, 5, 3, 2), 9, 9), "q4_0": ((8, 6, 3, 2), 8, 8), "q4_k": ((8, 5, 4, 2), 2, 5), diff --git a/tests/test_split_kernel_dispatch.py b/tests/test_split_kernel_dispatch.py index 20bbcd3..73bd449 100644 --- a/tests/test_split_kernel_dispatch.py +++ b/tests/test_split_kernel_dispatch.py @@ -109,7 +109,7 @@ def _pipelines(name): def test_split_codecs_found(): - assert {"q4_k", "q8_0", "q2_0", "pq2_0", "iq4_xs", "iq4_nl"} <= set(HOST) + assert {"q4_k", "q8_0", "q2_0", "pq2_0", "iq4_xs", "iq4_nl", "nvfp4"} <= set(HOST) # Every codec has the three split-K tiles and both NAX split-K tiles. for codec, names in HOST.items(): for frag in ( diff --git a/tests/test_switch_gemm.py b/tests/test_switch_gemm.py index 2c0c340..c49b4a1 100644 --- a/tests/test_switch_gemm.py +++ b/tests/test_switch_gemm.py @@ -292,6 +292,17 @@ def order(a): np.testing.assert_array_equal(order(g), order(r)) +def _nvfp4_wire(rng, experts, n_rows): + """nvfp4 wire with group scales of 1/64 to 1/32, the range of real + weights.""" + from kqref import synth_wire + + wire = synth_wire( + rng, "nvfp4", 36, experts * n_rows * (K // 64), nvfp4_scales=(0x08, 0x11) + ) + return wire.reshape(experts, n_rows, (K // 64) * 36) + + def _nax_ok(): import mlx_kquant as kq @@ -302,7 +313,7 @@ def _nax_ok(): bool(os.environ.get("KQUANT_FORCE_CPU")), reason="gather_qmm_seg is Metal-only.", ) -@pytest.mark.parametrize("codec", ["iq2_xs", "iq3_xxs", "q4_k", "q8_0"]) +@pytest.mark.parametrize("codec", ["iq2_xs", "iq3_xxs", "q4_k", "q8_0", "nvfp4"]) @pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16]) def test_gather_qmm_seg_nax_matches_loop(codec, dtype, monkeypatch): """The NAX seg kernel (taken by default on NAX GPUs) matches the @@ -345,6 +356,11 @@ def test_gather_qmm_seg_nax_matches_loop(codec, dtype, monkeypatch): num_experts=E, output_dims=N, input_dims=K, bias=False, codec=codec ) sw.weight = mx.array(wire) + elif codec == "nvfp4": + sw = KQuantSwitchLinear( + num_experts=E, output_dims=N, input_dims=K, bias=False, codec=codec + ) + sw.weight = mx.array(_nvfp4_wire(np.random.default_rng(12), E, N)) else: sw = _make_switch(codec) w, s = sw["weight"], sw["scales"] @@ -421,13 +437,14 @@ def test_gather_qmm_seg_nax_unaligned_n(monkeypatch): bool(os.environ.get("KQUANT_FORCE_CPU")), reason="gather_qmm_seg is Metal-only.", ) -@pytest.mark.parametrize("codec", ["iq2_xs", "iq3_xxs"]) +@pytest.mark.parametrize("codec", ["iq2_xs", "iq3_xxs", "nvfp4"]) @pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16]) def test_gather_qmm_seg_nax_grid_codec_unaligned_n(codec, dtype, monkeypatch): - """The IQ grid codecs of the prefill gathers on an N tail past the last - 32-column simdgroup band, with tiles of 1, 33 and 65 rows (one live - 16-row sub-band, a partial second band, a second map tile) and a full - tile, all match the per-expert loop.""" + """The IQ grid codecs of the prefill gathers and nvfp4, which runs the + same NAX loader, on an N tail past the last 32-column simdgroup band, + with tiles of 1, 33 and 65 rows (one live 16-row sub-band, a partial + second band, a second map tile) and a full tile, all match the + per-expert loop.""" import mlx_kquant as kq # Loop reference on the BM=64 NAX tile, as in the test above. @@ -438,11 +455,14 @@ def test_gather_qmm_seg_nax_grid_codec_unaligned_n(codec, dtype, monkeypatch): rng = np.random.default_rng(41) n_odd = 72 - bpb = CODEC_GEOMETRY[codec][2] - wire = rng.integers(0, 256, (4, n_odd, bytes_per_row(codec, K)), dtype=np.uint8) - d = np.frombuffer(np.float16(0.01).tobytes(), dtype=np.uint8) - for b in range(0, wire.shape[-1], bpb): - wire[..., b : b + 2] = d + if codec == "nvfp4": + wire = _nvfp4_wire(rng, 4, n_odd) + else: + bpb = CODEC_GEOMETRY[codec][2] + wire = rng.integers(0, 256, (4, n_odd, bytes_per_row(codec, K)), dtype=np.uint8) + d = np.frombuffer(np.float16(0.01).tobytes(), dtype=np.uint8) + for b in range(0, wire.shape[-1], bpb): + wire[..., b : b + 2] = d w = mx.array(wire) s = mx.zeros((1,), dtype=mx.uint8) counts = np.array([1, 33, 65, 64], dtype=np.int64) diff --git a/tests/test_weight_base.py b/tests/test_weight_base.py index 415853f..6a8430b 100644 --- a/tests/test_weight_base.py +++ b/tests/test_weight_base.py @@ -2,7 +2,7 @@ The GPU kernels need a packed weight to start on a 2-, 4-, 8- or 16-byte boundary, by codec, and the ops raise ValueError on a start off it instead -of returning wrong values. mxfp4 and nvfp4 run from any start, as do the +of returning wrong values. mxfp4 runs from any start, as do the CPU kernels of every codec. A weight is moved off its boundary by copying its bytes into a larger buffer at an offset and passing a view of them. @@ -31,8 +31,8 @@ "q4_1": 4, "q5_1": 4, "ptq1_0": 4, + "nvfp4": 4, "mxfp4": 1, - "nvfp4": 1, } # verify_nax's own start check, stricter than BASE_ALIGN for these codecs. VNAX_BASE = {"pq2_0": 4, "q2_0": 4, "q4_0": 4, "q8_0": 4}