From ad4b895d785286a0f70478decae7c5fb2408f14d Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:16:18 -0700 Subject: [PATCH 1/2] feat(route_shed): GPU-side routed-expert slot remap + residency shed op (gpu-dispatch Tier 2 front end) kq.route_shed(indices, scores, slot_table) -> (slots, mix, miss_ids, miss_scores): remaps routed ids to arena slots on GPU, sheds every non-resident expert (lazy graphs cannot demand-read disk), renormalizes kept gate weights mass-preserving, reports misses front-packed descending-score for between-token prestage. Shed entries reuse the row's first kept slot with zero mix weight so downstream gather kernels need no changes. CPU eval mirrors the kernel's f32 accumulation order bit-for-bit; GPU parity tests gated on Metal (deferred while a decode run holds the box). Claude-Session: https://claude.ai/code/session_01Szxi6DtQUJcrCgTFUmWcfR --- CMakeLists.txt | 2 + bindings.cpp | 30 +++ metal/kq_route_shed.metal | 5 + .../mlx/backend/metal/kernels/kq_route_shed.h | 91 +++++++ mlx_kquant/__init__.py | 2 + src/kquant.h | 42 ++++ src/kquant_route_shed.cpp | 236 ++++++++++++++++++ tests/test_route_shed.py | 175 +++++++++++++ 8 files changed, 583 insertions(+) create mode 100644 metal/kq_route_shed.metal create mode 100644 metal/mlx/backend/metal/kernels/kq_route_shed.h create mode 100644 src/kquant_route_shed.cpp create mode 100644 tests/test_route_shed.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 0197881..3a9402e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -127,6 +127,7 @@ target_sources( ${CMAKE_CURRENT_LIST_DIR}/src/kquant_dsa_qat.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_moe_glu.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_norm_fused.cpp + ${CMAKE_CURRENT_LIST_DIR}/src/kquant_route_shed.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_gather.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_encode.cpp ${CMAKE_CURRENT_LIST_DIR}/src/kquant_arena.cpp @@ -207,6 +208,7 @@ if(MLX_BUILD_METAL) ${CMAKE_CURRENT_LIST_DIR}/metal/kq_moe_glu.metal ${CMAKE_CURRENT_LIST_DIR}/metal/kq_moe_glu_kq.metal ${CMAKE_CURRENT_LIST_DIR}/metal/kq_norm_fused.metal + ${CMAKE_CURRENT_LIST_DIR}/metal/kq_route_shed.metal INCLUDE_DIRS ${CMAKE_CURRENT_LIST_DIR}/metal ${MLX_INCLUDE_DIRS} diff --git a/bindings.cpp b/bindings.cpp index 7e66d30..633a4e3 100644 --- a/bindings.cpp +++ b/bindings.cpp @@ -921,6 +921,36 @@ NB_MODULE(_ext, m) { array: same shape and dtype as a. )"); + m.def( + "route_shed", + &mlx_kquant::route_shed, + "indices"_a, + "scores"_a, + "slot_table"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"( + GPU-side routed-expert slot remap + residency shed for streamed MoE + decode. Remaps routed expert ids to arena slots via slot_table + (expert id -> slot, negative = non-resident), sheds every + non-resident expert (a lazy graph cannot demand-read disk), + renormalizes kept gate weights mass-preserving + (score * S_all / S_kept), and reports misses for between-token + prestaging. Shed entries keep a valid slot (the row's first kept + slot, else 0) with a zero mix weight, so downstream gather kernels + need no changes. All-miss rows return an all-zero mix. + + Args: + indices (array): expert indices [T, R], uint32. R <= 64. + scores (array): gate scores [T, R], float32. + slot_table (array): [n_experts], int32; slot index or negative. + + Returns: + tuple: (slots u32 [T, R], mix f32 [T, R], + miss_ids i32 [T, R] front-packed descending-score with -1 pad, + miss_scores f32 [T, R] aligned with 0 pad). + )"); + m.def( "gather_qmm", &mlx_kquant::gather_qmm, diff --git a/metal/kq_route_shed.metal b/metal/kq_route_shed.metal new file mode 100644 index 0000000..c1f32e1 --- /dev/null +++ b/metal/kq_route_shed.metal @@ -0,0 +1,5 @@ +// clang-format off +// Routed-expert slot remap + residency shed; see kq_route_shed.h. +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/kq_route_shed.h" +// clang-format on diff --git a/metal/mlx/backend/metal/kernels/kq_route_shed.h b/metal/mlx/backend/metal/kernels/kq_route_shed.h new file mode 100644 index 0000000..7eed7a3 --- /dev/null +++ b/metal/mlx/backend/metal/kernels/kq_route_shed.h @@ -0,0 +1,91 @@ +// Routed-expert slot remap + residency shed (gpu-dispatch front end). +// One thread per token row: R <= 64 entries of serial work per thread, so +// occupancy is irrelevant and launch cost dominates. Semantics and f32 +// accumulation order must stay bit-identical to KQuantRouteShed::eval_cpu. + +#define KQ_ROUTE_SHED_MAX_R 64 + +[[kernel]] void kq_route_shed( + const device uint32_t* indices [[buffer(0)]], + const device float* scores [[buffer(1)]], + const device int32_t* slot_table [[buffer(2)]], + device uint32_t* slots [[buffer(3)]], + device float* mix [[buffer(4)]], + device int32_t* miss_ids [[buffer(5)]], + device float* miss_scores [[buffer(6)]], + const constant int& R [[buffer(7)]], + const constant int& E [[buffer(8)]], + const constant int& T [[buffer(9)]], + uint tid [[thread_position_in_grid]]) { + if (int(tid) >= T) { + return; + } + const device uint32_t* id = indices + tid * R; + const device float* sc = scores + tid * R; + device uint32_t* sl = slots + tid * R; + device float* mo = mix + tid * R; + device int32_t* mi = miss_ids + tid * R; + device float* ms = miss_scores + tid * R; + + // Pass 1: residency, mass sums, first kept slot. Ascending-r f32 + // accumulation order is part of the op contract (CPU parity). + int32_t slot_of[KQ_ROUTE_SHED_MAX_R]; + float s_all = 0.0f; + float s_kept = 0.0f; + int32_t first_kept = 0; + bool have_kept = false; + for (int r = 0; r < R; r++) { + const int32_t e = int32_t(id[r]); + const int32_t slot = (e >= 0 && e < E) ? slot_table[e] : -1; + slot_of[r] = slot; + const float sv = sc[r]; + s_all += sv; + if (slot >= 0) { + s_kept += sv; + if (!have_kept) { + first_kept = slot; + have_kept = true; + } + } + } + const float renorm = s_kept > 0.0f ? s_all / s_kept : 0.0f; + + // Pass 2: slots + mix; collect misses into registers. + int32_t m_id[KQ_ROUTE_SHED_MAX_R]; + float m_sc[KQ_ROUTE_SHED_MAX_R]; + int n_miss = 0; + for (int r = 0; r < R; r++) { + if (slot_of[r] >= 0) { + sl[r] = uint32_t(slot_of[r]); + mo[r] = sc[r] * renorm; + } else { + sl[r] = uint32_t(first_kept); + mo[r] = 0.0f; + m_id[n_miss] = int32_t(id[r]); + m_sc[n_miss] = sc[r]; + n_miss++; + } + } + // Misses front-packed in descending score order (prestage priority); + // stable insertion sort, matching the CPU eval. + for (int i = 1; i < n_miss; i++) { + const int32_t vi = m_id[i]; + const float vs = m_sc[i]; + int j = i - 1; + while (j >= 0 && m_sc[j] < vs) { + m_id[j + 1] = m_id[j]; + m_sc[j + 1] = m_sc[j]; + j--; + } + m_id[j + 1] = vi; + m_sc[j + 1] = vs; + } + for (int r = 0; r < n_miss; r++) { + mi[r] = m_id[r]; + ms[r] = m_sc[r]; + } + for (int r = n_miss; r < R; r++) { + mi[r] = -1; + ms[r] = 0.0f; + } +} diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py index dd8721c..2fb0d67 100644 --- a/mlx_kquant/__init__.py +++ b/mlx_kquant/__init__.py @@ -62,6 +62,7 @@ quantized_matmul_qmv_bias, rmsnorm2_add, rmsnorm_multi3, + route_shed, sdpa_decode_gqa, sdpa_fa_verify, sdpa_vector, @@ -117,6 +118,7 @@ "quantized_matmul_qmv_bias", "rmsnorm2_add", "rmsnorm_multi3", + "route_shed", "sdpa_decode_gqa", "sdpa_fa_verify", "sdpa_vector", diff --git a/src/kquant.h b/src/kquant.h index c0a97cb..cea057e 100644 --- a/src/kquant.h +++ b/src/kquant.h @@ -514,6 +514,24 @@ mx::array rmsnorm2_add( float eps, mx::StreamOrDevice s = {}); +// GPU-side routed-expert slot remap + residency shed for streamed MoE decode +// (the gpu-dispatch autonomous-token front end; see kq_route_shed.h). +// slot_table maps expert id -> arena slot, negative = non-resident. Every +// non-resident routed expert is shed (a lazy graph cannot demand-read disk); +// kept gate weights are renormalized mass-preserving (score * S_all/S_kept). +// Shed entries keep a valid slot (the row's first kept slot, else 0) with a +// zero mix weight so downstream gather kernels stay exact. Misses come back +// packed to the front of miss_ids/miss_scores in descending score order +// (prestage priority), -1 / 0 padded; the host consumes them between tokens. +// Over-budget accounting vs a keep-mass P is host-side arithmetic on +// miss_scores. Returns {slots u32 [T,R], mix f32 [T,R], miss_ids i32 [T,R], +// miss_scores f32 [T,R]}. +std::vector route_shed( + mx::array indices, + mx::array scores, + mx::array slot_table, + mx::StreamOrDevice s = {}); + // ----------------------------- primitives ----------------------------- // Dequantize a single uint8 K-quant wire-byte tensor. Inference-only: @@ -1352,6 +1370,30 @@ class KQuantRMSNorm2Add : public mx::Primitive { float eps_; }; +// Routed-expert slot remap + residency shed (see route_shed). +// Inference-only. +class KQuantRouteShed : public mx::Primitive { + public: + explicit KQuantRouteShed(mx::Stream stream) : mx::Primitive(stream) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override; + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + std::vector output_shapes( + const std::vector& inputs) override; + + const char* name() const override { + return "KQuantRouteShed"; + } + bool is_equivalent(const mx::Primitive& other) const override { + return true; + } +}; + // Gather (MoE) quantized matmul. vjp implements only the gradient wrt x (a // per-expert gather with the transpose flipped, scatter-added back to the // source x rows) so LoRA can train on a frozen kquant base; jvp/vmap and the diff --git a/src/kquant_route_shed.cpp b/src/kquant_route_shed.cpp new file mode 100644 index 0000000..1bbab88 --- /dev/null +++ b/src/kquant_route_shed.cpp @@ -0,0 +1,236 @@ +// Routed-expert slot remap + residency shed for streamed MoE decode (the +// gpu-dispatch autonomous-token front end). One tiny dispatch per MoE layer +// replaces the per-layer host round-trip: the kernel remaps routed expert ids +// to arena slots via a GPU-resident slot table, sheds every non-resident +// expert (a lazy graph cannot demand-read disk), renormalizes the kept gate +// weights mass-preserving, and reports the misses so the host can prestage +// between tokens. The CPU eval mirrors the kernel's f32 arithmetic and +// ordering exactly (parity harness in tests/test_route_shed.py). +#include +#include + +#include "kquant.h" +#include "kquant_internal.h" + +#include "mlx/backend/cpu/encoder.h" +#include "mlx/ops.h" +#include "mlx/utils.h" + +#ifdef _METAL_ +#include "kquant_metal_internal.h" // kq_get_kernel +#include "mlx/backend/metal/device.h" +#endif + +namespace mx = mlx::core; + +namespace mlx_kquant { + +namespace { + +constexpr int KQ_ROUTE_SHED_MAX_R = 64; + +} // namespace + +#ifdef _METAL_ + +void KQuantRouteShed::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& s = stream(); + auto& d = mx::metal::device(s.device); + for (auto& out : outputs) { + out.set_data(mx::allocator::malloc(out.nbytes())); + } + + const auto& indices = inputs[0]; + const auto& scores = inputs[1]; + const auto& table = inputs[2]; + + int R = indices.shape(-1); + int T = int(indices.size() / R); + int E = table.shape(0); + + auto kernel = kq_get_kernel(d, "kq_route_shed"); + auto& ce = mx::metal::get_command_encoder(s); + ce.set_compute_pipeline_state(kernel); + ce.set_input_array(indices, 0); + ce.set_input_array(scores, 1); + ce.set_input_array(table, 2); + ce.set_output_array(outputs[0], 3); + ce.set_output_array(outputs[1], 4); + ce.set_output_array(outputs[2], 5); + ce.set_output_array(outputs[3], 6); + ce.set_bytes(R, 7); + ce.set_bytes(E, 8); + ce.set_bytes(T, 9); + // One thread per token row: R <= 64 entries of serial work, launch cost + // dominates. Grid = T threads in groups of 32. + int tg = T < 32 ? T : 32; + MTL::Size group_dims(tg, 1, 1); + MTL::Size grid_dims((T + tg - 1) / tg, 1, 1); + ce.dispatch_threadgroups(grid_dims, group_dims); +} + +#else // !_METAL_ + +void KQuantRouteShed::eval_gpu( + const std::vector&, + std::vector&) { + throw std::runtime_error( + "[mlx_kquant.route_shed] GPU eval requires a Metal build."); +} + +#endif // _METAL_ + +void KQuantRouteShed::eval_cpu( + const std::vector& inputs, + std::vector& outputs) { + for (auto& out : outputs) { + out.set_data(mx::allocator::malloc(out.nbytes())); + } + + auto& encoder = mx::cpu::get_command_encoder(stream()); + for (const auto& in : inputs) { + encoder.set_input_array(in); + } + for (auto& out : outputs) { + encoder.set_output_array(out); + } + encoder.dispatch( + [indices = mx::array::unsafe_weak_copy(inputs[0]), + scores = mx::array::unsafe_weak_copy(inputs[1]), + table = mx::array::unsafe_weak_copy(inputs[2]), + slots = mx::array::unsafe_weak_copy(outputs[0]), + mix = mx::array::unsafe_weak_copy(outputs[1]), + miss_ids = mx::array::unsafe_weak_copy(outputs[2]), + miss_scores = mx::array::unsafe_weak_copy(outputs[3])]() mutable { + const int R = indices.shape(-1); + const int64_t T = indices.size() / R; + const int E = table.shape(0); + const uint32_t* idp = indices.data(); + const float* scp = scores.data(); + const int32_t* tbp = table.data(); + uint32_t* slp = slots.data(); + float* mxp = mix.data(); + int32_t* mip = miss_ids.data(); + float* msp = miss_scores.data(); + + for (int64_t t = 0; t < T; t++) { + const uint32_t* id = idp + t * R; + const float* sc = scp + t * R; + uint32_t* sl = slp + t * R; + float* mo = mxp + t * R; + int32_t* mi = mip + t * R; + float* ms = msp + t * R; + + // Pass 1: residency, mass sums, first kept slot. Ascending-r f32 + // accumulation order is part of the op contract (GPU parity). + int32_t slot_of[KQ_ROUTE_SHED_MAX_R]; + float s_all = 0.0f; + float s_kept = 0.0f; + int32_t first_kept = 0; + bool have_kept = false; + for (int r = 0; r < R; r++) { + const int32_t e = int32_t(id[r]); + const int32_t slot = (e >= 0 && e < E) ? tbp[e] : -1; + slot_of[r] = slot; + const float sv = sc[r]; + s_all += sv; + if (slot >= 0) { + s_kept += sv; + if (!have_kept) { + first_kept = slot; + have_kept = true; + } + } + } + const float renorm = s_kept > 0.0f ? s_all / s_kept : 0.0f; + + // Pass 2: slots + mix; collect misses. + int n_miss = 0; + for (int r = 0; r < R; r++) { + if (slot_of[r] >= 0) { + sl[r] = uint32_t(slot_of[r]); + mo[r] = sc[r] * renorm; + } else { + sl[r] = uint32_t(first_kept); + mo[r] = 0.0f; + mi[n_miss] = int32_t(id[r]); + ms[n_miss] = sc[r]; + n_miss++; + } + } + // Misses front-packed in descending score order (prestage priority); + // stable insertion sort, matching the kernel. + for (int i = 1; i < n_miss; i++) { + const int32_t vi = mi[i]; + const float vs = ms[i]; + int j = i - 1; + while (j >= 0 && ms[j] < vs) { + mi[j + 1] = mi[j]; + ms[j + 1] = ms[j]; + j--; + } + mi[j + 1] = vi; + ms[j + 1] = vs; + } + for (int r = n_miss; r < R; r++) { + mi[r] = -1; + ms[r] = 0.0f; + } + } + }); +} + +std::vector KQuantRouteShed::output_shapes( + const std::vector& inputs) { + const auto& shape = inputs[0].shape(); + return {shape, shape, shape, shape}; +} + +std::vector route_shed( + mx::array indices, + mx::array scores, + mx::array slot_table, + mx::StreamOrDevice s_) { + auto s = mx::to_stream(s_); + const char* op = "[mlx_kquant.route_shed]"; + if (indices.ndim() < 1 || indices.shape(-1) < 1) { + throw std::invalid_argument( + std::string(op) + " indices must have a last axis."); + } + if (indices.shape(-1) > KQ_ROUTE_SHED_MAX_R) { + throw std::invalid_argument( + std::string(op) + " routed width R must be <= " + + std::to_string(KQ_ROUTE_SHED_MAX_R) + "."); + } + if (indices.dtype() != mx::uint32) { + throw std::invalid_argument(std::string(op) + " indices must be uint32."); + } + if (scores.shape() != indices.shape()) { + throw std::invalid_argument( + std::string(op) + " scores must match indices in shape."); + } + if (scores.dtype() != mx::float32) { + throw std::invalid_argument(std::string(op) + " scores must be float32."); + } + if (slot_table.ndim() != 1 || slot_table.shape(0) < 1) { + throw std::invalid_argument( + std::string(op) + " slot_table must be 1-D [n_experts]."); + } + if (slot_table.dtype() != mx::int32) { + throw std::invalid_argument(std::string(op) + " slot_table must be int32."); + } + + auto contig = [&](const mx::array& a) { + return a.flags().row_contiguous ? a : mx::contiguous(a, false, s); + }; + const auto& shape = indices.shape(); + return mx::array::make_arrays( + {shape, shape, shape, shape}, + {mx::uint32, mx::float32, mx::int32, mx::float32}, + std::make_shared(s), + {contig(indices), contig(scores), contig(slot_table)}); +} + +} // namespace mlx_kquant diff --git a/tests/test_route_shed.py b/tests/test_route_shed.py new file mode 100644 index 0000000..5560cf4 --- /dev/null +++ b/tests/test_route_shed.py @@ -0,0 +1,175 @@ +"""route_shed: slot remap + residency shed for streamed MoE decode. + +The op contract (shed all non-resident, mass-preserving renorm over kept, +misses front-packed descending-score) is pinned by a pure-python reference; +the CPU eval must match it bit-for-bit and the GPU kernel must match the CPU +eval bit-for-bit (same f32 accumulation order). +""" + +import mlx.core as mx +import numpy as np +import pytest + +import mlx_kquant as kq + +gpu = pytest.mark.skipif( + not mx.metal.is_available(), reason="GPU parity requires Metal" +) + + +def ref_route_shed(indices, scores, slot_table): + """Pure-python reference, mirroring the kernel's arithmetic order.""" + T, R = indices.shape + E = slot_table.shape[0] + slots = np.zeros((T, R), np.uint32) + mix = np.zeros((T, R), np.float32) + miss_ids = np.full((T, R), -1, np.int32) + miss_scores = np.zeros((T, R), np.float32) + for t in range(T): + slot_of = [int(slot_table[e]) if 0 <= e < E else -1 for e in indices[t]] + s_all = np.float32(0.0) + s_kept = np.float32(0.0) + first_kept = 0 + have_kept = False + for r in range(R): + sv = np.float32(scores[t, r]) + s_all = np.float32(s_all + sv) + if slot_of[r] >= 0: + s_kept = np.float32(s_kept + sv) + if not have_kept: + first_kept = slot_of[r] + have_kept = True + renorm = np.float32(s_all / s_kept) if s_kept > 0 else np.float32(0.0) + misses = [] + for r in range(R): + if slot_of[r] >= 0: + slots[t, r] = slot_of[r] + mix[t, r] = np.float32(np.float32(scores[t, r]) * renorm) + else: + slots[t, r] = first_kept + misses.append((int(indices[t, r]), np.float32(scores[t, r]))) + # Stable descending-score order. + misses.sort(key=lambda m: -m[1]) + for i, (e, sv) in enumerate(misses): + miss_ids[t, i] = e + miss_scores[t, i] = sv + return slots, mix, miss_ids, miss_scores + + +def make_case(T=4, R=8, E=256, resident_frac=0.7, seed=0): + rng = np.random.default_rng(seed) + indices = np.stack([rng.choice(E, size=R, replace=False) for _ in range(T)]).astype( + np.uint32 + ) + scores = rng.uniform(0.01, 1.0, size=(T, R)).astype(np.float32) + slot_table = np.full(E, -1, np.int32) + resident = rng.choice(E, size=int(E * resident_frac), replace=False) + slot_table[resident] = np.arange(len(resident), dtype=np.int32) + return indices, scores, slot_table + + +def run_op(indices, scores, slot_table, stream=None): + kwargs = {} if stream is None else {"stream": stream} + outs = kq.route_shed( + mx.array(indices), mx.array(scores), mx.array(slot_table), **kwargs + ) + mx.eval(*outs) + return [np.array(o) for o in outs] + + +@pytest.mark.parametrize("seed", range(5)) +@pytest.mark.parametrize("resident_frac", [0.0, 0.3, 0.7, 1.0]) +def test_cpu_matches_reference(seed, resident_frac): + indices, scores, slot_table = make_case(resident_frac=resident_frac, seed=seed) + got = run_op(indices, scores, slot_table, stream=mx.cpu) + want = ref_route_shed(indices, scores, slot_table) + for g, w in zip(got, want, strict=True): + np.testing.assert_array_equal(g, w) + + +def test_all_resident_is_identity(): + indices, scores, slot_table = make_case(resident_frac=1.0, seed=1) + slots, mix, miss_ids, miss_scores = run_op( + indices, scores, slot_table, stream=mx.cpu + ) + np.testing.assert_array_equal(slots, slot_table[indices.astype(np.int64)]) + np.testing.assert_allclose(mix, scores, rtol=1e-6) + assert (miss_ids == -1).all() + assert (miss_scores == 0).all() + + +def test_all_miss_row_zero_mix(): + indices, scores, slot_table = make_case(resident_frac=0.0, seed=2) + slots, mix, miss_ids, miss_scores = run_op( + indices, scores, slot_table, stream=mx.cpu + ) + assert (mix == 0).all() + assert (slots == 0).all() + assert (miss_ids >= 0).all() + # Descending score, all entries reported. + assert (np.diff(miss_scores, axis=-1) <= 0).all() + np.testing.assert_allclose( + np.sort(miss_scores, axis=-1), np.sort(scores, axis=-1), rtol=0 + ) + + +def test_mass_preserving_renorm(): + indices, scores, slot_table = make_case(resident_frac=0.5, seed=3) + _, mix, _, _ = run_op(indices, scores, slot_table, stream=mx.cpu) + np.testing.assert_allclose(mix.sum(axis=-1), scores.sum(axis=-1), rtol=1e-5) + + +def test_shed_entries_reuse_first_kept_slot(): + E = 16 + indices = np.array([[3, 5, 7, 9]], np.uint32) + scores = np.array([[0.4, 0.3, 0.2, 0.1]], np.float32) + slot_table = np.full(E, -1, np.int32) + slot_table[5] = 11 # only expert 5 resident + slots, mix, miss_ids, miss_scores = run_op( + indices, scores, slot_table, stream=mx.cpu + ) + np.testing.assert_array_equal(slots, [[11, 11, 11, 11]]) + np.testing.assert_allclose(mix[0, 1], 1.0, rtol=1e-6) # 0.3 * (1.0/0.3) + assert mix[0, 0] == mix[0, 2] == mix[0, 3] == 0 + np.testing.assert_array_equal(miss_ids[0], [3, 7, 9, -1]) + np.testing.assert_allclose(miss_scores[0], [0.4, 0.2, 0.1, 0.0]) + + +def test_input_validation(): + indices, scores, slot_table = make_case() + with pytest.raises(ValueError): + kq.route_shed( + mx.array(indices.astype(np.int32)), + mx.array(scores), + mx.array(slot_table), + ) + with pytest.raises(ValueError): + kq.route_shed( + mx.array(indices), + mx.array(scores.astype(np.float16)), + mx.array(slot_table), + ) + with pytest.raises(ValueError): + kq.route_shed( + mx.array(indices), + mx.array(scores), + mx.array(slot_table.astype(np.int16)), + ) + with pytest.raises(ValueError): + kq.route_shed( + mx.array(np.zeros((1, 65), np.uint32)), + mx.array(np.zeros((1, 65), np.float32)), + mx.array(slot_table), + ) + + +@gpu +@pytest.mark.parametrize("seed", range(3)) +@pytest.mark.parametrize("shape", [(1, 8), (4, 8), (33, 4), (2, 64)]) +def test_gpu_matches_cpu_bitwise(seed, shape): + T, R = shape + indices, scores, slot_table = make_case(T=T, R=R, resident_frac=0.6, seed=seed) + cpu = run_op(indices, scores, slot_table, stream=mx.cpu) + dev = run_op(indices, scores, slot_table, stream=mx.gpu) + for c, g in zip(cpu, dev, strict=True): + np.testing.assert_array_equal(c, g) From 74a93d96aa377f4344a5da6a4bbe1f46d785e6c8 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:58:40 -0700 Subject: [PATCH 2/2] changelog: gpu-dispatch entries under unreleased --- .github/workflows/ci.yml | 7 ------- CHANGELOG.md | 5 +++++ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 36294ae..e70f146 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,13 +59,6 @@ jobs: with: python-version: ${{ matrix.python }} - # The extension is a Metal C++ module - it compiles on macOS arm64 with - # the Xcode toolchain (no GPU needed to *build*). Pinned mlx wheel: the - # kernels include MLX's steel headers and link libmlx, so the ABI must - # match exactly. - # mlx-lm is the [tools] dep the model-level tests (loader round-trip) need; - # without a GPU those tests skip, but installing it lets them run on a - # GPU-bearing runner. - name: Build extension + test deps run: | python -m pip install --upgrade pip diff --git a/CHANGELOG.md b/CHANGELOG.md index 30b4976..a70634c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added +- `route_shed(indices, scores, slot_table)`: GPU-side routed-expert slot + remap plus residency shed for streamed MoE decode; non-resident experts + are shed and reported (miss ids and scores) without a host sync. + ## [0.3.7] Build and publish cp314 wheel