feat(kernel): deepseek4 hc glue ops, cb caps, faster iq gather decode

CHANGELOG.md

@@ -10,6 +10,19 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - `skinny_matmul`: x @ w.T at token widths 1..16 against small-N large-K
   nn.Linear-layout weights, 4-8x faster than the stock GEMM at widths
   2..16 (router gates, indexer projections at speculative verify widths).
+- `hc_front_reduce` / `hc_front_expand_reduce` / `hc_sinkhorn_collapse` /
+  `hc_expand`: fused deepseek4 hyper-connection glue for single-token
+  decode; replaces ~176 python kernel launches per step with 4 native ops.
+- `get_cb_caps` / `set_cb_caps`: runtime read/write of MLX's command
+  buffer split caps, so a server can run coarse buffers during decode and
+  fine buffers during deep prefill.
+
+### Changed
+- iq2_xxs / iq2_xs / iq2_s / iq3_s MoE gather decode is 9-12% faster per
+  call (hoisted block scale, byte-indexed grids); the ext mat-vec at
+  verify widths 2..8 gains 7-10% on the same codecs.
+- Score-mixed MoE down gather gains a slot-parallel kernel at decode
+  scale (bit-identical; KQ_MOE_SP forces either form).
 
 ## [0.3.9]
 

CMakeLists.txt

@@ -127,6 +127,8 @@ 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_hc_glue.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/src/kquant_cb_caps.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_route_shed.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_skinny_mv.cpp
          ${CMAKE_CURRENT_LIST_DIR}/src/kquant_gather.cpp
@@ -209,6 +211,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_hc_glue.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_route_shed.metal
     ${CMAKE_CURRENT_LIST_DIR}/metal/kq_skinny_mv.metal
     INCLUDE_DIRS

bindings.cpp

@@ -1,5 +1,6 @@
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/optional.h>
+#include <nanobind/stl/pair.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/variant.h>
 #include <nanobind/stl/vector.h>
@@ -1151,6 +1152,140 @@ NB_MODULE(_ext, m) {
             array: same shape and dtype as a.
       )");
 
