From d2262d16607933c809e0db109abcdbde2cde10fd Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:47:57 -0700 Subject: [PATCH 1/6] fix(gguf): zero-copy load for >2 GiB wire tensors --- CHANGELOG.md | 5 ++++ src/kquant_gguf.cpp | 57 ++++++++++++++++++++++++++++++++++++++------- tests/test_gguf.py | 43 ++++++++++++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eae38bd..e26ef06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Fixed +- Zero-copy GGUF load of tensors whose wire bytes exceed 2 GiB (e.g. the + expert stacks of a many-hundred-expert MoE): these silently fell back to + an eager per-tensor memcpy, exhausting memory at load on over-RAM models. + ## [0.3.8] Batched and shared-prefix decode attention: cascade, paged sparse, q8 KV diff --git a/src/kquant_gguf.cpp b/src/kquant_gguf.cpp index 5270857..a039e2f 100644 --- a/src/kquant_gguf.cpp +++ b/src/kquant_gguf.cpp @@ -106,8 +106,13 @@ const char* zc_dtype_name(mx::Dtype d) { // newBufferWithBytesNoCopy), then slices/reshapes to the tensor - all no-copy. // The window array's deleter releases that buffer and drops a ref to the // captured gguf_ctx, so the mmap survives exactly as long as some viewing array -// does. Returns nullopt when a no-copy wrap isn't possible (unaligned window or -// >INT32_MAX elements past the page-aligned base); the caller then memcpy's. +// does. A tensor with more elements than INT32_MAX at its own dtype (mx::Shape +// is int32 - e.g. a >2 GB uint8 expert wire stack of a many-hundred-expert +// MoE) is window-sliced at a wider integer dtype and view()ed back at the end; +// slice and view both stay buffer-sharing on the contiguous window, so that +// path is still no-copy. Returns nullopt when no wrap is possible (unaligned +// window, or the last dim / offsets don't divide at any wide dtype either); +// the caller then memcpy's. // // Alignment reasoning: gguflib mmaps at a page-aligned base and GGUF tensor // data sits at a 32-byte-aligned file offset, so `wd` is 32-aligned -> win_off @@ -136,12 +141,40 @@ std::optional try_zero_copy_array( } const size_t win_bytes = win_off + nbytes; - const size_t win_elems = win_bytes / isz; - const size_t off_elems = win_off / isz; - const size_t num_elems = nbytes / isz; - if (win_elems > static_cast(std::numeric_limits::max())) { - return std::nullopt; // mx::Shape elements are int32. + const size_t int_max = static_cast(std::numeric_limits::max()); + + // Dtype the 1-D window is built and sliced at: normally the tensor's own, + // widened when the element count would overflow int32 shape dims. + mx::Dtype win_dtype = dtype; + size_t win_isz = isz; + if (win_bytes / isz > int_max) { + const size_t last_bytes = + shape.empty() ? 0 : static_cast(shape.back()) * isz; + const std::pair wide[] = { + {8, mx::uint64}, {4, mx::uint32}, {2, mx::uint16}}; + bool widened = false; + for (const auto& [w, dt] : wide) { + if (w <= isz || last_bytes == 0) { + break; + } + if (win_off % w != 0 || nbytes % w != 0 || last_bytes % w != 0) { + continue; + } + if (win_bytes / w > int_max) { + continue; + } + win_dtype = dt; + win_isz = w; + widened = true; + break; + } + if (!widened) { + return std::nullopt; + } } + const size_t win_elems = win_bytes / win_isz; + const size_t off_elems = win_off / win_isz; + const size_t num_elems = nbytes / win_isz; mx::allocator::Buffer buf = mx::allocator::make_buffer(reinterpret_cast(win_base), win_bytes); @@ -154,12 +187,18 @@ std::optional try_zero_copy_array( zc_unregister(addr); mx::allocator::release(b); }; - mx::array window(buf, mx::Shape{static_cast(win_elems)}, dtype, del); + mx::array window(buf, mx::Shape{static_cast(win_elems)}, win_dtype, del); mx::array view = mx::slice( window, mx::Shape{static_cast(off_elems)}, mx::Shape{static_cast(off_elems + num_elems)}); - return mx::reshape(view, shape); + if (win_isz == isz) { + return mx::reshape(view, shape); + } + mx::Shape wide_shape = shape; + wide_shape.back() = static_cast( + static_cast(shape.back()) * isz / win_isz); + return mx::view(mx::reshape(view, wide_shape), dtype); } // Map a GGUF tensor type to one of the extension's supported codecs, or diff --git a/tests/test_gguf.py b/tests/test_gguf.py index e4ad6ff..a079364 100644 --- a/tests/test_gguf.py +++ b/tests/test_gguf.py @@ -141,3 +141,46 @@ def test_verify_zero_copy_views(tmp_path): # Unevaluated arrays are reported, not silently skipped. (prob,) = kq.verify_zero_copy_views([("z", a32 + 0.0)]) assert "unevaluated" in prob + + +@pytest.mark.skipif( + not __import__("os").environ.get("KQUANT_BIG_TESTS"), + reason="writes a >2 GiB GGUF; set KQUANT_BIG_TESTS=1 to run", +) +def test_load_gguf_wire_over_int32(tmp_path): + """A quantized tensor whose wire bytes exceed INT32_MAX must still load + zero-copy (wide-dtype window slice + view back), not fall back to the + eager per-tensor memcpy - which OOMs a many-hundred-expert MoE at load.""" + import resource + import sys + + rows, k = 247_000, 8192 # 34-byte q8_0 blocks -> 8704 B/row, ~2.15 GB + row_bytes = k // 32 * 34 + nbytes = rows * row_bytes + assert nbytes > 2**31 + + rng = np.random.default_rng(0) + wire = rng.integers(0, 256, size=nbytes, dtype=np.uint8) + path = str(tmp_path / "big.gguf") + w = GGUFWriter(path, "smoke") + # add_tensor takes the byte shape; the writer derives the logical shape. + w.add_tensor("big.q8", wire.reshape(rows, row_bytes), raw_dtype=GT.Q8_0) + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + + scale = 1 if sys.platform == "darwin" else 1024 # ru_maxrss units + rss0 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * scale + arrays, codecs, _meta, _shapes = kq.load_gguf(path, True) + rss1 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * scale + + a = arrays["big.q8"] + assert a.dtype == mx.uint8 and a.shape == (rows, row_bytes) + assert dict(codecs)["big.q8"] == "q8_0" + # The regression signature: the pre-fix path memcpy'd the wire eagerly. + assert rss1 - rss0 < 512 * 1024 * 1024, "load copied the wire bytes" + + for i in (0, 1, rows // 2, rows - 1): + got = np.array(a[i]) + assert np.array_equal(got, wire[i * row_bytes : (i + 1) * row_bytes]), i From eff7b41d19d353f0b2adda7be91b713244ce5466 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:06:10 -0700 Subject: [PATCH 2/6] feat(arena): itemsize parameter on arena_alloc for >2 GiB slots --- bindings.cpp | 28 ++++++++++++++++++++++++---- src/kquant.h | 4 +++- src/kquant_arena.cpp | 13 ++++++++++--- 3 files changed, 37 insertions(+), 8 deletions(-) diff --git a/bindings.cpp b/bindings.cpp index 7479374..22613fa 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -1392,9 +1392,26 @@ NB_MODULE(_ext, m) { m.def( "arena_alloc", - [](const std::vector& shape) { + [](const std::vector& shape, int itemsize) { + mlx::core::Dtype dt = mlx::core::uint8; + switch (itemsize) { + case 1: + break; + case 2: + dt = mlx::core::uint16; + break; + case 4: + dt = mlx::core::uint32; + break; + case 8: + dt = mlx::core::uint64; + break; + default: + throw std::invalid_argument( + "[mlx_kquant.arena_alloc] itemsize must be 1, 2, 4 or 8."); + } auto [arr, addr] = mlx_kquant::arena_alloc( - mlx::core::Shape(shape.begin(), shape.end())); + mlx::core::Shape(shape.begin(), shape.end()), dt); PyObject* mv = PyMemoryView_FromMemory( reinterpret_cast(addr), static_cast(arr.nbytes()), @@ -1405,11 +1422,14 @@ NB_MODULE(_ext, m) { return nb::make_tuple(arr, nb::steal(mv)); }, "shape"_a, + "itemsize"_a = 1, R"( Allocate a page-aligned host buffer wrapped zero-copy as a Metal - shared-storage uint8 array. + shared-storage unsigned-integer array (dtype uint8/16/32/64 per + ``itemsize``; wider itemsizes let a >2 GiB slot fit int32 shape dims). - Returns (array, memoryview): the same bytes seen from both sides. + Returns (array, memoryview): the same bytes seen from both sides, + the memoryview always byte-addressed over the full allocation. The writable memoryview is the CPU feeder's window (os.preadv into slices of it reads disk straight into GPU-visible memory); the array is what kernels consume. The memoryview is valid only while the array diff --git a/src/kquant.h b/src/kquant.h index 1315050..7d5c10e 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -1748,7 +1748,9 @@ bool shared_event_wait(uint64_t handle, uint64_t value, int64_t timeout_ms); // base address (for the Python binding's writable memoryview over the same // bytes; valid exactly as long as the array lives). The CPU-write -> GPU-read // ordering contract is the caller's, via the event ops below. -std::pair arena_alloc(const mx::Shape& shape); +std::pair arena_alloc( + const mx::Shape& shape, + mlx::core::Dtype dtype = mlx::core::uint8); // Stream side: identity ops on `x` that encode an MTLSharedEvent signal/wait // at their position in the graph's evaluation order. The returned array diff --git a/src/kquant_arena.cpp b/src/kquant_arena.cpp index c9f208b..54b3730 100644 --- a/src/kquant_arena.cpp +++ b/src/kquant_arena.cpp @@ -22,8 +22,15 @@ namespace mx = mlx::core; namespace mlx_kquant { -std::pair arena_alloc(const mx::Shape& shape) { - size_t nbytes = 1; +std::pair arena_alloc( + const mx::Shape& shape, + mx::Dtype dtype) { + if (dtype != mx::uint8 && dtype != mx::uint16 && dtype != mx::uint32 && + dtype != mx::uint64) { + throw std::invalid_argument( + "[mlx_kquant.arena_alloc] dtype must be an unsigned integer type."); + } + size_t nbytes = mx::size_of(dtype); for (auto d : shape) { if (d <= 0) { throw std::invalid_argument( @@ -50,7 +57,7 @@ std::pair arena_alloc(const mx::Shape& shape) { mx::allocator::release(b); std::free(ptr); }; - mx::array arr(buf, shape, mx::uint8, del); + mx::array arr(buf, shape, dtype, del); return {std::move(arr), reinterpret_cast(ptr)}; } From 3d75d0c06761129f579e8d34b4986ce25a0751f1 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 30 Jul 2026 17:49:23 -0700 Subject: [PATCH 3/6] feat(residency): wire chosen buffers into the Metal residency set --- bindings.cpp | 24 +++++++++++++++ mlx_kquant/__init__.py | 6 ++++ src/kquant.h | 12 ++++++++ src/kquant_arena.cpp | 67 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 109 insertions(+) diff --git a/bindings.cpp b/bindings.cpp index 22613fa..0b3a2f5 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -1438,6 +1438,30 @@ NB_MODULE(_ext, m) { (shared_event_set) after writing; nothing else orders them. )"); + m.def( + "residency_insert", + &mlx_kquant::residency_insert, + "a"_a, + "Stage ``a``'s underlying Metal buffer for the device residency set " + "(wired for the buffer's lifetime once residency_commit runs), so " + "command buffers stop re-wiring its pages on every use. The array " + "must have materialized data (evaluate first). False on non-Metal " + "builds or missing data."); + + m.def( + "residency_commit", + &mlx_kquant::residency_commit, + "Commit staged residency_insert additions and request residency. " + "False on non-Metal builds."); + + m.def( + "residency_erase", + &mlx_kquant::residency_erase, + "a"_a, + "Stage removal of ``a``'s buffer from the residency set (call before " + "dropping a member buffer; takes effect at the next commit). False " + "on non-Metal builds or missing data."); + // --- shared-event stream primitives (feeder loop) --- m.def( diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index 241b02c..e1c8d5f 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -60,6 +60,9 @@ quantize, quantized_matmul, quantized_matmul_qmv_bias, + residency_commit, + residency_erase, + residency_insert, rmsnorm2_add, rmsnorm_multi3, route_shed, @@ -126,6 +129,9 @@ "sdpa_decode_gqa_paged", "sdpa_fa_verify", "sdpa_vector", + "residency_commit", + "residency_erase", + "residency_insert", "shared_event_create", "shared_event_destroy", "shared_event_read", diff --git a/src/kquant.h b/src/kquant.h index 7d5c10e..3e223be 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -1752,6 +1752,18 @@ std::pair arena_alloc( const mx::Shape& shape, mlx::core::Dtype dtype = mlx::core::uint8); +// GPU residency: stage arrays' buffers into the Metal device residency +// set (insert per array, then one commit), so command buffers stop +// re-wiring those pages on every use (the cost that dominates streamed +// decode once weights are host-pinned). Membership lasts for the buffer's +// lifetime. Both return false on non-Metal builds; insert also returns +// false for an array with no materialized data. +bool residency_insert(const mx::array& a); +bool residency_commit(); +// Remove before freeing a set member (a freed buffer must not stay in the +// set); takes effect at the next commit. +bool residency_erase(const mx::array& a); + // Stream side: identity ops on `x` that encode an MTLSharedEvent signal/wait // at their position in the graph's evaluation order. The returned array // aliases x and MUST be threaded into downstream compute (or evaluated diff --git a/src/kquant_arena.cpp b/src/kquant_arena.cpp index 54b3730..65f13e9 100644 --- a/src/kquant_arena.cpp +++ b/src/kquant_arena.cpp @@ -18,6 +18,10 @@ #include "kquant.h" +#ifdef _METAL_ +#include "mlx/backend/metal/device.h" +#endif + namespace mx = mlx::core; namespace mlx_kquant { @@ -61,4 +65,67 @@ std::pair arena_alloc( return {std::move(arr), reinterpret_cast(ptr)}; } +#ifdef _METAL_ + +// MLX's ResidencySet wrapper methods are not exported from libmlx; its +// inline accessor hands back the raw MTL::ResidencySet (created at device +// init and attached to MLX's command queue), and metal-cpp is header-only, +// so the additions run entirely in this TU. +static MTL::ResidencySet* kq_residency_set() { + auto& d = mx::metal::device(mx::Device(mx::Device::gpu)); + return const_cast(d.residency_set().mtl_residency_set()); +} + +bool residency_insert(const mx::array& a) { + const void* ptr = a.buffer().ptr(); + if (ptr == nullptr) { + return false; + } + auto* rs = kq_residency_set(); + if (rs == nullptr) { + return false; + } + rs->addAllocation(static_cast(const_cast(ptr))); + return true; +} + +bool residency_commit() { + auto* rs = kq_residency_set(); + if (rs == nullptr) { + return false; + } + rs->commit(); + rs->requestResidency(); + return true; +} + +bool residency_erase(const mx::array& a) { + const void* ptr = a.buffer().ptr(); + if (ptr == nullptr) { + return false; + } + auto* rs = kq_residency_set(); + if (rs == nullptr) { + return false; + } + rs->removeAllocation(static_cast(const_cast(ptr))); + return true; +} + +#else + +bool residency_insert(const mx::array&) { + return false; +} + +bool residency_commit() { + return false; +} + +bool residency_erase(const mx::array&) { + return false; +} + +#endif + } // namespace mlx_kquant From efa2959a48c1dc4fc5f82de7f59a0e6c16fddf3a Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 30 Jul 2026 17:50:05 -0700 Subject: [PATCH 4/6] docs: changelog for arena itemsize + residency ops --- CHANGELOG.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e26ef06..d47f7ef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added +- `arena_alloc` accepts `itemsize` 2/4/8 so >2 GiB staging slots fit int32 + shape dims. +- `residency_insert` / `residency_commit` / `residency_erase`: wire chosen + buffers into the Metal residency set, ending per-command-buffer re-wiring + of large host-pinned weights. + ### Fixed - Zero-copy GGUF load of tensors whose wire bytes exceed 2 GiB (e.g. the expert stacks of a many-hundred-expert MoE): these silently fell back to From b7dff804e2ea4adb74531959977402c5f9c3b528 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:25:48 -0700 Subject: [PATCH 5/6] iq qmv: vectorized scale/grid unpack for iq1_m + iq2_xxs decode --- CHANGELOG.md | 4 ++ .../backend/metal/kernels/kq_quantized_iq.h | 62 +++++++++++++------ 2 files changed, 46 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d47f7ef..584fb26 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Changed +- iq1_m and iq2_xxs mat-vec decode is 1.2-1.6x faster per call (vectorized + scale and grid unpack, bit-exact). + ### Added - `arena_alloc` accepts `itemsize` 2/4/8 so >2 GiB staging slots fit int32 shape dims. diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index bd01b49..11adc37 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h @@ -2257,16 +2257,28 @@ METAL_FUNC void kq_iq2_xxs_qmv_impl( ib * KQ_IQ2_XXS_BLOCK_BYTES; const U d = U(float(*(const device half*)sb)); const device uint8_t* qs = sb + KQ_IQ2_XXS_QS_OFFSET + s * 8; - const uint signbits = uint(qs[4]) | (uint(qs[5]) << 8) | - (uint(qs[6]) << 16) | (uint(qs[7]) << 24); + // Blocks are 66 bytes (2-aligned), so the sign word is two ushort + // loads, not four byte loads. + const device ushort* qw = reinterpret_cast(qs); + const uint signbits = uint(qw[2]) | (uint(qw[3]) << 16); const U db = d * (U(0.5f) + U(signbits >> 28)) * U(0.25f); const uint8_t signs = ksigns_iq2xs[(signbits >> (7 * l)) & 127]; + // Reinterpret the grid word as two uchar4s; folding the sign into + // the integer-valued magnitude is exact, so the single rounding per + // fma is unchanged and results stay bit-exact. const uint64_t g = iq2xxs_grid[qs[l]]; + const uchar4 g_lo = as_type(uint(g & 0xffffffffu)); + const uchar4 g_hi = as_type(uint(g >> 32)); U partial = 0; #pragma unroll - for (int j = 0; j < 8; j++) { - partial += xt[j] * U((g >> (8 * j)) & 0xff) * - ((signs & kmask_iq2xs[j]) ? U(-1) : U(1)); + for (int j = 0; j < 4; j++) { + const U gv = (signs & kmask_iq2xs[j]) ? -U(g_lo[j]) : U(g_lo[j]); + partial += xt[j] * gv; + } +#pragma unroll + for (int j = 0; j < 4; j++) { + const U gv = (signs & kmask_iq2xs[4 + j]) ? -U(g_hi[j]) : U(g_hi[j]); + partial += xt[4 + j] * gv; } result[row] += db * partial; } @@ -4016,6 +4028,11 @@ METAL_FUNC void kq_iq1_m_qmv_impl( y += tid.x * out_vec_size; const int s = simd_lid / 4; // sub-block const int l = simd_lid % 4; // l-group (one 8-weight group) + // Per-lane constants: field shift within the scale word, qh nibble shift, + // sign-bit mask (invariant across superblocks). + const int shift0 = (l < 2) ? 0 : 3; + const int hshift = (l & 1) ? 4 : 8; + const uint8_t sign_mask = (l & 1) ? 0x80 : 0x08; U result[results_per_simdgroup] = {0}; for (int ib = 0; ib < nb; ib++) { U xt[vpt]; @@ -4027,29 +4044,34 @@ METAL_FUNC void kq_iq1_m_qmv_impl( const device uint8_t* sb = w + static_cast(out_row + row) * row_bytes + ib * KQ_IQ1_M_BLOCK_BYTES; - const device uint8_t* scp = sb + KQ_IQ1_M_SCALES_OFFSET; - const ushort sc0 = ushort(scp[0]) | (ushort(scp[1]) << 8); - const ushort sc1 = ushort(scp[2]) | (ushort(scp[3]) << 8); - const ushort sc2 = ushort(scp[4]) | (ushort(scp[5]) << 8); - const ushort sc3 = ushort(scp[6]) | (ushort(scp[7]) << 8); - const ushort scale_u16 = (sc0 >> 12) | ((sc1 >> 8) & 0x00f0) | - ((sc2 >> 4) & 0x0f00) | (sc3 & 0xf000); + // Blocks are 56 bytes and scales sit at +48, so the 8-byte scale + // block is always 8-aligned: one vector load replaces eight byte + // loads, and the per-half word is a select from the same vector. + const ushort4 scv = + *reinterpret_cast(sb + KQ_IQ1_M_SCALES_OFFSET); + const ushort scale_u16 = (scv.x >> 12) | ((scv.y >> 8) & 0x00f0) | + ((scv.z >> 4) & 0x0f00) | (scv.w & 0xf000); const U d = U(float(as_type(scale_u16))); - const device uint8_t* swp = scp + (s / 2) * 2; - const uint sc_word = uint(swp[0]) | (uint(swp[1]) << 8); - const int shift = 6 * (s & 1) + ((l < 2) ? 0 : 3); + const uint sc_word = scv[s / 2]; + const int shift = 6 * (s & 1) + shift0; const U dl = d * U(2 * int((sc_word >> shift) & 7) + 1); const uint8_t qh = sb[KQ_IQ1_M_QH_OFFSET + s * 2 + l / 2]; - const int hshift = (l & 1) ? 4 : 8; const uint idx = uint(sb[KQ_IQ1_M_QS_OFFSET + s * 4 + l]) | ((uint(qh) << hshift) & 0x700); - const U delta = (qh & ((l & 1) ? 0x80 : 0x08)) ? U(-0.125f) : U(0.125f); + const U delta = (qh & sign_mask) ? U(-0.125f) : U(0.125f); + // Reinterpret the signed grid word as two char4s: value-identical to + // the byte extract chain, same fma order, so results stay bit-exact. const uint64_t g = iq1s_grid[idx]; + const char4 g_lo = as_type(uint(g & 0xffffffffu)); + const char4 g_hi = as_type(uint(g >> 32)); U partial = 0; #pragma unroll - for (int j = 0; j < 8; j++) { - const int8_t gv = as_type(uint8_t((g >> (8 * j)) & 0xff)); - partial += xt[j] * (U(gv) + delta); + for (int j = 0; j < 4; j++) { + partial += xt[j] * (U(g_lo[j]) + delta); + } +#pragma unroll + for (int j = 0; j < 4; j++) { + partial += xt[4 + j] * (U(g_hi[j]) + delta); } result[row] += dl * partial; } From dbd7ccf34cebddeb9c9957da34d2135162e90176 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:35:21 -0700 Subject: [PATCH 6/6] iq qmv: extend vectorized unpack to iq1_s/iq2_xs/iq2_s/iq3_xxs/iq3_s --- CHANGELOG.md | 5 +- .../backend/metal/kernels/kq_quantized_iq.h | 105 ++++++++++++------ 2 files changed, 77 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 584fb26..9546709 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,8 +7,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] ### Changed -- iq1_m and iq2_xxs mat-vec decode is 1.2-1.6x faster per call (vectorized - scale and grid unpack, bit-exact). +- Grid-codec mat-vec decode is faster per call via vectorized scale and + grid unpack (bit-exact): iq1_m 1.2-1.6x, iq2_xxs 1.2x, iq2_xs 1.4x; + iq1_s, iq3_xxs, iq3_s small gains; iq2_s neutral. ### Added - `arena_alloc` accepts `itemsize` 2/4/8 so >2 GiB staging slots fit int32 diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index 11adc37..40b872f 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h @@ -1361,20 +1361,25 @@ METAL_FUNC void kq_iq3_xxs_qmv_impl( ib * KQ_IQ3_XXS_BLOCK_BYTES; const U d = U(float(*(const device half*)sb)); const device uint8_t* qs = sb + KQ_IQ3_XXS_QS_OFFSET + s * 8; - const device uint8_t* gas = sb + KQ_IQ3_XXS_GAS_OFFSET + s * 4; - const uint aux32 = uint(gas[0]) | (uint(gas[1]) << 8) | - (uint(gas[2]) << 16) | (uint(gas[3]) << 24); + // Blocks are 98 bytes (2-aligned), so the aux word is two ushort + // loads, not four byte loads. + const device ushort* gw = reinterpret_cast( + sb + KQ_IQ3_XXS_GAS_OFFSET + s * 4); + const uint aux32 = uint(gw[0]) | (uint(gw[1]) << 16); const U db = d * (U(0.5f) + U(aux32 >> 28)) * U(0.5f); const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * l)) & 127]; - const uint g1 = iq3xxs_grid[qs[2 * l]]; - const uint g2 = iq3xxs_grid[qs[2 * l + 1]]; + // Reinterpret each grid word as a uchar4; folding the sign into the + // integer-valued magnitude is exact, so the single rounding per fma + // is unchanged. + const uchar4 g1 = as_type(iq3xxs_grid[qs[2 * l]]); + const uchar4 g2 = as_type(iq3xxs_grid[qs[2 * l + 1]]); U partial = 0; #pragma unroll for (int j = 0; j < 4; j++) { - partial += xt[j] * U((g1 >> (8 * j)) & 0xff) * - ((signs & kmask_iq2xs[j]) ? U(-1) : U(1)); - partial += xt[j + 4] * U((g2 >> (8 * j)) & 0xff) * - ((signs & kmask_iq2xs[j + 4]) ? U(-1) : U(1)); + const U gv1 = (signs & kmask_iq2xs[j]) ? -U(g1[j]) : U(g1[j]); + partial += xt[j] * gv1; + const U gv2 = (signs & kmask_iq2xs[4 + j]) ? -U(g2[j]) : U(g2[j]); + partial += xt[j + 4] * gv2; } result[row] += db * partial; } @@ -1808,19 +1813,25 @@ METAL_FUNC void kq_iq3_s_qmv_impl( const device uint8_t* scales = sb + KQ_IQ3_S_SCALES_OFFSET; const U db = d * U(1 + 2 * ((scales[s / 2] >> (4 * (s & 1))) & 0xf)); const uint qh = sb[KQ_IQ3_S_QH_OFFSET + s]; - const device uint8_t* qs = sb + KQ_IQ3_S_QS_OFFSET + s * 8; + // The qs pair is 2-aligned: one ushort load instead of two byte + // loads. + const uint qpair = uint(*reinterpret_cast( + sb + KQ_IQ3_S_QS_OFFSET + s * 8 + 2 * l)); const uint8_t signs = sb[KQ_IQ3_S_SIGNS_OFFSET + s * 4 + l]; - const uint i1 = qs[2 * l] | ((qh << (8 - 2 * l)) & 256); - const uint i2 = qs[2 * l + 1] | ((qh << (7 - 2 * l)) & 256); - const uint g1 = iq3s_grid[i1]; - const uint g2 = iq3s_grid[i2]; + const uint i1 = (qpair & 0xff) | ((qh << (8 - 2 * l)) & 256); + const uint i2 = (qpair >> 8) | ((qh << (7 - 2 * l)) & 256); + // Reinterpret each grid word as a uchar4; folding the sign into the + // integer-valued magnitude is exact, so the single rounding per fma + // is unchanged. + const uchar4 g1 = as_type(iq3s_grid[i1]); + const uchar4 g2 = as_type(iq3s_grid[i2]); U partial = 0; #pragma unroll for (int j = 0; j < 4; j++) { - partial += xt[j] * U((g1 >> (8 * j)) & 0xff) * - ((signs & kmask_iq2xs[j]) ? U(-1) : U(1)); - partial += xt[j + 4] * U((g2 >> (8 * j)) & 0xff) * - ((signs & kmask_iq2xs[j + 4]) ? U(-1) : U(1)); + const U gv1 = (signs & kmask_iq2xs[j]) ? -U(g1[j]) : U(g1[j]); + partial += xt[j] * gv1; + const U gv2 = (signs & kmask_iq2xs[4 + j]) ? -U(g2[j]) : U(g2[j]); + partial += xt[j + 4] * gv2; } result[row] += db * partial; } @@ -2708,18 +2719,30 @@ METAL_FUNC void kq_iq2_xs_qmv_impl( static_cast(out_row + row) * row_bytes + ib * KQ_IQ2_XS_BLOCK_BYTES; const U d = U(float(*(const device half*)sb)); - const device uint8_t* qp = sb + KQ_IQ2_XS_QS_OFFSET + s * 8 + l * 2; - const uint q = uint(qp[0]) | (uint(qp[1]) << 8); + // The qs entry is a 2-aligned uint16: one ushort load instead of + // two byte loads. + const uint q = uint(*reinterpret_cast( + sb + KQ_IQ2_XS_QS_OFFSET + s * 8 + l * 2)); const uint8_t sc = sb[KQ_IQ2_XS_SCALES_OFFSET + s]; const int sc_nib = (l < 2) ? (sc & 0xf) : (sc >> 4); const U db = d * (U(0.5f) + U(sc_nib)) * U(0.25f); const uint8_t signs = ksigns_iq2xs[q >> 9]; + // Reinterpret the grid word as two uchar4s; folding the sign into + // the integer-valued magnitude is exact, so the single rounding per + // fma is unchanged. const uint64_t g = iq2xs_grid[q & 511]; + const uchar4 g_lo = as_type(uint(g & 0xffffffffu)); + const uchar4 g_hi = as_type(uint(g >> 32)); U partial = 0; #pragma unroll - for (int j = 0; j < 8; j++) { - partial += xt[j] * U((g >> (8 * j)) & 0xff) * - ((signs & kmask_iq2xs[j]) ? U(-1) : U(1)); + for (int j = 0; j < 4; j++) { + const U gv = (signs & kmask_iq2xs[j]) ? -U(g_lo[j]) : U(g_lo[j]); + partial += xt[j] * gv; + } +#pragma unroll + for (int j = 0; j < 4; j++) { + const U gv = (signs & kmask_iq2xs[4 + j]) ? -U(g_hi[j]) : U(g_hi[j]); + partial += xt[4 + j] * gv; } result[row] += db * partial; } @@ -3159,12 +3182,23 @@ METAL_FUNC void kq_iq2_s_qmv_impl( const U db = d * (U(0.5f) + U(sc_nib)) * U(0.25f); const uint idx = qs[l] | ((qh << (8 - 2 * l)) & 0x300); const uint8_t signs_byte = sg[l]; + // Reinterpret the grid word as two uchar4s; folding the sign into + // the integer-valued magnitude is exact, so the single rounding per + // fma is unchanged. const uint64_t g = iq2s_grid[idx]; + const uchar4 g_lo = as_type(uint(g & 0xffffffffu)); + const uchar4 g_hi = as_type(uint(g >> 32)); U partial = 0; #pragma unroll - for (int j = 0; j < 8; j++) { - partial += xt[j] * U((g >> (8 * j)) & 0xff) * - ((signs_byte & kmask_iq2xs[j]) ? U(-1) : U(1)); + for (int j = 0; j < 4; j++) { + const U gv = (signs_byte & kmask_iq2xs[j]) ? -U(g_lo[j]) : U(g_lo[j]); + partial += xt[j] * gv; + } +#pragma unroll + for (int j = 0; j < 4; j++) { + const U gv = + (signs_byte & kmask_iq2xs[4 + j]) ? -U(g_hi[j]) : U(g_hi[j]); + partial += xt[4 + j] * gv; } result[row] += db * partial; } @@ -3598,18 +3632,27 @@ METAL_FUNC void kq_iq1_s_qmv_impl( static_cast(out_row + row) * row_bytes + ib * KQ_IQ1_S_BLOCK_BYTES; const U d = U(float(*(const device half*)sb)); - const device uint8_t* qhp = sb + KQ_IQ1_S_QH_OFFSET + s * 2; - const uint qh = uint(qhp[0]) | (uint(qhp[1]) << 8); + // The qh entry is a 2-aligned uint16: one ushort load instead of + // two byte loads. + const uint qh = uint(*reinterpret_cast( + sb + KQ_IQ1_S_QH_OFFSET + s * 2)); const uint8_t qs = sb[KQ_IQ1_S_QS_OFFSET + s * 4 + l]; const U dl = d * U(2 * int((qh >> 12) & 7) + 1); const U delta = (qh & 0x8000) ? U(-0.125f) : U(0.125f); const uint idx = uint(qs) | (((qh >> (3 * l)) & 7) << 8); + // Reinterpret the signed grid word as two char4s: value-identical to + // the byte extract chain, same fma order, so results stay bit-exact. const uint64_t g = iq1s_grid[idx]; + const char4 g_lo = as_type(uint(g & 0xffffffffu)); + const char4 g_hi = as_type(uint(g >> 32)); U partial = 0; #pragma unroll - for (int j = 0; j < 8; j++) { - const int8_t gv = as_type(uint8_t((g >> (8 * j)) & 0xff)); - partial += xt[j] * (U(gv) + delta); + for (int j = 0; j < 4; j++) { + partial += xt[j] * (U(g_lo[j]) + delta); + } +#pragma unroll + for (int j = 0; j < 4; j++) { + partial += xt[4 + j] * (U(g_hi[j]) + delta); } result[row] += dl * partial; }