From db154f98088ea3135741197d77ca4c2b1cf0bb65 Mon Sep 17 00:00:00 2001 From: jonathan308 Date: Mon, 3 Aug 2026 22:24:37 -0700 Subject: [PATCH 001/338] fix(cache): copy pooling-cache deltas out of their parent buffer (#2500) Detach retained pooling-cache deltas from cumulative parent buffers to prevent quadratic memory retention. --- omlx/cache/pooling_delta.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/omlx/cache/pooling_delta.py b/omlx/cache/pooling_delta.py index 392440d73c..48178de5ca 100644 --- a/omlx/cache/pooling_delta.py +++ b/omlx/cache/pooling_delta.py @@ -6,6 +6,13 @@ import logging from typing import Any +try: + import mlx.core as mx + + HAS_MLX = True +except ImportError: + HAS_MLX = False + logger = logging.getLogger(__name__) POOLING_CACHE_DELTA_CLASS = "PoolingCacheDelta" @@ -88,7 +95,16 @@ def compact_pooling_cache_snapshot( ) continue + # mx.contiguous, not a bare slice: a retained view keeps its + # whole parent buffer alive, so an in-memory snapshot would pin + # the entire cumulative pooled tensor instead of this block's + # delta -- the compaction saves nothing and total retention goes + # quadratic in context length. The SSD path escapes this only + # because serialization copies the bytes out. (PoolingCache + # .extract() copies for the same reason, cache_extras.py:930.) delta = pooled[:, expected_start:expected_end] + if HAS_MLX and hasattr(mx, "contiguous"): + delta = mx.contiguous(delta) compacted_states[sub_idx] = ( state[0], state[1], From 5aaf006e7f58c8ffdc23dc6a391c2394ddbbae3c Mon Sep 17 00:00:00 2001 From: jonathan308 Date: Mon, 3 Aug 2026 22:25:03 -0700 Subject: [PATCH 002/338] fix(deepseek_v4): make retained-reasoning prompts append-only (#2501) Preserve historical thinking markers and retained reasoning when drop_thinking is disabled so cached prompt prefixes remain reusable. --- omlx/patches/deepseek_v4/chat_template_v4.py | 66 +++++++++++++---- .../test_deepseek_v4_template_append_only.py | 73 +++++++++++++++++++ 2 files changed, 126 insertions(+), 13 deletions(-) create mode 100644 tests/test_deepseek_v4_template_append_only.py diff --git a/omlx/patches/deepseek_v4/chat_template_v4.py b/omlx/patches/deepseek_v4/chat_template_v4.py index e3bfdc3de7..7b7e1d94f7 100644 --- a/omlx/patches/deepseek_v4/chat_template_v4.py +++ b/omlx/patches/deepseek_v4/chat_template_v4.py @@ -178,11 +178,38 @@ def find_last_user_index(messages: List[Dict[str, Any]]) -> int: return last_user_index +def _user_thinking_suffix( + index: int, + last_user_idx: int, + thinking_mode: str, + drop_thinking: bool, +) -> str: + """Marker that follows a user/developer message. + + The current turn opens thinking. Historical turns normally close it, + because their reasoning is dropped from the transcript — but that makes + rendering position-dependent: when a new user message arrives, the + previous user message's suffix flips from ```` to ```` + and every cached prefix from that point on is invalidated. + + When reasoning is retained (``drop_thinking=False``), the historical + turn still carries its reasoning block, so keeping ```` here + reproduces exactly what was rendered while that turn was current and + the transcript grows append-only — which is what a prefix cache needs. + """ + if thinking_mode != "thinking": + return thinking_end_token + if index == last_user_idx or not drop_thinking: + return thinking_start_token + return thinking_end_token + + def render_message( index: int, messages: List[Dict[str, Any]], thinking_mode: str, tools: Any = None, + drop_thinking: bool = True, ) -> str: assert 0 <= index < len(messages) assert thinking_mode in [ @@ -228,18 +255,16 @@ def render_message( content_developer += "\n\n# The user's message is: {}".format(content) prompt += user_msg_template.format(content=content_developer) - if index == last_user_idx and thinking_mode == "thinking": - prompt += thinking_start_token - else: - prompt += thinking_end_token + prompt += _user_thinking_suffix( + index, last_user_idx, thinking_mode, drop_thinking + ) elif role == "user": prompt += user_msg_template.format(content=content) - if index == last_user_idx and thinking_mode == "thinking": - prompt += thinking_start_token - else: - prompt += thinking_end_token + prompt += _user_thinking_suffix( + index, last_user_idx, thinking_mode, drop_thinking + ) elif role == "tool": prev_assistant_idx = index - 1 @@ -268,7 +293,14 @@ def render_message( if tool_call_order == len(assistant_tool_calls): prompt += "\n" - if index >= last_user_idx and thinking_mode == "thinking": + # Same append-only rule as the user/developer suffix: when + # reasoning is retained, a historical tool result keeps the + # opener it was rendered with at generation time (the + # following assistant message renders `reasoning...` + # with no opener of its own, so the pair stays consistent). + if thinking_mode == "thinking" and ( + index >= last_user_idx or not drop_thinking + ): prompt += "\n\n" + thinking_start_token else: prompt += "\n\n" + thinking_end_token @@ -293,10 +325,17 @@ def render_message( summary_content = content or "" - if thinking_mode == "thinking" and index > last_user_idx: - assert ( - reasoning_content or tool_calls - ), f"ThinkingMode: {thinking_mode}, invalid message without reasoning_content/tool_calls `{msg}` after last user message" + # Historical assistant turns render their reasoning too when it is + # retained, so the transcript stays byte-identical to what was + # rendered while that turn was current (append-only prefix). + render_thinking = thinking_mode == "thinking" and ( + index > last_user_idx or not drop_thinking + ) + if render_thinking: + if index > last_user_idx: + assert ( + reasoning_content or tool_calls + ), f"ThinkingMode: {thinking_mode}, invalid message without reasoning_content/tool_calls `{msg}` after last user message" thinking_part = ( thinking_template.format(reasoning_content=reasoning_content or "") + thinking_end_token @@ -371,6 +410,7 @@ def encode_messages( full_messages, thinking_mode=thinking_mode, tools=tools, + drop_thinking=drop_thinking, ) return prompt diff --git a/tests/test_deepseek_v4_template_append_only.py b/tests/test_deepseek_v4_template_append_only.py new file mode 100644 index 0000000000..b351ebbd98 --- /dev/null +++ b/tests/test_deepseek_v4_template_append_only.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Append-only rendering for the DeepSeek V4 DSML chat template. + +With drop_thinking=False (reasoning retained), the rendered transcript must +be append-only across new user turns so prefix caches stay valid: a turn's +rendering never changes once it is historical. Default rendering +(drop_thinking=True) must be byte-identical to the previous behavior. +""" +import pytest + +from omlx.patches.deepseek_v4 import chat_template_v4 as tmpl + +SYSTEM = {"role": "system", "content": "You are a coding agent."} +U1 = {"role": "user", "content": "Refactor the parser module."} +A1 = { + "role": "assistant", + "content": "Done: extracted tokenize().", + "reasoning_content": "Plan the split, then extract.", +} +U2 = {"role": "user", "content": "Now add tests."} +A_TOOL = { + "role": "assistant", + "content": "", + "reasoning_content": "Need to inspect the file.", + "tool_calls": [ + {"function": {"name": "read_file", "arguments": '{"path": "a.py"}'}} + ], +} +TOOL = {"role": "tool", "content": "file contents"} + + +def render(messages, **kwargs): + return tmpl.encode_messages(messages, thinking_mode="thinking", **kwargs) + + +class TestAppendOnlyRendering: + def test_new_user_turn_is_append_only_when_reasoning_retained(self): + turn1 = render([SYSTEM, U1], drop_thinking=False) + turn2 = render([SYSTEM, U1, A1, U2], drop_thinking=False) + assert turn2.startswith(turn1) + + def test_tool_loop_then_user_turn_is_append_only(self): + loop = render([SYSTEM, U1, A_TOOL, TOOL], drop_thinking=False) + follow = render([SYSTEM, U1, A_TOOL, TOOL, A1, U2], drop_thinking=False) + assert follow.startswith(loop) + + def test_historical_assistant_reasoning_is_rendered_when_retained(self): + out = render([SYSTEM, U1, A1, U2], drop_thinking=False) + assert A1["reasoning_content"] in out + + def test_default_rendering_flips_previous_user_marker(self): + # Documents the existing default behavior this feature works around: + # with drop_thinking=True the previous user suffix flips to , + # so the rendering is NOT append-only. + turn1 = render([SYSTEM, U1]) + turn2 = render([SYSTEM, U1, A1, U2]) + assert not turn2.startswith(turn1) + + @pytest.mark.parametrize("mode", ["thinking", "chat"]) + def test_default_rendering_unchanged_across_cases(self, mode): + # drop_thinking=True is the default; passing it explicitly must be + # byte-identical to omitting it for every case/mode combination. + cases = [ + [SYSTEM, U1], + [SYSTEM, U1, A1, U2], + [SYSTEM, U1, A_TOOL, TOOL], + [SYSTEM, U1, A_TOOL, TOOL, A1, U2], + [U1, A1, U2], + ] + for msgs in cases: + assert tmpl.encode_messages( + msgs, thinking_mode=mode + ) == tmpl.encode_messages(msgs, thinking_mode=mode, drop_thinking=True) From 75987b9f66d154f58ecdcdf43747b60ea797c2e6 Mon Sep 17 00:00:00 2001 From: jonathan308 Date: Mon, 3 Aug 2026 22:25:34 -0700 Subject: [PATCH 003/338] perf(deepseek_v4): guard indexer fallback past int32 limits (#2502) Tile oversized MLX indexer fallback operations, fuse the head reduction, and warn once when the native kernel is unavailable. --- omlx/patches/deepseek_v4/deepseek_v4_model.py | 81 +++++++++++++++++-- tests/test_deepseek_v4_patch.py | 72 +++++++++++++++++ 2 files changed, 145 insertions(+), 8 deletions(-) diff --git a/omlx/patches/deepseek_v4/deepseek_v4_model.py b/omlx/patches/deepseek_v4/deepseek_v4_model.py index ae9736ee2c..1344f1cb63 100644 --- a/omlx/patches/deepseek_v4/deepseek_v4_model.py +++ b/omlx/patches/deepseek_v4/deepseek_v4_model.py @@ -1,5 +1,6 @@ # Copyright © 2026 Apple Inc. +import logging import math from dataclasses import dataclass, field from functools import lru_cache, partial @@ -459,6 +460,33 @@ def _overlap_compress_kv(kv, gate, ape, head_dim): return (kv * weights).sum(axis=-2) +# Pooled-axis tile for the MLX indexer fallback (used when the native +# glm_moe_dsa extension is not built). The native dsa_indexer_scores kernel +# keeps the (heads, L, P) score tensor in registers; the fallback has to +# materialize it, and at ratio 4 with a 512-token chunk P = ctx/4, so the +# intermediate is (1, 64, 512, ctx/4) fp32 — 8.3 GiB at 273k context, read +# five more times. Worse, 64*512*P crosses 2**31 elements at ctx = 256k, the +# boundary where mlx's int32 kernel indexing silently zeros the tail and +# corrupts top-k selection. Tiling the pooled axis bounds the live +# intermediate at 64*512*TILE and keeps every matmul under 2**31. +_INDEXER_POOL_TILE = 16384 +# mlx kernels index with int32, so a tensor at or past 2**31 elements has its +# tail silently zeroed. Stay a factor of 2 below that. For a 512-token chunk +# with 64 index heads this is reached at P = ctx/4 ~= 32768, i.e. ctx ~= 128k. +_INDEXER_MAX_ELEMS = 2**30 +_DEEPSEEK_V4_INDEXER_FALLBACK_WARNED = False + + +@partial(mx.compile, shapeless=True) +def _indexer_head_reduce(scores, weights, scale): + """relu -> scale -> head-weight -> head-sum, fused over `scores`. + + `scores` is (B, H, L, P_tile); the reduction is over the head axis (1), + entirely within each pooled tile, so tiling P never crosses the reduce. + """ + return (mx.maximum(scores, 0) * scale * weights).sum(axis=1) + + @partial(mx.compile, shapeless=True) def _split_softmax(log_normalizer, logits_a, logits_b, sinks=None): if sinks is not None: @@ -1173,9 +1201,26 @@ def __call__( try: from omlx.custom_kernels.glm_moe_dsa import fast as glm_fast - if glm_fast.has_symbol("dsa_indexer_scores") and glm_fast.has_symbol( - "dsa_topk_indices" - ): + _have_native = glm_fast.has_symbol( + "dsa_indexer_scores" + ) and glm_fast.has_symbol("dsa_topk_indices") + if not _have_native: + # Warn only when the kernels are genuinely missing -- the + # shape predicates above legitimately skip small/early + # chunks, so warning outside this branch cries wolf. The + # miss is otherwise completely silent, and the MLX + # fallback's prefill slope is ~4x worse at long context. + global _DEEPSEEK_V4_INDEXER_FALLBACK_WARNED + if not _DEEPSEEK_V4_INDEXER_FALLBACK_WARNED: + _DEEPSEEK_V4_INDEXER_FALLBACK_WARNED = True + logging.getLogger(__name__).warning( + "deepseek_v4: native dsa_indexer_scores/" + "dsa_topk_indices unavailable (glm_moe_dsa " + "extension not built); falling back to MLX. " + "Long-context prefill is several times slower " + "(rebuild with OMLX_WITH_CUSTOM_KERNEL=1)." + ) + if _have_native: weights = ( self.weights_proj(x) if projected_weights is None @@ -1207,14 +1252,34 @@ def __call__( except Exception: _DEEPSEEK_V4_INDEXER_NATIVE_DISABLED = True - scores = q.astype(mx.float32) @ pooled[:, None].swapaxes(-1, -2).astype( - mx.float32 - ) - scores = mx.maximum(scores, 0) * self.scale weights = ( self.weights_proj(x) if projected_weights is None else projected_weights ).astype(mx.float32) * (self.n_heads**-0.5) - scores = (scores * weights.swapaxes(-1, -2)[..., None]).sum(axis=1) + weights = weights.swapaxes(-1, -2)[..., None] # (B, H, L, 1) + qf = q.astype(mx.float32) + kf = pooled[:, None].swapaxes(-1, -2).astype(mx.float32) # (B, 1, Dh, P) + n_pool = pooled.shape[1] + n_elems = qf.shape[0] * qf.shape[1] * qf.shape[2] * n_pool + if n_elems < _INDEXER_MAX_ELEMS: + # Single matmul: fastest, and the intermediate is safely indexable. + scores = _indexer_head_reduce(qf @ kf, weights, self.scale) + else: + # Only past the int32 indexing limit is tiling worth its cost: + # splitting the pooled axis adds matmul launches and a full-width + # concatenate (measured: ~1.2x slower slope at 273k), but an + # intermediate over 2**31 elements is silently zeroed by mlx's + # int32 kernel indexing, which corrupts top-k selection. Choose + # the largest tile that stays under the limit so the split is as + # coarse as correctness allows. + per_pool = max(1, qf.shape[0] * qf.shape[1] * qf.shape[2]) + tile = max(1024, min(_INDEXER_POOL_TILE, _INDEXER_MAX_ELEMS // per_pool)) + scores = mx.concatenate( + [ + _indexer_head_reduce(qf @ kf[..., s : s + tile], weights, self.scale) + for s in range(0, n_pool, tile) + ], + axis=-1, + ) if pmask is not None: scores = mx.where( pmask if pmask.ndim == 3 else pmask[None], diff --git a/tests/test_deepseek_v4_patch.py b/tests/test_deepseek_v4_patch.py index c66d5cac66..e896dc7fc9 100644 --- a/tests/test_deepseek_v4_patch.py +++ b/tests/test_deepseek_v4_patch.py @@ -1460,3 +1460,75 @@ def test_native_block_kind_short_circuits_on_nax_prefill(self, monkeypatch): gated = linear._native_block_kind(decode_x, True) monkeypatch.setattr(sl, "_nax_prefers_stock", lambda n: False) assert gated == linear._native_block_kind(decode_x, True) + + +class TestIndexerFallbackTiling: + """The MLX indexer fallback (used when the native glm_moe_dsa kernel is + not built) tiles the pooled axis so its (B, heads, L, P) intermediate + never crosses 2**31 elements — the boundary where mlx int32 kernel + indexing silently zeroes the tail and corrupts top-k selection at + >256k context — while keeping top-k selection identical to the untiled + reduction.""" + + def _reduce_and_ref(self): + # The patch registers deepseek_v4_model.py as mlx_lm.models.deepseek_v4 + # (its relative `.base` import resolves there); import it by that name. + import sys + + import mlx.core as mx + + dm = sys.modules["mlx_lm.models.deepseek_v4"] + return mx, dm + + def test_head_reduce_matches_naive(self, applied_patch): + mx, dm = self._reduce_and_ref() + mx.random.seed(0) + scores = mx.random.normal((1, 8, 16, 64)) + weights = mx.random.normal((1, 8, 16, 1)) + got = dm._indexer_head_reduce(scores, weights, 0.125) + ref = (mx.maximum(scores, 0) * 0.125 * weights).sum(axis=1) + assert float(mx.abs(got - ref).max()) < 1e-5 + + def test_tiling_selects_identical_topk(self, applied_patch): + # Split a non-reduced axis: the top-k indices must be bit-stable + # regardless of tile size, at pooled counts that straddle the tile. + mx, dm = self._reduce_and_ref() + mx.random.seed(1) + H, L, Dh, topk = 64, 32, 128, 64 + scale = Dh ** -0.5 + + def untiled(q, pooled, w): + wf = (w.astype(mx.float32)).swapaxes(-1, -2)[..., None] + s = q.astype(mx.float32) @ pooled[:, None].swapaxes(-1, -2).astype(mx.float32) + return dm._indexer_head_reduce(s, wf, scale) + + def tiled(q, pooled, w, tile): + wf = (w.astype(mx.float32)).swapaxes(-1, -2)[..., None] + qf = q.astype(mx.float32) + kf = pooled[:, None].swapaxes(-1, -2).astype(mx.float32) + P = pooled.shape[1] + if P <= tile: + return dm._indexer_head_reduce(qf @ kf, wf, scale) + return mx.concatenate( + [dm._indexer_head_reduce(qf @ kf[..., s:s + tile], wf, scale) + for s in range(0, P, tile)], + axis=-1, + ) + + for P in (255, 256, 257, 800): + q = mx.random.normal((1, H, L, Dh)) + pooled = mx.random.normal((1, P, Dh)) + w = mx.random.normal((1, L, H)) + a = untiled(q, pooled, w) + b = tiled(q, pooled, w, tile=256) + k = min(topk, P) + ia = mx.sort(mx.argpartition(-a, k - 1, axis=-1)[..., :k], axis=-1) + ib = mx.sort(mx.argpartition(-b, k - 1, axis=-1)[..., :k], axis=-1) + assert int((ia != ib).sum()) == 0, f"top-k differs at P={P}" + + def test_tile_stays_under_int32_index_limit(self, applied_patch): + # The prefill chunk is 512 and index heads are 64; the tiled matmul + # output must stay below 2**31 elements at any context length. + _, dm = self._reduce_and_ref() + assert 64 * 512 * dm._INDEXER_POOL_TILE < 2 ** 31 + assert dm._INDEXER_MAX_ELEMS < 2 ** 31 From 7cd9586efeacefcd4fbbf15e060b6f0476dd4428 Mon Sep 17 00:00:00 2001 From: jundot Date: Tue, 4 Aug 2026 14:29:36 +0900 Subject: [PATCH 004/338] test: cover DeepSeek V4 indexer fallback warning --- omlx/patches/deepseek_v4/deepseek_v4_model.py | 4 +- tests/test_deepseek_v4_patch.py | 92 ++++++++++++++++--- 2 files changed, 80 insertions(+), 16 deletions(-) diff --git a/omlx/patches/deepseek_v4/deepseek_v4_model.py b/omlx/patches/deepseek_v4/deepseek_v4_model.py index 1344f1cb63..24a0b3e79c 100644 --- a/omlx/patches/deepseek_v4/deepseek_v4_model.py +++ b/omlx/patches/deepseek_v4/deepseek_v4_model.py @@ -1275,7 +1275,9 @@ def __call__( tile = max(1024, min(_INDEXER_POOL_TILE, _INDEXER_MAX_ELEMS // per_pool)) scores = mx.concatenate( [ - _indexer_head_reduce(qf @ kf[..., s : s + tile], weights, self.scale) + _indexer_head_reduce( + qf @ kf[..., s : s + tile], weights, self.scale + ) for s in range(0, n_pool, tile) ], axis=-1, diff --git a/tests/test_deepseek_v4_patch.py b/tests/test_deepseek_v4_patch.py index e896dc7fc9..bb4dbf8cea 100644 --- a/tests/test_deepseek_v4_patch.py +++ b/tests/test_deepseek_v4_patch.py @@ -1489,46 +1489,108 @@ def test_head_reduce_matches_naive(self, applied_patch): ref = (mx.maximum(scores, 0) * 0.125 * weights).sum(axis=1) assert float(mx.abs(got - ref).max()) < 1e-5 + def test_missing_native_warning_fires_once( + self, applied_patch, caplog, monkeypatch + ): + import logging + import sys + + import mlx.core as mx + + from omlx.custom_kernels.glm_moe_dsa import fast as glm_fast + + dm = sys.modules["mlx_lm.models.deepseek_v4"] + config = dm.ModelArgs( + hidden_size=16, + q_lora_rank=16, + qk_rope_head_dim=2, + num_hidden_layers=1, + compress_ratios=[4], + index_n_heads=32, + index_head_dim=128, + index_topk=8, + ) + indexer = dm.Indexer(config, compress_ratio=4) + pooled = mx.zeros((1, 64, 128), dtype=mx.float16) + monkeypatch.setattr( + dm.Compressor, + "__call__", + lambda self, x, pool_cache, offset: pooled, + ) + monkeypatch.setattr(glm_fast, "has_symbol", lambda name: False) + monkeypatch.setattr(dm, "_DEEPSEEK_V4_INDEXER_NATIVE_DISABLED", False) + monkeypatch.setattr(dm, "_DEEPSEEK_V4_INDEXER_FALLBACK_WARNED", False) + + x = mx.zeros((1, 64, 16), dtype=mx.float16) + projected_q = mx.zeros((1, 32, 64, 128), dtype=mx.float16) + projected_weights = mx.zeros((1, 64, 32), dtype=mx.float16) + with caplog.at_level(logging.WARNING, logger=dm.__name__): + for _ in range(2): + result = indexer( + x, + q_residual=x, + position_rope=None, + pool_cache=None, + offset=0, + projected_q=projected_q, + projected_weights=projected_weights, + ) + mx.eval(result) + + fallback_warnings = [ + record + for record in caplog.records + if "native dsa_indexer_scores/dsa_topk_indices unavailable" + in record.getMessage() + ] + assert len(fallback_warnings) == 1 + def test_tiling_selects_identical_topk(self, applied_patch): # Split a non-reduced axis: the top-k indices must be bit-stable # regardless of tile size, at pooled counts that straddle the tile. mx, dm = self._reduce_and_ref() mx.random.seed(1) - H, L, Dh, topk = 64, 32, 128, 64 - scale = Dh ** -0.5 + heads, length, head_dim, topk = 64, 32, 128, 64 + scale = head_dim**-0.5 def untiled(q, pooled, w): wf = (w.astype(mx.float32)).swapaxes(-1, -2)[..., None] - s = q.astype(mx.float32) @ pooled[:, None].swapaxes(-1, -2).astype(mx.float32) + s = q.astype(mx.float32) @ pooled[:, None].swapaxes(-1, -2).astype( + mx.float32 + ) return dm._indexer_head_reduce(s, wf, scale) def tiled(q, pooled, w, tile): wf = (w.astype(mx.float32)).swapaxes(-1, -2)[..., None] qf = q.astype(mx.float32) kf = pooled[:, None].swapaxes(-1, -2).astype(mx.float32) - P = pooled.shape[1] - if P <= tile: + pool_count = pooled.shape[1] + if pool_count <= tile: return dm._indexer_head_reduce(qf @ kf, wf, scale) return mx.concatenate( - [dm._indexer_head_reduce(qf @ kf[..., s:s + tile], wf, scale) - for s in range(0, P, tile)], + [ + dm._indexer_head_reduce(qf @ kf[..., s : s + tile], wf, scale) + for s in range(0, pool_count, tile) + ], axis=-1, ) - for P in (255, 256, 257, 800): - q = mx.random.normal((1, H, L, Dh)) - pooled = mx.random.normal((1, P, Dh)) - w = mx.random.normal((1, L, H)) + for pool_count in (255, 256, 257, 800): + q = mx.random.normal((1, heads, length, head_dim)) + pooled = mx.random.normal((1, pool_count, head_dim)) + w = mx.random.normal((1, length, heads)) a = untiled(q, pooled, w) b = tiled(q, pooled, w, tile=256) - k = min(topk, P) + k = min(topk, pool_count) ia = mx.sort(mx.argpartition(-a, k - 1, axis=-1)[..., :k], axis=-1) ib = mx.sort(mx.argpartition(-b, k - 1, axis=-1)[..., :k], axis=-1) - assert int((ia != ib).sum()) == 0, f"top-k differs at P={P}" + assert ( + int((ia != ib).sum()) == 0 + ), f"top-k differs at pool_count={pool_count}" def test_tile_stays_under_int32_index_limit(self, applied_patch): # The prefill chunk is 512 and index heads are 64; the tiled matmul # output must stay below 2**31 elements at any context length. _, dm = self._reduce_and_ref() - assert 64 * 512 * dm._INDEXER_POOL_TILE < 2 ** 31 - assert dm._INDEXER_MAX_ELEMS < 2 ** 31 + assert 64 * 512 * dm._INDEXER_POOL_TILE < 2**31 + assert dm._INDEXER_MAX_ELEMS < 2**31 From ed9c5e7397bf203ac0318829a072138f77789906 Mon Sep 17 00:00:00 2001 From: jundot Date: Tue, 4 Aug 2026 15:01:55 +0900 Subject: [PATCH 005/338] fix: preserve CacheList signatures when stripping rotating tips (#2493) Re-wrap loaded CacheList payloads before re-saving superseded rotating tips so sub-cache composition metadata survives the round trip. --- omlx/cache/prefix_cache.py | 6 ++ tests/test_prefix_cache_rotating_tip_strip.py | 65 +++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/omlx/cache/prefix_cache.py b/omlx/cache/prefix_cache.py index 264fd18039..423e7972c2 100644 --- a/omlx/cache/prefix_cache.py +++ b/omlx/cache/prefix_cache.py @@ -1225,6 +1225,12 @@ def _strip_rotating_payload(self, block_hash: bytes) -> bool: ): new_data.append((mx.zeros((1,)), mx.zeros((1,)))) stripped += 1 + elif type_name == "CacheList" and isinstance(layer, list): + # load_block_with_metadata() exposes CacheList payloads as + # their legacy list shape. Restore the storage marker + # before re-saving so save_block keeps sub_count metadata + # and can stamp the same cachelist_subtypes signature. + new_data.append(("__cache_list__", layer)) else: new_data.append(layer) if stripped == 0: diff --git a/tests/test_prefix_cache_rotating_tip_strip.py b/tests/test_prefix_cache_rotating_tip_strip.py index dfc9a79f6d..4037188f06 100644 --- a/tests/test_prefix_cache_rotating_tip_strip.py +++ b/tests/test_prefix_cache_rotating_tip_strip.py @@ -40,6 +40,7 @@ from omlx.cache._rotating_subclass import PrefillReadyRotatingKVCache from omlx.cache.paged_cache import BlockTable, PagedCacheManager from omlx.cache.paged_ssd_cache import PagedSSDCacheManager, SharedHotCacheBudget +from omlx.cache.pooling_delta import POOLING_CACHE_DELTA_CLASS from omlx.cache.prefix_cache import BlockAwarePrefixCache from omlx.cache.type_registry import CacheTypeRegistry @@ -130,6 +131,37 @@ def _kvcache_only_data(seq_len): ] +def _v4_shaped_storage_data(): + """Serialized V4 shape: top-level rotating plus mixed CacheList.""" + rotating = ( + mx.ones((1, 1, WINDOW, 2), dtype=mx.float32), + mx.ones((1, 1, WINDOW, 2), dtype=mx.float32), + ) + pooling = ( + "__nstate__", + POOLING_CACHE_DELTA_CLASS, + [ + None, + None, + mx.ones((1, 1, 2), dtype=mx.float32), + mx.ones((1, 1, 2), dtype=mx.float32), + mx.ones((1, 1, 2), dtype=mx.float32), + mx.array([0, 1], dtype=mx.int64), + ], + ) + return ( + [rotating, ("__cache_list__", [rotating, pooling, pooling])], + ["RotatingKVCache", "CacheList"], + [ + (0, WINDOW, WINDOW, WINDOW), + ( + ["RotatingKVCache", "PoolingCache", "PoolingCache"], + [None, 4, 4], + ), + ], + ) + + def _store_turn(cache, turn, num_blocks, data_fn=_hybrid_cache_data): """Simulate one conversation turn: store a chain of num_blocks blocks. @@ -229,6 +261,39 @@ def test_stripped_block_keeps_sliceable_layers(tmp_path): assert cache._is_placeholder_state(data[rotating_idx]) +def test_stripped_block_preserves_v4_cachelist_signature(tmp_path): + """Re-saving a stripped V4 tip must preserve CacheList composition.""" + from omlx.cache.paged_ssd_cache import _signature_cachelist_subtypes + + cache, ssd = _make_cache(tmp_path) + block_hash = b"\x24" * 32 + data, layer_types, layer_meta = _v4_shaped_storage_data() + expected = {"1": ["RotatingKVCache", "PoolingCache:5", "PoolingCache:5"]} + + assert ssd.save_block( + block_hash=block_hash, + cache_data=data, + token_count=BLOCK_SIZE, + model_name="test-model", + layer_cache_types=layer_types, + layer_meta_states=layer_meta, + ) + _, before_meta = ssd.load_block_with_metadata(block_hash) + assert before_meta is not None + before_signature = before_meta["cache_signature"] + assert _signature_cachelist_subtypes(before_signature) == expected + + assert cache._strip_rotating_payload(block_hash) + + after_data, after_meta = ssd.load_block_with_metadata(block_hash) + assert after_data is not None and after_meta is not None + assert after_meta["cache_signature"] == before_signature + assert _signature_cachelist_subtypes(after_meta["cache_signature"]) == expected + assert cache._is_placeholder_state(after_data[0]) + assert isinstance(after_data[1], list) + assert len(after_data[1]) == 3 + + def test_hot_cache_byte_counter_consistent(tmp_path): """The stripped entry shrinks and the byte counter stays exact.""" cache, ssd = _make_cache(tmp_path) From 51327ccc64e5df9b7d097b0f58fd8f1ed899cb5a Mon Sep 17 00:00:00 2001 From: jundot Date: Tue, 4 Aug 2026 15:02:06 +0900 Subject: [PATCH 006/338] chore: bump version to 0.5.6 --- omlx/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/omlx/_version.py b/omlx/_version.py index 86716a713a..a779a44262 100644 --- a/omlx/_version.py +++ b/omlx/_version.py @@ -1 +1 @@ -__version__ = "0.5.5" +__version__ = "0.5.6" From d5be00e7f7c99f851a03b917950d07db29102e89 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 4 Aug 2026 06:24:36 +0000 Subject: [PATCH 007/338] formula: bump to 0.5.6 --- Formula/omlx.rb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Formula/omlx.rb b/Formula/omlx.rb index 54343779b2..5f01d83725 100644 --- a/Formula/omlx.rb +++ b/Formula/omlx.rb @@ -3,8 +3,8 @@ class Omlx < Formula desc "LLM inference server optimized for Apple Silicon" homepage "https://github.com/jundot/omlx" - url "https://github.com/jundot/omlx/archive/refs/tags/v0.5.5.tar.gz" - sha256 "d77b58c007b3f1d3b5463ac66ddcd9923db5839f82213d9b221e0f68b867ff3c" + url "https://github.com/jundot/omlx/archive/refs/tags/v0.5.6.tar.gz" + sha256 "abcdff98302c8e1063ecceae272fc7b6278e9abda5ae3aff330813b4f0fba7dd" license "Apache-2.0" head "https://github.com/jundot/omlx.git", branch: "main" From 4450eca406e58cfb334ec0dca500c5aad6c7882e Mon Sep 17 00:00:00 2001 From: jundot Date: Tue, 4 Aug 2026 17:25:55 +0900 Subject: [PATCH 008/338] fix: align DeepSeek V4 template with official encoder --- omlx/api/utils.py | 25 +- omlx/patches/deepseek_v4/chat_template_v4.py | 891 ++++++++++++++----- omlx/patches/deepseek_v4/tokenizer_patch.py | 7 + tests/test_api_utils.py | 61 ++ tests/test_deepseek_v4_patch.py | 147 ++- 5 files changed, 876 insertions(+), 255 deletions(-) diff --git a/omlx/api/utils.py b/omlx/api/utils.py index fa6f11a1a8..4416d7d4ee 100644 --- a/omlx/api/utils.py +++ b/omlx/api/utils.py @@ -374,6 +374,13 @@ def chat_template_preserves_mid_system( if placement not in {"tail", "between"}: return False + explicit_capability = getattr(tokenizer, "_omlx_supports_mid_system_messages", None) + if explicit_capability is None: + template = getattr(tokenizer, "_chat_template", None) + explicit_capability = getattr(template, "supports_mid_system_messages", None) + if explicit_capability is False: + return False + has_tools = bool(tools) cache_key = _mid_system_probe_cache_key( tokenizer, @@ -656,8 +663,9 @@ def prepare_system_messages_for_template( ) -> list[dict]: """Preserve cache-friendly mid-system turns when the template supports them. - Unsupported placements or templates fall back to the historical behavior: - all system messages are consolidated at the front. + Model-specific templates may first relocate supported system turns to a + native reminder role. Unsupported placements or templates then use the + configured strict or user-note fallback. """ messages = [dict(msg) for msg in messages] if unsupported_mid_system_policy not in {"strict", "user_note_safe"}: @@ -682,6 +690,19 @@ def unsupported_fallback() -> list[dict]: return prepared return strict_fallback() + if not is_partial and has_nonleading_system_message(messages): + relocator = getattr(tokenizer, "_omlx_relocate_mid_system_messages", None) + if relocator is None: + template = getattr(tokenizer, "_chat_template", None) + relocator = getattr(template, "relocate_mid_system_messages", None) + if callable(relocator): + try: + relocated = relocator(messages) + except Exception: + relocated = None + if relocated is not None: + messages = [dict(msg) for msg in relocated] + placements = _mid_system_placement_kinds(messages) if not placements: if placements is None: diff --git a/omlx/patches/deepseek_v4/chat_template_v4.py b/omlx/patches/deepseek_v4/chat_template_v4.py index 7b7e1d94f7..766b77f74c 100644 --- a/omlx/patches/deepseek_v4/chat_template_v4.py +++ b/omlx/patches/deepseek_v4/chat_template_v4.py @@ -1,72 +1,56 @@ -# Copyright © 2025 Apple Inc. -# SPDX-License-Identifier: Apache-2.0 -"""DeepSeek V4 DSML chat template (derived from mlx-lm deepseek_v32). - -This file is a near-verbatim copy of -``mlx_lm/chat_templates/deepseek_v32.py`` (Apple Inc., Apache 2.0). The -only edit is the outer DSML marker name: V3.2 wraps tool calls in -``<｜DSML｜function_calls>...`` while V4 uses -``<｜DSML｜tool_calls>...`` (per vllm's -``DeepSeekV4ToolParser`` which subclasses ``DeepSeekV32ToolParser`` -overriding only those two tokens). The inner ``<｜DSML｜invoke>`` / -``<｜DSML｜parameter>`` grammar is identical between V3.2 and V4. - -omlx registers this module as ``mlx_lm.chat_templates.deepseek_v4`` so -mlx-lm's tokenizer_config ``chat_template_type`` lookup picks it up -transparently. +# Copyright (c) 2023 DeepSeek +# SPDX-License-Identifier: MIT +# ruff: noqa: N806, SIM114, UP006, UP007, UP035, UP045 +"""DeepSeek V4 chat encoding based on the official 0731 reference. + +The encoding core is vendored from ``deepseek-ai/DeepSeek-V4-Flash-0731``. +The small adapter at the end exposes mlx-lm's ``apply_chat_template`` API. """ import copy import json -import re -from inspect import isfunction from typing import Any, Dict, List, Optional, Tuple, Union -from transformers.utils.chat_template_utils import get_json_schema - -TOOLS_SYSTEM_TEMPLATE = """## Tools - -You have access to a set of tools you can use to answer the user's question. -You can invoke functions by writing a "<{dsml_token}tool_calls>" block like the following as part of your reply to the user: -<{dsml_token}tool_calls> -<{dsml_token}invoke name="$FUNCTION_NAME"> -<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE -... - -<{dsml_token}invoke name="$FUNCTION_NAME2"> -... - - - -String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects). - -If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example: - -<{dsml_token}tool_calls> -... - - - -... - - -{thinking_start_token}...thinking about results{thinking_end_token} - -Here are the functions available in JSONSchema format: - -{tool_schemas} - -""" +# ============================================================ +# Special Tokens +# ============================================================ bos_token: str = "<｜begin▁of▁sentence｜>" eos_token: str = "<｜end▁of▁sentence｜>" thinking_start_token: str = "" thinking_end_token: str = "" dsml_token: str = "｜DSML｜" + +USER_SP_TOKEN = "<｜User｜>" +ASSISTANT_SP_TOKEN = "<｜Assistant｜>" +LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>" + +# The open-weight 0731 encoding has no role delimiter for a system message +# after the conversation starts. oMLX exposes this capability explicitly so +# callers do not mistake raw marker ordering for model-level support. +supports_mid_system_messages = False + +# Task special tokens for internal classification tasks +DS_TASK_SP_TOKENS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} +VALID_TASKS = set(DS_TASK_SP_TOKENS.keys()) + +# ============================================================ +# Templates +# ============================================================ + system_msg_template: str = "{content}" -user_msg_template: str = "<｜User｜>{content}<｜Assistant｜>" -assistant_msg_template: str = "{reasoning}{content}{tool_calls}<｜end▁of▁sentence｜>" -thinking_template = "{reasoning_content}" +user_msg_template: str = "{content}" +latest_reminder_msg_template: str = "{content}" +assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token +assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}" +thinking_template: str = "{reasoning_content}" response_format_template: str = ( "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" @@ -75,29 +59,77 @@ '<{dsml_token}invoke name="{name}">\n{arguments}\n' ) tool_calls_template = ( - "<{dsml_token}tool_calls>\n{tool_calls}\n" + "<{dsml_token}{tc_block_name}>\n{tool_calls}\n" ) +tool_calls_block_name: str = "tool_calls" + +tool_output_template: str = "{content}" + +# Reasoning effort levels. In thinking mode, the prompt for the selected level is +# prepended at the very beginning of the conversation. `low` is the default and +# adds nothing. +REASONING_EFFORT_PROMPTS: Dict[str, str] = { + "low": "", + "high": ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n" + ), + "max": ( + "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n" + "You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n" + "Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n" + ), +} +DEFAULT_REASONING_EFFORT = "low" + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + -tool_output_template: str = "\n{content}" +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + +# ============================================================ +# Utility Functions +# ============================================================ def to_json(value: Any) -> str: + """Serialize a value to JSON string.""" try: return json.dumps(value, ensure_ascii=False) - except: + except (TypeError, ValueError): return json.dumps(value, ensure_ascii=True) def tools_from_openai_format(tools): - def normalize_tool(tool): - if isfunction(tool): - return get_json_schema(tool) - return tool["function"] - - return [normalize_tool(tool) for tool in tools] + """Extract function definitions from OpenAI-format tool list.""" + return [tool["function"] for tool in tools] def tool_calls_from_openai_format(tool_calls): + """Convert OpenAI-format tool calls to internal format.""" return [ { "name": tool_call["function"]["name"], @@ -107,25 +139,41 @@ def tool_calls_from_openai_format(tool_calls): ] +def tool_calls_to_openai_format(tool_calls): + """Convert internal tool calls to OpenAI format.""" + return [ + { + "type": "function", + "function": { + "name": tool_call["name"], + "arguments": tool_call["arguments"], + }, + } + for tool_call in tool_calls + ] + + def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str: - p_dsml_template = """<{dsml_token}parameter name="{key}" string="{is_str}">{value}""" + """ + Encode tool call arguments into DSML parameter format. + + Args: + tool_call: Dict with "name" and "arguments" (JSON string) keys. + + Returns: + DSML-formatted parameter string. + """ + p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}' P_dsml_strs = [] - # OpenAI tool_calls store arguments as a JSON string; the omlx - # Anthropic adapter (api/anthropic_utils.py:198) decodes ``input`` - # into a dict before storing it on assistant messages. Accept both - # so multi-turn conversations whose history was authored from - # either side render without raising. - raw_args = tool_call["arguments"] - if isinstance(raw_args, str): - arguments = json.loads(raw_args) - elif isinstance(raw_args, dict): - arguments = raw_args + raw_arguments = tool_call["arguments"] + if isinstance(raw_arguments, dict): + arguments = raw_arguments else: - raise TypeError( - f"tool_call['arguments'] must be str or dict, got " - f"{type(raw_args).__name__}" - ) + try: + arguments = json.loads(raw_arguments) + except (TypeError, ValueError, json.JSONDecodeError): + arguments = {"arguments": raw_arguments} for k, v in arguments.items(): p_dsml_str = p_dsml_template.format( @@ -134,7 +182,6 @@ def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str: is_str="true" if isinstance(v, str) else "false", value=v if isinstance(v, str) else to_json(v), ) - P_dsml_strs.append(p_dsml_str) return "\n".join(P_dsml_strs) @@ -143,6 +190,17 @@ def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str: def decode_dsml_to_arguments( tool_name: str, tool_args: Dict[str, Tuple[str, str]] ) -> Dict[str, str]: + """ + Decode DSML parameters back to a tool call dict. + + Args: + tool_name: Name of the tool. + tool_args: Dict mapping param_name -> (value, is_string_flag). + + Returns: + Dict with "name" and "arguments" (JSON string) keys. + """ + def _decode_value(key: str, value: str, string: str): if string == "true": value = to_json(value) @@ -159,9 +217,18 @@ def _decode_value(key: str, value: str, string: str): def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str: + """ + Render tool schemas into the system prompt format. + + Args: + tools: List of tool schema dicts (each with name, description, parameters). + + Returns: + Formatted tools section string. + """ tools_json = [to_json(t) for t in tools] - return TOOLS_SYSTEM_TEMPLATE.format( + return TOOLS_TEMPLATE.format( tool_schemas="\n".join(tools_json), dsml_token=dsml_token, thinking_start_token=thinking_start_token, @@ -170,6 +237,7 @@ def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str: def find_last_user_index(messages: List[Dict[str, Any]]) -> int: + """Find the index of the last user/developer message.""" last_user_index = -1 for idx in range(len(messages) - 1, -1, -1): if messages[idx].get("role") in ["user", "developer"]: @@ -178,39 +246,35 @@ def find_last_user_index(messages: List[Dict[str, Any]]) -> int: return last_user_index -def _user_thinking_suffix( - index: int, - last_user_idx: int, - thinking_mode: str, - drop_thinking: bool, -) -> str: - """Marker that follows a user/developer message. - - The current turn opens thinking. Historical turns normally close it, - because their reasoning is dropped from the transcript — but that makes - rendering position-dependent: when a new user message arrives, the - previous user message's suffix flips from ```` to ```` - and every cached prefix from that point on is invalidated. - - When reasoning is retained (``drop_thinking=False``), the historical - turn still carries its reasoning block, so keeping ```` here - reproduces exactly what was rendered while that turn was current and - the transcript grows append-only — which is what a prefix cache needs. - """ - if thinking_mode != "thinking": - return thinking_end_token - if index == last_user_idx or not drop_thinking: - return thinking_start_token - return thinking_end_token +# ============================================================ +# Message Rendering +# ============================================================ def render_message( index: int, messages: List[Dict[str, Any]], thinking_mode: str, - tools: Any = None, drop_thinking: bool = True, + reasoning_effort: Optional[str] = None, ) -> str: + """ + Render a single message at the given index into its encoded string form. + + This is the core function that converts each message in the conversation + into the DeepSeek-V4 format. + + Args: + index: Index of the message to render. + messages: Full list of messages in the conversation. + thinking_mode: Either "chat" or "thinking". + drop_thinking: Whether to drop reasoning content from earlier turns. + reasoning_effort: Reasoning effort level, one of "low", "high", "max". + None is treated as "low". + + Returns: + Encoded string for this message. + """ assert 0 <= index < len(messages) assert thinking_mode in [ "chat", @@ -223,19 +287,29 @@ def render_message( role = msg.get("role") content = msg.get("content") - tools = tools or msg.get("tools") + tools = msg.get("tools") response_format = msg.get("response_format") tool_calls = msg.get("tool_calls") reasoning_content = msg.get("reasoning_content") + wo_eos = msg.get("wo_eos", False) + if tools: + tools = tools_from_openai_format(tools) if tool_calls: tool_calls = tool_calls_from_openai_format(tool_calls) + # Reasoning effort prefix (only at index 0 in thinking mode; "low" adds nothing) + reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT + assert ( + reasoning_effort in REASONING_EFFORT_PROMPTS + ), f"Invalid reasoning effort: {reasoning_effort}, expected one of {list(REASONING_EFFORT_PROMPTS)}" + if index == 0 and thinking_mode == "thinking": + prompt += REASONING_EFFORT_PROMPTS[reasoning_effort] + if role == "system": prompt += system_msg_template.format(content=content or "") if tools: - prompt += "\n\n" + render_tools(tools_from_openai_format(tools)) - + prompt += "\n\n" + render_tools(tools) if response_format: prompt += "\n\n" + response_format_template.format( schema=to_json(response_format) @@ -243,187 +317,503 @@ def render_message( elif role == "developer": assert content, f"Invalid message for role `{role}`: {msg}" - content_developer = "" - if tools: - content_developer += "\n\n" + render_tools(tools_from_openai_format(tools)) + content_developer = USER_SP_TOKEN + content_developer += content + + if tools: + content_developer += "\n\n" + render_tools(tools) if response_format: content_developer += "\n\n" + response_format_template.format( schema=to_json(response_format) ) - content_developer += "\n\n# The user's message is: {}".format(content) - prompt += user_msg_template.format(content=content_developer) - prompt += _user_thinking_suffix( - index, last_user_idx, thinking_mode, drop_thinking - ) elif role == "user": - prompt += user_msg_template.format(content=content) - - prompt += _user_thinking_suffix( - index, last_user_idx, thinking_mode, drop_thinking + prompt += USER_SP_TOKEN + + # Handle content blocks (tool results mixed with text) + content_blocks = msg.get("content_blocks") + if content_blocks: + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(block.get("text", "")) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [] + for b in tool_content: + if b.get("type") == "text": + text_parts.append(b.get("text", "")) + else: + text_parts.append(f"[Unsupported {b.get('type')}]") + tool_content = "\n\n".join(text_parts) + parts.append(tool_output_template.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + prompt += "\n\n".join(parts) + else: + prompt += content or "" + + elif role == "latest_reminder": + prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format( + content=content ) elif role == "tool": - prev_assistant_idx = index - 1 - assistant_msg = messages[prev_assistant_idx] - while prev_assistant_idx >= 0 and assistant_msg.get("role") == "tool": - prev_assistant_idx -= 1 - assistant_msg = messages[prev_assistant_idx] - - assert ( - index == 0 - or prev_assistant_idx >= 0 - and assistant_msg.get("role") == "assistant" - ), f"Invalid messages at {index}:\n{assistant_msg}" - - tool_call_order = index - prev_assistant_idx - assistant_tool_calls = assistant_msg.get("tool_calls") - assert ( - assistant_tool_calls and len(assistant_tool_calls) >= tool_call_order - ), "No tool calls but found tool output" - - if tool_call_order == 1: - prompt += "\n\n" - - prompt += tool_output_template.format(content=content) - - if tool_call_order == len(assistant_tool_calls): - prompt += "\n" - - # Same append-only rule as the user/developer suffix: when - # reasoning is retained, a historical tool result keeps the - # opener it was rendered with at generation time (the - # following assistant message renders `reasoning...` - # with no opener of its own, so the pair stays consistent). - if thinking_mode == "thinking" and ( - index >= last_user_idx or not drop_thinking - ): - prompt += "\n\n" + thinking_start_token - else: - prompt += "\n\n" + thinking_end_token + raise NotImplementedError( + "deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()" + ) elif role == "assistant": - prev_assistant_idx = index thinking_part = "" + tc_content = "" - tool_calls_content = "" if tool_calls: - tool_calls = [ + tc_list = [ tool_call_template.format( dsml_token=dsml_token, - name=tool_call.get("name"), - arguments=encode_arguments_to_dsml(tool_call), + name=tc.get("name"), + arguments=encode_arguments_to_dsml(tc), ) - for tool_call in tool_calls + for tc in tool_calls ] - tool_calls_content += "\n\n" + tool_calls_template.format( - dsml_token=dsml_token, tool_calls="\n".join(tool_calls) + tc_content += "\n\n" + tool_calls_template.format( + dsml_token=dsml_token, + tool_calls="\n".join(tc_list), + tc_block_name=tool_calls_block_name, ) summary_content = content or "" + rc = reasoning_content or "" - # Historical assistant turns render their reasoning too when it is - # retained, so the transcript stays byte-identical to what was - # rendered while that turn was current (append-only prefix). - render_thinking = thinking_mode == "thinking" and ( - index > last_user_idx or not drop_thinking - ) - if render_thinking: - if index > last_user_idx: - assert ( - reasoning_content or tool_calls - ), f"ThinkingMode: {thinking_mode}, invalid message without reasoning_content/tool_calls `{msg}` after last user message" - thinking_part = ( - thinking_template.format(reasoning_content=reasoning_content or "") - + thinking_end_token - ) + # Check if previous message has a task - if so, this is a task output (no thinking) + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None - prompt += assistant_msg_template.format( - reasoning=thinking_part, - content=summary_content, - tool_calls=tool_calls_content, - ) + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_idx: + thinking_part = ( + thinking_template.format(reasoning_content=rc) + thinking_end_token + ) + else: + thinking_part = "" + + if wo_eos: + prompt += assistant_msg_wo_eos_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + prompt += assistant_msg_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) else: raise NotImplementedError(f"Unknown role: {role}") + # Append transition tokens based on what follows + if index + 1 < len(messages) and messages[index + 1].get("role") not in [ + "assistant", + "latest_reminder", + ]: + return prompt + + task = messages[index].get("task") + if task is not None: + # Task special token for internal classification tasks + assert ( + task in VALID_TASKS + ), f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}" + task_sp_token = DS_TASK_SP_TOKENS[task] + + if task != "action": + # Non-action tasks: append task sp token directly after the message + prompt += task_sp_token + else: + # Action task: append Assistant + thinking token + action sp token + prompt += ASSISTANT_SP_TOKEN + prompt += ( + thinking_end_token + if thinking_mode != "thinking" + else thinking_start_token + ) + prompt += task_sp_token + + elif messages[index].get("role") in ["user", "developer"]: + # Normal generation: append Assistant + thinking token + prompt += ASSISTANT_SP_TOKEN + if not drop_thinking and thinking_mode == "thinking": + prompt += thinking_start_token + elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx: + prompt += thinking_start_token + else: + prompt += thinking_end_token + return prompt -def drop_thinking_messages( - messages: List[Dict[str, Any]], last_user_idx: Optional[int] = None +# ============================================================ +# Preprocessing +# ============================================================ + + +def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Merge tool messages into the preceding user message using content_blocks format. + + DeepSeek-V4 does not have a standalone "tool" role; instead, tool results + are encoded as blocks within user messages. + + This function converts a standard OpenAI-format conversation (with separate + "tool" role messages) into V4 format where tool results are merged into + user messages. + + Args: + messages: List of message dicts in OpenAI format. + + Returns: + Processed message list with tool messages merged into user messages. + """ + merged: List[Dict[str, Any]] = [] + + for msg in messages: + msg = copy.deepcopy(msg) + role = msg.get("role") + + if role == "tool": + # Convert tool message to a user message with tool_result block + tool_block = { + "type": "tool_result", + "tool_use_id": msg.get("tool_call_id", ""), + "content": msg.get("content", ""), + } + # Merge into previous message if it's already a user (merged tool) + if ( + merged + and merged[-1].get("role") == "user" + and "content_blocks" in merged[-1] + ): + merged[-1]["content_blocks"].append(tool_block) + else: + merged.append( + { + "role": "user", + "content_blocks": [tool_block], + } + ) + elif role == "user": + text_block = {"type": "text", "text": msg.get("content", "")} + if ( + merged + and merged[-1].get("role") == "user" + and "content_blocks" in merged[-1] + and merged[-1].get("task") is None + ): + merged[-1]["content_blocks"].append(text_block) + else: + new_msg = { + "role": "user", + "content": msg.get("content", ""), + "content_blocks": [text_block], + } + # Preserve extra fields (task, wo_eos, mask, etc.) + for key in ("task", "wo_eos", "mask"): + if key in msg: + new_msg[key] = msg[key] + merged.append(new_msg) + else: + merged.append(msg) + + return merged + + +def sort_tool_results_by_call_order( + messages: List[Dict[str, Any]], ) -> List[Dict[str, Any]]: - messages_wo_thinking: List[Dict[str, Any]] = [] - last_user_idx = ( - find_last_user_index(messages) if last_user_idx is None else last_user_idx - ) - for idx, msg in enumerate(messages): + """ + Sort tool_result blocks within user messages by the order of tool_calls + in the preceding assistant message. + + Args: + messages: Preprocessed message list (after merge_tool_messages). + + Returns: + Message list with sorted tool result blocks. + """ + last_tool_call_order: Dict[str, int] = {} + + for msg in messages: role = msg.get("role") - if role in ["user", "system", "tool"] or idx >= last_user_idx: - messages_wo_thinking.append(msg) - continue + if role == "assistant" and msg.get("tool_calls"): + last_tool_call_order = {} + for idx, tc in enumerate(msg["tool_calls"]): + tc_id = tc.get("id") or tc.get("function", {}).get("id", "") + if tc_id: + last_tool_call_order[tc_id] = idx + + elif role == "user" and msg.get("content_blocks"): + tool_blocks = [ + b for b in msg["content_blocks"] if b.get("type") == "tool_result" + ] + if len(tool_blocks) > 1 and last_tool_call_order: + sorted_blocks = sorted( + tool_blocks, + key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0), + ) + sorted_idx = 0 + new_blocks = [] + for block in msg["content_blocks"]: + if block.get("type") == "tool_result": + new_blocks.append(sorted_blocks[sorted_idx]) + sorted_idx += 1 + else: + new_blocks.append(block) + msg["content_blocks"] = new_blocks + + return messages - elif role == "assistant": - msg_wo_thinking = copy.copy(msg) - msg_wo_thinking.pop("reasoning_content", None) - messages_wo_thinking.append(msg_wo_thinking) - return messages_wo_thinking +# ============================================================ +# Main Encoding Function +# ============================================================ def encode_messages( messages: List[Dict[str, Any]], - thinking_mode: str = "thinking", + thinking_mode: str, context: Optional[List[Dict[str, Any]]] = None, drop_thinking: bool = True, add_default_bos_token: bool = True, - tools: Any = None, + reasoning_effort: Optional[str] = None, ) -> str: + """ + Encode a list of messages into the DeepSeek-V4 prompt format. + + This is the main entry point for encoding conversations. It handles: + - BOS token insertion + - Thinking mode with optional reasoning content dropping + - Tool message merging into user messages + - Multi-turn conversation context + + Args: + messages: List of message dicts to encode. + thinking_mode: Either "chat" or "thinking". + context: Optional preceding context messages (already encoded prefix). + drop_thinking: If True, drop reasoning_content from earlier assistant turns + (only keep reasoning for messages after the last user message). + add_default_bos_token: Whether to prepend BOS token at conversation start. + reasoning_effort: Reasoning effort level, one of "low", "high", "max". + Only takes effect in thinking mode. None is treated as "low". + + Returns: + The encoded prompt string. + """ context = context if context else [] - # render_message only injects the DSML tools block on system / developer - # roles (chat_template_v4.py:194-207). When the first message is a - # plain user (e.g. OpenAI request without a system message, or an - # Anthropic request whose system field was empty) the tools schema - # never reaches the model and it cannot emit a tool_calls block. - # Prepend an empty synthetic system message so render_tools fires - # without otherwise altering the conversation. - if ( - tools - and messages - and messages[0].get("role") not in ("system", "developer") - and not (context and context[0].get("role") in ("system", "developer")) - ): - messages = [{"role": "system", "content": ""}, *messages] + # Preprocess: merge tool messages and sort tool results + messages = merge_tool_messages(messages) + messages = sort_tool_results_by_call_order(context + messages)[len(context) :] + if context: + context = merge_tool_messages(context) + context = sort_tool_results_by_call_order(context) full_messages = context + messages + prompt = bos_token if add_default_bos_token and len(context) == 0 else "" - if thinking_mode == "thinking" and drop_thinking: - full_messages = drop_thinking_messages(full_messages) + # Resolve drop_thinking: if any message has tools defined, don't drop thinking + effective_drop_thinking = drop_thinking + if any(m.get("tools") for m in full_messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + full_messages = _drop_thinking_messages(full_messages) + # After dropping, recalculate how many messages to render + # (context may have shrunk too) + num_to_render = len(full_messages) - len(_drop_thinking_messages(context)) + context_len = len(full_messages) - num_to_render + else: + num_to_render = len(messages) + context_len = len(context) - for idx in range(len(messages)): + for idx in range(num_to_render): prompt += render_message( - idx + len(context), + idx + context_len, full_messages, thinking_mode=thinking_mode, - tools=tools, - drop_thinking=drop_thinking, + drop_thinking=effective_drop_thinking, + reasoning_effort=reasoning_effort, ) return prompt +def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Drop reasoning_content and non-essential messages before the last user message. + + Behavior: + - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept. + - Messages at or after the last user index are always kept. + - Assistant messages before the last user get reasoning_content removed. + - Developer messages before the last user are dropped entirely. + """ + last_user_idx = find_last_user_index(messages) + result = [] + keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"} + + for idx, msg in enumerate(messages): + role = msg.get("role") + if role in keep_roles or idx >= last_user_idx: + result.append(msg) + elif role == "assistant": + msg = copy.copy(msg) + msg.pop("reasoning_content", None) + result.append(msg) + # developer and other roles before last_user_idx are dropped + + return result + + +# ============================================================ +# mlx-lm Adapter +# ============================================================ + + +def _attach_request_tools( + messages: List[Dict[str, Any]], + tools: Any, + context: Optional[List[Dict[str, Any]]] = None, +) -> tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]: + """Attach request-level tools where the reference encoder expects them.""" + copied_messages = [copy.deepcopy(message) for message in messages] + copied_context = ( + [copy.deepcopy(message) for message in context] if context else context + ) + if not tools: + return copied_messages, copied_context + + # A context prefix that starts with system/developer is assumed to have + # already encoded the tool definitions. Re-injecting them into the delta + # would duplicate the schema and break append-only prompt caching. + if copied_context and any( + message.get("role") in {"system", "developer"} for message in copied_context + ): + return copied_messages, copied_context + + for message in copied_messages: + if message.get("role") in {"system", "developer"}: + message.setdefault("tools", tools) + return copied_messages, copied_context + + copied_messages.insert( + 0, + {"role": "system", "content": "", "tools": copy.deepcopy(tools)}, + ) + return copied_messages, copied_context + + +def _latest_reminder_content_as_text(content: Any) -> Optional[str]: + """Return text-only system content, or None for unsupported content.""" + if content is None: + return "" + if isinstance(content, str): + return content + if not isinstance(content, list): + return None + + parts = [] + for block in content: + if not isinstance(block, dict) or block.get("type", "text") != "text": + return None + text = block.get("text", "") + if text is not None and not isinstance(text, str): + return None + if text: + parts.append(text) + return "\n".join(parts) + + +def relocate_mid_system_messages( + messages: List[Dict[str, Any]], +) -> Optional[List[Dict[str, Any]]]: + """Move supported inline system runs to V4 ``latest_reminder`` context. + + Claude Code appends a volatile system message immediately after the user + message it qualifies. The V4 reference encoder has no delimiter for that + raw placement. Reclassify the system run as a new ``latest_reminder`` + context immediately before the same user, preserving append-only prompt + caching across requests. + + Only the observed ``user -> system -> (assistant | end)`` shape is + rewritten. Other placements and non-leading developer messages return + None so oMLX can use its configured strict or user-note fallback. + """ + source = [copy.deepcopy(message) for message in messages] + relocated: List[Dict[str, Any]] = [] + seen_non_system = False + index = 0 + + while index < len(source): + message = source[index] + if message.get("role") != "system": + relocated.append(message) + seen_non_system = True + index += 1 + continue + + start = index + parts = [] + while index < len(source) and source[index].get("role") == "system": + text = _latest_reminder_content_as_text(source[index].get("content")) + if text is None: + return None + if text: + parts.append(text) + index += 1 + + if not seen_non_system: + relocated.extend(source[start:index]) + continue + + next_role = source[index].get("role") if index < len(source) else None + if ( + not relocated + or relocated[-1].get("role") != "user" + or next_role not in {None, "assistant"} + ): + return None + + associated_user = relocated.pop() + if parts: + relocated.append( + { + "role": "latest_reminder", + "content": "\n\n".join(parts), + } + ) + relocated.append(associated_user) + + return relocated + + def apply_chat_template( - messages, continue_final_message=False, add_generation_prompt=False, **kwargs + messages, + continue_final_message: bool = False, + add_generation_prompt: bool = False, + **kwargs, ): - # mlx-lm and the omlx server forward an ``enable_thinking`` boolean - # kwarg through ``tokenizer.apply_chat_template``. The V3.2-derived - # ``encode_messages`` signature only knows ``thinking_mode`` ("chat" - # | "thinking"). Translate here so the caller's kwarg shape is - # preserved without leaking the rename downstream. + """Expose the official encoder through mlx-lm's template callable API.""" + if continue_final_message and add_generation_prompt: + raise ValueError( + "Only one of continue_final_message or add_generation_prompt can be True" + ) + if "enable_thinking" in kwargs and "thinking_mode" not in kwargs: kwargs["thinking_mode"] = ( "thinking" if kwargs.pop("enable_thinking") else "chat" @@ -431,25 +821,38 @@ def apply_chat_template( else: kwargs.pop("enable_thinking", None) - # Drop unknown kwargs that some API frontends inject but - # encode_messages does not consume — keeps the wrapper resilient - # against future template_kwargs additions. - _accepted = { + kwargs.setdefault("thinking_mode", "thinking") + tools = kwargs.pop("tools", None) + context = kwargs.get("context") + prepared_messages, prepared_context = _attach_request_tools( + messages, + tools, + context=context, + ) + if context is not None: + kwargs["context"] = prepared_context + + accepted = { "thinking_mode", "context", "drop_thinking", "add_default_bos_token", - "tools", + "reasoning_effort", } - kwargs = {k: v for k, v in kwargs.items() if k in _accepted} + encode_kwargs = {key: value for key, value in kwargs.items() if key in accepted} + out = encode_messages(prepared_messages, **encode_kwargs) - out = encode_messages(messages, **kwargs) - if continue_final_message and add_generation_prompt: - raise ValueError( - "Only one of continue_final_message or add_generation_prompt can be True" - ) - if not add_generation_prompt and messages[-1]["role"] == "user": - out = out.removesuffix("<｜Assistant｜>") - if continue_final_message and messages[-1]["role"] == "assistant": + if not add_generation_prompt: + out = out.removesuffix(ASSISTANT_SP_TOKEN + thinking_start_token) + out = out.removesuffix(ASSISTANT_SP_TOKEN + thinking_end_token) + if continue_final_message and messages and messages[-1].get("role") == "assistant": out = out.removesuffix(eos_token) return out + + +apply_chat_template.supports_mid_system_messages = ( # type: ignore[attr-defined] + supports_mid_system_messages +) +apply_chat_template.relocate_mid_system_messages = ( # type: ignore[attr-defined] + relocate_mid_system_messages +) diff --git a/omlx/patches/deepseek_v4/tokenizer_patch.py b/omlx/patches/deepseek_v4/tokenizer_patch.py index 2851755473..db06c1a5f3 100644 --- a/omlx/patches/deepseek_v4/tokenizer_patch.py +++ b/omlx/patches/deepseek_v4/tokenizer_patch.py @@ -218,6 +218,13 @@ def patched_load(model_path, tokenizer_config_extra=None, eos_token_ids=None): if wrapper._chat_template is None: wrapper._chat_template = _ct.apply_chat_template wrapper.has_chat_template = True + if wrapper._chat_template is _ct.apply_chat_template: + wrapper._omlx_supports_mid_system_messages = ( + _ct.supports_mid_system_messages + ) + wrapper._omlx_relocate_mid_system_messages = ( + _ct.relocate_mid_system_messages + ) if wrapper._tool_parser is None: wrapper._tool_parser = _tp.parse_tool_call wrapper._tool_call_start = _tp.tool_call_start diff --git a/tests/test_api_utils.py b/tests/test_api_utils.py index eab6ef2002..6939f28082 100644 --- a/tests/test_api_utils.py +++ b/tests/test_api_utils.py @@ -2178,6 +2178,17 @@ def apply_chat_template(self, messages, **kwargs): ) return "user:__OMLX_MID_SYSTEM_PROBE_USER__" + class ExplicitlyUnsupportedTokenizer(PreserveTokenizer): + _omlx_supports_mid_system_messages = False + + class RelocatingTokenizer(ExplicitlyUnsupportedTokenizer): + @staticmethod + def _omlx_relocate_mid_system_messages(messages): + return [ + {"role": "latest_reminder", "content": messages[1]["content"]}, + messages[0], + ] + def test_preserves_tail_system_when_template_keeps_position(self): messages = [ {"role": "user", "content": "Hello"}, @@ -2191,6 +2202,56 @@ def test_preserves_tail_system_when_template_keeps_position(self): assert [m["role"] for m in result] == ["user", "system"] assert result[1]["content"] == "Plan mode" + def test_explicit_capability_overrides_marker_order_false_positive(self): + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Plan mode"}, + ] + + result = prepare_system_messages_for_template( + messages, + self.ExplicitlyUnsupportedTokenizer(), + unsupported_mid_system_policy="user_note_safe", + ) + + assert [message["role"] for message in result] == ["user"] + assert result[0]["content"].endswith("[System note]\nPlan mode\n[/System note]") + + def test_model_specific_relocation_precedes_capability_fallback(self): + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Plan mode"}, + ] + + result = prepare_system_messages_for_template( + messages, + self.RelocatingTokenizer(), + unsupported_mid_system_policy="user_note_safe", + ) + + assert result == [ + {"role": "latest_reminder", "content": "Plan mode"}, + {"role": "user", "content": "Hello"}, + ] + + def test_partial_mode_does_not_use_model_specific_relocation(self): + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Plan mode"}, + ] + + result = prepare_system_messages_for_template( + messages, + self.RelocatingTokenizer(), + is_partial=True, + unsupported_mid_system_policy="user_note_safe", + ) + + assert result == [ + {"role": "system", "content": "Plan mode"}, + {"role": "user", "content": "Hello"}, + ] + def test_preserves_between_turn_system_when_template_keeps_position(self): messages = [ {"role": "user", "content": "Hello"}, diff --git a/tests/test_deepseek_v4_patch.py b/tests/test_deepseek_v4_patch.py index bb4dbf8cea..1b2aa57880 100644 --- a/tests/test_deepseek_v4_patch.py +++ b/tests/test_deepseek_v4_patch.py @@ -367,17 +367,14 @@ def test_outer_markers_exposed(self, applied_patch): class TestChatTemplateV4: - """chat_template_v4 — DSML system prompt + tool_calls render.""" + """Official DeepSeek V4 0731 encoding plus the mlx-lm adapter.""" def test_outer_marker_uses_tool_calls_not_function_calls(self, applied_patch): from omlx.patches.deepseek_v4 import chat_template_v4 as ct - # vllm's DeepSeekV4ToolParser overrides only the outer marker - # name (tool_calls vs V3.2's function_calls). Verify our copy - # made that one edit. assert "function_calls" not in ct.tool_calls_template assert "tool_calls" in ct.tool_calls_template - assert "function_calls" not in ct.TOOLS_SYSTEM_TEMPLATE + assert "function_calls" not in ct.TOOLS_TEMPLATE def test_inner_grammar_unchanged_from_v32(self, applied_patch): from omlx.patches.deepseek_v4 import chat_template_v4 as ct @@ -400,7 +397,9 @@ def test_round_trip_encode_then_parse(self, applied_patch): dsml_token=ct.dsml_token, name="f", arguments=encoded_args ) block = ct.tool_calls_template.format( - dsml_token=ct.dsml_token, tool_calls=invoke + dsml_token=ct.dsml_token, + tool_calls=invoke, + tc_block_name=ct.tool_calls_block_name, ) # Strip the outer markers as TokenizerWrapper would. inner = ( @@ -441,7 +440,7 @@ def test_user_only_request_with_tools_injects_dsml(self, applied_patch): tools=tools, add_generation_prompt=True, ) - assert "" in prompt + assert "### Available Tool Schemas" in prompt assert "get_weather" in prompt assert ct.dsml_token in prompt @@ -476,7 +475,7 @@ def test_system_user_request_with_tools_unchanged(self, applied_patch): ) assert "You are a helpful assistant." in prompt # Only one tools block — no double-injection from synthetic prepend. - assert prompt.count("") == 1 + assert prompt.count("### Available Tool Schemas") == 1 def test_user_only_no_tools_no_prepend(self, applied_patch): """No tools → no synthetic system. Plain user-only request renders @@ -487,9 +486,139 @@ def test_user_only_no_tools_no_prepend(self, applied_patch): [{"role": "user", "content": "Hi"}], add_generation_prompt=True, ) - assert "" not in prompt + assert "### Available Tool Schemas" not in prompt assert "## Tools" not in prompt + def test_official_basic_thinking_prompt(self, applied_patch): + from omlx.patches.deepseek_v4 import chat_template_v4 as ct + + prompt = ct.apply_chat_template( + [ + {"role": "system", "content": "Be helpful."}, + {"role": "user", "content": "Hello"}, + ], + add_generation_prompt=True, + ) + + assert prompt == ( + "<｜begin▁of▁sentence｜>Be helpful." "<｜User｜>Hello<｜Assistant｜>" + ) + + def test_official_latest_reminder_before_user(self, applied_patch): + from omlx.patches.deepseek_v4 import chat_template_v4 as ct + + prompt = ct.apply_chat_template( + [ + {"role": "system", "content": "Be helpful."}, + {"role": "latest_reminder", "content": "2026-08-04,Seoul"}, + {"role": "user", "content": "Hello"}, + ], + add_generation_prompt=True, + ) + + assert prompt == ( + "<｜begin▁of▁sentence｜>Be helpful." + "<｜latest_reminder｜>2026-08-04,Seoul" + "<｜User｜>Hello<｜Assistant｜>" + ) + + def test_official_tool_result_is_merged_into_user_turn(self, applied_patch): + from omlx.patches.deepseek_v4 import chat_template_v4 as ct + + prompt = ct.apply_chat_template( + [ + {"role": "user", "content": "Look it up"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lookup", + "arguments": {"query": "oMLX"}, + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "result", + }, + ], + add_generation_prompt=True, + ) + + assert "<｜User｜>result" in prompt + assert prompt.endswith("<｜Assistant｜>") + + def test_declares_generic_mid_system_unsupported(self, applied_patch): + from omlx.patches.deepseek_v4 import chat_template_v4 as ct + + assert ct.supports_mid_system_messages is False + assert ct.apply_chat_template.supports_mid_system_messages is False + + def test_relocates_claude_tail_system_before_its_user(self, applied_patch): + from omlx.patches.deepseek_v4 import chat_template_v4 as ct + + messages = [ + {"role": "system", "content": "Be helpful."}, + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Plan mode"}, + ] + + relocated = ct.relocate_mid_system_messages(messages) + + assert relocated == [ + {"role": "system", "content": "Be helpful."}, + {"role": "latest_reminder", "content": "Plan mode"}, + {"role": "user", "content": "Hello"}, + ] + assert messages[1]["role"] == "user" + prompt = ct.apply_chat_template(relocated, add_generation_prompt=True) + assert prompt == ( + "<｜begin▁of▁sentence｜>Be helpful." + "<｜latest_reminder｜>Plan mode" + "<｜User｜>Hello<｜Assistant｜>" + ) + + def test_relocation_merges_system_run_before_same_user(self, applied_patch): + from omlx.patches.deepseek_v4 import chat_template_v4 as ct + + relocated = ct.relocate_mid_system_messages( + [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Plan mode"}, + {"role": "system", "content": "Hook context"}, + {"role": "assistant", "content": "OK"}, + ] + ) + + assert relocated == [ + { + "role": "latest_reminder", + "content": "Plan mode\n\nHook context", + }, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "OK"}, + ] + + def test_relocation_refuses_ambiguous_system_placement(self, applied_patch): + from omlx.patches.deepseek_v4 import chat_template_v4 as ct + + assert ( + ct.relocate_mid_system_messages( + [ + {"role": "user", "content": "First"}, + {"role": "system", "content": "Ambiguous"}, + {"role": "user", "content": "Second"}, + ] + ) + is None + ) + def test_encode_arguments_accepts_dict(self, applied_patch): """Anthropic /v1/messages history stores tool_call arguments as a dict (anthropic_utils.py decodes the input before saving). From f24198e94794e6cabaf5b47e13b7f964f79ec464 Mon Sep 17 00:00:00 2001 From: jundot Date: Tue, 4 Aug 2026 18:24:20 +0900 Subject: [PATCH 009/338] chore: bump version to 0.5.7 --- omlx/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/omlx/_version.py b/omlx/_version.py index a779a44262..1cc82e6b87 100644 --- a/omlx/_version.py +++ b/omlx/_version.py @@ -1 +1 @@ -__version__ = "0.5.6" +__version__ = "0.5.7" From 50846648273591a621bb96cb1b3956c45d43efed Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 4 Aug 2026 10:01:02 +0000 Subject: [PATCH 010/338] formula: bump to 0.5.7 --- Formula/omlx.rb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Formula/omlx.rb b/Formula/omlx.rb index 5f01d83725..66ea925cba 100644 --- a/Formula/omlx.rb +++ b/Formula/omlx.rb @@ -3,8 +3,8 @@ class Omlx < Formula desc "LLM inference server optimized for Apple Silicon" homepage "https://github.com/jundot/omlx" - url "https://github.com/jundot/omlx/archive/refs/tags/v0.5.6.tar.gz" - sha256 "abcdff98302c8e1063ecceae272fc7b6278e9abda5ae3aff330813b4f0fba7dd" + url "https://github.com/jundot/omlx/archive/refs/tags/v0.5.7.tar.gz" + sha256 "6a69cdffa7ddeb7d2fbee235971ee025407bf0cc019a3a481f278320d31a1189" license "Apache-2.0" head "https://github.com/jundot/omlx.git", branch: "main" From c62446352814b8b4ca8bd4d435b97f32bbddfb2d Mon Sep 17 00:00:00 2001 From: Jun Kim <64250138+jundot@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:22:28 +0900 Subject: [PATCH 011/338] feat: complete Ling 3.0 Flash support (#2526) Add the vendored bailing_hybrid model, recurrent cache handling, FP8 checkpoint conversion, and oQ/oQe integration. Fix variable-length batch state updates and verify SSD prefix-cache correctness. Based on the original Ling 3.0 Flash implementation in #2524 by Mike Wallio. Co-authored-by: Mike Wallio --- omlx/cache/prefix_cache.py | 13 +- omlx/cache/type_handlers.py | 11 + omlx/oq.py | 37 +- omlx/patches/bailing_hybrid/__init__.py | 99 +++ .../bailing_hybrid/bailing_hybrid_model.py | 825 ++++++++++++++++++ omlx/scheduler.py | 23 + omlx/utils/model_loading.py | 34 + tests/test_bailing_hybrid_patch.py | 591 +++++++++++++ tests/test_cache_type_handlers.py | 13 + tests/test_prefix_cache.py | 48 + 10 files changed, 1686 insertions(+), 8 deletions(-) create mode 100644 omlx/patches/bailing_hybrid/__init__.py create mode 100644 omlx/patches/bailing_hybrid/bailing_hybrid_model.py create mode 100644 tests/test_bailing_hybrid_patch.py diff --git a/omlx/cache/prefix_cache.py b/omlx/cache/prefix_cache.py index 423e7972c2..11d40e861f 100644 --- a/omlx/cache/prefix_cache.py +++ b/omlx/cache/prefix_cache.py @@ -2898,9 +2898,16 @@ def _is_non_sliceable_sub_class(class_name: str) -> bool: last_block_meta_states ): meta_state = last_block_meta_states[layer_idx] - cache = marker_handler.deserialize_state( - tuple(elements), meta_state - ) + if marker_handler.is_variable_length_state(): + cache = marker_handler.reconstruct_cache( + {"states": list(elements)}, + meta_state, + token_count=valid_token_count, + ) + else: + cache = marker_handler.deserialize_state( + tuple(elements), meta_state + ) if cache is None: logger.error( f"Layer {layer_idx}: failed to reconstruct {marker_class}" diff --git a/omlx/cache/type_handlers.py b/omlx/cache/type_handlers.py index 9830068e7f..9c36bd3491 100644 --- a/omlx/cache/type_handlers.py +++ b/omlx/cache/type_handlers.py @@ -813,6 +813,17 @@ def concatenate_states( # Use latest state return states[-1] + def deserialize_state( + self, + elements: tuple[Any, ...], + meta_state: Any | None = None, + ) -> Any: + """Reconstruct every element of a variable-length ArraysCache state.""" + return self.reconstruct_cache( + {"states": list(elements)}, + meta_state, + ) + def reconstruct_cache( self, state: dict[str, Any], diff --git a/omlx/oq.py b/omlx/oq.py index 65b0321a38..e0a97d6fc6 100644 --- a/omlx/oq.py +++ b/omlx/oq.py @@ -154,7 +154,10 @@ def _uses_quantized_source_sensitivity(config: dict) -> bool: quant_method = str(quantization_config.get("quant_method", "")).lower() if quant_method == "mxfp8": return True - return quant_method == "fp8" and _is_deepseek_v4_config(config) + return quant_method == "fp8" and ( + _is_deepseek_v4_config(config) + or config.get("model_type") == "bailing_hybrid" + ) def _is_minimax_m3_config(config: dict) -> bool: @@ -3687,6 +3690,14 @@ def __getattr__(self, name): except Exception as patch_err: logger.debug(f"mimo_v2 patch not applied: {patch_err}") + if config.get("model_type") == "bailing_hybrid": + try: + from omlx.patches.bailing_hybrid import apply_bailing_hybrid_patch + + apply_bailing_hybrid_patch() + except Exception as patch_err: + logger.debug(f"bailing_hybrid patch not applied: {patch_err}") + # Apply mlx-lm MTP patch so the patched __init__/sanitize handle # mtp.* tensors correctly. Idempotent — apply() is a no-op once # patched. @@ -6248,6 +6259,9 @@ def _find_model_layers(model): if hasattr(model, "model") and hasattr(model.model, "embed_tokens"): embed_fn = model.model.embed_tokens layers = model.model.layers + elif hasattr(model, "model") and hasattr(model.model, "word_embeddings"): + embed_fn = model.model.word_embeddings + layers = model.model.layers elif hasattr(model, "language_model") and hasattr(model.language_model, "model"): lm = model.language_model.model if hasattr(lm, "embed_tokens"): @@ -6529,14 +6543,20 @@ def _forward_layer(block, inputs, mask, position_ids): def _layer_masks_for_model(model, layers, inputs): """Build the per-layer mask schedule used by the original model.""" if hasattr(model, "make_cache") and any( - hasattr(layer, "is_linear") for layer in layers + hasattr(layer, "is_linear") or hasattr(layer, "is_global") + for layer in layers ): try: from mlx_lm.models.base import create_attention_mask, create_ssm_mask cache = model.make_cache() - fa_idx = getattr(getattr(model, "model", model), "fa_idx", 0) - ssm_idx = getattr(getattr(model, "model", model), "ssm_idx", 0) + model_core = getattr(model, "model", model) + fa_idx = getattr( + model_core, "fa_idx", getattr(model_core, "_attn_idx", 0) + ) + ssm_idx = getattr( + model_core, "ssm_idx", getattr(model_core, "_gla_idx", 0) + ) fa_cache = cache[fa_idx] if fa_idx < len(cache) else None ssm_cache = cache[ssm_idx] if ssm_idx < len(cache) else None try: @@ -6556,8 +6576,15 @@ def _layer_masks_for_model(model, layers, inputs): # SSM layers (GatedDeltaNet) expect (B, S) boolean mask, not # (S, S) causal mask. During calibration there is no padding, # so None is the correct mask for SSM layers. + def _is_ssm_layer(layer): + if hasattr(layer, "is_linear"): + return bool(layer.is_linear) + if hasattr(layer, "is_global"): + return not bool(layer.is_global) + return False + return [ - ssm_mask if getattr(layer, "is_linear", False) else fa_mask + ssm_mask if _is_ssm_layer(layer) else fa_mask for layer in layers ] except (ImportError, AttributeError): diff --git a/omlx/patches/bailing_hybrid/__init__.py b/omlx/patches/bailing_hybrid/__init__.py new file mode 100644 index 0000000000..20af41c127 --- /dev/null +++ b/omlx/patches/bailing_hybrid/__init__.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Ling 3.0 Flash support for the pinned mlx-lm dependency. + +Vendors ``mlx_lm.models.bailing_hybrid`` from scaryrawr/mlx-lm's +``ling-3.0-flash`` branch without changing oMLX's official mlx-lm pin. The +branch is based on that exact pinned revision and adds the mixed MLA/KDA model +used by Ling 3.0 Flash. + +MLX-LM resolves architectures by importing modules under its own namespace. +Register the vendored implementation there only while upstream lacks it. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import logging +import sys +from pathlib import Path + +logger = logging.getLogger(__name__) + +BRANCH_HEAD_SHA = "d719464ff754e65d9dec496ef3fea27bddefd79c" +SOURCE_URL = ( + "https://github.com/scaryrawr/mlx-lm/blob/ling-3.0-flash/" + "mlx_lm/models/bailing_hybrid.py" +) + +_MODULE_NAME = "mlx_lm.models.bailing_hybrid" +_APPLIED = False + + +def _register_module() -> None: + if _MODULE_NAME in sys.modules: + return + + file_path = Path(__file__).parent / "bailing_hybrid_model.py" + spec = importlib.util.spec_from_file_location(_MODULE_NAME, str(file_path)) + if spec is None or spec.loader is None: + raise ImportError(f"Could not create spec for {_MODULE_NAME} from {file_path}") + + module = importlib.util.module_from_spec(spec) + module.__package__ = "mlx_lm.models" + sys.modules[_MODULE_NAME] = module + try: + spec.loader.exec_module(module) + models_pkg = importlib.import_module("mlx_lm.models") + models_pkg.bailing_hybrid = module + except BaseException: + if sys.modules.get(_MODULE_NAME) is module: + sys.modules.pop(_MODULE_NAME) + raise + + logger.info("Registered %s from %s", _MODULE_NAME, file_path.name) + + +def apply_bailing_hybrid_patch() -> bool: + """Register Ling's ``bailing_hybrid`` model when upstream lacks it.""" + global _APPLIED + if _APPLIED: + return False + + try: + module = importlib.import_module(_MODULE_NAME) + except ModuleNotFoundError as error: + if error.name == "mlx_lm": + logger.debug("mlx_lm not importable - bailing_hybrid patch skipped") + return False + if error.name != _MODULE_NAME: + raise + _register_module() + applied = True + else: + models_pkg = importlib.import_module("mlx_lm.models") + models_pkg.bailing_hybrid = module + applied = False + + _APPLIED = True + if applied: + logger.info( + "Ling 3.0 Flash mlx-lm patch applied (branch head %s)", + BRANCH_HEAD_SHA[:8], + ) + return True + + logger.debug("mlx_lm.models.bailing_hybrid already available upstream") + return False + + +def is_applied() -> bool: + return _APPLIED + + +__all__ = [ + "BRANCH_HEAD_SHA", + "SOURCE_URL", + "apply_bailing_hybrid_patch", + "is_applied", +] diff --git a/omlx/patches/bailing_hybrid/bailing_hybrid_model.py b/omlx/patches/bailing_hybrid/bailing_hybrid_model.py new file mode 100644 index 0000000000..6e18a0f73c --- /dev/null +++ b/omlx/patches/bailing_hybrid/bailing_hybrid_model.py @@ -0,0 +1,825 @@ +# SPDX-License-Identifier: MIT +# ruff: noqa +# Copyright © 2026 Apple Inc. + +import math +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union + +import mlx.core as mx +import mlx.nn as nn + +from .activations import swiglu +from .base import ( + BaseModelArgs, + create_attention_mask, + create_ssm_mask, + scaled_dot_product_attention, +) +from .cache import ArraysCache, KVCache +from .gated_delta import gated_delta_kernel, gated_delta_ops +from .mla import MultiLinear +from .rope_utils import initialize_rope +from .switch_layers import SwitchGLU + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + intermediate_size: int + moe_intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + num_experts: int + num_experts_per_tok: int + num_shared_experts: int + n_group: int + topk_group: int + first_k_dense_replace: int + layer_group_size: int + group_norm_size: int + vocab_size: int + rms_norm_eps: float + rope_theta: float + max_position_embeddings: int + routed_scaling_factor: float + head_dim: int + kv_lora_rank: int + qk_rope_head_dim: int + qk_nope_head_dim: int + v_head_dim: int + q_lora_rank: Optional[int] = None + rope_interleave: bool = True + partial_rotary_factor: float = 0.5 + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + use_qkv_bias: bool = False + use_bias: bool = False + use_qk_norm: bool = True + score_function: str = "sigmoid" + norm_topk_prob: bool = True + moe_router_enable_expert_bias: bool = True + tie_word_embeddings: bool = False + num_nextn_predict_layers: int = 0 + gated_attention_proj_granularity_type: Optional[str] = None + no_kda_lora: bool = True + kda_safe_gate: bool = False + kda_lower_bound: Optional[float] = None + short_conv_kernel_size: int = 4 + quantization_config: Optional[Dict[str, Any]] = None + + +def recurrent_gla( + q: mx.array, + k: mx.array, + v: mx.array, + g: mx.array, + scale: float, + h: Optional[mx.array] = None, +) -> Tuple[mx.array, mx.array]: + L = q.shape[2] + in_dtype = q.dtype + # Keep the recurrent state in float32; precision loss here compounds over + # long prompts through repeated multiplication by exp(g) < 1. + exp_g = mx.exp(g)[:, None, None].astype(mx.float32) + q = (q * scale).astype(mx.float32) + k = k.astype(mx.float32) + v = v.astype(mx.float32) + if h is not None: + h = h.astype(mx.float32) + outputs = [] + for t in range(L): + q_t = q[:, :, t : t + 1] + k_t = k[:, :, t : t + 1] + v_t = v[:, :, t : t + 1] + h_up = k_t.transpose(0, 1, 3, 2) @ v_t + h = h_up if h is None else h * exp_g + h_up + outputs.append(q_t @ h) + return mx.concatenate(outputs, axis=2).astype(in_dtype), h + + +class GroupRMSNorm(nn.Module): + def __init__(self, dims: int, eps: float = 1e-5, groups: int = 1): + super().__init__() + self.weight = mx.ones((dims,)) + self.groups = groups + self.eps = eps + + def __call__(self, x: mx.array) -> mx.array: + x = mx.unflatten(x, axis=-1, shape=(self.groups, -1)) + x = mx.fast.rms_norm(x, weight=None, eps=self.eps) + return self.weight * mx.flatten(x, -2) + + +class MLP(nn.Module): + def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None): + super().__init__() + dim = intermediate_size if intermediate_size is not None else args.intermediate_size + self.gate_proj = nn.Linear(args.hidden_size, dim, bias=args.use_bias) + self.up_proj = nn.Linear(args.hidden_size, dim, bias=args.use_bias) + self.down_proj = nn.Linear(dim, args.hidden_size, bias=args.use_bias) + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class MultiLatentAttention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.num_heads = args.num_attention_heads + self.q_lora_rank = args.q_lora_rank + self.qk_rope_head_dim = args.qk_rope_head_dim + self.kv_lora_rank = args.kv_lora_rank + self.v_head_dim = args.v_head_dim + self.qk_nope_head_dim = args.qk_nope_head_dim + self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim + self.gated_attention_proj_granularity_type = ( + args.gated_attention_proj_granularity_type + ) + + self.scale = self.qk_head_dim**-0.5 + + if self.q_lora_rank is None: + self.q_proj = nn.Linear( + args.hidden_size, self.num_heads * self.qk_head_dim, bias=False + ) + else: + self.q_a_proj = nn.Linear( + args.hidden_size, self.q_lora_rank, bias=args.use_qkv_bias + ) + self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=args.rms_norm_eps) + self.q_b_proj = nn.Linear( + self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False + ) + + self.kv_a_proj_with_mqa = nn.Linear( + args.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=args.use_qkv_bias, + ) + self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=args.rms_norm_eps) + + self.embed_q = MultiLinear( + self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads + ) + self.unembed_out = MultiLinear( + self.kv_lora_rank, self.v_head_dim, self.num_heads + ) + + self.dense = nn.Linear( + self.num_heads * self.v_head_dim, + args.hidden_size, + bias=args.use_qkv_bias, + ) + if args.gated_attention_proj_granularity_type == "head_wise": + self.g_proj = nn.Linear(args.hidden_size, self.num_heads, bias=False) + elif args.gated_attention_proj_granularity_type == "element_wise": + self.g_proj = nn.Linear( + args.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + ) + else: + self.g_proj = None + + if args.rope_scaling is not None: + mscale_all_dim = args.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = args.rope_scaling.get("factor", 1) + if mscale_all_dim and scaling_factor > 1: + s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0 + self.scale = self.scale * s * s + + self.rope = initialize_rope( + dims=self.qk_rope_head_dim, + base=args.rope_theta, + traditional=args.rope_interleave, + max_position_embeddings=args.max_position_embeddings, + scaling_config=args.rope_scaling, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + + if self.q_lora_rank is None: + q = self.q_proj(x) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x))) + + q = q.reshape(B, L, self.num_heads, self.qk_head_dim).transpose(0, 2, 1, 3) + q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1) + + compressed_kv = self.kv_a_proj_with_mqa(x) + compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1) + k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3) + kv_latent = self.kv_a_layernorm(compressed_kv) + + offset = cache.offset if cache is not None else 0 + q_pe = self.rope(q_pe, offset) + k_pe = self.rope(k_pe, offset) + + kv_latent = mx.expand_dims(kv_latent, axis=1) + + if cache is not None: + kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe) + + pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2) + if mask is not None: + pe_scores = mx.where( + mask, + pe_scores, + mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype), + ) + + if L == 1: + q_nope = self.embed_q(q_nope) + k = v = kv_latent + else: + k = self.embed_q(kv_latent, transpose=False) + v = self.unembed_out(kv_latent) + + output = scaled_dot_product_attention( + q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores + ) + if L == 1: + output = self.unembed_out(output) + + if self.g_proj is not None: + gate = mx.sigmoid(self.g_proj(x).astype(mx.float32)).astype(output.dtype) + if self.gated_attention_proj_granularity_type == "head_wise": + output = output * gate.transpose(0, 2, 1)[..., None] + else: + output = output * gate.reshape( + B, L, self.num_heads, self.v_head_dim + ).transpose(0, 2, 1, 3) + + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.dense(output) + + +class DepthwiseConv1d(nn.Module): + def __init__(self, channels: int, kernel_size: int): + super().__init__() + scale = math.sqrt(1.0 / channels) + self.weight = mx.random.uniform( + low=-scale, + high=scale, + shape=(channels, 1, kernel_size), + ) + + def __call__( + self, + x: mx.array, + cache: Optional[mx.array] = None, + mask: Optional[mx.array] = None, + lengths: Optional[mx.array] = None, + ) -> Tuple[mx.array, mx.array]: + batch, length, channels = x.shape + kernel_size = self.weight.shape[-1] + if cache is None: + cache = mx.zeros((batch, channels, kernel_size), dtype=x.dtype) + + if mask is not None: + x = mx.where(mask[..., None], x, 0) + + history = cache[:, :, -kernel_size + 1 :].transpose(0, 2, 1) + conv_input = mx.concatenate([history, x], axis=1) + weight = self.weight.moveaxis(2, 1).astype(x.dtype) + output = nn.silu(mx.conv_general(conv_input, weight, groups=channels)) + + cache_input = mx.concatenate([cache, x.transpose(0, 2, 1)], axis=2) + if lengths is not None: + ends = mx.clip(lengths, 0, length) + positions = (ends[:, None] + mx.arange(kernel_size))[..., None] + next_cache = mx.take_along_axis( + cache_input.transpose(0, 2, 1), positions, axis=1 + ).transpose(0, 2, 1) + else: + next_cache = mx.contiguous(cache_input[:, :, -kernel_size:]) + return output, mx.contiguous(next_cache) + + +class GatedRMSNorm(nn.Module): + def __init__(self, head_dim: int, eps: float): + super().__init__() + self.weight = mx.ones((head_dim,)) + self.eps = eps + + def __call__(self, x: mx.array, gate: mx.array) -> mx.array: + x = mx.fast.rms_norm(x, weight=self.weight, eps=self.eps) + return x * mx.sigmoid(gate.astype(mx.float32)).astype(x.dtype) + + +def recurrent_kda( + q: mx.array, + k: mx.array, + v: mx.array, + g: mx.array, + beta: mx.array, + a_log: mx.array, + dt_bias: mx.array, + h: Optional[mx.array] = None, + safe_gate: bool = True, + lower_bound: Optional[float] = -5.0, + mask: Optional[mx.array] = None, +) -> Tuple[mx.array, mx.array]: + """Run Ling's KDA recurrence through mlx-lm's fused gated-delta kernel.""" + batch, length, heads, head_dim = q.shape + if h is None: + state = mx.zeros((batch, heads, head_dim, head_dim), dtype=mx.float32) + else: + # Ling stores the reference state as [key_dim, value_dim], while + # mlx-lm's fused kernel uses [value_dim, key_dim]. + state = h.swapaxes(-1, -2).astype(mx.float32) + + inv_scale = head_dim**-0.5 + q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6) + k = inv_scale * mx.fast.rms_norm(k, None, 1e-6) + + gate_input = g.astype(mx.float32) + dt_bias.reshape( + 1, 1, heads, head_dim + ).astype(mx.float32) + a_log = a_log.reshape(1, 1, heads, 1).astype(mx.float32) + if safe_gate and lower_bound is not None: + log_decay = lower_bound * mx.sigmoid(mx.exp(a_log) * gate_input) + else: + log_decay = -mx.exp(a_log) * nn.softplus(gate_input) + decay = mx.exp(log_decay) + beta = mx.sigmoid(beta) + + if ( + head_dim % 32 == 0 + and mx.default_device() == mx.gpu + and mx.metal.is_available() + ): + output, state = gated_delta_kernel(q, k, v, decay, beta, state, mask) + else: + output, state = gated_delta_ops(q, k, v, decay, beta, state, mask) + + return output, state.swapaxes(-1, -2) + + +class LinearAttention(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.num_attention_heads = args.num_attention_heads + self.head_dim = args.head_dim + self.safe_gate = args.kda_safe_gate + self.lower_bound = args.kda_lower_bound + + projection_size = self.num_attention_heads * self.head_dim + self.q_proj = nn.Linear(args.hidden_size, projection_size, bias=False) + self.k_proj = nn.Linear(args.hidden_size, projection_size, bias=False) + self.v_proj = nn.Linear(args.hidden_size, projection_size, bias=False) + self.q_conv1d = DepthwiseConv1d( + projection_size, args.short_conv_kernel_size + ) + self.k_conv1d = DepthwiseConv1d( + projection_size, args.short_conv_kernel_size + ) + self.v_conv1d = DepthwiseConv1d( + projection_size, args.short_conv_kernel_size + ) + self.A_log = mx.zeros((self.num_attention_heads,), dtype=mx.float32) + self.dt_bias = mx.zeros((projection_size,), dtype=mx.float32) + self.f_proj = nn.Linear(args.hidden_size, projection_size, bias=False) + self.b_proj = nn.Linear(args.hidden_size, self.num_attention_heads, bias=False) + self.g_proj = nn.Linear(args.hidden_size, projection_size, bias=False) + self.o_norm = GatedRMSNorm(self.head_dim, args.rms_norm_eps) + self.o_proj = nn.Linear(projection_size, args.hidden_size, bias=False) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + offset: int = 0, + ) -> mx.array: + B, L, _ = x.shape + + # Keep each recurrent tensor in its own ArraysCache slot. The source + # branch stores one tuple in one slot, but pinned mlx-lm batches by + # indexing/merging slots individually; a tuple there breaks every + # BatchGenerator prompt-to-decode transition. Prefixes created by the + # source branch may still restore as a one-slot cache, so migrate that + # legacy layout in place before external prefill reuses it. + state = None + if cache is not None: + slots = list(cache.state) + if len(slots) == 1: + legacy_state = slots[0] + if legacy_state is None: + slots = [None] * 4 + elif isinstance(legacy_state, (list, tuple)) and len(legacy_state) == 4: + slots = list(legacy_state) + else: + raise ValueError( + "Invalid bailing_hybrid recurrent cache: expected a " + "four-tensor legacy state" + ) + cache.state = slots + if len(slots) != 4: + raise ValueError( + "Invalid bailing_hybrid recurrent cache: expected 4 slots, " + f"got {len(slots)}" + ) + state = tuple(slots) + if state is None or all(value is None for value in state): + recurrent_state = conv_q = conv_k = conv_v = None + else: + recurrent_state, conv_q, conv_k, conv_v = state + + lengths = cache.lengths if cache is not None else None + if lengths is not None: + # ``ArraysCache.merge`` initializes ``left_padding`` even for an + # empty cache, and its generic ``make_mask`` checks that field + # before ``lengths``. During right-padded prompt batches the + # resulting all-true mask would advance the recurrent state over + # padding. The current chunk lengths are authoritative here. + mask = mx.arange(L)[None, :] < lengths[:, None] + q, conv_q = self.q_conv1d(self.q_proj(x), conv_q, mask, lengths) + k, conv_k = self.k_conv1d(self.k_proj(x), conv_k, mask, lengths) + v, conv_v = self.v_conv1d(self.v_proj(x), conv_v, mask, lengths) + q = q.reshape(B, L, self.num_attention_heads, self.head_dim) + k = k.reshape(B, L, self.num_attention_heads, self.head_dim) + v = v.reshape(B, L, self.num_attention_heads, self.head_dim) + f = self.f_proj(x).reshape( + B, L, self.num_attention_heads, self.head_dim + ) + beta = self.b_proj(x) + + output, recurrent_state = recurrent_kda( + q, + k, + v, + f, + beta, + self.A_log, + self.dt_bias, + recurrent_state, + safe_gate=self.safe_gate, + lower_bound=self.lower_bound, + mask=mask, + ) + if cache is not None: + cache[0] = recurrent_state + cache[1] = conv_q + cache[2] = conv_k + cache[3] = conv_v + cache.advance(L) + + gate = self.g_proj(x).reshape( + B, L, self.num_attention_heads, self.head_dim + ) + output = self.o_norm(output, gate) + return self.o_proj(output.reshape(B, L, -1)) + + +def group_expert_select( + gates: mx.array, + e_score_correction_bias: Optional[mx.array], + top_k: int, + n_group: int, + topk_group: int, + routed_scaling_factor: float, + norm_topk_prob: bool, + score_function: str, +) -> Tuple[mx.array, mx.array]: + in_type = gates.dtype + if score_function == "sigmoid": + scores = mx.sigmoid(gates.astype(mx.float32)) + else: + scores = mx.softmax(gates.astype(mx.float32), axis=-1) + orig_scores = scores + if e_score_correction_bias is not None: + scores = scores + e_score_correction_bias + if n_group > 1: + scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1)) + group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True) + k = n_group - topk_group + group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :] + scores = mx.put_along_axis( + scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2 + ) + scores = mx.flatten(scores, -2, -1) + + inds = mx.argpartition(-scores, kth=top_k - 1, axis=-1)[..., :top_k] + scores = mx.take_along_axis(orig_scores, inds, axis=-1) + if top_k > 1 and norm_topk_prob: + scores = scores / (scores.sum(axis=-1, keepdims=True) + 1e-20) + scores = scores * routed_scaling_factor + return inds, scores.astype(in_type) + + +class Gate(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.top_k = args.num_experts_per_tok + self.n_group = args.n_group + self.topk_group = args.topk_group + self.norm_topk_prob = args.norm_topk_prob + self.routed_scaling_factor = args.routed_scaling_factor + self.score_function = args.score_function + + self.gate_proj = nn.Linear(args.hidden_size, args.num_experts, bias=False) + self.expert_bias = ( + mx.zeros((args.num_experts,)) + if args.moe_router_enable_expert_bias + else None + ) + + def __call__(self, x: mx.array) -> Tuple[mx.array, mx.array]: + return group_expert_select( + self.gate_proj(x), + self.expert_bias, + self.top_k, + self.n_group, + self.topk_group, + self.routed_scaling_factor, + self.norm_topk_prob, + self.score_function, + ) + + +class SparseMoeBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.num_experts_per_tok = args.num_experts_per_tok + self.switch_mlp = SwitchGLU( + args.hidden_size, + args.moe_intermediate_size, + args.num_experts, + bias=args.use_bias, + ) + self.gate = Gate(args) + self.shared_experts = ( + MLP(args, intermediate_size=args.moe_intermediate_size * args.num_shared_experts) + if args.num_shared_experts > 0 + else None + ) + + def __call__(self, x: mx.array) -> mx.array: + topk_idx, topk_weight = self.gate(x) + out = self.switch_mlp(x, topk_idx) + out = (out * topk_weight[..., None]).sum(axis=-2) + if self.shared_experts is not None: + out = out + self.shared_experts(x) + return out + + +class DecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + n_layers = args.num_hidden_layers + group = args.layer_group_size + self.is_global = ( + (layer_idx + 1) % group == 0 or layer_idx >= (n_layers // group) * group + ) + + if self.is_global: + self.attention = MultiLatentAttention(args) + else: + self.attention = LinearAttention(args, layer_idx=layer_idx) + + if args.num_experts is not None and layer_idx >= args.first_k_dense_replace: + self.mlp = SparseMoeBlock(args) + else: + self.mlp = MLP(args) + + self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + offset: int = 0, + ) -> mx.array: + if self.is_global: + r = self.attention(self.input_layernorm(x), mask, cache) + else: + r = self.attention(self.input_layernorm(x), mask, cache, offset=offset) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + return h + r + + +class LanguageModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + DecoderLayer(args, layer_idx=i) for i in range(args.num_hidden_layers) + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + # Find a representative attention layer index for offset/mask sizing. + self._attn_idx = next( + (i for i, l in enumerate(self.layers) if l.is_global), 0 + ) + self._gla_idx = next( + (i for i, l in enumerate(self.layers) if not l.is_global), 0 + ) + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + ) -> mx.array: + h = self.word_embeddings(inputs) + + if cache is None: + cache = [None] * len(self.layers) + + attn_mask = create_attention_mask(h, cache[self._attn_idx], return_array=True) + gla_mask = create_ssm_mask(h, cache[self._gla_idx]) + offset = ( + cache[self._attn_idx].offset if cache[self._attn_idx] is not None else 0 + ) + if hasattr(offset, "dtype"): + offset = offset + 0 + + for layer, c in zip(self.layers, cache): + mask = attn_mask if layer.is_global else gla_mask + h = layer(h, mask, c, offset=offset) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = LanguageModel(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + ) -> mx.array: + out = self.model(inputs, cache) + if self.args.tie_word_embeddings: + return self.model.word_embeddings.as_linear(out) + return self.lm_head(out) + + def sanitize(self, weights): + n_layers = self.args.num_hidden_layers + + # Drop MTP and any extra non-base layers (Ling 2.6 has 1 MTP layer + # appended after num_hidden_layers; deletes weights for those). + weights = { + k: v + for k, v in weights.items() + if not ( + k.startswith("model.layers.") + and k.split(".")[2].isdigit() + and int(k.split(".")[2]) >= n_layers + ) + } + + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + + for l in range(n_layers): + prefix = f"model.layers.{l}" + + # MoE expert stacking + gate remap + if l >= self.args.first_k_dense_replace: + for m in ["gate_proj", "down_proj", "up_proj"]: + for k in ["weight", "scales", "biases", "weight_scale_inv"]: + if f"{prefix}.mlp.experts.0.{m}.{k}" in weights: + stacked = [ + weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}") + for e in range(self.args.num_experts) + ] + weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack( + stacked + ) + + for suffix in ( + "weight", + "bias", + "scales", + "biases", + "weight_scale_inv", + ): + source = f"{prefix}.mlp.gate.{suffix}" + if source in weights: + weights[f"{prefix}.mlp.gate.gate_proj.{suffix}"] = weights.pop( + source + ) + + # MLA kv_b_proj split for global attention layers. + kv_b_key = f"{prefix}.attention.kv_b_proj.weight" + if kv_b_key in weights: + v = weights.pop(kv_b_key) + head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim + num_heads = self.args.num_attention_heads + v = v.reshape(num_heads, head_dim, -1) + wk = mx.contiguous( + v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2) + ) + wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :]) + weights[f"{prefix}.attention.embed_q.weight"] = wk + weights[f"{prefix}.attention.unembed_out.weight"] = wv + + return self._convert_fp8_block_weights(weights) + + def _convert_fp8_block_weights(self, weights): + """Convert Ling's block-scaled E4M3 tensors to 8-bit affine MLX. + + Published FP8 checkpoints use a float32 ``weight_scale_inv`` grid, + normally with 128x128 blocks. Metal cannot multiply that layout + directly. Converting one stacked tensor at a time keeps peak memory + bounded while preserving the checkpoint's compact runtime footprint. + """ + quantization_config = self.args.quantization_config or {} + block_size = quantization_config.get("weight_block_size", (128, 128)) + if not isinstance(block_size, (list, tuple)) or len(block_size) != 2: + block_size = (128, 128) + block_rows, block_cols = (int(block_size[0]), int(block_size[1])) + + scale_keys = [key for key in weights if key.endswith(".weight_scale_inv")] + for scale_key in scale_keys: + weight_key = scale_key[: -len("_scale_inv")] + if weight_key not in weights or weights[weight_key].dtype != mx.uint8: + continue + + scale = weights.pop(scale_key).astype(mx.float32) + weight = mx.from_fp8(weights.pop(weight_key), dtype=mx.float32) + out_dim, in_dim = weight.shape[-2:] + target_out = scale.shape[-2] * block_rows + target_in = scale.shape[-1] * block_cols + if target_out < out_dim or target_in < in_dim: + raise ValueError( + f"Invalid FP8 block scale for {weight_key}: weight " + f"{weight.shape}, scale {scale.shape}, block {block_size}" + ) + + pad_out = target_out - out_dim + pad_in = target_in - in_dim + if pad_out or pad_in: + padding = [(0, 0)] * (weight.ndim - 2) + padding.extend(((0, pad_out), (0, pad_in))) + weight = mx.pad(weight, padding) + + lead = weight.shape[:-2] + weight = weight.reshape( + *lead, + scale.shape[-2], + block_rows, + scale.shape[-1], + block_cols, + ) + weight = weight * scale[..., :, None, :, None] + weight = weight.reshape(*lead, target_out, target_in) + weight = weight[..., :out_dim, :in_dim].astype(mx.bfloat16) + + quantized, scales, biases = mx.quantize(weight, group_size=64, bits=8) + weights[weight_key] = quantized + base = weight_key[: -len("weight")] + weights[f"{base}scales"] = scales + weights[f"{base}biases"] = biases + mx.eval(quantized, scales, biases) + mx.clear_cache() + + return weights + + @property + def quant_predicate(self): + def predicate(path, _): + if path.endswith("mlp.gate.gate_proj"): + return {"group_size": 64, "bits": 8} + return True + + return predicate + + @property + def cast_predicate(self): + def predicate(k): + return "expert_bias" not in k + + return predicate + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + caches = [] + for l in self.layers: + if l.is_global: + caches.append(KVCache()) + else: + caches.append(ArraysCache(size=4)) + return caches diff --git a/omlx/scheduler.py b/omlx/scheduler.py index 4adf3d9293..26dac2974b 100644 --- a/omlx/scheduler.py +++ b/omlx/scheduler.py @@ -6288,6 +6288,29 @@ def _validate_cache(self, cache: Any) -> bool: if isinstance(cache, list): if len(cache) == 0: return False + + # Variable-state caches must keep the arity declared by the live + # model. Older model implementations can persist an ArraysCache + # with fewer (or zero) slots under the same class name; accepting + # it reaches the model with missing recurrent state and either + # crashes or silently corrupts a cached continuation. + try: + expected_cache = make_prompt_cache(self.model) + except Exception: + expected_cache = None + if isinstance(expected_cache, (list, tuple)) and len( + expected_cache + ) == len(cache): + arrays_names = {"ArraysCache", "SizedArraysCache"} + for layer_cache, expected_layer in zip(cache, expected_cache): + if ( + type(expected_layer).__name__ == "ArraysCache" + and type(layer_cache).__name__ in arrays_names + and len(getattr(layer_cache, "state", ())) + != len(getattr(expected_layer, "state", ())) + ): + return False + # Check each layer for layer_cache in cache: if layer_cache is None: diff --git a/omlx/utils/model_loading.py b/omlx/utils/model_loading.py index 73d59a4823..084bb94595 100644 --- a/omlx/utils/model_loading.py +++ b/omlx/utils/model_loading.py @@ -179,6 +179,30 @@ def normalize_laguna_compressed_quant(cfg: dict) -> dict: return cfg +def normalize_bailing_hybrid_fp8_quant(cfg: dict) -> dict: + """Map Ling block-scaled FP8 checkpoints to an MLX runtime format. + + Ling 3.0 Flash FP8 checkpoints store E4M3 weights with float32 + ``weight_scale_inv`` tensors on a 128x128 block grid. MLX has no native + matmul for that layout, so the vendored model sanitizer dequantizes each + block and requantizes it to 8-bit affine. Declaring the matching runtime + quantization here makes ``mlx_lm.utils.load_model`` construct + ``QuantizedLinear`` modules for the generated ``scales`` sidecars. + + Mutates *cfg* in place and returns it for convenience. + """ + if cfg.get("model_type") != "bailing_hybrid": + return cfg + if isinstance(cfg.get("quantization"), dict): + return cfg + qc = cfg.get("quantization_config") + if not isinstance(qc, dict) or qc.get("quant_method") != "fp8": + return cfg + + cfg["quantization"] = {"group_size": 64, "bits": 8} + return cfg + + def _patch_mlx_lm_load_config() -> None: """Wrap ``mlx_lm.utils.load_config`` to expand per-layer quant keys.""" global _MLX_LM_LOAD_CONFIG_PATCHED @@ -197,6 +221,7 @@ def _patched(model_path, *args, **kwargs): expand_per_layer_quant_keys(cfg) expand_glm_moe_dsa_fused_quant_keys(cfg) normalize_laguna_compressed_quant(cfg) + normalize_bailing_hybrid_fp8_quant(cfg) return cfg _lu.load_config = _patched @@ -219,6 +244,9 @@ def maybe_apply_pre_load_patches( - MiMo V2.5 text backbone (PR 1219) when ``config.json`` declares ``model_type == "mimo_v2"``. The vendored model intentionally ignores the base checkpoint's vision, audio, speech, and MTP weights. + - Ling 3.0 Flash mixed MLA/KDA model when ``config.json`` declares + ``model_type == "bailing_hybrid"``. The vendored module is registered + as ``mlx_lm.models.bailing_hybrid`` before mlx-lm resolves its classes. - Llama 4 attention offset patch when ``config.json`` declares ``model_type == "llama4"`` directly or under ``text_config``. - GLM-5.2 ``glm_moe_dsa`` patch (mlx-lm PR 1410) when ``config.json`` @@ -339,6 +367,12 @@ def maybe_apply_pre_load_patches( if apply_mimo_v2_patch(): logger.info("MiMo V2.5 text pre-load patch applied for %s", model_name) + if model_type == "bailing_hybrid": + from ..patches.bailing_hybrid import apply_bailing_hybrid_patch + + if apply_bailing_hybrid_patch(): + logger.info("Ling 3.0 Flash pre-load patch applied for %s", model_name) + if model_type == "laguna": # MLX-LM dynamically imports the architecture and tokenizer-configured # parser during ``lm_load_compat``; register both before that load starts. diff --git a/tests/test_bailing_hybrid_patch.py b/tests/test_bailing_hybrid_patch.py new file mode 100644 index 0000000000..55b7f503c5 --- /dev/null +++ b/tests/test_bailing_hybrid_patch.py @@ -0,0 +1,591 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the Ling 3.0 Flash ``bailing_hybrid`` mlx-lm patch.""" + +import importlib +import json +import sys +from types import SimpleNamespace + +import mlx.core as mx +import pytest + + +def _minimal_config(**overrides): + config = { + "model_type": "bailing_hybrid", + "architectures": ["BailingHybridForCausalLM"], + "hidden_size": 32, + "intermediate_size": 64, + "moe_intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "num_experts": 2, + "num_experts_per_tok": 1, + "num_shared_experts": 0, + "n_group": 1, + "topk_group": 1, + "first_k_dense_replace": 1, + "layer_group_size": 2, + "group_norm_size": 1, + "vocab_size": 128, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "max_position_embeddings": 256, + "routed_scaling_factor": 1.0, + "head_dim": 8, + "kv_lora_rank": 8, + "qk_rope_head_dim": 4, + "qk_nope_head_dim": 4, + "v_head_dim": 4, + "short_conv_kernel_size": 3, + } + config.update(overrides) + return config + + +def _load_patch_module(): + from omlx.patches.bailing_hybrid import apply_bailing_hybrid_patch + + apply_bailing_hybrid_patch() + return importlib.import_module("mlx_lm.models.bailing_hybrid") + + +def test_apply_registers_bailing_hybrid_module(): + module = _load_patch_module() + + assert module.__package__ == "mlx_lm.models" + assert sys.modules["mlx_lm.models.bailing_hybrid"] is module + + import mlx_lm.models as models_pkg + + assert models_pkg.bailing_hybrid is module + + +def test_apply_is_idempotent(): + from omlx.patches.bailing_hybrid import ( + apply_bailing_hybrid_patch, + is_applied, + ) + + first = apply_bailing_hybrid_patch() + second = apply_bailing_hybrid_patch() + + assert is_applied() is True + assert second is False + assert first in (True, False) + + +def test_apply_prefers_upstream_module(monkeypatch): + from omlx.patches import bailing_hybrid + + upstream = object() + models_pkg = SimpleNamespace() + + def fake_import(name): + if name == "mlx_lm.models.bailing_hybrid": + return upstream + if name == "mlx_lm.models": + return models_pkg + raise AssertionError(f"unexpected import: {name}") + + monkeypatch.setattr(bailing_hybrid, "_APPLIED", False) + monkeypatch.setattr(bailing_hybrid.importlib, "import_module", fake_import) + monkeypatch.setattr( + bailing_hybrid, + "_register_module", + lambda: (_ for _ in ()).throw(AssertionError("vendored module used")), + ) + + assert bailing_hybrid.apply_bailing_hybrid_patch() is False + assert models_pkg.bailing_hybrid is upstream + + +def test_get_classes_resolves_bailing_hybrid(): + _load_patch_module() + + from mlx_lm.utils import _get_classes + + model_cls, args_cls = _get_classes(_minimal_config()) + + assert model_cls.__name__ == "Model" + assert args_cls.__name__ == "ModelArgs" + + +def test_mixed_global_and_linear_attention_cache_forward(): + bailing_hybrid = _load_patch_module() + from mlx_lm.generate import BatchGenerator + from mlx_lm.models.cache import ArraysCache, KVCache + + model = bailing_hybrid.Model( + bailing_hybrid.ModelArgs.from_dict(_minimal_config()) + ) + cache = model.make_cache() + + assert type(cache[0]) is ArraysCache + assert type(cache[1]) is KVCache + + prefill = model(mx.array([[1, 2, 3]], dtype=mx.int32), cache=cache) + decode = model(mx.array([[4]], dtype=mx.int32), cache=cache) + mx.eval(prefill, decode) + + assert prefill.shape == (1, 3, 128) + assert decode.shape == (1, 1, 128) + assert cache[0][0] is not None + assert cache[1].offset == 4 + + generator = BatchGenerator( + model, + max_tokens=2, + prefill_batch_size=2, + completion_batch_size=2, + sampler=lambda logits: mx.argmax(logits, axis=-1), + ) + uids = generator.insert([[1, 2, 3], [4, 5]], max_tokens=[2, 2]) + finished = [] + for _ in range(8): + _, responses = generator.next() + finished.extend(r for r in responses if r.finish_reason is not None) + if len(finished) == 2: + break + + assert uids == [0, 1] + assert {response.uid for response in finished} == {0, 1} + assert all(response.finish_reason == "length" for response in finished) + + +def _batch_greedy_tokens(model, prompts, max_tokens=6): + from mlx_lm.generate import BatchGenerator + + generator = BatchGenerator( + model, + max_tokens=max_tokens, + prefill_batch_size=len(prompts), + completion_batch_size=len(prompts), + sampler=lambda logits: mx.argmax(logits, axis=-1), + ) + uids = generator.insert(prompts, max_tokens=[max_tokens] * len(prompts)) + tokens = {uid: [] for uid in uids} + for _ in range(max_tokens + 4): + _, responses = generator.next() + for response in responses: + tokens[response.uid].append(response.token) + if all(len(output) == max_tokens for output in tokens.values()): + break + return [tokens[uid] for uid in uids] + + +def test_variable_length_batch_matches_single_request_greedy_tokens(): + bailing_hybrid = _load_patch_module() + mx.random.seed(7) + model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(_minimal_config())) + + short_prompt = [4, 5] + long_prompt = [7, 8, 9, 10, 11, 12] + single = _batch_greedy_tokens(model, [short_prompt])[0] + batched = _batch_greedy_tokens(model, [short_prompt, long_prompt])[0] + + assert batched == single + + +def test_depthwise_conv_matches_token_loop_reference(): + bailing_hybrid = _load_patch_module() + conv = bailing_hybrid.DepthwiseConv1d(channels=4, kernel_size=3) + conv.weight = mx.arange(12, dtype=mx.float32).reshape(4, 1, 3) / 12 + x = mx.arange(32, dtype=mx.float32).reshape(2, 4, 4) / 32 + initial_cache = mx.arange(24, dtype=mx.float32).reshape(2, 4, 3) / 24 + + expected_cache = initial_cache + expected_outputs = [] + weight = conv.weight[:, 0, :] + for token_idx in range(x.shape[1]): + current = x[:, token_idx : token_idx + 1, :].transpose(0, 2, 1) + expected_cache = mx.concatenate( + [expected_cache[:, :, 1:], current], + axis=2, + ) + value = (expected_cache * weight[None, :, :]).sum(axis=2) + expected_outputs.append(mx.sigmoid(value) * value) + expected = mx.stack(expected_outputs, axis=1) + + actual, actual_cache = conv(x, initial_cache) + mx.eval(expected, expected_cache, actual, actual_cache) + + assert mx.allclose(actual, expected, rtol=1e-5, atol=1e-6) + assert mx.allclose(actual_cache, expected_cache) + + +def test_depthwise_conv_uses_lengths_for_right_padded_cache_state(): + bailing_hybrid = _load_patch_module() + conv = bailing_hybrid.DepthwiseConv1d(channels=4, kernel_size=3) + conv.weight = mx.arange(12, dtype=mx.float32).reshape(4, 1, 3) / 12 + x = mx.arange(32, dtype=mx.float32).reshape(2, 4, 4) / 32 + initial_cache = mx.arange(24, dtype=mx.float32).reshape(2, 4, 3) / 24 + mask = mx.array( + [[True, True, False, False], [True, True, True, True]], + dtype=mx.bool_, + ) + + batch_output, batch_cache = conv( + x, + initial_cache, + mask=mask, + lengths=mx.array([2, 4]), + ) + single_output, single_cache = conv(x[:1, :2], initial_cache[:1]) + mx.eval(batch_output, batch_cache, single_output, single_cache) + + assert mx.allclose(batch_output[0, :2], single_output[0]) + assert mx.allclose(batch_cache[0], single_cache[0]) + + +@pytest.mark.parametrize("safe_gate", [False, True]) +def test_fused_kda_matches_reference(safe_gate): + bailing_hybrid = _load_patch_module() + batch, length, heads, head_dim = 1, 5, 2, 8 + q = mx.arange(batch * length * heads * head_dim, dtype=mx.float32).reshape( + batch, length, heads, head_dim + ) + q = q / 100 + k = q + 0.1 + v = q + 0.2 + g = q + 0.3 + beta = mx.arange(batch * length * heads, dtype=mx.float32).reshape( + batch, length, heads + ) + beta = beta / 10 + a_log = mx.array([-0.2, 0.3], dtype=mx.float32) + dt_bias = mx.arange(heads * head_dim, dtype=mx.float32) / 50 + initial_state = mx.arange( + batch * heads * head_dim * head_dim, + dtype=mx.float32, + ).reshape(batch, heads, head_dim, head_dim) + initial_state = initial_state / 1000 + + reference_state = initial_state + reference_outputs = [] + for token_idx in range(length): + q_t = q[:, token_idx] + k_t = k[:, token_idx] + v_t = v[:, token_idx] + q_t = q_t / mx.sqrt(mx.sum(q_t * q_t, axis=-1, keepdims=True) + 1e-6) + k_t = k_t / mx.sqrt(mx.sum(k_t * k_t, axis=-1, keepdims=True) + 1e-6) + gate_input = g[:, token_idx] + dt_bias.reshape(heads, head_dim) + if safe_gate: + log_decay = -5.0 * mx.sigmoid( + mx.exp(a_log)[None, :, None] * gate_input + ) + else: + log_decay = -mx.exp(a_log)[None, :, None] * mx.logaddexp( + gate_input, + mx.array(0.0), + ) + reference_state = reference_state * mx.exp(log_decay)[..., None] + delta = v_t - mx.sum(reference_state * k_t[..., None], axis=2) + delta = delta * mx.sigmoid(beta[:, token_idx])[..., None] + reference_state = reference_state + k_t[..., None] * delta[..., None, :] + reference_outputs.append( + mx.sum(reference_state * q_t[..., None], axis=2) * (head_dim**-0.5) + ) + expected = mx.stack(reference_outputs, axis=1) + + actual, actual_state = bailing_hybrid.recurrent_kda( + q, + k, + v, + g, + beta, + a_log, + dt_bias, + initial_state, + safe_gate=safe_gate, + lower_bound=-5.0, + ) + mx.eval(expected, reference_state, actual, actual_state) + + assert mx.allclose(actual, expected, rtol=2e-4, atol=2e-5) + assert mx.allclose(actual_state, reference_state, rtol=2e-4, atol=2e-5) + + +def test_external_prefill_upgrades_legacy_one_slot_cache(): + bailing_hybrid = _load_patch_module() + from mlx_lm.models.cache import ArraysCache + + from omlx.request import Request, SamplingParams + from omlx.scheduler import Scheduler + + model = bailing_hybrid.Model( + bailing_hybrid.ModelArgs.from_dict(_minimal_config()) + ) + source_cache = model.make_cache() + prefix_logits = model( + mx.array([[1, 2]], dtype=mx.int32), + cache=source_cache, + ) + mx.eval(prefix_logits) + + legacy_cache = ArraysCache(size=1) + legacy_cache[0] = tuple(source_cache[0].state) + cache = [legacy_cache, source_cache[1]] + request = Request( + request_id="ling-legacy-prefill", + prompt=[3, 4], + sampling_params=SamplingParams(max_tokens=1), + ) + request.prompt_token_ids = [3, 4] + request.num_prompt_tokens = 2 + + tokenizer = SimpleNamespace( + encode=lambda _text: [0], + eos_token_id=127, + all_special_ids=[127], + ) + scheduler = Scheduler(model=model, tokenizer=tokenizer) + prefilled_cache, last_token = scheduler._do_external_prefill( + request, + request.prompt_token_ids, + cache, + ) + + assert prefilled_cache is cache + assert last_token == [4] + assert len(legacy_cache.state) == 4 + assert all(state is not None for state in legacy_cache.state) + + +def test_scheduler_rejects_legacy_zero_slot_cache(): + bailing_hybrid = _load_patch_module() + from mlx_lm.models.cache import ArraysCache + + from omlx.scheduler import Scheduler + + model = bailing_hybrid.Model( + bailing_hybrid.ModelArgs.from_dict(_minimal_config()) + ) + source_cache = model.make_cache() + logits = model(mx.array([[1, 2]], dtype=mx.int32), cache=source_cache) + mx.eval(logits) + + tokenizer = SimpleNamespace( + encode=lambda _text: [0], + eos_token_id=127, + all_special_ids=[127], + ) + scheduler = Scheduler(model=model, tokenizer=tokenizer) + + assert scheduler._validate_cache([ArraysCache(size=0), source_cache[1]]) is False + assert scheduler._validate_cache(source_cache) is True + + +def test_sanitize_remaps_moe_and_mla_weights(): + bailing_hybrid = _load_patch_module() + model = bailing_hybrid.Model( + bailing_hybrid.ModelArgs.from_dict(_minimal_config()) + ) + + weights = { + "model.layers.1.mlp.gate.weight": mx.ones((2, 32)), + "model.layers.1.mlp.gate.bias": mx.ones((2,)), + "model.layers.1.attention.kv_b_proj.weight": mx.arange(128).reshape(16, 8), + "model.layers.2.mtp.weight": mx.ones((1,)), + } + for projection, shape in ( + ("gate_proj", (16, 32)), + ("up_proj", (16, 32)), + ("down_proj", (32, 16)), + ): + for expert in range(2): + weights[f"model.layers.1.mlp.experts.{expert}.{projection}.weight"] = ( + mx.full(shape, expert + 1) + ) + + sanitized = model.sanitize(weights) + + assert "model.layers.1.mlp.gate.weight" not in sanitized + assert "model.layers.1.mlp.gate.bias" not in sanitized + assert sanitized["model.layers.1.mlp.gate.gate_proj.weight"].shape == (2, 32) + assert sanitized["model.layers.1.mlp.gate.gate_proj.bias"].shape == (2,) + assert sanitized["model.layers.1.mlp.switch_mlp.gate_proj.weight"].shape == ( + 2, + 16, + 32, + ) + assert sanitized["model.layers.1.mlp.switch_mlp.up_proj.weight"].shape == ( + 2, + 16, + 32, + ) + assert sanitized["model.layers.1.mlp.switch_mlp.down_proj.weight"].shape == ( + 2, + 32, + 16, + ) + assert sanitized["model.layers.1.attention.embed_q.weight"].shape == (2, 8, 4) + assert sanitized["model.layers.1.attention.unembed_out.weight"].shape == ( + 2, + 4, + 8, + ) + assert "model.layers.1.attention.kv_b_proj.weight" not in sanitized + assert "model.layers.2.mtp.weight" not in sanitized + + +def test_sanitize_converts_block_fp8_weights_to_affine_runtime_layout(): + bailing_hybrid = _load_patch_module() + config = _minimal_config( + hidden_size=64, + intermediate_size=128, + quantization_config={ + "quant_method": "fp8", + "weight_block_size": [128, 128], + }, + ) + model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(config)) + source = mx.linspace(-1.0, 1.0, 16 * 64).reshape(16, 64) + fp8 = mx.to_fp8(source) + weight_key = "model.layers.0.attention.q_proj.weight" + scale_key = f"{weight_key}_scale_inv" + + sanitized = model.sanitize( + { + weight_key: fp8, + scale_key: mx.array([[0.5]], dtype=mx.float32), + } + ) + restored = mx.dequantize( + sanitized[weight_key], + sanitized[weight_key.replace("weight", "scales")], + sanitized[weight_key.replace("weight", "biases")], + group_size=64, + bits=8, + ) + expected = mx.from_fp8(fp8, dtype=mx.bfloat16) * 0.5 + mx.eval(restored, expected) + + assert scale_key not in sanitized + assert sanitized[weight_key].dtype == mx.uint32 + assert mx.allclose(restored, expected, rtol=2e-2, atol=5e-3) + + +def test_sanitize_stacks_fp8_expert_weights_and_sidecars(): + bailing_hybrid = _load_patch_module() + config = _minimal_config( + hidden_size=64, + intermediate_size=128, + quantization_config={ + "quant_method": "fp8", + "weight_block_size": [128, 128], + }, + ) + model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(config)) + weights = {} + for expert in range(2): + prefix = f"model.layers.1.mlp.experts.{expert}.gate_proj" + source = mx.full((16, 64), 0.25 * (expert + 1), dtype=mx.float32) + weights[f"{prefix}.weight"] = mx.to_fp8(source) + weights[f"{prefix}.weight_scale_inv"] = mx.ones((1, 1)) + + sanitized = model.sanitize(weights) + prefix = "model.layers.1.mlp.switch_mlp.gate_proj" + + assert sanitized[f"{prefix}.weight"].shape == (2, 16, 16) + assert sanitized[f"{prefix}.scales"].shape == (2, 16, 1) + assert sanitized[f"{prefix}.biases"].shape == (2, 16, 1) + assert not any(key.endswith("weight_scale_inv") for key in sanitized) + + +def test_bailing_fp8_config_normalizes_to_affine_runtime_quantization(): + from omlx.utils.model_loading import normalize_bailing_hybrid_fp8_quant + + config = _minimal_config( + quantization_config={ + "quant_method": "fp8", + "weight_block_size": [128, 128], + } + ) + + assert normalize_bailing_hybrid_fp8_quant(config) is config + assert config["quantization"] == {"group_size": 64, "bits": 8} + + +def test_fp8_checkpoint_loads_strictly_as_quantized_model(tmp_path): + bailing_hybrid = _load_patch_module() + import mlx.nn as nn + from mlx.utils import tree_flatten + + config = _minimal_config( + hidden_size=64, + intermediate_size=128, + quantization_config={ + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [128, 128], + }, + ) + source_model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(config)) + weights = dict(tree_flatten(source_model.parameters())) + weight_key = "model.layers.0.attention.q_proj.weight" + source_weight = weights[weight_key] + weights[weight_key] = mx.to_fp8(source_weight.astype(mx.float32)) + weights[f"{weight_key}_scale_inv"] = mx.ones((1, 1), dtype=mx.float32) + mx.save_safetensors(str(tmp_path / "model.safetensors"), weights) + (tmp_path / "config.json").write_text(json.dumps(config)) + + from mlx_lm.utils import load_model + + from omlx.utils.model_loading import maybe_apply_pre_load_patches + + maybe_apply_pre_load_patches(str(tmp_path)) + loaded, loaded_config = load_model(tmp_path, strict=True) + logits = loaded(mx.array([[1, 2, 3]], dtype=mx.int32)) + mx.eval(logits) + + assert loaded_config["quantization"] == {"group_size": 64, "bits": 8} + assert isinstance(loaded.model.layers[0].attention.q_proj, nn.QuantizedLinear) + assert logits.shape == (1, 3, config["vocab_size"]) + + +def test_oq_discovers_ling_embeddings_and_hybrid_layer_masks(): + bailing_hybrid = _load_patch_module() + from omlx.oq import ( + _find_model_layers, + _layer_masks_for_model, + _uses_quantized_source_sensitivity, + ) + + config = _minimal_config( + quantization_config={"quant_method": "fp8"}, + ) + model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(config)) + embed_fn, layers = _find_model_layers(model) + hidden = embed_fn(mx.array([[1, 2, 3]], dtype=mx.int32)) + masks = _layer_masks_for_model(model, layers, hidden) + + assert embed_fn is model.model.word_embeddings + assert layers is model.model.layers + assert masks[0] is None + assert masks[1] is not None + assert _uses_quantized_source_sensitivity(config) is True + + +def test_pre_load_dispatch_calls_bailing_hybrid_patch(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr( + "omlx.patches.bailing_hybrid.apply_bailing_hybrid_patch", + lambda: calls.append(True) or True, + ) + (tmp_path / "config.json").write_text(json.dumps(_minimal_config())) + + from omlx.utils.model_loading import maybe_apply_pre_load_patches + + maybe_apply_pre_load_patches(str(tmp_path)) + + assert calls == [True] + + +def test_bailing_hybrid_is_discovered_as_llm(tmp_path): + from omlx.model_discovery import detect_model_type + + (tmp_path / "config.json").write_text(json.dumps(_minimal_config())) + + assert detect_model_type(tmp_path) == "llm" diff --git a/tests/test_cache_type_handlers.py b/tests/test_cache_type_handlers.py index 0a9d58a646..b088ed2771 100644 --- a/tests/test_cache_type_handlers.py +++ b/tests/test_cache_type_handlers.py @@ -718,6 +718,19 @@ def test_concatenate_states(self, handler): assert result["states"] == [3, 4, 5] + def test_deserialize_state_preserves_all_slots(self, handler): + """Variable-length state must not be dropped by fixed-axis decoding.""" + import mlx.core as mx + + elements = tuple(mx.full((1,), i) for i in range(4)) + + restored = handler.deserialize_state(elements) + + assert isinstance(restored, SizedArraysCache) + assert len(restored.state) == 4 + for expected, actual in zip(elements, restored.state): + assert mx.array_equal(expected, actual).item() + def test_state_keys(self, handler): """Test state keys.""" assert handler._get_state_keys() == ("states",) diff --git a/tests/test_prefix_cache.py b/tests/test_prefix_cache.py index 1964b9cdf6..a76e140268 100644 --- a/tests/test_prefix_cache.py +++ b/tests/test_prefix_cache.py @@ -3105,6 +3105,54 @@ def test_reconstruct_accepts_sized_arrays_metadata_with_turboquant(self, mx): assert block_table.num_tokens == 8 mock_ssd.forget_block.assert_not_called() + def test_reconstruct_preserves_variable_length_arrays_state(self, mx): + """N-tuple ArraysCache markers restore every recurrent state slot.""" + from omlx.cache.paged_ssd_cache import PagedSSDCacheManager + from omlx.cache.type_handlers import SizedArraysCache + + mock_ssd = MagicMock(spec=PagedSSDCacheManager) + paged_cache = PagedCacheManager( + block_size=4, + max_blocks=100, + model_name="test-model", + initial_blocks=100, + ) + cache = BlockAwarePrefixCache( + model=MockModel(num_layers=1), + paged_cache_manager=paged_cache, + paged_ssd_cache_manager=mock_ssd, + ) + block = paged_cache.allocate_block() + block.block_hash = b"arrays-four-state" + block.token_count = 4 + block.ref_count = 2 + paged_cache.cached_block_hash_to_block.insert(block.block_hash, block) + block_table = BlockTable( + request_id="req-arrays-four-state", + block_ids=[block.block_id], + num_tokens=4, + ) + states = [mx.full((1, 2, 3), i) for i in range(4)] + mock_ssd.load_block_with_metadata.return_value = ( + [("__nstate__", "ArraysCache", states)], + { + "model_name": "test-model", + "num_layers": 1, + "block_size": 4, + "layer_cache_types": ["ArraysCache"], + "layer_meta_states": [()], + }, + ) + + result = cache.reconstruct_cache(block_table) + + assert result is not None + assert isinstance(result[0], SizedArraysCache) + assert result[0].size() == 4 + assert len(result[0].state) == 4 + for expected, actual in zip(states, result[0].state): + assert mx.array_equal(expected, actual).item() + class TestPerBlockMetaStates: """Tests for per-block meta_states in store_cache with boundary snapshots. From d4adcc357e1653574d7fc281123ff7323dee755e Mon Sep 17 00:00:00 2001 From: Mike Wallio Date: Wed, 5 Aug 2026 20:57:39 -0400 Subject: [PATCH 012/338] feat: support Ling 3.0 Flash FP4 checkpoints (#2534) Load routed expert tensors through MLX's native MXFP4 layout while preserving the existing FP8 conversion path for remaining projections. Decode E8M0 scales and add strict mixed-checkpoint coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1daf7029-ff16-4d57-9002-a2ad1d7bb2ee --- .../bailing_hybrid/bailing_hybrid_model.py | 36 +++- omlx/utils/model_loading.py | 24 ++- tests/test_bailing_hybrid_patch.py | 182 ++++++++++++++++++ 3 files changed, 237 insertions(+), 5 deletions(-) diff --git a/omlx/patches/bailing_hybrid/bailing_hybrid_model.py b/omlx/patches/bailing_hybrid/bailing_hybrid_model.py index 6e18a0f73c..0135800bc7 100644 --- a/omlx/patches/bailing_hybrid/bailing_hybrid_model.py +++ b/omlx/patches/bailing_hybrid/bailing_hybrid_model.py @@ -736,12 +736,17 @@ def sanitize(self, weights): return self._convert_fp8_block_weights(weights) def _convert_fp8_block_weights(self, weights): - """Convert Ling's block-scaled E4M3 tensors to 8-bit affine MLX. + """Convert Ling's published FP8 and MXFP4 tensor layouts for MLX. Published FP8 checkpoints use a float32 ``weight_scale_inv`` grid, normally with 128x128 blocks. Metal cannot multiply that layout directly. Converting one stacked tensor at a time keeps peak memory bounded while preserving the checkpoint's compact runtime footprint. + + The mixed FP4 checkpoint stores routed experts as two packed E2M1 + values per int8 byte with one E8M0 scale per 32 logical values. That is + MLX's native MXFP4 layout, so only reinterpret the packed bytes and + rename the scale sidecar instead of dequantizing the experts. """ quantization_config = self.args.quantization_config or {} block_size = quantization_config.get("weight_block_size", (128, 128)) @@ -752,10 +757,35 @@ def _convert_fp8_block_weights(self, weights): scale_keys = [key for key in weights if key.endswith(".weight_scale_inv")] for scale_key in scale_keys: weight_key = scale_key[: -len("_scale_inv")] - if weight_key not in weights or weights[weight_key].dtype != mx.uint8: + if weight_key not in weights: + continue + + source_weight = weights[weight_key] + is_routed_mxfp4 = ( + quantization_config.get("routed_experts_quant_method") == "mxfp4" + and ".mlp.switch_mlp." in weight_key + and source_weight.dtype == mx.int8 + ) + if is_routed_mxfp4: + scale = weights.pop(scale_key) + packed = weights.pop(weight_key).view(mx.uint32) + base = weight_key[: -len("weight")] + weights[weight_key] = packed + weights[f"{base}scales"] = scale + mx.eval(packed, scale) continue - scale = weights.pop(scale_key).astype(mx.float32) + if source_weight.dtype != mx.uint8: + continue + + scale = weights.pop(scale_key) + if scale.dtype == mx.uint8: + scale = mx.power( + mx.array(2.0, dtype=mx.float32), + scale.astype(mx.float32) - 127.0, + ) + else: + scale = scale.astype(mx.float32) weight = mx.from_fp8(weights.pop(weight_key), dtype=mx.float32) out_dim, in_dim = weight.shape[-2:] target_out = scale.shape[-2] * block_rows diff --git a/omlx/utils/model_loading.py b/omlx/utils/model_loading.py index 084bb94595..1e00fb0c76 100644 --- a/omlx/utils/model_loading.py +++ b/omlx/utils/model_loading.py @@ -180,7 +180,7 @@ def normalize_laguna_compressed_quant(cfg: dict) -> dict: def normalize_bailing_hybrid_fp8_quant(cfg: dict) -> dict: - """Map Ling block-scaled FP8 checkpoints to an MLX runtime format. + """Map Ling mixed FP8/MXFP4 checkpoints to MLX runtime formats. Ling 3.0 Flash FP8 checkpoints store E4M3 weights with float32 ``weight_scale_inv`` tensors on a 128x128 block grid. MLX has no native @@ -189,6 +189,11 @@ def normalize_bailing_hybrid_fp8_quant(cfg: dict) -> dict: quantization here makes ``mlx_lm.utils.load_model`` construct ``QuantizedLinear`` modules for the generated ``scales`` sidecars. + The FP4 release keeps non-expert projections in that FP8 layout, but stores + routed expert projections as packed MXFP4 with E8M0 scales. Those tensors + already match MLX's native MXFP4 representation after a byte reinterpret, + so add per-module overrides for the runtime ``SwitchGLU`` paths. + Mutates *cfg* in place and returns it for convenience. """ if cfg.get("model_type") != "bailing_hybrid": @@ -199,7 +204,22 @@ def normalize_bailing_hybrid_fp8_quant(cfg: dict) -> dict: if not isinstance(qc, dict) or qc.get("quant_method") != "fp8": return cfg - cfg["quantization"] = {"group_size": 64, "bits": 8} + quantization: dict[str, Any] = {"group_size": 64, "bits": 8} + if qc.get("routed_experts_quant_method") == "mxfp4": + group_size = int(qc.get("routed_experts_group_size", 32)) + first_sparse_layer = int(cfg.get("first_k_dense_replace", 0)) + num_hidden_layers = int(cfg.get("num_hidden_layers", 0)) + expert_quantization = { + "group_size": group_size, + "bits": 4, + "mode": "mxfp4", + } + for layer_idx in range(first_sparse_layer, num_hidden_layers): + base = f"model.layers.{layer_idx}.mlp.switch_mlp" + for projection in ("gate_proj", "up_proj", "down_proj"): + quantization[f"{base}.{projection}"] = dict(expert_quantization) + + cfg["quantization"] = quantization return cfg diff --git a/tests/test_bailing_hybrid_patch.py b/tests/test_bailing_hybrid_patch.py index 55b7f503c5..57514c53b1 100644 --- a/tests/test_bailing_hybrid_patch.py +++ b/tests/test_bailing_hybrid_patch.py @@ -494,6 +494,98 @@ def test_sanitize_stacks_fp8_expert_weights_and_sidecars(): assert not any(key.endswith("weight_scale_inv") for key in sanitized) +def test_sanitize_preserves_packed_mxfp4_expert_weights(): + bailing_hybrid = _load_patch_module() + config = _minimal_config( + hidden_size=64, + intermediate_size=128, + quantization_config={ + "quant_method": "fp8", + "weight_block_size": [128, 128], + "routed_experts_quant_method": "mxfp4", + "routed_experts_group_size": 32, + }, + ) + model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(config)) + weights = {} + expected = [] + for expert in range(2): + prefix = f"model.layers.1.mlp.experts.{expert}.gate_proj" + source = mx.linspace(-1.0, 1.0, 16 * 64).reshape(16, 64) * (expert + 1) + packed, scales = mx.quantize( + source, + group_size=32, + bits=4, + mode="mxfp4", + ) + weights[f"{prefix}.weight"] = packed.view(mx.int8) + weights[f"{prefix}.weight_scale_inv"] = scales + expected.append( + mx.dequantize( + packed, + scales, + None, + group_size=32, + bits=4, + mode="mxfp4", + ) + ) + + sanitized = model.sanitize(weights) + prefix = "model.layers.1.mlp.switch_mlp.gate_proj" + restored = mx.dequantize( + sanitized[f"{prefix}.weight"], + sanitized[f"{prefix}.scales"], + None, + group_size=32, + bits=4, + mode="mxfp4", + ) + expected = mx.stack(expected) + mx.eval(restored, expected) + + assert sanitized[f"{prefix}.weight"].shape == (2, 16, 8) + assert sanitized[f"{prefix}.weight"].dtype == mx.uint32 + assert sanitized[f"{prefix}.scales"].shape == (2, 16, 2) + assert sanitized[f"{prefix}.scales"].dtype == mx.uint8 + assert not any(key.endswith("weight_scale_inv") for key in sanitized) + assert mx.array_equal(restored, expected) + + +def test_sanitize_decodes_e8m0_fp8_block_scales(): + bailing_hybrid = _load_patch_module() + config = _minimal_config( + hidden_size=64, + intermediate_size=128, + quantization_config={ + "quant_method": "fp8", + "weight_block_size": [128, 128], + }, + ) + model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(config)) + source = mx.linspace(-1.0, 1.0, 16 * 64).reshape(16, 64) + fp8 = mx.to_fp8(source) + weight_key = "model.layers.0.attention.q_proj.weight" + + sanitized = model.sanitize( + { + weight_key: fp8, + f"{weight_key}_scale_inv": mx.array([[126]], dtype=mx.uint8), + } + ) + restored = mx.dequantize( + sanitized[weight_key], + sanitized[weight_key.replace("weight", "scales")], + sanitized[weight_key.replace("weight", "biases")], + group_size=64, + bits=8, + ) + expected = mx.from_fp8(fp8, dtype=mx.bfloat16) * 0.5 + mx.eval(restored, expected) + + assert mx.allclose(restored, expected, rtol=2e-2, atol=5e-3) + + def test_bailing_fp8_config_normalizes_to_affine_runtime_quantization(): from omlx.utils.model_loading import normalize_bailing_hybrid_fp8_quant @@ -508,6 +600,36 @@ def test_bailing_fp8_config_normalizes_to_affine_runtime_quantization(): assert config["quantization"] == {"group_size": 64, "bits": 8} +def test_bailing_mixed_fp4_config_adds_routed_expert_overrides(): + from omlx.utils.model_loading import normalize_bailing_hybrid_fp8_quant + + config = _minimal_config( + num_hidden_layers=3, + first_k_dense_replace=1, + quantization_config={ + "quant_method": "fp8", + "weight_block_size": [128, 128], + "routed_experts_quant_method": "mxfp4", + "routed_experts_group_size": 32, + }, + ) + + assert normalize_bailing_hybrid_fp8_quant(config) is config + quantization = config["quantization"] + assert quantization["group_size"] == 64 + assert quantization["bits"] == 8 + assert "model.layers.0.mlp.switch_mlp.gate_proj" not in quantization + expected = {"group_size": 32, "bits": 4, "mode": "mxfp4"} + for layer_idx in (1, 2): + for projection in ("gate_proj", "up_proj", "down_proj"): + assert ( + quantization[ + f"model.layers.{layer_idx}.mlp.switch_mlp.{projection}" + ] + == expected + ) + + def test_fp8_checkpoint_loads_strictly_as_quantized_model(tmp_path): bailing_hybrid = _load_patch_module() import mlx.nn as nn @@ -545,6 +667,66 @@ def test_fp8_checkpoint_loads_strictly_as_quantized_model(tmp_path): assert logits.shape == (1, 3, config["vocab_size"]) +def test_mixed_fp4_checkpoint_loads_strictly(tmp_path): + bailing_hybrid = _load_patch_module() + from mlx.utils import tree_flatten + from mlx_lm.models.switch_layers import QuantizedSwitchLinear + + config = _minimal_config( + hidden_size=64, + intermediate_size=128, + moe_intermediate_size=32, + quantization_config={ + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [128, 128], + "routed_experts_quant_method": "mxfp4", + "routed_experts_group_size": 32, + }, + ) + source_model = bailing_hybrid.Model(bailing_hybrid.ModelArgs.from_dict(config)) + weights = dict(tree_flatten(source_model.parameters())) + for projection in ("gate_proj", "up_proj", "down_proj"): + runtime_key = f"model.layers.1.mlp.switch_mlp.{projection}.weight" + expert_weights = weights.pop(runtime_key) + for expert, expert_weight in enumerate(expert_weights): + packed, scales = mx.quantize( + expert_weight, + group_size=32, + bits=4, + mode="mxfp4", + ) + checkpoint_prefix = ( + f"model.layers.1.mlp.experts.{expert}.{projection}" + ) + weights[f"{checkpoint_prefix}.weight"] = packed.view(mx.int8) + weights[f"{checkpoint_prefix}.weight_scale_inv"] = scales + + mx.save_safetensors(str(tmp_path / "model.safetensors"), weights) + (tmp_path / "config.json").write_text(json.dumps(config)) + + from mlx_lm.utils import load_model + + from omlx.utils.model_loading import maybe_apply_pre_load_patches + + maybe_apply_pre_load_patches(str(tmp_path)) + loaded, loaded_config = load_model(tmp_path, strict=True) + logits = loaded(mx.array([[1, 2, 3]], dtype=mx.int32)) + mx.eval(logits) + + quantization = loaded_config["quantization"] + expected = {"group_size": 32, "bits": 4, "mode": "mxfp4"} + assert ( + quantization["model.layers.1.mlp.switch_mlp.gate_proj"] == expected + ) + assert isinstance( + loaded.model.layers[1].mlp.switch_mlp.gate_proj, + QuantizedSwitchLinear, + ) + assert loaded.model.layers[1].mlp.switch_mlp.gate_proj.mode == "mxfp4" + assert logits.shape == (1, 3, config["vocab_size"]) + + def test_oq_discovers_ling_embeddings_and_hybrid_layer_masks(): bailing_hybrid = _load_patch_module() from omlx.oq import ( From e1acb0bcca9f292bb0bbe00100e5c4159f607f51 Mon Sep 17 00:00:00 2001 From: Maxim Mazurok Date: Thu, 6 Aug 2026 11:00:07 +1000 Subject: [PATCH 013/338] feat(vlm_mtp): enforce thinking budget inside MTP speculative decoding (#2456) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(vlm_mtp): enforce thinking budget inside MTP speculative decoding The #2399 fix routes any request carrying per-request logits processors away from vlm_mtp to BatchGenerator, making thinking_budget and vlm_mtp_enabled mutually exclusive. For the thinking budget specifically the fallback is unnecessary: mlx-vlm's verify walk already exposes a positioned sampling hook (sampler.sample_target(logprobs, row_ids, positions)), and two properties of the MTP round loop make processor application there exact — every emitted token is the target-side sample at its position, and a verify slot's sample is only used when the draft prefix it was conditioned on was accepted. MTPProcessingSampler threads snapshot-capable processors (today: ThinkingBudgetProcessor, which gains snapshot_state/restore_state) into that hook, checkpointing processor state per position and rewinding on draft rejection. The forced close composes with speculation naturally: a forced token mismatches the draft, is emitted through the normal rejection path, and costs ~1 token/round only for the 3-4 close tokens. Acceptance rates during normal decode are untouched because the budget processor is a no-op on logits until the boundary. Grammar and penalty processors still fall back to BatchGenerator (the gate is now capability-based); grammar in particular defeats speculation by construction. Settings/UI exclusivity is relaxed for thinking budget only. Non-English i18n strings for the two vlm_mtp conflict keys still mention thinking budget and need a follow-up retranslation. Co-Authored-By: Claude Fable 5 * fix(vlm_mtp): gate positioned sampling on the round-loop view, not the inner LM The eligibility check for routing budget-carrying requests through vlm_mtp probed speculative_logits_from_hidden on the inner language model. That is not the object mlx-vlm's round loop sees: for mRoPE adapters (Qwen VLMs) _VLMAdapterMTPProxy intentionally hides the inner model's speculative_* fast paths to preserve mRoPE position handling, so the gate passed while the verify walk fell back to plain vectorized sampling and silently dropped the thinking budget — the exact failure mode #2399 exists to prevent. vlm_mtp_positioned_sampling_available() now mirrors the proxy's visibility rules (adapter attr first; speculative_* blocked for mRoPE adapters; fall-through to the inner model otherwise) and the scheduler gate uses it, declining to BatchGenerator when the hook is not visible. An adapter that later grows an mRoPE-safe positioned hook re-opens the route automatically. Tests: an equivalence sweep pins the helper against the real proxy's attribute resolution for every (mrope, adapter-hook, lm-hook) combination, plus regressions for the reported case (inner LM has the hook, mRoPE adapter hides it -> decline) and the non-mRoPE pass-through. Co-Authored-By: Claude Fable 5 --------- Co-authored-by: Claude Fable 5 --- .../AppView/Screens/ModelSettingsScreen.swift | 10 +- .../ViewModels/ModelSettingsScreenVM.swift | 10 +- omlx/admin/i18n/en.json | 2 +- omlx/admin/static/js/dashboard.js | 3 +- .../dashboard/_modal_model_settings.html | 10 +- omlx/api/thinking.py | 49 ++ omlx/model_settings.py | 21 +- omlx/scheduler.py | 93 ++- omlx/speculative/processing_sampler.py | 216 +++++++ omlx/speculative/vlm_mtp.py | 37 ++ tests/test_model_settings.py | 11 +- tests/test_vlm_mtp_thinking_budget.py | 605 ++++++++++++++++++ 12 files changed, 1020 insertions(+), 47 deletions(-) create mode 100644 omlx/speculative/processing_sampler.py create mode 100644 tests/test_vlm_mtp_thinking_budget.py diff --git a/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift b/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift index f131dc97b5..bd9f05e571 100644 --- a/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift +++ b/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift @@ -714,11 +714,9 @@ private struct AdvancedTab: View { Row(label: String(localized: "settings.advanced.thinking_budget.label", defaultValue: "Thinking Budget", comment: "Row label for the thinking budget field"), - sublabel: vm.vlmMtpEnabled - ? vm.vlmMtpProcessorLockedReason - : String(localized: "settings.advanced.thinking_budget.sub", - defaultValue: "Limit thinking tokens for reasoning models. Forces end of thinking when exceeded.", - comment: "Sublabel for the thinking budget field")) { + sublabel: String(localized: "settings.advanced.thinking_budget.sub", + defaultValue: "Limit thinking tokens for reasoning models. Forces end of thinking when exceeded.", + comment: "Sublabel for the thinking budget field")) { HStack(spacing: 8) { if vm.thinkingBudgetEnabled { TextInput(text: vm.bindProfile($vm.thinkingBudgetTokens), @@ -727,8 +725,6 @@ private struct AdvancedTab: View { Toggle("", isOn: vm.bindProfile($vm.thinkingBudgetEnabled)) .labelsHidden().toggleStyle(.switch) } - .disabled(vm.vlmMtpEnabled) - .help(vm.vlmMtpEnabled ? vm.vlmMtpProcessorLockedReason : "") } Row(label: String(localized: "settings.advanced.tool_result_limit.label", defaultValue: "Limit Tool Result Tokens", diff --git a/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift b/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift index 600d69bb9a..060002c112 100644 --- a/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift +++ b/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift @@ -786,8 +786,8 @@ final class ModelSettingsScreenVM { } if vlmMtpProcessorConflict { return String(localized: "settings.vlm_mtp.conflict.processors", - defaultValue: "Unset repetition / presence penalty and thinking budget before enabling VLM MTP.", - comment: "Tooltip / sublabel shown when VLM MTP can't be enabled because penalty or thinking-budget settings are set") + defaultValue: "Unset repetition / presence penalty before enabling VLM MTP.", + comment: "Tooltip / sublabel shown when VLM MTP can't be enabled because penalty settings are set") } return nil } @@ -795,11 +795,13 @@ final class ModelSettingsScreenVM { /// Settings that materialize as per-request logits processors, which the /// vlm_mtp decode path cannot apply (#2399). Mirrors /// vlm_mtp_processor_conflicts() in model_settings.py; neutral values - /// (repetition 1.0, presence 0.0) do not conflict. + /// (repetition 1.0, presence 0.0) do not conflict. Thinking budget is + /// exempt: it is applied on the vlm_mtp path at verify time + /// (MTPProcessingSampler). var vlmMtpProcessorConflict: Bool { if let rep = Double(repetitionPenalty), rep != 1.0 { return true } if let pres = Double(presencePenalty), pres != 0.0 { return true } - return thinkingBudgetEnabled + return false } /// Sublabel / tooltip for the sampling rows locked while VLM MTP is on. diff --git a/omlx/admin/i18n/en.json b/omlx/admin/i18n/en.json index 0848e3bb0a..c8a39d8418 100644 --- a/omlx/admin/i18n/en.json +++ b/omlx/admin/i18n/en.json @@ -548,7 +548,7 @@ "modal.model_settings.vlm_mtp": "VLM MTP", "modal.model_settings.vlm_mtp_hint": "Speculative decoding via an external MTP drafter model (Gemma 4 assistant or Qwen MTP). ~1.3–1.6× single-request speedup; concurrent eligible requests fall back to standard batching.", "modal.model_settings.vlm_mtp_conflict": "Disable DFlash / SpecPrefill / MTP / TurboQuant KV first — VLM MTP owns the speculative decode path.", - "modal.model_settings.vlm_mtp_processor_conflict": "Unset repetition / presence penalty, thinking budget and guided grammar first — VLM MTP cannot apply per-request logits processors.", + "modal.model_settings.vlm_mtp_processor_conflict": "Unset repetition / presence penalty and guided grammar first — VLM MTP cannot apply these per-request logits processors. Thinking budget is supported.", "modal.model_settings.vlm_mtp_processor_locked": "Locked while VLM MTP is enabled — this setting needs per-request logits processors, which VLM MTP cannot apply.", "modal.model_settings.vlm_mtp_draft_model": "Drafter model", "modal.model_settings.vlm_mtp_draft_model_placeholder": "Select an assistant or MTP drafter…", diff --git a/omlx/admin/static/js/dashboard.js b/omlx/admin/static/js/dashboard.js index 6e1cf54484..2afb9e610b 100644 --- a/omlx/admin/static/js/dashboard.js +++ b/omlx/admin/static/js/dashboard.js @@ -1437,6 +1437,8 @@ // which the VLM MTP decode path cannot apply (#2399). Mirrors // vlm_mtp_processor_conflicts() in model_settings.py; neutral // values (repetition 1.0, presence 0.0) do not conflict. + // Thinking budget is exempt: it is applied on the vlm_mtp path + // at verify time (MTPProcessingSampler). vlmMtpProcessorConflict() { const ms = this.modelSettings; if (!ms) return false; @@ -1445,7 +1447,6 @@ const pres = num(ms.presence_penalty); return (rep !== null && rep !== 1.0) || (pres !== null && pres !== 0.0) - || !!ms.enableThinkingBudget || !!ms.guided_grammar_enabled; }, diff --git a/omlx/admin/templates/dashboard/_modal_model_settings.html b/omlx/admin/templates/dashboard/_modal_model_settings.html index 626182deb8..88c1cd1ed3 100644 --- a/omlx/admin/templates/dashboard/_modal_model_settings.html +++ b/omlx/admin/templates/dashboard/_modal_model_settings.html @@ -458,15 +458,9 @@

{{
{{ t('modal.model_settings.thinking_budget') }}

{{ t('modal.model_settings.thinking_budget_hint') }}

-

{{ t('modal.model_settings.vlm_mtp_processor_locked') }}

-
- Clear memory cache? + {{ t('status.runtime_cache.clear_memory_confirm') }}
- SSD + {{ t('status.runtime_cache.tier_ssd') }}
@@ -445,12 +445,12 @@

{{ t('status.head