+  m.def(
+      "hc_front_reduce",
+      &mlx_kquant::hc_front_reduce,
+      "x"_a,
+      "fn"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        Hyper-connection front reduction for the fused M=1 decode route:
+        the 24 mix dots of x against fn plus the row sum of squares
+        (deferred rms factor). hc_mult 4 only.
+
+        Args:
+            x (array): [..., 4, D] streams, float16/bfloat16, D % 8 == 0,
+                D <= 8192.
+            fn (array): [24, 4 * D] float32 mix matrix.
+
+        Returns:
+            tuple: (mixes_raw f32 [..., 24], sumsq f32 [..., 1]).
+      )");
+
+  m.def(
+      "hc_front_expand_reduce",
+      &mlx_kquant::hc_front_expand_reduce,
+      "x_sub"_a,
+      "resid"_a,
+      "post"_a,
+      "comb"_a,
+      "fn"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        The previous cycle's hc_expand fused ahead of the front reduction:
+        one dispatch expands (x_sub, resid, post, comb) to h, writes it,
+        and reduces the mix dots and sum of squares of h.
+
+        Args:
+            x_sub (array): [..., D] sublayer output.
+            resid (array): [..., 4, D] residual streams.
+            post (array): [..., 4] float32.
+            comb (array): [..., 4, 4] float32.
+            fn (array): [24, 4 * D] float32 mix matrix.
+
+        Returns:
+            tuple: (h [..., 4, D], mixes_raw f32 [..., 24],
+            sumsq f32 [..., 1]); h is bit-identical to the unfused expand.
+      )");
+
+  m.def(
+      "hc_sinkhorn_collapse",
+      &mlx_kquant::hc_sinkhorn_collapse,
+      "x"_a,
+      "mixes_raw"_a,
+      "sumsq"_a,
+      "scale"_a,
+      "base"_a,
+      "w"_a,
+      "iters"_a,
+      "hc_eps"_a,
+      "norm_eps"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        Sinkhorn mix normalization plus stream collapse with the sublayer
+        RMSNorm folded into the single output rounding. The deferred front
+        rms factor enters through sumsq and multiplies the three scales.
+
+        Args:
+            x (array): [..., 4, D] streams, float16/bfloat16.
+            mixes_raw (array): [..., 24] float32 from the front reduction.
+            sumsq (array): [..., 1] float32 row sum of squares.
+            scale (array): [3] float32 pre/post/comb scales.
+            base (array): [24] float32 mix biases.
+            w (array): [D] sublayer norm weight, same dtype as x.
+            iters (int): sinkhorn iterations.
+            hc_eps (float): sinkhorn epsilon.
+            norm_eps (float): rms_norm epsilon.
+
+        Returns:
+            tuple: (collapsed [..., D], post f32 [..., 4],
+            comb f32 [..., 4, 4]).
+      )");
+
+  m.def(
+      "hc_expand",
+      &mlx_kquant::hc_expand,
+      "x"_a,
+      "resid"_a,
+      "post"_a,
+      "comb"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        Expand the sublayer output back to four streams:
+        out[i] = post[i] * x + sum_j comb[j][i] * resid[j].
+
+        Args:
+            x (array): [..., D] sublayer output.
+            resid (array): [..., 4, D] residual streams.
+            post (array): [..., 4] float32.
+            comb (array): [..., 4, 4] float32.
+
+        Returns:
+            array: [..., 4, D], dtype of resid.
+      )");
+
+  m.def(
+      "get_cb_caps",
+      &mlx_kquant::get_cb_caps,
+      R"(
+        Read MLX's live command-buffer split caps.
+
+        Returns:
+            tuple: (max_ops_per_buffer, max_mb_per_buffer).
+      )");
+
+  m.def(
+      "set_cb_caps",
+      &mlx_kquant::set_cb_caps,
+      "max_ops"_a,
+      "max_mb"_a,
+      R"(
+        Set MLX's command-buffer split caps at runtime. The env knobs
+        latch at device init; decode wants coarse buffers, deep prefill
+        fine ones, so servers flip these per phase.
+
+        Args:
+            max_ops (int): ops per command buffer, in [1, 2^30].
+            max_mb (int): MB per command buffer, in [1, 2^30].
+
+        Returns:
+            tuple: the previous (max_ops, max_mb).
+      )");
+
   m.def(
       "skinny_matmul",
       &mlx_kquant::skinny_matmul,

metal/kq_hc_glue.metal

@@ -0,0 +1,14 @@
+// clang-format off
+// Hyper-connection M=1 glue kernel instantiations; see kq_hc_glue.h.
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/kq_hc_glue.h"
+
+instantiate_kernel("kq_hc_front_reduce_bfloat16_t", kq_hc_front_reduce, bfloat16_t)
+instantiate_kernel("kq_hc_front_reduce_float16_t", kq_hc_front_reduce, float16_t)
+instantiate_kernel("kq_hc_front_expand_reduce_bfloat16_t", kq_hc_front_expand_reduce, bfloat16_t)
+instantiate_kernel("kq_hc_front_expand_reduce_float16_t", kq_hc_front_expand_reduce, float16_t)
+instantiate_kernel("kq_hc_sinkhorn_collapse_bfloat16_t", kq_hc_sinkhorn_collapse, bfloat16_t)
+instantiate_kernel("kq_hc_sinkhorn_collapse_float16_t", kq_hc_sinkhorn_collapse, float16_t)
+instantiate_kernel("kq_hc_expand_bfloat16_t", kq_hc_expand, bfloat16_t)
+instantiate_kernel("kq_hc_expand_float16_t", kq_hc_expand, float16_t)
+    // clang-format on

metal/kq_moe_glu_kq.metal

@@ -129,10 +129,18 @@ instantiate_kq_moe_glu_kq_fine(q8_0, float16_t)
       "kq_" #codec "_gather_qmv_mix_ns" sfx "_float16_t",                     \
       kq_ext_gather_qmv_mix_ns, float16_t, traits, nx)
 
+// The slot-parallel variant is NX = 8 only (wide K-lanes measured
+// flat-to-negative; sp multiplies threads without shortening K-chains).
 #define instantiate_kq_ext_mix_ns(codec, traits)                              \
   instantiate_kq_ext_mix_ns_nx(codec, traits, 8, "")                           \
   instantiate_kq_ext_mix_ns_nx(codec, traits, 16, "_nx16")                     \
-  instantiate_kq_ext_mix_ns_nx(codec, traits, 32, "_nx32")
+  instantiate_kq_ext_mix_ns_nx(codec, traits, 32, "_nx32")                     \
+  instantiate_kernel(                                                         \
+      "kq_" #codec "_gather_qmv_mix_ns_sp_bfloat16_t",                        \
+      kq_ext_gather_qmv_mix_ns_sp, bfloat16_t, traits, 8)                      \
+  instantiate_kernel(                                                         \
+      "kq_" #codec "_gather_qmv_mix_ns_sp_float16_t",                         \
+      kq_ext_gather_qmv_mix_ns_sp, float16_t, traits, 8)
 
 // Biased experts (gpt-oss): per-(expert, out_dim) f32 biases fused into the
 // GLU epilogue / qmv store. Only the clamped-SwiGLU epilogue is emitted --