feat(streaming): GPU-side routed-expert slot remap + residency shed

.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

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

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}

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,

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

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;
+  }
+}

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",

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<mx::array> 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<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+  void eval_gpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+
+  std::vector<mx::Shape> output_shapes(
+      const std::vector<mx::array>& 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