diff --git a/CHANGELOG.md b/CHANGELOG.md index f6ff726..2cee7e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). `KQ_QMV_SPLITK=0` restores the previous kernels. - The iq2_s fused MoE gate/up and down gathers at one activation row take less time on every GPU. `KQ_MOE_C32=0` restores the previous kernels. +- On M5 GPUs the iq2_s MoE gathers at one activation row take a few percent + less time again. `KQ_MOE_NX=8` restores the earlier width. ## [0.5.1] diff --git a/README.md b/README.md index a8b2206..4c4a44d 100644 --- a/README.md +++ b/README.md @@ -305,8 +305,9 @@ All optional; the defaults are right for normal use. A/B lever. - `KQ_MOE_C32` - on Metal, the 32-weight decode step of the iq2_s fused MoE gate/up and down gathers. By default it runs at one activation row on every GPU. `0` restores the 16-weight - kernels, and `1` also runs it at wider rows where the dedupe kernels do not apply. `KQ_MOE_NX=16` - selects its 16-lane form, and `KQ_MOE_NX=32` keeps the 16-weight kernels. Read per call. The + kernels, and `1` also runs it at wider rows where the dedupe kernels do not apply. It runs 16 + K-lanes per row on M5 and later GPUs and 8 on earlier ones. `KQ_MOE_NX=8` or `16` forces a + width, and `KQ_MOE_NX=32` keeps the 16-weight kernels. Read per call. The outputs differ from the 16-weight kernels only in summation order. A/B lever. - `KQ_SDPA_IDX_NAX=0` - run `sdpa_fa_indexed` on its simdgroup kernel even where the NAX tile diff --git a/bindings.cpp b/bindings.cpp index f15632d..96a0940 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -127,6 +127,13 @@ NB_MODULE(_ext, m) { "codec, rows of x, N and K, as [min_m, min_k, max_k, min_n, max_n]; " "empty where the qmm_t kernel runs."); + m.def( + "_moe_c32_nx", + &mlx_kquant::moe_c32_nx, + "arch"_a, + "The K-lanes per output row (8 or 16) of the iq2_s c32 MoE gathers " + "for an architecture string, unless KQ_MOE_NX forces one."); + m.def( "_dequantize_route_takes", &mlx_kquant::dequantize_route_takes, diff --git a/docs/kernels.md b/docs/kernels.md index f139eee..f68c0da 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -1041,8 +1041,9 @@ term; requires the expert biases and is instantiated for `mxfp4`/`nvfp4` only). At one activation row the iq2_s gate/up and down gathers decode 32 weights per lane step, one 32-weight group of a block, where the other codecs decode 16. The block scale, the high-bit byte and the scale byte then load once per 32 weights, and the grid indices and sign bytes load as -16-bit pairs. Each lane covers one 32-weight group per step at 8 K-lanes per output row, or 16 -under `KQ_MOE_NX=16`. Outputs differ from the 16-weight kernels only in summation order. +16-bit pairs. Each lane covers one 32-weight group per step, with 16 K-lanes per output row on +M5 and later GPUs and 8 on earlier ones. `KQ_MOE_NX=8` or `16` forces a width. Outputs differ from +the 16-weight kernels only in summation order. `KQ_MOE_C32=0` restores the 16-weight kernels. `KQ_MOE_C32=1` also runs the 32-weight kernels at wider rows where the dedupe kernels do not apply. A forced `KQ_MOE_NX=32` keeps the 16-weight kernels. @@ -1052,6 +1053,9 @@ On M3 Max at E=256, N=K=4096 and top 8, back to back, the gate/up gather took 0. step of eight Muse Spark UD-IQ2_S layers, the MoE blocks took 8.3% less time. With the `q8_0` split-K kernel, the step went from 1.414 to 1.335 ms per layer (-5.6%). +On M5 Max the 16-lane form is faster. Back to back at the same shapes, a MoE block replica with +both gathers at 16 K-lanes took 2.3% to 5.2% less time than at 8, in both arm orders. + The grid lookup is the next lever. With the grid read removed (wrong output), the 32-weight gathers took about 25% less time, and 37% less than the 16-weight gathers. That bounds the gain of a cheaper lookup. Two other grid sources gave no gain. A 64-bit grid in a device buffer matched the diff --git a/src/kquant.h b/src/kquant.h index 7379f7b..930ccdd 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -84,6 +84,11 @@ bool dequantize_route_takes( int N, int K); +// The K-lanes per output row (8 or 16) of the iq2_s c32 MoE gathers on a GPU +// with this architecture string (as get_architecture gives it), when +// KQ_MOE_NX does not force one. +int moe_c32_nx(const std::string& arch); + // Every kernel name the split-K dispatchers can build (qmm_splitk, // qmm_nax_splitk, verify_mma, verify_nax and their fold), for the kernel // name tests; empty on non-Metal builds. diff --git a/src/kquant_moe_glu.cpp b/src/kquant_moe_glu.cpp index 3a5c779..253ed72 100644 --- a/src/kquant_moe_glu.cpp +++ b/src/kquant_moe_glu.cpp @@ -5,6 +5,7 @@ // K-quant wire bytes: moe_glu_gather_kq (gate + up + GLU epilogue; // optional biases with act "swiglu_clamp", mxfp4/nvfp4 only) and // gather_qmv_kq, per-codec kernels (full codec matrix). +#include #include #include #include @@ -358,8 +359,19 @@ inline bool kq_moe_c32(const std::string& codec, int T) { return v == -1 ? T == 1 : v != 0; } +// A forced KQ_MOE_NX of 8 or 16, else this GPU's class width (moe_c32_nx). inline int kq_moe_c32_nx() { - return kq_moe_nx_env() == 16 ? 16 : 8; + const int v = kq_moe_nx_env(); + if (v == 8 || v == 16) { + return v; + } +#ifdef _METAL_ + static const int here = + moe_c32_nx(mx::metal::device(mx::Device::gpu).get_architecture()); + return here; +#else + return 8; +#endif } inline bool kq_gather_fine(int64_t coarse_tgs) { @@ -1317,6 +1329,26 @@ std::vector KQuantGatherQMVMixBias::output_shapes( return {{inputs[4].shape(0), inputs[0].shape(1)}}; } +// applegpu_g, two digit gen, a lower case suffix other than p: +// NX 16 from gen 17 (M5), where it beat NX 8 by 2 to 6% on a Muse Spark MoE +// block (M5 Max, E=256, top 8, N=K=4096). NX 8 elsewhere. +int moe_c32_nx(const std::string& arch) { + const std::string head = "applegpu_g"; + if (arch.size() != head.size() + 3 || arch.compare(0, head.size(), head)) { + return 8; + } + const char d0 = arch[head.size()]; + const char d1 = arch[head.size() + 1]; + const char suffix = arch[head.size() + 2]; + if (!std::isdigit(static_cast(d0)) || + !std::isdigit(static_cast(d1)) || + !std::islower(static_cast(suffix)) || suffix == 'p') { + return 8; + } + const int gen = (d0 - '0') * 10 + (d1 - '0'); + return gen >= 17 ? 16 : 8; +} + // Codecs with the fused kq GLU/gather kernels wired (kq_moe_glu_kq.h): // tuned kernels for q6_k/q8_0, generic Ext-trait kernels for the rest of the // codec matrix. Unsupported codecs must fall back to the stock path diff --git a/tests/test_moe_c32_class.py b/tests/test_moe_c32_class.py new file mode 100644 index 0000000..d4d2baa --- /dev/null +++ b/tests/test_moe_c32_class.py @@ -0,0 +1,31 @@ +"""The K-lanes per row of the iq2_s c32 MoE gathers by GPU class. + +``_moe_c32_nx`` answers for any architecture string, so these cases run on +every machine, CPU builds included. +""" + +import pytest +from mlx_kquant._ext import _moe_c32_nx + +WIDE = [f"applegpu_g{g}{c}" for g in (17, 18, 19) for c in "gsdc"] +NARROW = [f"applegpu_g{g}{c}" for g in (13, 14, 15, 16) for c in "gsdc"] + [ + "applegpu_g17p", + "applegpu_g18p", + "applegpu_g12s", + "applegpu_g17", + "applegpu_g17S", + "applegpu_g1xs", + "applegpu_g170s", + "gpu_g17s", + "", +] + + +@pytest.mark.parametrize("arch", WIDE) +def test_m5_and_later_take_nx16(arch): + assert _moe_c32_nx(arch) == 16 + + +@pytest.mark.parametrize("arch", NARROW) +def test_earlier_phones_and_malformed_keep_nx8(arch): + assert _moe_c32_nx(arch) == 8 diff --git a/tests/test_moe_glu_variants.py b/tests/test_moe_glu_variants.py index ad7f17e..c1b5dc1 100644 --- a/tests/test_moe_glu_variants.py +++ b/tests/test_moe_glu_variants.py @@ -552,24 +552,37 @@ def _c32_run(T, dtype, env, tmp_path, arm): return np.load(f), r.stderr +def _c32_class_suffix(): + import mlx.core as mx + from mlx_kquant._ext import _moe_c32_nx + + nx = _moe_c32_nx(mx.device_info().get("architecture", "")) + return "" if nx == 8 else f"_nx{nx}" + + @pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) def test_c32_dispatch(dtype, tmp_path): """iq2_s gate/up and down gathers take the 32-weight-step kernels at one - row by default, at NX = 8 or a forced 16. KQ_MOE_C32=0 and a forced - KQ_MOE_NX=32 keep the 16-weight kernels. Two rows keep them too unless - KQ_MOE_C32=1. The forms differ only in summation order.""" - want_glu = f"kq_iq2_s_moe_glu_gather_silu_c32_{dtype}_t" - want_qmv = f"kq_iq2_s_gather_qmv_c32_{dtype}_t" + row by default, at this GPU class's NX (8, or 16 from M5), or a forced 8 + or 16. KQ_MOE_C32=0 and a forced KQ_MOE_NX=32 keep the 16-weight kernels. + Two rows keep them too unless KQ_MOE_C32=1. The forms differ only in + summation order.""" + sfx = _c32_class_suffix() + want_glu = f"kq_iq2_s_moe_glu_gather_silu_c32{sfx}_{dtype}_t" + want_qmv = f"kq_iq2_s_gather_qmv_c32{sfx}_{dtype}_t" new, names = _c32_run(1, dtype, {}, tmp_path, "default") assert want_glu in names and want_qmv in names, names old, names = _c32_run(1, dtype, {"KQ_MOE_C32": "0"}, tmp_path, "old") assert "_c32" not in names, names + narrow, names = _c32_run(1, dtype, {"KQ_MOE_NX": "8"}, tmp_path, "nx8") + assert f"kq_iq2_s_moe_glu_gather_silu_c32_{dtype}_t" in names, names + assert f"kq_iq2_s_gather_qmv_c32_{dtype}_t" in names, names wide, names = _c32_run(1, dtype, {"KQ_MOE_NX": "16"}, tmp_path, "nx16") assert "kq_iq2_s_moe_glu_gather_silu_c32_nx16_" in names, names assert "kq_iq2_s_gather_qmv_c32_nx16_" in names, names for key in ("glu", "qmv"): scale = np.abs(old[key]).max() - for got in (new, wide): + for got in (new, narrow, wide): rel = np.abs(got[key] - old[key]).max() / scale assert rel < 2e-3, (key, rel) _, names = _c32_run(1, dtype, {"KQ_MOE_NX": "32"}, tmp_path, "nx32") @@ -610,8 +623,9 @@ def test_c32_reads_live(): {"KQ_MOE_NX_LOG": "1"}, "c32 live read", ) + sfx = _c32_class_suffix() assert "kq_iq2_s_gather_qmv_bfloat16_t" in r.stderr, r.stderr - assert "kq_iq2_s_gather_qmv_c32_bfloat16_t" in r.stderr, r.stderr + assert f"kq_iq2_s_gather_qmv_c32{sfx}_bfloat16_t" in r.stderr, r.stderr _C32_REF_SNIPPET = r"""