perf(deepseek_v4): faster exact DeepSeek-V4 prefill (wsdpa fusion + scheduler/cache amortization)

omlx/custom_kernels/glm_moe_dsa/csrc/bindings.cpp

@@ -39,6 +39,8 @@ NB_MODULE(_ext, m) {
       "unused_causal_prefix_topk"_a = 0,
       "skip_causal_future_store"_a = false,
       "causal_q_offset"_a = -1,
+      "mask_ratio"_a = 0,
+      "mask_q_offset"_a = 0,
       "stream"_a = nb::none());
   m.def(
       "dsa_topk_indices",

omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.cpp

@@ -54,13 +54,17 @@ class DSAIndexerScoresPrimitive : public Primitive {
       bool weights_lh,
       int unused_causal_prefix_topk,
       bool skip_causal_future_store,
-      int causal_q_offset)
+      int causal_q_offset,
+      int mask_ratio,
+      int mask_q_offset)
       : Primitive(stream),
         causal_(causal),
         weights_lh_(weights_lh),
         unused_causal_prefix_topk_(unused_causal_prefix_topk),
         skip_causal_future_store_(skip_causal_future_store),
-        causal_q_offset_(causal_q_offset) {}
+        causal_q_offset_(causal_q_offset),
+        mask_ratio_(mask_ratio),
+        mask_q_offset_(mask_q_offset) {}
 
   static bool unsupported(
       const array& q,
@@ -128,16 +132,26 @@ class DSAIndexerScoresPrimitive : public Primitive {
     out.set_data(allocator::malloc(out.nbytes()));
 
     constexpr int bm = 64;
-    constexpr int bn = 64;
     constexpr int bk = 16;
-    constexpr int wm = 2;
-    constexpr int wn = 2;
 
     const int B = q.shape(0);
     const int H = q.shape(1);
     const int M = q.shape(2);
     const int N = k.shape(2);
     const int D = q.shape(3);
+
+    // bm/bn/wm/wn do not enter the per-element K-reduction order (bk=16 and
+    // the MMA fragment K-layout are unchanged), so the tile config only
+    // affects scheduling. bn=128 (paired with wm=2,wn=4) was measured on
+    // M3 Ultra (L=2048, bf16, H=64): it ties bn=64 at P=25k but is ~9%
+    // slower at P=125k and ~11% slower at P=2.5k — the kernel is
+    // compute/barrier-bound (Q and pooled-K panels are largely
+    // L2-resident), so the traffic reduction does not pay. bn=64/wm2/wn2
+    // is the fixed configuration.
+    const int bn = 64;
+    const int wm = 2;
+    const int wn = 2;
+
     const int tiles_m = (M + bm - 1) / bm;
     const int tiles_n = (N + bn - 1) / bn;
 
@@ -203,6 +217,8 @@ class DSAIndexerScoresPrimitive : public Primitive {
     compute_encoder.set_bytes(unused_causal_prefix_topk_, 6);
     compute_encoder.set_bytes(skip_causal_future_store_, 7);
     compute_encoder.set_bytes(causal_q_offset_, 8);
+    compute_encoder.set_bytes(mask_ratio_, 9);
+    compute_encoder.set_bytes(mask_q_offset_, 10);
 
     MTL::Size group_dims = MTL::Size(wm * wn * 32, 1, 1);
     MTL::Size grid_dims = MTL::Size(tiles_n, tiles_m, B);
@@ -216,15 +232,19 @@ class DSAIndexerScoresPrimitive : public Primitive {
     return causal_ == rhs.causal_ && weights_lh_ == rhs.weights_lh_ &&
         unused_causal_prefix_topk_ == rhs.unused_causal_prefix_topk_ &&
         skip_causal_future_store_ == rhs.skip_causal_future_store_ &&
-        causal_q_offset_ == rhs.causal_q_offset_;
+        causal_q_offset_ == rhs.causal_q_offset_ &&
+        mask_ratio_ == rhs.mask_ratio_ &&
+        mask_q_offset_ == rhs.mask_q_offset_;
   }
   auto state() const {
     return std::make_tuple(
         causal_,
         weights_lh_,
         unused_causal_prefix_topk_,
         skip_causal_future_store_,
-        causal_q_offset_);
+        causal_q_offset_,
+        mask_ratio_,
+        mask_q_offset_);
   }
 
  private:
@@ -233,6 +253,8 @@ class DSAIndexerScoresPrimitive : public Primitive {
   int unused_causal_prefix_topk_;
   bool skip_causal_future_store_;
   int causal_q_offset_;
+  int mask_ratio_;
+  int mask_q_offset_;
 };
 
 class DSATopKIndicesPrimitive : public Primitive {
@@ -572,6 +594,8 @@ array dsa_indexer_scores(
     int unused_causal_prefix_topk,
     bool skip_causal_future_store,
     int causal_q_offset,
+    int mask_ratio,
+    int mask_q_offset,
     StreamOrDevice s) {
   if (queries.ndim() != 4 || keys.ndim() != 4 ||
       (weights.ndim() != 3 && weights.ndim() != 4)) {
@@ -625,6 +649,19 @@ array dsa_indexer_scores(
         << "-1 or non-negative, got " << causal_q_offset << ".";
     throw std::invalid_argument(msg.str());
   }
+  if (mask_ratio < 0) {
+    std::ostringstream msg;
+    msg << "[omlx_glm_kernels.dsa_indexer_scores] mask_ratio must be "
+        << "non-negative (0 disables the fused pooled-causal mask), got "
+        << mask_ratio << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (mask_ratio > 0 && mask_q_offset < 0) {
+    std::ostringstream msg;
+    msg << "[omlx_glm_kernels.dsa_indexer_scores] mask_q_offset must be "
+        << "non-negative when mask_ratio > 0, got " << mask_q_offset << ".";
+    throw std::invalid_argument(msg.str());
+  }
 
   auto stream = to_stream(s);
   auto q = ensure_row_contiguous(astype(queries, final_type, stream), stream);
@@ -647,7 +684,9 @@ array dsa_indexer_scores(
           weights_lh,
           unused_causal_prefix_topk,
           skip_causal_future_store,
-          causal_q_offset),
+          causal_q_offset,
+          mask_ratio,
+          mask_q_offset),
       std::move(inputs));
 }
 

omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.h

@@ -16,6 +16,13 @@ mx::array dsa_indexer_scores(
     int unused_causal_prefix_topk = 0,
     bool skip_causal_future_store = false,
     int causal_q_offset = -1,
+    // Optional fused pooled-ratio causal mask (0 = disabled, the historical
+    // behavior). When mask_ratio > 0, pooled column c is masked for query row
+    // r iff c >= (mask_q_offset + r + 1) / mask_ratio, and masked positions
+    // are written as finfo(dtype).min in the kernel epilogue — bit-identical
+    // to the mx.where pass it replaces.
+    int mask_ratio = 0,
+    int mask_q_offset = 0,
     mx::StreamOrDevice s = {});
 
 mx::array dsa_topk_indices(

omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.metal

@@ -34,8 +34,20 @@ struct OMLXDSATopKParams {
       topk,                                                             \
       threads)
 
+// bn=64,wm=2,wn=2 (128 threads): the historical config, default for all
+// shapes. bn=128,wm=2,wn=4 (256 threads): halves A-tile re-reads and barrier
+// count per output element for N % 128 == 0 shapes; wn=4 keeps per-thread
+// accumulator/fragment register pressure identical to bn=64 (wn=2 at bn=128
+// doubles it and measurably regresses). bm/bn/wm/wn never enter the
+// per-element K-reduction order (bk=16 and the MMA fragment K-layout are
+// unchanged), so all configs are bit-identical.
+// NOTE: a K-panel hoist (resident B tile across heads) was tried and
+// REVERTED: the 27KB threadgroup footprint collapses occupancy to 1
+// threadgroup/core (32KB budget) and costs ~18% despite the traffic saving.
 instantiate_dsa_indexer_score(float16, half, 64, 64, 16, 2, 2);
 instantiate_dsa_indexer_score(bfloat16, bfloat16_t, 64, 64, 16, 2, 2);
+instantiate_dsa_indexer_score(float16, half, 64, 128, 16, 2, 4);
+instantiate_dsa_indexer_score(bfloat16, bfloat16_t, 64, 128, 16, 2, 4);
 
 instantiate_dsa_topk_indices(float16, half, 2048, 1024);
 instantiate_dsa_topk_indices(bfloat16, bfloat16_t, 2048, 1024);

omlx/custom_kernels/glm_moe_dsa/csrc/kernels/steel_dsa_indexer_score.h

@@ -20,6 +20,16 @@ METAL_FUNC uint dsa_ordered_key_16_bits(ushort bits) {
   return (bits & 0x8000) ? uint((~bits) & 0xffff) : uint(bits | 0x8000);
 }
 
+// finfo(T).min as an exact bit pattern: bf16 0xFF7F (-3.3895313892515355e38),
+// fp16 0xFBFF (-65504). This is the sentinel the call site's mx.where pass
+// writes (mx.finfo(dtype).min) — finite, and distinct from the -inf the
+// kernel's own causal path uses. Bit-exact by construction.
+template <typename T>
+METAL_FUNC T dsa_finfo_min() {
+  return as_type<T>(
+      ushort(metal::is_same<T, bfloat16_t>::value ? 0xFF7F : 0xFBFF));
+}
+
 template <typename T, typename O, int TOPK, int THREADS>
 [[kernel, max_total_threads_per_threadgroup(THREADS)]] void dsa_topk_indices_16bit(
     const device T* scores [[buffer(0)]],
@@ -358,6 +368,8 @@ dsa_indexer_score(
     const constant int& unused_causal_prefix_topk [[buffer(6)]],
     const constant bool& skip_causal_future_store [[buffer(7)]],
     const constant int& causal_q_offset [[buffer(8)]],
+    const constant int& mask_ratio [[buffer(9)]],
+    const constant int& mask_q_offset [[buffer(10)]],
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]],
     uint3 tid [[threadgroup_position_in_grid]],
@@ -483,6 +495,7 @@ dsa_indexer_score(
     }
   }
 
+  const T pooled_sentinel = dsa_finfo_min<T>();
   device T* Dst = O + size_t(mma_op.sm) * params->ldd + mma_op.sn;
   short ai = 0;
   STEEL_PRAGMA_UNROLL
@@ -498,8 +511,21 @@ dsa_indexer_score(
       for (short e = 0; e < decltype(mma_op.Ctile)::kElemsPerFrag; ++e) {
         const int col = col_base + e;
         const bool future = do_causal && col > q_offset + row;
+        // ── Pooled-ratio causal mask (lossless opt 3) ─────────────────────
+        // Folds the call site's separate mx.where pass into the epilogue.
+        // Validity rule (cache_extras.py PoolingCache.make_mask): pooled
+        // column `col` is visible to query row `row` iff
+        //   col < (mask_q_offset + row + 1) // mask_ratio
+        // so masked iff col >= that bound (int operands are non-negative,
+        // C++ truncation == Python floor). Masked positions receive the SAME
+        // sentinel the where pass wrote — finfo(T).min, not -inf — so the
+        // post-mask output is bit-identical to the old two-step path.
+        // mask_ratio == 0 disables the mode (pmask == None callers).
+        const bool pooled_masked = mask_ratio > 0 &&
+            col >= (mask_q_offset + row + 1) / mask_ratio;
         const T value = future ? static_cast<T>(-INFINITY)
-                               : static_cast<T>(accum[ai]);
+            : pooled_masked  ? pooled_sentinel
+                             : static_cast<T>(accum[ai]);
         Dst[out_base + e] = value;
         ai++;
       }

omlx/custom_kernels/glm_moe_dsa/fast.py

@@ -69,6 +69,27 @@ def _verify_abi(ext, import_error):
 _ext, _IMPORT_ERROR = _verify_abi(_ext, _IMPORT_ERROR)
 
 
+def _probe_mask_fold(ext) -> bool:
+    """True iff the built extension accepts the mask-fold kwargs.
+
+    ``dsa_indexer_scores`` grew ``mask_ratio``/``mask_q_offset`` after the
+    first native builds shipped. An older ``_ext`` parses no such kwargs, so
+    passing them unconditionally raises ``TypeError`` for every caller —
+    including GLM-5.2's historical unmasked path. Nanobind renders named
+    args into ``__doc__``, so probe the signature once at import; callers
+    on older builds keep the historical call signature and the mask is
+    applied in a second pass with identical sentinel semantics.
+    """
+    fn = getattr(ext, "dsa_indexer_scores", None)
+    if fn is None:
+        return False
+    doc = getattr(fn, "__doc__", None) or ""
+    return "mask_ratio" in doc and "mask_q_offset" in doc
+
+
+_EXT_MASK_FOLD = _probe_mask_fold(_ext)
+
+
 NATIVE_SYMBOLS = (
     "dsa_decode_scores",
     "dsa_indexer_scores",
@@ -128,10 +149,23 @@ def dsa_indexer_scores(
     unused_causal_prefix_topk: int = 0,
     skip_causal_future_store: bool = False,
     causal_q_offset: int = -1,
+    mask_ratio: int = 0,
+    mask_q_offset: int = 0,
     *,
     stream=None,
 ) -> mx.array:
-    if _ext is not None:
+    """Head-summed DSA indexer scores.
+
+    ``mask_ratio > 0`` folds the pooled-ratio causal mask into the kernel
+    epilogue: pooled column ``c`` is masked for query row ``r`` iff
+    ``c >= (mask_q_offset + r + 1) // mask_ratio`` and receives the
+    ``finfo(dtype).min`` sentinel — bit-identical to applying
+    ``mx.where(mask, scores, finfo.min)`` in a second pass. ``mask_ratio=0``
+    (default) is the historical unmasked behavior. On extension builds
+    predating the fold kwargs, the historical call signature is kept and
+    the same mask is applied in a second pass with identical semantics.
+    """
+    if _ext is not None and _EXT_MASK_FOLD:
         return _ext.dsa_indexer_scores(
             queries,
             keys,
@@ -140,18 +174,46 @@ def dsa_indexer_scores(
             unused_causal_prefix_topk=unused_causal_prefix_topk,
             skip_causal_future_store=skip_causal_future_store,
             causal_q_offset=causal_q_offset,
+            mask_ratio=mask_ratio,
+            mask_q_offset=mask_q_offset,
             **_native_stream_kwargs(stream),
         )
-    return mx.fast.dsa_indexer_scores(
-        queries,
-        keys,
-        weights,
-        causal=causal,
-        unused_causal_prefix_topk=unused_causal_prefix_topk,
-        skip_causal_future_store=skip_causal_future_store,
-        causal_q_offset=causal_q_offset,
-        stream=stream or mx.gpu,
-    )
+    if _ext is not None:
+        # Older build without the mask-fold kwargs: keep the historical
+        # call signature; the mask is applied in a second pass below.
+        scores = _ext.dsa_indexer_scores(
+            queries,
+            keys,
+            weights,
+            causal=causal,
+            unused_causal_prefix_topk=unused_causal_prefix_topk,
+            skip_causal_future_store=skip_causal_future_store,
+            causal_q_offset=causal_q_offset,
+            **_native_stream_kwargs(stream),
+        )
+    else:
+        scores = mx.fast.dsa_indexer_scores(
+            queries,
+            keys,
+            weights,
+            causal=causal,
+            unused_causal_prefix_topk=unused_causal_prefix_topk,
+            skip_causal_future_store=skip_causal_future_store,
+            causal_q_offset=causal_q_offset,
+            stream=stream or mx.gpu,
+        )
+    if mask_ratio > 0:
+        # Preserve the fused kernel's exact sentinel semantics on the
+        # non-fused paths (same validity rule, same finfo.min sentinel).
+        L = queries.shape[2]
+        P = keys.shape[2]
+        pool_idx = mx.arange(P)
+        query_idx = mx.arange(mask_q_offset + 1, mask_q_offset + L + 1)
+        mask = pool_idx < query_idx[:, None] // mask_ratio
+        scores = mx.where(
+            mask[None, None], scores, mx.finfo(scores.dtype).min
+        )
+    return scores
 
 
 def dsa_decode_scores(

omlx/patches/deepseek_v4/__init__.py

@@ -129,6 +129,37 @@ def _register_model_type_aliases() -> None:
         sys.modules.setdefault("mlx_lm.models.deepseek_v4_mtp", base_module)
 
 
+def _probe_native_indexer_kernels() -> None:
+    """One-time startup probe for the native indexer kernels.
+
+    The per-chunk gate in ``Indexer.__call__`` only warns the first time a
+    native-eligible chunk actually hits the fallback, which can be minutes
+    into a long prefill. Surface availability (and the rebuild fix) once at
+    patch time instead.
+    """
+    try:
+        from omlx.custom_kernels.glm_moe_dsa import fast as glm_fast
+
+        have = glm_fast.has_symbol("dsa_indexer_scores") and glm_fast.has_symbol(
+            "dsa_topk_indices"
+        )
+    except Exception:
+        have = False
+    if have:
+        logger.info(
+            "deepseek_v4: native indexer kernels available "
+            "(dsa_indexer_scores/dsa_topk_indices)"
+        )
+    else:
+        logger.warning(
+            "deepseek_v4: native indexer kernels dsa_indexer_scores/"
+            "dsa_topk_indices unavailable (glm_moe_dsa extension not built); "
+            "the indexer will use the MLX fallback and long-context prefill "
+            "will be several times slower. Rebuild with "
+            "OMLX_WITH_CUSTOM_KERNEL=1."
+        )
+
+
 def apply_deepseek_v4_patch() -> bool:
     """Apply the DeepSeek V4 patch to mlx-lm. Idempotent.
 
@@ -184,6 +215,9 @@ def apply_deepseek_v4_patch() -> bool:
     # 8. Register omlx-side cache handlers.
     _register_cache_handlers()
 
+    # 9. One-time native indexer kernel probe (INFO/WARNING only).
+    _probe_native_indexer_kernels()
+
     _APPLIED = True
     logger.info("DeepSeek V4 patch applied (PR 1192 head %s)", PR_HEAD_SHA[:8])
     return True