diff --git a/omlx/custom_kernels/glm_moe_dsa/csrc/bindings.cpp b/omlx/custom_kernels/glm_moe_dsa/csrc/bindings.cpp index 491f21773..6231bf78b 100644 --- a/omlx/custom_kernels/glm_moe_dsa/csrc/bindings.cpp +++ b/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", diff --git a/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.cpp b/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.cpp index 795c80780..5cbeb55ae 100644 --- a/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.cpp +++ b/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,7 +232,9 @@ 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( @@ -224,7 +242,9 @@ class DSAIndexerScoresPrimitive : public Primitive { 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)); } diff --git a/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.h b/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.h index 0a58e096e..30d22910e 100644 --- a/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.h +++ b/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( diff --git a/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.metal b/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.metal index ba387870c..a2a1a8f4c 100644 --- a/omlx/custom_kernels/glm_moe_dsa/csrc/dsa_indexer.metal +++ b/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); diff --git a/omlx/custom_kernels/glm_moe_dsa/csrc/kernels/steel_dsa_indexer_score.h b/omlx/custom_kernels/glm_moe_dsa/csrc/kernels/steel_dsa_indexer_score.h index 3b31c70b6..90ca8a69a 100644 --- a/omlx/custom_kernels/glm_moe_dsa/csrc/kernels/steel_dsa_indexer_score.h +++ b/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 +METAL_FUNC T dsa_finfo_min() { + return as_type( + ushort(metal::is_same::value ? 0xFF7F : 0xFBFF)); +} + template [[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(); 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(-INFINITY) - : static_cast(accum[ai]); + : pooled_masked ? pooled_sentinel + : static_cast(accum[ai]); Dst[out_base + e] = value; ai++; } diff --git a/omlx/custom_kernels/glm_moe_dsa/fast.py b/omlx/custom_kernels/glm_moe_dsa/fast.py index 4a589496b..add0b10dc 100644 --- a/omlx/custom_kernels/glm_moe_dsa/fast.py +++ b/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( diff --git a/omlx/patches/deepseek_v4/__init__.py b/omlx/patches/deepseek_v4/__init__.py index 029149984..683f80db8 100644 --- a/omlx/patches/deepseek_v4/__init__.py +++ b/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 diff --git a/omlx/patches/deepseek_v4/cache_extras.py b/omlx/patches/deepseek_v4/cache_extras.py index 14a3f99da..cb1fe8f52 100644 --- a/omlx/patches/deepseek_v4/cache_extras.py +++ b/omlx/patches/deepseek_v4/cache_extras.py @@ -32,7 +32,15 @@ def __init__(self, ratio: int): self.buf_gate = None self.remainder = 0 - self.pooled = None + # Append-in-place pooled storage. ``_pool_buf`` is the backing + # allocation (capacity >= logical length); ``_pool_len`` is the + # logical row count that the old step-allocated ``pooled`` tensor + # reported as ``pooled.shape[1]``. Appends only ever write rows at + # [old_len, new_len), so a view captured before an append keeps + # reading the same bytes afterwards; regrowth allocates a fresh + # buffer and leaves any outstanding views on the old one intact. + self._pool_buf = None + self._pool_len = 0 self._undo = None self._undo_chain = False @@ -52,9 +60,49 @@ def __init__(self, ratio: int): # prefix crosses a compression boundary. self._mtp_cross_boundary_rollback = True + @property + def pooled(self): + """Logical pooled tensor: a view of the backing buffer's first + ``_pool_len`` rows, with exactly the shape/contents the old + step-allocated ``self.pooled`` array had. + + Long-lived consumers that must survive later appends must copy + (see ``state`` and ``BatchPoolingCache.extract``); rows below the + current logical length are never rewritten, so views used within + the current chunk/step graph stay correct. + """ + if self._pool_buf is None: + return None + return self._pool_buf[:, : self._pool_len] + + @pooled.setter + def pooled(self, v): + # Full-tensor rebinding (state restore, merge/extract targets, + # trim restore). Capacity collapses to the logical length; the next + # append regrows geometrically. + if v is None: + self._pool_buf = None + self._pool_len = 0 + else: + self._pool_buf = v + self._pool_len = v.shape[1] + @property def offset(self): - return 0 if self.pooled is None else self.pooled.shape[1] + return self._pool_len + + def _grow_pool(self, needed: int) -> None: + """Ensure backing capacity for ``needed`` rows (geometric growth). + + Copies only the logical region; rows outside it were never visible + through ``pooled``. Outstanding views keep referencing the old + buffer, whose committed bytes stay valid. + """ + old = self._pool_buf + capacity = max(needed, 2 * old.shape[1]) + new = mx.zeros((old.shape[0], capacity, old.shape[2]), dtype=old.dtype) + new[:, : self._pool_len] = old[:, : self._pool_len] + self._pool_buf = new def accumulate_windows(self, kv: mx.array, gate: mx.array, offset): B, L, D1 = kv.shape @@ -160,14 +208,22 @@ def accumulate_windows(self, kv: mx.array, gate: mx.array, offset): def update_and_fetch(self, px: mx.array): if px.shape[1] == 0: - if self.pooled is None: + if self._pool_buf is None: return mx.zeros((px.shape[0], 0, px.shape[-1]), dtype=px.dtype) return self.pooled - if self.pooled is None: - self.pooled = px + n = px.shape[1] + if self._pool_buf is None: + # First append: adopt the compressor output as the backing + # buffer (exact fit, no copy) — identical to the old + # ``self.pooled = px`` rebinding. + self._pool_buf = px + self._pool_len = n else: - self.pooled = mx.concatenate([self.pooled, px], axis=1) + if self._pool_len + n > self._pool_buf.shape[1]: + self._grow_pool(self._pool_len + n) + self._pool_buf[:, self._pool_len : self._pool_len + n] = px + self._pool_len += n return self.pooled def make_mask(self, L: int = 1, offset: int = 0): @@ -253,7 +309,6 @@ def trim(self, n): buf_kv, buf_gate, rem_prev, pooled_prev, kv, gate, prev_kv, prev_gate = ( self._undo ) - pooled_after = self.pooled self._undo = None self._undo_chain = False k = kv.shape[1] - n @@ -266,13 +321,23 @@ def trim(self, n): completed = prefix_kv.shape[1] // self.ratio previous_pooled = 0 if pooled_prev is None else pooled_prev.shape[1] if completed == 0: - self.pooled = pooled_prev + # Restore the pre-update logical length. The appended rows were + # only ever written at [previous_pooled, ...), so dropping the + # logical length discards exactly what the old + # ``self.pooled = pooled_prev`` rebinding discarded. + if previous_pooled == 0: + self._pool_buf = None + self._pool_len = 0 + else: + self._pool_len = previous_pooled self.prev_win_kv = prev_kv self.prev_win_gate = prev_gate else: # The full verify already computed these prefix windows. Keep # their exact rows instead of recompressing them during rollback. - self.pooled = pooled_after[:, : previous_pooled + completed] + # Old code rebound ``self.pooled = pooled_after[:, :N]``; the + # buffer's first N rows are exactly that slice. + self._pool_len = previous_pooled + completed end = completed * self.ratio start = end - self.ratio self.prev_win_kv = prefix_kv[:, start:end, :][:, None] @@ -306,18 +371,19 @@ def store_prev(self, kv, gate, dropped): self.prev_win_gate = gate[:, -1:] def size(self): - return 0 if self.pooled is None else self.pooled.shape[1] + return self._pool_len def empty(self): - return self.pooled is None and self.remainder == 0 + return self._pool_buf is None and self.remainder == 0 @property def nbytes(self): total = 0 if self.buf_kv is not None: total += self.buf_kv.nbytes + self.buf_gate.nbytes - if self.pooled is not None: - total += self.pooled.nbytes + if self._pool_buf is not None: + # Resident allocation (capacity), not just the logical view. + total += self._pool_buf.nbytes if self.prev_win_kv is not None: total += self.prev_win_kv.nbytes + self.prev_win_gate.nbytes return total @@ -342,7 +408,15 @@ def __init__(self, ratio: int, left_padding: List[int]): self.buf_gate = None self.remainder = [0] * batch_size - self.pooled = None + # Append-in-place pooled storage (see PoolingCache). ``_pool_buf`` + # is the backing allocation; ``_pool_lengths`` (already tracked for + # offset/mask bookkeeping) is the per-row logical length; + # ``_pool_extent`` reproduces the physical ``pooled.shape[1]`` the + # old code exposed, which could overshoot max(_pool_lengths) when + # the longest row and the row completing windows differed + # (old: max(lengths_before) + max_new). + self._pool_buf = None + self._pool_extent = 0 self._pool_lengths = [0] * batch_size self._lengths = [2**31] * batch_size @@ -363,6 +437,31 @@ def __init__(self, ratio: int, left_padding: List[int]): self._prev_valid = [False] * batch_size self._last_usable = [0] * batch_size + @property + def pooled(self): + """Logical pooled tensor: view of the backing buffer's first + ``_pool_extent`` columns, matching the old physical tensor's + shape/contents exactly (per-row validity lives in + ``_pool_lengths`` and is applied via ``make_mask``). + + Rows/columns below the current extent are never rewritten, so + views used inside the current chunk/step graph stay correct; + long-lived consumers must copy. + """ + if self._pool_buf is None: + return None + return self._pool_buf[:, : self._pool_extent] + + @pooled.setter + def pooled(self, v): + # Full-tensor rebinding (state restore, filter/extend/merge). + if v is None: + self._pool_buf = None + self._pool_extent = 0 + else: + self._pool_buf = v + self._pool_extent = v.shape[1] + @property def offset(self): return mx.array(self._pool_lengths, dtype=mx.int32) @@ -512,7 +611,7 @@ def update_and_fetch(self, px: mx.array): B, N, D = px.shape if N == 0: - if self.pooled is None: + if self._pool_buf is None: return mx.zeros((B, 0, D), dtype=px.dtype) return self.pooled @@ -524,39 +623,34 @@ def update_and_fetch(self, px: mx.array): ] max_new = max(new_counts) if max_new == 0: - if self.pooled is None: + if self._pool_buf is None: return mx.zeros((B, 0, D), dtype=px.dtype) return self.pooled - # The singleton path is the common decode/prefill case. Build a - # fresh logical value instead of mutating a zero-filled allocation - # that the same lazy graph immediately consumes in attention. - if B == 1: - count = new_counts[0] - current = self._pool_lengths[0] - new_rows = px[:, :count] - if self.pooled is None or current == 0: - self.pooled = new_rows - else: - self.pooled = mx.concatenate( - [self.pooled[:, :current], new_rows], axis=1 - ) - self._pool_lengths[0] = current + count - return self.pooled - + # Physical extent exactly as the old code computed it, including the + # overshoot when the longest row is not the one completing windows. max_pool = max(self._pool_lengths) + max_new - if self.pooled is None: - self.pooled = mx.zeros((B, max_pool, D), dtype=px.dtype) - elif self.pooled.shape[1] < max_pool: - pad = mx.zeros((B, max_pool - self.pooled.shape[1], D), dtype=px.dtype) - self.pooled = mx.concatenate([self.pooled, pad], axis=1) - + if self._pool_buf is None: + self._pool_buf = mx.zeros((B, max_pool, D), dtype=px.dtype) + elif self._pool_buf.shape[1] < max_pool: + # Geometric regrowth; copy only the visible region so columns + # beyond the extent stay zero-filled like the old pad path. + capacity = max(max_pool, 2 * self._pool_buf.shape[1]) + new_buf = mx.zeros((B, capacity, D), dtype=px.dtype) + new_buf[:, : self._pool_extent] = self._pool_buf[:, : self._pool_extent] + self._pool_buf = new_buf + self._pool_extent = max(self._pool_extent, max_pool) + + # Append in place. The old singleton path rebound + # ``self.pooled = concatenate([self.pooled[:, :current], new_rows])``; + # writing rows [pl, pl+nc) into the buffer is value-identical, and + # the B > 1 path already used exactly this in-place scheme. for i in range(B): nc = new_counts[i] if nc > 0: pl = self._pool_lengths[i] - self.pooled[i, pl : pl + nc] = px[i, :nc] + self._pool_buf[i, pl : pl + nc] = px[i, :nc] self._pool_lengths[i] = pl + nc return self.pooled @@ -660,7 +754,6 @@ def trim(self, n): prev_valid, ) = self._undo if self._mtp_cross_boundary_rollback: - pooled_after = self.pooled self._undo = None self._undo_chain = False k = kv.shape[1] - n @@ -674,9 +767,14 @@ def trim(self, n): ) completed = prefix_kv.shape[1] // self.ratio next_pool_length = pool_lengths[0] + completed - self.pooled = ( - pooled_after[:, :next_pool_length] if next_pool_length else None - ) + # Old code rebound ``self.pooled = pooled_after[:, :N]``; the + # buffer's first N columns are exactly that slice (appends only + # wrote at [pool_lengths[0], ...)). + if next_pool_length: + self._pool_extent = next_pool_length + else: + self._pool_buf = None + self._pool_extent = 0 self._pool_lengths = [next_pool_length] self._processed = [processed[0] + k] @@ -734,11 +832,13 @@ def trim(self, n): def _truncate_pooled_tail(self): """Drop pooled rows written by a rejected speculative suffix.""" - if self.pooled is None: + if self._pool_buf is None: return logical_size = max(self._pool_lengths, default=0) - if self.pooled.shape[1] > logical_size: - self.pooled = self.pooled[:, :logical_size] + if self._pool_extent > logical_size: + # Logical truncation only; the buffer's first logical_size + # columns are exactly the old ``pooled[:, :logical_size]``. + self._pool_extent = logical_size def prev_for_prepend(self): """Per-row previous window with invalid rows masked via -inf gates. @@ -776,18 +876,19 @@ def store_prev(self, kv, gate, dropped): self._prev_valid = [v or n > 0 for v, n in zip(self._prev_valid, n_new)] def size(self): - return 0 if self.pooled is None else self.pooled.shape[1] + return 0 if self._pool_buf is None else self._pool_extent def empty(self): - return self.pooled is None and all(r == 0 for r in self.remainder) + return self._pool_buf is None and all(r == 0 for r in self.remainder) @property def nbytes(self): total = 0 if self.buf_kv is not None: total += self.buf_kv.nbytes + self.buf_gate.nbytes - if self.pooled is not None: - total += self.pooled.nbytes + if self._pool_buf is not None: + # Resident allocation (capacity), not just the logical view. + total += self._pool_buf.nbytes return total def filter(self, batch_indices): diff --git a/omlx/patches/deepseek_v4/deepseek_v4_model.py b/omlx/patches/deepseek_v4/deepseek_v4_model.py index ca080805f..3fd62997e 100644 --- a/omlx/patches/deepseek_v4/deepseek_v4_model.py +++ b/omlx/patches/deepseek_v4/deepseek_v4_model.py @@ -17,6 +17,7 @@ from .mla import MultiLinear from .pipeline import PipelineMixin from omlx.patches.deepseek_v4.switch_layers import SwitchGLU +from omlx.patches.deepseek_v4.wsdpa_attention import wsdpa_prefill, wsdpa_topk_prefill from omlx.patches.deepseek_v4.decode_consistency import ( is_armed as is_dspark_verify_armed, ) @@ -637,6 +638,7 @@ def _sparse_pooled_attention( local_window: Optional[int] = None, decode_consistent: bool = False, native_only: bool = False, + _standard_mask: bool = False, ) -> Optional[mx.array]: global _DEEPSEEK_V4_SPARSE_ATTENTION_NATIVE_DISABLED @@ -661,6 +663,20 @@ def _sparse_pooled_attention( and pooled.shape[-1] == D and topk.ndim == 3 ): + if _standard_mask and B == 1 and q.dtype == mx.bfloat16 and topk.shape[1] == L: + out = wsdpa_topk_prefill( + q, + local_kv, + pooled, + topk, + sinks, + scale, + int(q_offset), + int(local_window), + int(compress_ratio), + ) + if out is not None: + return out try: from omlx.custom_kernels.glm_moe_dsa import fast as glm_fast @@ -1246,13 +1262,31 @@ def __call__( if projected_weights is None else projected_weights ).astype(q.dtype) * ((self.n_heads**-0.5) * self.scale) + # Fused pooled-ratio causal mask (lossless): the kernel + # epilogue writes the finfo(bf16).min sentinel itself + # when the mask is the plain 2-D PoolingCache ratio mask, + # bit-identical to the mx.where pass it replaces. + # Batched 3-D masks keep the where pass; pmask is None + # stays unmasked as before. + _mask_ratio = 0 + _mask_q_offset = 0 + if ( + pmask is not None + and pmask.ndim == 2 + and isinstance(offset, int) + and type(pool_cache).__name__ == "PoolingCache" + ): + _mask_ratio = int(pool_cache.ratio) + _mask_q_offset = int(offset) scores4 = glm_fast.dsa_indexer_scores( q, pooled[:, None], weights, causal=False, + mask_ratio=_mask_ratio, + mask_q_offset=_mask_q_offset, ) - if pmask is not None: + if _mask_ratio == 0 and pmask is not None: scores4 = mx.where( (pmask[:, None] if pmask.ndim == 3 else pmask[None, None]), scores4, @@ -1475,6 +1509,8 @@ def __call__( x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, + *, + _standard_mask: bool = False, ) -> mx.array: B, L, _ = x.shape offset = cache.offset if cache is not None else 0 @@ -1502,15 +1538,28 @@ def __call__( if self.dspark and B == 1 and L == 1: out = exact_attention(q, [kv], self.scale, sinks) else: - out = scaled_dot_product_attention( - q, - kv, - kv, - cache=cache, - scale=self.scale, - mask=mask, - sinks=sinks, - ) + out = None + if _standard_mask and B == 1 and L > 1: + out = wsdpa_prefill( + q, + kv, + None, + sinks, + self.scale, + offset, + self.config.sliding_window, + 1, + ) + if out is None: + out = scaled_dot_product_attention( + q, + kv, + kv, + cache=cache, + scale=self.scale, + mask=mask, + sinks=sinks, + ) out = _project_attention_output(self, out, offset) if self.sharding_group is not None: @@ -1620,10 +1669,22 @@ def __call__( pooled_mask = ( pool_cache.make_mask(L, offset) if pool_cache is not None else None ) - # The native kernel reconstructs the model's causal/sliding masks - # from offsets; direct callers with custom masks stay on dense SDPA. + # The wsdpa and native kernels reconstruct the model's causal/sliding + # masks from offsets; direct callers with custom masks stay on the + # reference path. out = None - if ( + if _standard_mask and B == 1 and L > 1: + out = wsdpa_prefill( + q, + kv, + pooled if pooled.shape[1] > 0 else None, + sinks, + self.scale, + offset, + self.config.sliding_window, + self.compress_ratio, + ) + if out is None and ( self.config.use_native_ratio128_attention and self.compress_ratio == 128 and _standard_mask @@ -1650,6 +1711,7 @@ def __call__( local_window=self.config.sliding_window, decode_consistent=self.dspark, native_only=True, + _standard_mask=_standard_mask, ) if out is None: if pooled.shape[1] > 0: @@ -1726,6 +1788,8 @@ def __call__( x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, + *, + _standard_mask: bool = False, ) -> mx.array: B, L, _ = x.shape local_cache = cache[0] if cache is not None else None @@ -1884,30 +1948,57 @@ def __call__( if self.dspark and B == 1 and L == 1: out = exact_attention(q, [kv], self.scale, sinks) else: - out = scaled_dot_product_attention( - q, - kv, - kv, - cache=local_cache, - scale=self.scale, - mask=mask, - sinks=sinks, - ) + out = None + if _standard_mask and B == 1 and L > 1: + out = wsdpa_prefill( + q, + kv, + None, + sinks, + self.scale, + offset, + self.config.sliding_window, + self.compress_ratio, + ) + if out is None: + out = scaled_dot_product_attention( + q, + kv, + kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=sinks, + ) elif pooled.shape[1] <= self.indexer.index_topk: - full_kv = mx.concatenate([kv, pooled[:, None]], axis=2) - mask = _extend_mask(mask, pmask, full_kv.shape[2]) if self.dspark and B == 1 and L == 1: + full_kv = mx.concatenate([kv, pooled[:, None]], axis=2) out = exact_attention(q, [full_kv], self.scale, sinks) else: - out = scaled_dot_product_attention( - q, - full_kv, - full_kv, - cache=local_cache, - scale=self.scale, - mask=mask, - sinks=sinks, - ) + out = None + if _standard_mask and B == 1 and L > 1: + out = wsdpa_prefill( + q, + kv, + pooled, + sinks, + self.scale, + offset, + self.config.sliding_window, + self.compress_ratio, + ) + if out is None: + full_kv = mx.concatenate([kv, pooled[:, None]], axis=2) + mask = _extend_mask(mask, pmask, full_kv.shape[2]) + out = scaled_dot_product_attention( + q, + full_kv, + full_kv, + cache=local_cache, + scale=self.scale, + mask=mask, + sinks=sinks, + ) else: out = _sparse_pooled_attention( q, @@ -1922,6 +2013,7 @@ def __call__( compress_ratio=self.compress_ratio, local_window=self.config.sliding_window, decode_consistent=self.dspark, + _standard_mask=_standard_mask, ) out = _project_attention_output(self, out, offset) @@ -1964,15 +2056,12 @@ def __call__( residual = h x, post, comb = self.attn_hc(h) attn_input = self.attn_norm(x) - if isinstance(self.attn, CompressedAttention): - x = self.attn( - attn_input, - mask=mask, - cache=cache, - _standard_mask=_standard_mask, - ) - else: - x = self.attn(attn_input, mask=mask, cache=cache) + x = self.attn( + attn_input, + mask=mask, + cache=cache, + _standard_mask=_standard_mask, + ) h = hc_expand(x, residual, post, comb) residual = h diff --git a/omlx/patches/deepseek_v4/wsdpa_attention.py b/omlx/patches/deepseek_v4/wsdpa_attention.py new file mode 100644 index 000000000..eaf6c8bb3 --- /dev/null +++ b/omlx/patches/deepseek_v4/wsdpa_attention.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused windowed + pooled-sparse prefill attention for DeepSeek-V4. + +During chunked prefill the stock path runs full N^2 scaled_dot_product_attention +with a materialized boolean mask, even though each query can only see + local rows j: (p - window, p] (causal + sliding window) + pool rows i: i < (p + 1) // ratio (compressed KV, ratio in {4, 128}) +plus a per-head attention sink. At ratio-4 chunk boundaries that is ~4x the +necessary FLOPs (measured 39-48 ms per layer-call at L=2048, head_dim=512). + +This module computes the identical math with a single fused Metal kernel that +visits only visible rows (fp32 online softmax in fixed row order, sink in the +denominator). Any setup/runtime failure permanently falls back to the stock +path; OMLX_DSV4_WSDPA=0 disables the kernel outright. +""" + +import logging +import os + +import mlx.core as mx + +logger = logging.getLogger(__name__) + +_ENABLED = os.environ.get("OMLX_DSV4_WSDPA", "1") == "1" +_kernel = None +_broken = False + +_HEADER = """ +#include +using namespace metal; +""" + +_SOURCE = """ + // q: [H, L, D] bf16 (contiguous) + // kv: [S, D] bf16 (local rows, linear layout) + // pooled: [P, D] bf16 (compressed rows; dummy row when P == 0) + // sinks: [H] bf16 + // params: int32 [6] = {offset, window, ratio, P, S, L} + // scalep: fp32 [1] = scale + // out: [H, L, D] bf16 + // + // grid (16, L): gx = head-group (4 heads), gy = query row. + // threadgroup 128 threads = 4 simdgroups, one head per simdgroup. + // No threadgroup staging: rows are read directly (L2-resident at these + // sizes); zero barriers measurably beats staged variants (13.7 vs 17.3ms). + const uint t = threadgroup_position_in_grid.y; + const uint tid = thread_index_in_threadgroup; + const uint head = threadgroup_position_in_grid.x * 4 + (tid / 32); + const uint lane = tid % 32; + + constant int *prm = (constant int *)¶ms[0]; + const int offset = prm[0]; + const int window = prm[1]; + const int ratio = prm[2]; + const int P = prm[3]; + const int S = prm[4]; + const int L = prm[5]; + const float scale = ((constant float *)&scalep[0])[0]; + + const uint D = 512; + const uint p = (uint)offset + t; + + device const bfloat4 *qh = (device const bfloat4 *)(q + ((uint64_t)head * L + t) * D); + float4 qa0 = float4(qh[lane + 0]); + float4 qa1 = float4(qh[lane + 32]); + float4 qa2 = float4(qh[lane + 64]); + float4 qa3 = float4(qh[lane + 96]); + + float m = -INFINITY, s = 0.0f; + float4 a0 = 0.0f, a1 = 0.0f, a2 = 0.0f, a3 = 0.0f; + + device const bfloat4 *kvr = (device const bfloat4 *)kv; + device const bfloat4 *pol = (device const bfloat4 *)pooled; + + // ---- local rows (p-window, p], then pooled prefix i < (p+1)/ratio + // The local cache is a RotatingKVCache that trims to the last S rows, so + // buffer row 0 sits at absolute position base = offset + L - S. Translate + // absolute window bounds into buffer indices. + const int base = offset + L - S; + const int j0 = max(base, (int)p - window + 1) - base; + const int j1 = (int)p - base; + const int pvis = min((int)((p + 1) / (uint)ratio), P); + const int n_local = j1 - j0 + 1; + const int total = n_local + pvis; + + for (int r = 0; r < total; r++) { + const bool is_pool = r >= n_local; + const uint idx = is_pool ? (uint)(r - n_local) : (uint)(j0 + r); + device const bfloat4 *row = (is_pool ? pol : kvr) + (uint64_t)idx * 128; + float4 k0 = float4(row[lane + 0]), k1 = float4(row[lane + 32]); + float4 k2 = float4(row[lane + 64]), k3 = float4(row[lane + 96]); + float d = dot(qa0, k0) + dot(qa1, k1) + dot(qa2, k2) + dot(qa3, k3); + d = simd_sum(d) * scale; + float mn = max(m, d); + float c = (m == -INFINITY) ? 0.0f : exp(m - mn); + float w = exp(d - mn); + m = mn; s = s * c + w; + a0 = a0 * c + w * k0; a1 = a1 * c + w * k1; + a2 = a2 * c + w * k2; a3 = a3 * c + w * k3; + } + + // ---- attention sink (denominator only) + s += exp(float(((device const bfloat *)sinks)[head]) - m); + + const float inv = (s == 0.0f) ? 0.0f : 1.0f / s; + device bfloat4 *oh = (device bfloat4 *)(out + ((uint64_t)head * L + t) * D); + oh[lane + 0] = bfloat4(a0 * inv); + oh[lane + 32] = bfloat4(a1 * inv); + oh[lane + 64] = bfloat4(a2 * inv); + oh[lane + 96] = bfloat4(a3 * inv); +""" + + +def _get_kernel(): + global _kernel, _broken + if _broken or not _ENABLED: + return None + if _kernel is None: + try: + _kernel = mx.fast.metal_kernel( + name="dsv4_wsdpa", + input_names=["q", "kv", "pooled", "sinks", "params", "scalep"], + output_names=["out"], + source=_SOURCE, + header=_HEADER, + ) + except Exception: + _broken = True + logger.warning( + "DSV4 wsdpa kernel setup failed; using stock SDPA", exc_info=True + ) + return None + return _kernel + + +_KERNEL_TOPK = None +_TOPK_ENABLED = os.environ.get("OMLX_DSV4_WSDPA_TOPK", "1") == "1" + +_SOURCE_TOPK = """ + // q: [H, L, D] bf16 (contiguous) + // kv: [S, D] bf16 (local rows, linear absolute layout) + // pooled: [P, D] bf16 (compressed rows) + // topk: [L, K] uint32 (temporally sorted pooled indices) + // sinks: [H] bf16 + // params: int32 [7] = {offset, window, ratio, P, S, L, K} + // scalep: fp32 [1] = scale + // out: [H, L, D] bf16 + const uint t = threadgroup_position_in_grid.y; + const uint tid = thread_index_in_threadgroup; + const uint head = threadgroup_position_in_grid.x * 4 + (tid / 32); + const uint lane = tid % 32; + + constant int *prm = (constant int *)¶ms[0]; + const int offset = prm[0]; + const int window = prm[1]; + const int ratio = prm[2]; + const int P = prm[3]; + const int S = prm[4]; + const int L = prm[5]; + const int K = prm[6]; + const float scale = ((constant float *)&scalep[0])[0]; + + const uint D = 512; + const uint p = (uint)offset + t; + + device const bfloat4 *qh = (device const bfloat4 *)(q + ((uint64_t)head * L + t) * D); + float4 qa0 = float4(qh[lane + 0]); + float4 qa1 = float4(qh[lane + 32]); + float4 qa2 = float4(qh[lane + 64]); + float4 qa3 = float4(qh[lane + 96]); + + float m = -INFINITY, s = 0.0f; + float4 a0 = 0.0f, a1 = 0.0f, a2 = 0.0f, a3 = 0.0f; + + device const bfloat4 *kvr = (device const bfloat4 *)kv; + device const bfloat4 *pol = (device const bfloat4 *)pooled; + + const int base = offset + L - S; + const int j0 = max(base, (int)p - window + 1) - base; + const int j1 = (int)p - base; + const int pvis = min((int)((p + 1) / (uint)ratio), P); + const int n_local = j1 - j0 + 1; + + for (int r = 0; r < n_local; r++) { + const uint idx = (uint)(j0 + r); + device const bfloat4 *row = kvr + (uint64_t)idx * 128; + float4 k0 = float4(row[lane + 0]), k1 = float4(row[lane + 32]); + float4 k2 = float4(row[lane + 64]), k3 = float4(row[lane + 96]); + float d = dot(qa0, k0) + dot(qa1, k1) + dot(qa2, k2) + dot(qa3, k3); + d = simd_sum(d) * scale; + float mn = max(m, d); + float c = (m == -INFINITY) ? 0.0f : exp(m - mn); + float w = exp(d - mn); + m = mn; s = s * c + w; + a0 = a0 * c + w * k0; a1 = a1 * c + w * k1; + a2 = a2 * c + w * k2; a3 = a3 * c + w * k3; + } + + device const uint *tk = (device const uint *)topk; + const uint64_t tk_base = (uint64_t)t * (uint)K; + for (int r = 0; r < K; r++) { + const uint idx = tk[tk_base + (uint)r]; + if ((int)idx >= pvis) break; // indexer topk rows are temporally sorted + device const bfloat4 *row = pol + (uint64_t)idx * 128; + float4 k0 = float4(row[lane + 0]), k1 = float4(row[lane + 32]); + float4 k2 = float4(row[lane + 64]), k3 = float4(row[lane + 96]); + float d = dot(qa0, k0) + dot(qa1, k1) + dot(qa2, k2) + dot(qa3, k3); + d = simd_sum(d) * scale; + float mn = max(m, d); + float c = (m == -INFINITY) ? 0.0f : exp(m - mn); + float w = exp(d - mn); + m = mn; s = s * c + w; + a0 = a0 * c + w * k0; a1 = a1 * c + w * k1; + a2 = a2 * c + w * k2; a3 = a3 * c + w * k3; + } + + s += exp(float(((device const bfloat *)sinks)[head]) - m); + + const float inv = (s == 0.0f) ? 0.0f : 1.0f / s; + device bfloat4 *oh = (device bfloat4 *)(out + ((uint64_t)head * L + t) * D); + oh[lane + 0] = bfloat4(a0 * inv); + oh[lane + 32] = bfloat4(a1 * inv); + oh[lane + 64] = bfloat4(a2 * inv); + oh[lane + 96] = bfloat4(a3 * inv); +""" + + +def _get_topk_kernel(): + global _KERNEL_TOPK, _broken + if _broken or not _ENABLED or not _TOPK_ENABLED: + return None + if _KERNEL_TOPK is None: + try: + _KERNEL_TOPK = mx.fast.metal_kernel( + name="dsv4_wsdpa_topk", + input_names=["q", "kv", "pooled", "topk", "sinks", "params", "scalep"], + output_names=["out"], + source=_SOURCE_TOPK, + header=_HEADER, + ) + except Exception: + _broken = True + logger.warning( + "DSV4 wsdpa topk kernel setup failed; using stock sparse attention", + exc_info=True, + ) + return None + return _KERNEL_TOPK + + +def wsdpa_prefill( + q: mx.array, + kv: mx.array, + pooled: mx.array | None, + sinks: mx.array, + scale: float, + offset: int, + window: int, + ratio: int, +) -> mx.array | None: + """Fused window+pool attention; returns [B, H, L, D] or None (fallback). + + q: [B, H, L, D] (B == 1), kv: [B, 1, S, D], pooled: [B, P, D] or None. + Only valid for prefill shapes (L > 1); callers keep decode paths. + """ + global _broken + if isinstance(offset, mx.array): + return None # should not happen in prefill; stay safe + if ( + q.dtype != mx.bfloat16 + or q.shape[0] != 1 + or q.shape[1] != 64 + or q.shape[3] != 512 + ): + return None + kfn = _get_kernel() + if kfn is None: + return None + try: + heads, q_len, head_dim = q.shape[1], q.shape[2], q.shape[3] + kv_len = kv.shape[2] + pooled_len = 0 if pooled is None else pooled.shape[1] + qc = mx.contiguous(q[0]) # [H, L, D] + kvc = mx.contiguous(kv[0, 0]) # [S, D] + pol = ( + mx.contiguous(pooled[0]) + if pooled_len + else mx.zeros((1, head_dim), dtype=q.dtype) # constant-space guard + ) + params = mx.array( + [offset, window, ratio, pooled_len, kv_len, q_len], dtype=mx.int32 + ) + scalep = mx.array([scale], dtype=mx.float32) + out = kfn( + inputs=[qc, kvc, pol, sinks, params, scalep], + grid=(16 * 128, q_len, 1), + threadgroup=(128, 1, 1), + output_shapes=[(heads, q_len, head_dim)], + output_dtypes=[mx.bfloat16], + )[0] + return out[None] + except Exception: + _broken = True + logger.warning("DSV4 wsdpa kernel disabled; using stock SDPA", exc_info=True) + return None + + +def wsdpa_topk_prefill( + q: mx.array, + kv: mx.array, + pooled: mx.array, + topk: mx.array, + sinks: mx.array, + scale: float, + offset: int, + window: int, + ratio: int, +) -> mx.array | None: + """Fused window + indexer-topk pooled attention for ratio-4 prefill. + + Requires the Indexer's temporally sorted topk rows. Returns None to keep the + native/stock path whenever shapes or dtypes are not the exact prefill case. + """ + global _broken + if isinstance(offset, mx.array) or not _TOPK_ENABLED: + return None + if ( + q.dtype != mx.bfloat16 + or q.shape[0] != 1 + or q.shape[1] != 64 + or q.shape[3] != 512 + or topk.dtype != mx.uint32 + or topk.ndim != 3 + or topk.shape[0] != 1 + or topk.shape[1] != q.shape[2] + or pooled is None + or pooled.shape[1] == 0 + ): + return None + kfn = _get_topk_kernel() + if kfn is None: + return None + try: + heads, q_len, head_dim = q.shape[1], q.shape[2], q.shape[3] + params = mx.array( + [ + offset, + window, + ratio, + pooled.shape[1], + kv.shape[2], + q_len, + topk.shape[2], + ], + dtype=mx.int32, + ) + scalep = mx.array([scale], dtype=mx.float32) + out = kfn( + inputs=[ + mx.contiguous(q[0]), + mx.contiguous(kv[0, 0]), + mx.contiguous(pooled[0]), + mx.contiguous(topk[0]), + sinks, + params, + scalep, + ], + grid=(16 * 128, q_len, 1), + threadgroup=(128, 1, 1), + output_shapes=[(heads, q_len, head_dim)], + output_dtypes=[mx.bfloat16], + )[0] + return out[None] + except Exception: + _broken = True + logger.warning( + "DSV4 wsdpa topk kernel disabled; using stock sparse attention", + exc_info=True, + ) + return None diff --git a/omlx/patches/mlx_lm_mtp/deepseek_v4_model.py b/omlx/patches/mlx_lm_mtp/deepseek_v4_model.py index e19697467..3c7a8f277 100644 --- a/omlx/patches/mlx_lm_mtp/deepseek_v4_model.py +++ b/omlx/patches/mlx_lm_mtp/deepseek_v4_model.py @@ -350,7 +350,33 @@ def __call__( cache=None, return_hidden: bool = False, n_confirmed: int = 0, + skip_lm_head: bool = False, ): + if skip_lm_head: + # Chunked prefill discards per-chunk logits (the prompt's final + # token is scored by the first decode step instead). Run the + # hidden pass and preserve capture side effects, but skip the + # full-vocabulary projection for every chunk. + if ( + getattr(self, "_omlx_dspark_decode_enabled", False) + and not n_confirmed + and cache is not None + ): + h, h_aux = self.model(inputs, cache, return_dspark_hidden=True) + try: + deepseek_v4_dspark.capture_prompt(self, inputs, h_aux, cache) + except Exception: + logger.debug("DeepSeek DSpark prompt capture failed", exc_info=True) + return None + if not n_confirmed and prompt_priming.capture_eligible(self, cache): + h, h_raw = self.model(inputs, cache, return_raw_hidden=True) + try: + prompt_priming.maybe_capture(self, inputs, h_raw, cache) + except Exception: + logger.debug("MTP prompt-priming capture failed", exc_info=True) + return None + self.model(inputs, cache) + return None # ``n_confirmed`` is part of the patched-backbone interface: # batch_generator._call_backbone passes n_confirmed=1 during MTP # verify cycles. It only matters for models with module-level diff --git a/omlx/scheduler.py b/omlx/scheduler.py index a26d6244a..75e078208 100644 --- a/omlx/scheduler.py +++ b/omlx/scheduler.py @@ -3366,6 +3366,8 @@ def _do_external_prefill( model_kwargs["vlm_extra_kwargs"] = _slice_vlm_extra( extra_kwargs, n_to_process ) + if self._supports_skip_lm_head(): + model_kwargs["skip_lm_head"] = True self.model( input_arr[:, :n_to_process], cache=prompt_cache, @@ -4431,6 +4433,36 @@ def _record_chunk_transient( self._prefill_transient_tracker.samples, ) + def _supports_skip_lm_head(self) -> bool: + """Whether the loaded model accepts ``skip_lm_head=True``. + + Chunked prefill discards every chunk's logits (the prompt's final + token is scored by the first decode step), so models whose patched + ``__call__`` accepts the flag can skip the full-vocabulary + projection for every prefill chunk — for a 129k-vocab model that + GEMM is the single largest per-chunk matmul and was pure waste. + Detected once per scheduler; unknown models keep stock behavior. + """ + supported = getattr(self, "_skip_lm_head_supported", None) + if supported is None: + try: + import inspect + + call = getattr(type(self.model), "__call__", None) + supported = bool( + call is not None + and "skip_lm_head" in inspect.signature(call).parameters + ) + except Exception: + supported = False + self._skip_lm_head_supported = supported + if supported: + logger.info( + "Prefill lm_head skip enabled: chunk logits are discarded, " + "vocabulary projection deferred to first decode step." + ) + return supported + def _maybe_record_fixed_state_bytes(self, cache_list: Any) -> None: """Measure the GDN/Mamba fixed recurrent-state footprint once. @@ -4650,7 +4682,10 @@ def _step_prefill_chunk(self, state: _PrefillState) -> bool: with mx.stream(self._stream): chunk = state.tokens_remaining[:, :n] state.tokens_remaining = state.tokens_remaining[:, n:] - self.model(chunk, cache=state.cache) + if self._supports_skip_lm_head(): + self.model(chunk, cache=state.cache, skip_lm_head=True) + else: + self.model(chunk, cache=state.cache) mx.eval([c.state for c in state.cache]) _throttle_post = get_phys_footprint() self._record_chunk_transient( diff --git a/setup.py b/setup.py index 1e1e25cc5..22146dc54 100644 --- a/setup.py +++ b/setup.py @@ -32,6 +32,19 @@ def _custom_kernel_build_kwargs() -> dict: os.environ["CMAKE_ARGS"] = ( f"{cmake_args} {target_arg}".strip() if cmake_args else target_arg ) + cmake_args = os.environ["CMAKE_ARGS"] + + # CMake otherwise chooses the first framework Python on PATH, which can + # differ from the interpreter running pip (and lack nanobind / MLX). The + # extensions must use the active environment's ABI and CMake packages. + python_args = " ".join( + ( + f"-DPython_EXECUTABLE={sys.executable}", + f"-DPython3_EXECUTABLE={sys.executable}", + ) + ) + if "Python_EXECUTABLE" not in cmake_args: + os.environ["CMAKE_ARGS"] = f"{cmake_args} {python_args}".strip() from mlx import extension diff --git a/tests/test_custom_kernel_abi_probe.py b/tests/test_custom_kernel_abi_probe.py index 04b924f05..a021287e7 100644 --- a/tests/test_custom_kernel_abi_probe.py +++ b/tests/test_custom_kernel_abi_probe.py @@ -74,4 +74,119 @@ def test_local_build_probe_is_healthy(fast): pytest.skip(f"{fast.__name__} native build unavailable") import mlx.core as mx - assert fast._ext.abi_probe(mx.zeros((3,))) == 3 \ No newline at end of file + assert fast._ext.abi_probe(mx.zeros((3,))) == 3 + + +class _FoldAwareExt: + """New build: nanobind-style doc includes the mask-fold kwargs.""" + + def dsa_indexer_scores(self, *args, **kwargs): + raise AssertionError("probe must not call the kernel") + + dsa_indexer_scores.__doc__ = ( + "dsa_indexer_scores(queries: array, keys: array, weights: array, " + "causal: bool = True, 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 = None)" + ) + + +class _PreFoldExt: + """Old build: same symbol, but without the mask-fold kwargs.""" + + def dsa_indexer_scores(self, *args, **kwargs): + raise AssertionError("probe must not call the kernel") + + dsa_indexer_scores.__doc__ = ( + "dsa_indexer_scores(queries: array, keys: array, weights: array, " + "causal: bool = True, unused_causal_prefix_topk: int = 0, " + "skip_causal_future_store: bool = False, causal_q_offset: int = -1, " + "stream: None = None)" + ) + + +class _NoScoresExt: + """A build without dsa_indexer_scores at all.""" + + +def test_mask_fold_probe_detects_fold_aware_build(): + assert glm_fast._probe_mask_fold(_FoldAwareExt()) is True + + +def test_mask_fold_probe_rejects_pre_fold_build(): + assert glm_fast._probe_mask_fold(_PreFoldExt()) is False + + +def test_mask_fold_probe_handles_missing_symbol_and_ext(): + assert glm_fast._probe_mask_fold(_NoScoresExt()) is False + assert glm_fast._probe_mask_fold(None) is False + + +def test_pre_fold_build_keeps_historical_call_signature(monkeypatch): + """An old _ext must receive no mask kwargs and still get exact masking. + + Regression for the unconditional-kwargs break: GLM-5.2's native path + raised TypeError on every call, and the V4 indexer silently fell back + while the startup probe still reported the kernels as available. + """ + import mlx.core as mx + + calls = [] + + def old_scores(queries, keys, weights, **kwargs): + assert "mask_ratio" not in kwargs + assert "mask_q_offset" not in kwargs + calls.append(kwargs) + B, H, L, D = queries.shape + P = keys.shape[2] + return mx.zeros((B, H, L, P), dtype=queries.dtype) + + monkeypatch.setattr(glm_fast, "_ext", type("E", (), {"dsa_indexer_scores": staticmethod(old_scores)})()) + monkeypatch.setattr(glm_fast, "_EXT_MASK_FOLD", False) + + H, D, L, P = 64, 128, 64, 512 + q = mx.zeros((1, H, L, D), dtype=mx.bfloat16) + keys = mx.zeros((1, 1, P, D), dtype=mx.bfloat16) + weights = mx.zeros((1, L, H), dtype=mx.bfloat16) + + ratio, q_off = 4, 256 + out = glm_fast.dsa_indexer_scores( + q, keys, weights, causal=False, mask_ratio=ratio, mask_q_offset=q_off + ) + assert len(calls) == 1 + + rows = mx.arange(L)[:, None] + cols = mx.arange(P)[None, :] + expected = mx.where( + (cols < ((q_off + rows + 1) // ratio))[None, None], + mx.zeros((1, H, L, P), dtype=mx.bfloat16), + mx.finfo(mx.bfloat16).min, + ) + mx.eval(out, expected) + assert bool(mx.array_equal(out.view(mx.uint16), expected.view(mx.uint16))) + + +def test_fold_aware_build_receives_mask_kwargs(monkeypatch): + import mlx.core as mx + + seen = {} + + def new_scores(queries, keys, weights, **kwargs): + seen.update(kwargs) + B, H, L, _ = queries.shape + P = keys.shape[2] + return mx.zeros((B, H, L, P), dtype=queries.dtype) + + monkeypatch.setattr(glm_fast, "_ext", type("E", (), {"dsa_indexer_scores": staticmethod(new_scores)})()) + monkeypatch.setattr(glm_fast, "_EXT_MASK_FOLD", True) + + H, D, L, P = 64, 128, 64, 512 + q = mx.zeros((1, H, L, D), dtype=mx.bfloat16) + keys = mx.zeros((1, 1, P, D), dtype=mx.bfloat16) + weights = mx.zeros((1, L, H), dtype=mx.bfloat16) + + glm_fast.dsa_indexer_scores( + q, keys, weights, causal=False, mask_ratio=4, mask_q_offset=256 + ) + assert seen.get("mask_ratio") == 4 + assert seen.get("mask_q_offset") == 256 \ No newline at end of file diff --git a/tests/test_deepseek_v4_wsdpa.py b/tests/test_deepseek_v4_wsdpa.py new file mode 100644 index 000000000..9a19c4bd0 --- /dev/null +++ b/tests/test_deepseek_v4_wsdpa.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for DeepSeek-V4 fused windowed + pooled prefill attention.""" + +import math + +import mlx.core as mx +import pytest + +requires_metal = pytest.mark.skipif( + not mx.metal.is_available(), reason="Metal is required" +) + + +def _max_abs(a, b): + return mx.max(mx.abs(a.astype(mx.float32) - b.astype(mx.float32))).item() + + +def _reference_attention( + q, + kv, + pooled, + sinks, + scale, + offset, + window, + ratio, + topk=None, +): + """Explicit fp32 reference for the exact rows visited by the kernels. + + ``kv`` may be a trimmed RotatingKVCache buffer holding only the last + ``kv.shape[2]`` rows; buffer row 0 then maps to absolute position + ``offset + q_len - kv_len``. + """ + base = offset + q.shape[2] - kv.shape[2] + head_outputs = [] + for head in range(q.shape[1]): + row_outputs = [] + for row in range(q.shape[2]): + position = offset + row + local_start = max(base, position - window + 1) - base + local_end = position - base + key_parts = [kv[0, 0, local_start : local_end + 1].astype(mx.float32)] + + pooled_length = 0 if pooled is None else pooled.shape[1] + visible_pool = min((position + 1) // ratio, pooled_length) + if topk is None: + if visible_pool: + key_parts.append(pooled[0, :visible_pool].astype(mx.float32)) + else: + indices = [] + for index in topk[0, row].tolist(): + if index >= visible_pool: + break + indices.append(index) + if indices: + key_parts.append(pooled[0, indices].astype(mx.float32)) + + keys = mx.concatenate(key_parts, axis=0) + scores = (keys @ q[0, head, row].astype(mx.float32)) * scale + normalizer = mx.logsumexp(scores, axis=-1) + normalizer = mx.logaddexp(normalizer, sinks[head].astype(mx.float32)) + weights = mx.exp(scores - normalizer) + row_outputs.append((weights[:, None] * keys).sum(axis=0)) + head_outputs.append(mx.stack(row_outputs)) + return mx.stack(head_outputs)[None].astype(mx.bfloat16) + + +def _inputs(q_len, offset, window, pooled_len, ratio, trim=0): + mx.random.seed(7) + kv_len = offset + q_len - trim + q = (mx.random.normal((1, 64, q_len, 512)) * 0.25).astype(mx.bfloat16) + kv = (mx.random.normal((1, 1, kv_len, 512)) * 0.25).astype(mx.bfloat16) + pooled = (mx.random.normal((1, pooled_len, 512)) * 0.25).astype(mx.bfloat16) + sinks = (mx.random.normal((64,)) * 0.1).astype(mx.bfloat16) + scale = 1.0 / math.sqrt(512) + mx.eval(q, kv, pooled, sinks) + return q, kv, pooled, sinks, scale, offset, window, ratio + + +def _reset_wsdpa(monkeypatch): + from omlx.patches.deepseek_v4 import wsdpa_attention as wsdpa + + monkeypatch.setattr(wsdpa, "_ENABLED", True) + monkeypatch.setattr(wsdpa, "_TOPK_ENABLED", True) + monkeypatch.setattr(wsdpa, "_broken", False) + return wsdpa + + +@requires_metal +def test_wsdpa_prefill_matches_explicit_reference(monkeypatch): + wsdpa = _reset_wsdpa(monkeypatch) + args = _inputs(q_len=5, offset=7, window=4, pooled_len=3, ratio=4) + + out = wsdpa.wsdpa_prefill(*args) + ref = _reference_attention(*args) + + assert out is not None + mx.eval(out, ref) + assert out.shape == ref.shape + assert out.dtype == mx.bfloat16 + assert _max_abs(out, ref) < 8e-3 + + +@requires_metal +def test_wsdpa_topk_prefill_matches_explicit_reference(monkeypatch): + wsdpa = _reset_wsdpa(monkeypatch) + q, kv, pooled, sinks, scale, offset, window, ratio = _inputs( + q_len=6, + offset=11, + window=5, + pooled_len=5, + ratio=4, + ) + topk = mx.array( + [ + [0, 1, 2], + [0, 1, 2], + [0, 1, 2], + [0, 1, 3], + [1, 2, 3], + [1, 2, 4], + ], + dtype=mx.uint32, + )[None] + + out = wsdpa.wsdpa_topk_prefill( + q, kv, pooled, topk, sinks, scale, offset, window, ratio + ) + ref = _reference_attention( + q, kv, pooled, sinks, scale, offset, window, ratio, topk=topk + ) + + assert out is not None + mx.eval(out, ref) + assert out.shape == ref.shape + assert out.dtype == mx.bfloat16 + assert _max_abs(out, ref) < 8e-3 + + +def test_wsdpa_prefill_rejects_non_deepseek_v4_head_count(monkeypatch): + wsdpa = _reset_wsdpa(monkeypatch) + q = mx.zeros((1, 16, 4, 512), dtype=mx.bfloat16) + kv = mx.zeros((1, 1, 4, 512), dtype=mx.bfloat16) + sinks = mx.zeros((16,), dtype=mx.bfloat16) + + assert wsdpa.wsdpa_prefill(q, kv, None, sinks, 1.0, 0, 128, 1) is None + + +@requires_metal +def test_wsdpa_prefill_matches_reference_with_trimmed_rotating_cache(monkeypatch): + """RotatingKVCache trims the local buffer to the last W + L - 1 rows, so + during later prefill chunks buffer row 0 is at absolute position + base = offset + L - S > 0. The kernel must translate window bounds.""" + wsdpa = _reset_wsdpa(monkeypatch) + args = _inputs(q_len=8, offset=15, window=6, pooled_len=4, ratio=4, trim=5) + + out = wsdpa.wsdpa_prefill(*args) + ref = _reference_attention(*args) + + assert out is not None + mx.eval(out, ref) + assert out.shape == ref.shape + assert _max_abs(out, ref) < 8e-3 + + +@requires_metal +def test_wsdpa_topk_prefill_matches_reference_with_trimmed_rotating_cache( + monkeypatch, +): + wsdpa = _reset_wsdpa(monkeypatch) + q, kv, pooled, sinks, scale, offset, window, ratio = _inputs( + q_len=6, + offset=11, + window=5, + pooled_len=5, + ratio=4, + trim=3, + ) + topk = mx.array( + [ + [0, 1, 2], + [0, 1, 2], + [0, 1, 2], + [0, 1, 3], + [1, 2, 3], + [1, 2, 4], + ], + dtype=mx.uint32, + )[None] + + out = wsdpa.wsdpa_topk_prefill( + q, kv, pooled, topk, sinks, scale, offset, window, ratio + ) + ref = _reference_attention( + q, kv, pooled, sinks, scale, offset, window, ratio, topk=topk + ) + + assert out is not None + mx.eval(out, ref) + assert out.shape == ref.shape + assert _max_abs(out, ref) < 8e-3 diff --git a/tests/test_dsa_indexer_fused_mask.py b/tests/test_dsa_indexer_fused_mask.py new file mode 100644 index 000000000..ce8b93a4f --- /dev/null +++ b/tests/test_dsa_indexer_fused_mask.py @@ -0,0 +1,89 @@ +"""Regression tests for the fused pooled-ratio mask in dsa_indexer_scores. + +The kernel epilogue can apply the PoolingCache causal mask directly +(mask_ratio/mask_q_offset), replacing a separate mx.where pass. The fused +path must be BIT-IDENTICAL to the unfused reference (mask_ratio=0 followed +by mx.where with finfo.min), including -0.0/+0.0 behavior, and top-k +indices must match exactly. +""" + +import mlx.core as mx +import pytest + +from omlx.custom_kernels.glm_moe_dsa import fast as glm_fast + +pytestmark = pytest.mark.skipif( + not ( + glm_fast.has_symbol("dsa_indexer_scores") + and glm_fast.has_symbol("dsa_topk_indices") + ), + reason="glm_moe_dsa native extension not built", +) + + +def _reference_mask(L, P, ratio, q_offset): + """PoolingCache.make_mask semantics: visible iff col < (q_offset+row+1)//ratio.""" + rows = mx.arange(L)[:, None] + cols = mx.arange(P)[None, :] + return cols < ((q_offset + rows + 1) // ratio) + + +def _bit_equal(a, b): + mx.eval(a, b) + return bool(mx.array_equal(a.view(mx.uint16), b.view(mx.uint16))) + + +@pytest.mark.parametrize( + "L,P,ratio,q_offset,dtype", + [ + (128, 1088, 4, 256, mx.bfloat16), # offset > 0 + (64, 2048, 4, 0, mx.bfloat16), # offset 0 + (128, 2560, 128, 1024, mx.bfloat16), # ratio-128 style + (128, 1088, 4, 256, mx.float16), # fp16 + ], +) +def test_fused_mask_bit_identical(L, P, ratio, q_offset, dtype): + mx.random.seed(42) + H, D = 64, 128 + q = mx.random.normal((1, H, L, D)).astype(dtype) + pooled = mx.random.normal((1, P, D)).astype(dtype) + weights = mx.random.normal((1, L, H)).astype(dtype) # [B, L, H] convention + + mask = _reference_mask(L, P, ratio, q_offset) + + # Reference: unfused scores + mx.where pass (the old call-site flow). + ref = glm_fast.dsa_indexer_scores(q, pooled[:, None], weights, causal=False) + ref = mx.where(mask[None, None], ref, mx.finfo(ref.dtype).min) + + # Fused: kernel epilogue applies the same mask with the same sentinel. + fused = glm_fast.dsa_indexer_scores( + q, + pooled[:, None], + weights, + causal=False, + mask_ratio=ratio, + mask_q_offset=q_offset, + ) + + assert fused.shape == ref.shape + assert _bit_equal(fused, ref), "fused mask output differs bitwise from reference" + + k = min(512, P) + idx_ref = glm_fast.dsa_topk_indices(ref, k, bucketed=False) + idx_fused = glm_fast.dsa_topk_indices(fused, k, bucketed=False) + mx.eval(idx_ref, idx_fused) + assert bool(mx.array_equal(idx_ref, idx_fused)), "top-k indices differ" + + +def test_mask_ratio_zero_matches_unmasked(): + mx.random.seed(7) + H, D, L, P = 64, 128, 64, 512 + q = mx.random.normal((1, H, L, D)).astype(mx.bfloat16) + pooled = mx.random.normal((1, P, D)).astype(mx.bfloat16) + weights = mx.random.normal((1, L, H)).astype(mx.bfloat16) # [B, L, H] + + plain = glm_fast.dsa_indexer_scores(q, pooled[:, None], weights, causal=False) + zero_ratio = glm_fast.dsa_indexer_scores( + q, pooled[:, None], weights, causal=False, mask_ratio=0, mask_q_offset=0 + ) + assert _bit_equal(plain, zero_ratio) diff --git a/tests/test_pooling_cache_append_inplace.py b/tests/test_pooling_cache_append_inplace.py new file mode 100644 index 000000000..7a92b4c13 --- /dev/null +++ b/tests/test_pooling_cache_append_inplace.py @@ -0,0 +1,466 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Losslessness tests for the append-in-place PoolingCache rework. + +The caches in ``omlx/patches/deepseek_v4/cache_extras.py`` used to rebuild +``self.pooled`` with ``mx.concatenate`` on every chunk; they now append into +a preallocated backing buffer with geometric regrowth and expose the logical +tensor as a view. These tests pin the exact old observable behavior: +contents, shapes, offset/size bookkeeping, snapshot/delta immunity, and +trim/rollback semantics. +""" + +from __future__ import annotations + +import mlx.core as mx +import pytest + +from omlx.patches.deepseek_v4.cache_extras import ( + BatchPoolingCache, + PoolingCache, +) + + +def _rows(start: int, count: int, D: int, B: int = 1) -> mx.array: + """Deterministic distinct values per append so mis-ordering shows up.""" + vals = mx.arange(start * D * B, (start + count) * D * B, dtype=mx.float32) + return (vals.reshape(B, count, D) % 997) / 997.0 + + +class _RefSingle: + """Old concatenate semantics for PoolingCache.update_and_fetch.""" + + def __init__(self): + self.pooled = None + + def update_and_fetch(self, px: mx.array): + if px.shape[1] == 0: + return self.pooled + if self.pooled is None: + self.pooled = px + else: + self.pooled = mx.concatenate([self.pooled, px], axis=1) + return self.pooled + + +def _assert_same(actual, expected): + mx.eval(actual, expected) + assert actual.shape == expected.shape + assert bool(mx.array_equal(actual, expected)) + + +# --------------------------------------------------------------------------- +# PoolingCache (single sequence) +# --------------------------------------------------------------------------- + + +def test_single_varied_appends_match_concatenate_reference(): + cache = PoolingCache(4) + ref = _RefSingle() + # Include regrowth-forcing big appends, single rows, and zero-row calls. + sizes = [1, 3, 2, 8, 1, 1, 16, 5, 0, 33, 2, 0, 1, 64] + start = 0 + for n in sizes: + px = _rows(start, n, 8) + start += n + got = cache.update_and_fetch(px) + want = ref.update_and_fetch(px) + if want is None: + assert got.shape[1] == 0 + continue + _assert_same(got, want) + _assert_same(cache.pooled, want) + assert cache.offset == want.shape[1] + assert cache.size() == want.shape[1] + # Geometric capacity: backing buffer never smaller than the logical view. + assert cache._pool_buf.shape[1] >= cache._pool_len == ref.pooled.shape[1] + + +def test_single_capacity_regrowth_preserves_data(): + cache = PoolingCache(4) + ref = _RefSingle() + capacities = [] + for i in range(40): + px = _rows(i * 2, 2, 4) + cache.update_and_fetch(px) + ref.update_and_fetch(px) + mx.eval(cache.pooled) + capacities.append(cache._pool_buf.shape[1]) + _assert_same(cache.pooled, ref.pooled) + # Growth happened and was geometric (never grows by less than double). + assert capacities[-1] >= 80 + for prev, cur in zip(capacities, capacities[1:]): + assert cur == prev or cur >= 2 * prev + + +def test_single_snapshot_immune_to_later_appends(): + cache = PoolingCache(4) + ref = _RefSingle() + for i in range(6): + px = _rows(i, 1 + (i % 3), 8) + cache.update_and_fetch(px) + ref.update_and_fetch(px) + + # Materialized snapshot of the logical region (what pooling_delta does: + # slice + mx.contiguous, then the scheduler mx.evals the delta). + prefix_len = ref.pooled.shape[1] + snap = mx.contiguous(cache.pooled[:, :prefix_len]) + mx.eval(snap) + + # The state view, evaluated in place like the per-chunk fence does. + state_view = cache.state[2] + mx.eval(state_view) + + for i in range(10): + cache.update_and_fetch(_rows(100 + i * 4, 4, 8)) + mx.eval(cache.pooled) + + _assert_same(snap, ref.pooled) + # Rows below the snapshot length are never rewritten, so even the + # evaluated view still reads the same values. + _assert_same(state_view, ref.pooled) + assert cache.pooled.shape[1] == prefix_len + 40 + + +def test_single_state_setter_roundtrip(): + cache = PoolingCache(4) + ref = _RefSingle() + for i in range(5): + px = _rows(i, 3, 8) + cache.update_and_fetch(px) + ref.update_and_fetch(px) + + restored = PoolingCache(4) + restored.state = cache.state + _assert_same(restored.pooled, ref.pooled) + assert restored.offset == ref.pooled.shape[1] + + # Appends continue seamlessly after a restore (regrowth from exact fit). + more = _rows(50, 7, 8) + restored.update_and_fetch(more) + ref.update_and_fetch(more) + _assert_same(restored.pooled, ref.pooled) + + +def test_single_zero_row_append_on_empty_cache(): + cache = PoolingCache(4) + got = cache.update_and_fetch(mx.zeros((1, 0, 8), dtype=mx.float32)) + assert got.shape == (1, 0, 8) + assert cache.pooled is None + assert cache.offset == 0 + assert cache.empty() + + +def test_single_trim_within_remainder_keeps_pooled(): + cache = PoolingCache(4) + D1, D2 = 8, 8 + # Prompt of 5 tokens: completes one window, remainder 1. + kv = _rows(0, 5, D1) + gate = _rows(10, 5, D2) + r_kv, r_gate, _ = cache.accumulate_windows(kv, gate, 0) + assert r_kv.shape[1] == 4 + px = _rows(20, 1, 8) + cache.update_and_fetch(px) + mx.eval(cache.pooled) + assert cache.remainder == 1 + assert cache.offset == 1 + + assert cache.trim(1) == 1 + assert cache.remainder == 0 + # Pooled rows are untouched by a remainder trim. + _assert_same(cache.pooled, px) + + +def test_single_undo_trim_restores_pre_update_rows(): + """MTP draft rejection: a decode-sized update that completed a window is + rolled back through the one-update undo log; pooled must return to the + exact pre-update logical contents.""" + from omlx.patches.mlx_lm_mtp import cache_rollback + + cache_rollback.set_undo_armed(True) + try: + cache = PoolingCache(4) + ref = _RefSingle() + for i in range(3): + px = _rows(i * 2, 2, 8) + cache.update_and_fetch(px) + ref.update_and_fetch(px) + mx.eval(cache.pooled) + pre_update = mx.contiguous(cache.pooled) + mx.eval(pre_update) + + # Decode-sized update (L=1) that completes a window: 3 tokens sat in + # the remainder, so one more token produces a pooled row. + cache.remainder = 3 + cache.buf_kv = mx.zeros((1, 4, 8)) + cache.buf_gate = mx.zeros((1, 4, 8)) + kv = _rows(60, 1, 8) + gate = _rows(70, 1, 8) + r_kv, r_gate, _ = cache.accumulate_windows(kv, gate, 24) + assert r_kv.shape[1] == 4 # window completed + new_row = _rows(80, 1, 8) + cache.update_and_fetch(new_row) + mx.eval(cache.pooled) + assert cache.pooled.shape[1] == pre_update.shape[1] + 1 + + assert cache.is_trimmable() + assert cache.trim(1) == 1 + mx.eval(cache.pooled) + _assert_same(cache.pooled, pre_update) + assert cache.offset == pre_update.shape[1] + + # Appending after the rollback rewrites the trimmed slot; the + # pre-update snapshot taken before must stay immune. + cache.update_and_fetch(_rows(90, 2, 8)) + mx.eval(cache.pooled) + _assert_same(pre_update, ref.pooled) + finally: + cache_rollback.set_undo_armed(False) + + +# --------------------------------------------------------------------------- +# BatchPoolingCache +# --------------------------------------------------------------------------- + + +def _old_batch_update(state, px, ratio): + """Verbatim old (pre-rework) BatchPoolingCache.update_and_fetch semantics. + + ``state`` is a dict with keys pooled, pool_lengths, processed, remainder. + Returns the new pooled tensor; mutates pool_lengths in place. + """ + B, N, D = px.shape + pooled = state["pooled"] + pool_lengths = state["pool_lengths"] + + if N == 0: + return pooled + + new_counts = [ + (state["processed"][i] - state["remainder"][i]) // ratio - pool_lengths[i] + for i in range(B) + ] + max_new = max(new_counts) + if max_new == 0: + return pooled + + if B == 1: + count = new_counts[0] + current = pool_lengths[0] + new_rows = px[:, :count] + if pooled is None or current == 0: + pooled = new_rows + else: + pooled = mx.concatenate([pooled[:, :current], new_rows], axis=1) + pool_lengths[0] = current + count + return pooled + + max_pool = max(pool_lengths) + max_new + if pooled is None: + pooled = mx.zeros((B, max_pool, D), dtype=px.dtype) + elif pooled.shape[1] < max_pool: + pad = mx.zeros((B, max_pool - pooled.shape[1], D), dtype=px.dtype) + pooled = mx.concatenate([pooled, pad], axis=1) + + for i in range(B): + nc = new_counts[i] + if nc > 0: + pl = pool_lengths[i] + pooled[i, pl : pl + nc] = px[i, :nc] + pool_lengths[i] = pl + nc + return pooled + + +@pytest.mark.parametrize("B", [1, 2, 3]) +def test_batch_varied_appends_match_old_semantics(B): + ratio = 4 + cache = BatchPoolingCache(ratio, [0] * B) + ref = {"pooled": None, "pool_lengths": [0] * B} + + # Each step: every row consumes `step + i` tokens (some completing + # windows, some only filling remainders). + step = 0 + for tokens in ([4, 8, 3, 12, 5, 16, 1, 7, 20, 2, 9, 6],): + for t in tokens: + L = t + kv = _rows(step, L, 8, B) + gate = _rows(1000 + step, L, 8, B) + step += L + cache.prepare(lengths=[L] * B) + r_kv, r_gate, _ = cache.accumulate_windows(kv, gate, 0) + n_rows = r_kv.shape[1] // ratio + px = _rows(2000 + step, n_rows, 8, B) if n_rows else mx.zeros( + (B, 0, 8), dtype=mx.float32 + ) + # Reference bookkeeping mirrors the real cache's fields. + ref["processed"] = list(cache._processed) + ref["remainder"] = list(cache.remainder) + cache.update_and_fetch(px) + ref["pooled"] = _old_batch_update(ref, px, ratio) + ref["pool_lengths"] = list(cache._pool_lengths) + if ref["pooled"] is None: + assert cache.pooled is None or cache.pooled.shape[1] == 0 + continue + _assert_same(cache.pooled, ref["pooled"]) + assert cache.pooled.shape[1] == ref["pooled"].shape[1] + assert cache.size() == ref["pooled"].shape[1] + + +def test_batch_extent_overshoot_matches_old_shape(): + """Old physical shape could overshoot max(_pool_lengths) when the longest + row was not the row completing windows (max(lengths)+max_new).""" + ratio = 4 + B = 2 + cache = BatchPoolingCache(ratio, [0] * B) + ref = {"pooled": None, "pool_lengths": [0] * B} + + # Step 1: row 0 completes 10 windows (40 tokens), row 1 completes 1 + # (4 valid tokens; per-row valid lengths come from prepare()). + cache.prepare(lengths=[40, 4]) + kv = _rows(0, 40, 8, B) + gate = _rows(100, 40, 8, B) + cache.accumulate_windows(kv, gate, 0) + ref["processed"] = list(cache._processed) + ref["remainder"] = list(cache.remainder) + px = _rows(200, 10, 8, B) # row 0 -> 10 rows, row 1 -> 1 row + cache.update_and_fetch(px) + ref["pooled"] = _old_batch_update(ref, px, ratio) + ref["pool_lengths"] = list(cache._pool_lengths) + _assert_same(cache.pooled, ref["pooled"]) + assert cache._pool_lengths == [10, 1] + + # Step 2: row 0 completes nothing (3 tokens), row 1 completes 2 windows. + # Old max_pool overshoots: max(lengths)=10 + max_new=2 -> 12 while the + # new lengths are [10, 3]. + cache.prepare(lengths=[3, 8]) + kv = _rows(300, 8, 8, B) + gate = _rows(400, 8, 8, B) + cache.accumulate_windows(kv, gate, 0) + ref["processed"] = list(cache._processed) + ref["remainder"] = list(cache.remainder) + px = _rows(500, 2, 8, B) + cache.update_and_fetch(px) + ref["pooled"] = _old_batch_update(ref, px, ratio) + ref["pool_lengths"] = list(cache._pool_lengths) + + _assert_same(cache.pooled, ref["pooled"]) + assert ref["pooled"].shape[1] == 12 # overshoot really happened + assert cache.pooled.shape[1] == 12 + assert cache._pool_lengths == [10, 3] + # Overshoot columns stay zero-filled exactly like the old pad path. + tail = cache.pooled[:, 10:] + mx.eval(tail) + assert float(mx.abs(tail).max()) == 0.0 + + +def test_batch_snapshot_and_extract_immune_to_later_appends(): + ratio = 4 + cache = BatchPoolingCache(ratio, [0, 0]) + ref = {"pooled": None, "pool_lengths": [0] * 2} + + for step, t in enumerate([8, 8, 8, 8]): + cache.prepare(lengths=[t] * 2) + kv = _rows(step * 10, t, 8, 2) + gate = _rows(500 + step * 10, t, 8, 2) + cache.accumulate_windows(kv, gate, 0) + ref["processed"] = list(cache._processed) + ref["remainder"] = list(cache.remainder) + px = _rows(900 + step * 4, 2, 8, 2) + cache.update_and_fetch(px) + ref["pooled"] = _old_batch_update(ref, px, ratio) + ref["pool_lengths"] = list(cache._pool_lengths) + + mx.eval(cache.pooled) + snap = mx.contiguous(cache.pooled) + mx.eval(snap) + extracted = cache.extract(1) + mx.eval(extracted.pooled) + + # Keep appending (forces regrowth) and re-verify both snapshots. + for step, t in enumerate([12, 8, 16]): + cache.prepare(lengths=[t] * 2) + kv = _rows(2000 + step * 10, t, 8, 2) + gate = _rows(3000 + step * 10, t, 8, 2) + cache.accumulate_windows(kv, gate, 0) + ref["processed"] = list(cache._processed) + ref["remainder"] = list(cache.remainder) + n_rows = t // ratio + px = _rows(4000 + step * 4, n_rows, 8, 2) + cache.update_and_fetch(px) + ref["pooled"] = _old_batch_update(ref, px, ratio) + ref["pool_lengths"] = list(cache._pool_lengths) + + _assert_same(snap, ref["pooled"][:, : snap.shape[1]]) + # extract() holds row 1's first pl rows as an independent copy; each of + # the 4 pre-snapshot steps completed 2 windows per row, so pl == 8. + pl = 8 + _assert_same(extracted.pooled, ref["pooled"][1:2, :pl]) + assert isinstance(extracted, PoolingCache) + assert extracted.offset == pl + + +def test_batch_truncate_pooled_tail_matches_old_slice(): + ratio = 4 + cache = BatchPoolingCache(ratio, [0, 0]) + cache.prepare(lengths=[8, 8]) + kv = _rows(0, 8, 8, 2) + gate = _rows(100, 8, 8, 2) + cache.accumulate_windows(kv, gate, 0) + px = _rows(200, 2, 8, 2) + cache.update_and_fetch(px) + mx.eval(cache.pooled) + assert cache.pooled.shape[1] == 2 + + # Simulate a rejected speculative suffix on row 1 only: old code sliced + # pooled to max(_pool_lengths). + cache._pool_lengths[1] = 1 + cache._truncate_pooled_tail() + assert cache.pooled.shape[1] == 2 # max length still 2 (row 0) + cache._pool_lengths[0] = 1 + cache._truncate_pooled_tail() + assert cache.pooled.shape[1] == 1 + _assert_same(cache.pooled, px[:, :1]) + + +def test_batch_state_setter_and_filter(): + ratio = 4 + cache = BatchPoolingCache(ratio, [0, 0]) + cache.prepare(lengths=[8, 8]) + kv = _rows(0, 8, 8, 2) + gate = _rows(100, 8, 8, 2) + cache.accumulate_windows(kv, gate, 0) + px = _rows(200, 2, 8, 2) + cache.update_and_fetch(px) + mx.eval(cache.pooled) + + restored = BatchPoolingCache(ratio, [0, 0]) + restored.state = cache.state + _assert_same(restored.pooled, cache.pooled) + + filtered = BatchPoolingCache(ratio, [0, 0]) + filtered.state = cache.state + filtered._pool_lengths = list(cache._pool_lengths) + filtered.filter([1]) + _assert_same(filtered.pooled, cache.pooled[1:2]) + assert filtered._pool_lengths == [cache._pool_lengths[1]] + + +def test_merge_single_caches_preserves_contents(): + caches = [] + refs = [] + for b in range(3): + c = PoolingCache(4) + ref = _RefSingle() + for i in range(b + 2): + px = _rows(10 * b + i, 1 + i, 8) + c.update_and_fetch(px) + ref.update_and_fetch(px) + mx.eval(c.pooled) + caches.append(c) + refs.append(ref.pooled) + + batch = PoolingCache.merge(caches) + assert isinstance(batch, BatchPoolingCache) + max_pool = max(r.shape[1] for r in refs) + assert batch.pooled.shape == (3, max_pool, 8) + for i, r in enumerate(refs): + _assert_same(batch.pooled[i : i + 1, : r.shape[1]], r) diff --git a/tests/test_prefill_oom_graceful.py b/tests/test_prefill_oom_graceful.py index 95c30ebee..557340ae7 100644 --- a/tests/test_prefill_oom_graceful.py +++ b/tests/test_prefill_oom_graceful.py @@ -812,6 +812,7 @@ def test_step_prefill_reclaims_before_first_guard(): _memory_limit_bytes=0, _glm_dsa_adaptive_prefill=None, model=lambda *args, **kwargs: events.append("model"), + _supports_skip_lm_head=lambda: False, _adaptive_chunk_size=lambda n, **kwargs: events.append("adaptive") or n, _guard_prefill_chunk=lambda n, **kwargs: events.append("guard") or n, _record_chunk_transient=MagicMock(), diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 775df7324..5a5cd9f51 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -6061,3 +6061,43 @@ def test_prompt_progress_uses_model_name_not_basename( progress_none = tracker.get_model_progress("def456") assert len(progress_none) == 0 tracker.clear() + + +class TestSupportsSkipLmHead: + """Regression coverage for Scheduler._supports_skip_lm_head. + + Chunked prefill discards every chunk's logits, so patched DeepSeek-V4 + models accept ``skip_lm_head=True`` to skip the full-vocabulary + projection. Unknown models must keep stock behavior. + """ + + def _scheduler_with_model(self, model): + scheduler = Scheduler.__new__(Scheduler) + scheduler.model = model + return scheduler + + def test_detects_support(self): + class PatchedModel: + def __call__(self, inputs, cache=None, skip_lm_head=False): + return None + + scheduler = self._scheduler_with_model(PatchedModel()) + assert scheduler._supports_skip_lm_head() is True + # Result is cached on the instance. + assert scheduler._skip_lm_head_supported is True + + def test_rejects_stock_model(self): + class StockModel: + def __call__(self, inputs, cache=None): + return None + + scheduler = self._scheduler_with_model(StockModel()) + assert scheduler._supports_skip_lm_head() is False + assert scheduler._skip_lm_head_supported is False + + def test_rejects_missing_call(self): + class NoCall: + pass + + scheduler = self._scheduler_with_model(NoCall()) + assert scheduler._supports_skip_lm_head() is False