From a614cf2b9d13bd8c33465454c026d717e750a194 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:44:46 -0700 Subject: [PATCH 01/17] feat(arch): muse-glimmer text decoder (vendored model, remap, synth, arch row, profiles) --- gmlx/arch_table.py | 5 + gmlx/config_synth.py | 66 +++++++++ gmlx/loader.py | 8 ++ gmlx/muse_glimmer_model.py | 276 +++++++++++++++++++++++++++++++++++++ gmlx/profiles.py | 24 ++++ gmlx/remap.py | 24 ++++ 6 files changed, 403 insertions(+) create mode 100644 gmlx/muse_glimmer_model.py diff --git a/gmlx/arch_table.py b/gmlx/arch_table.py index adbe582..5349b06 100644 --- a/gmlx/arch_table.py +++ b/gmlx/arch_table.py @@ -88,6 +88,7 @@ class ArchEntry: "hunyuan-moe": ("hunyuan", "Tencent Hunyuan-A13B; softmax-gated fine-grained MoE + per-layer shared expert, per-head qk-norm (named query/key_layernorm), NTK-alpha rope. NEOX rope (no qk-permute); router -> mlp.gate.wg, shared expert -> mlp.shared_mlp; the GGUF materializes k/v on every layer so use_cla=False; rope alpha defaults to 1.0 (folded into freq_base); top-k router scores renormalized at load (norm_topk_prob - upstream mlx-lm omits it and degenerates)"), "hy_v3": ("hunyuan", "Tencent Hy3 (299B-A21B, llama.cpp PR #25395); sigmoid-gated fine-grained MoE (192 experts top-8) with selection-only expert bias (exp_probs_b, stored suffix-less) + top-k renorm x expert_weights_scale + one ungated shared expert, single leading dense layer (derived from tensor presence - no KV), per-head qk-norm, plain NEOX rope theta 11.16M (no qk-permute). Native MTP/NextN block past the trunk (stripped from the trunk on load; drafts via HyV3MTPDrafter - single-depth head, block_size 2, GMLX_HY3_MTP_BLOCK raises it). Router gate + expert bias pinned fp32 (llama.cpp routes fp32). HF enable_lm_head_fp32 pinned off (llama.cpp, the parity oracle, also computes the head in compute dtype). Early community GGUFs with arch 'hy-v3' (dash) are not mapped - reconvert. Model class vendored from mlx-lm PR #1485 with the MTP hidden-state wiring fixed to the vLLM-verified post-final-norm form"), "kimi-k3": ("kimi", "Moonshot Kimi-K3 (2.8T-A50B, llama.cpp PR #26185); hybrid KDA linear attention (per-key-channel-decay delta rule via the fused gated_delta kernels, short conv1d Q/K/V, full-rank sigmoid output gate) + nope-only MLA every attn_res-th-ish layer (per-layer head_count_kv array, 0 = KDA; absorbed embed_q/unembed_out, sigmoid output gate, NO rope anywhere), cross-layer residual attention (softmax-scored convex mix over banked residual checkpoints, restart on bank layers), latent MoE (896 experts top-16 at routed_expert_hidden_size behind routed_down/routed_up, sigmoid gating + correction bias, shared experts at full width), and situ activation replacing SwiGLU everywhere. ssm_a arrives folded (-exp(A_log)) and stays folded. Native MXFP4 experts (gpt-oss codec). XTML chat format (<|open|>tag<|sep|>), tiktoken kimi-k2 pretokenizer, EOS 163586 <|end_of_msg|>. Hybrid cache is non-trimmable -> chat re-prefills on trim. Model class vendored (no upstream mlx-lm class; kimi_linear is the nearest relative)"), + "muse-glimmer": ("muse", "Meta Muse Glimmer 30B (llama.cpp 'muse-glimmer'); dense 52-layer sandwich-norm decoder - pre/post norms around both attention and FFN, at two epsilons (1e-5 pre/final, 1e-8 post) - with an afmoe-shaped attention output gate (sigmoid(x_norm @ W_gate) applied before o_proj), per-head qk-norm whose weights are synthesized at conversion to absorb qk_scale_factor (q_norm a uniform 3.87, k_norm ones), an unweighted RMSNorm on the token embeddings before layer 0, and a logit scale (output_multiplier 0.196) + gemma-style tanh softcap at 20. Attention runs a 3-of-4 sliding window (2048) with RoPE, and full attention with NoPE on every 4th layer - the inverse of the usual arrangement, and the reason the model card says '131072+': the largest positional offset ever resolved is the 2048 window, so there is no rope extrapolation ceiling and past 131k the only limit is KV memory at ~13 KB/token (just the 13 global layers grow; the 39 sliding ones are pinned by their rotating cache). The four per-layer norms carry a baked +1 that a plain RMSNorm consumes as-is (no gemma unbake). NORM rope with Q/K un-permuted at conversion => traditional=True, no qk_permute. Model class vendored (no upstream mlx-lm class; afmoe is the nearest relative). Pairs with the Muse Glimmer mmproj (--mmproj) for vision and with the DFlash drafter GGUF (--draft-gguf) for speculative decoding"), "granitehybrid": ("granite", "IBM Granite 4.x hybrid (H-Micro/H-Tiny/H-Small); alternating Mamba2 + attention (layer_types from per-layer head_count_kv==0), softmax MoE + fused-input shared MLP (loader pre-fuses ffn_{gate,up}_shexp -> input_linear), granite runtime multipliers, NoPE via rope.scaling.finetuned=false. NORM rope => qk_permute on attention layers"), "falcon-h1": ("falcon", "TII Falcon-H1 (0.5B-34B); parallel attention + Mamba2 in every layer (one input_layernorm feeds both, outputs summed), dense gated MLP under feed_forward.*, explicit head_dim from key_length. NEOX rope (no qk-permute). The muP multiplier zoo is folded into the wire weights at convert => synth pins every multiplier neutral; ffn_norm/ssm_a/ssm_d stored with no .weight suffix"), "qwen3next": ("qwen", "Qwen3-Next-80B-A3B; gated-DeltaNet linear attention (3 of every 4 layers) + gated full attention (gate fused in attn_q), every-layer 512-expert MoE + shared expert. NEOX rope (no qk-permute, not qwen35's packing). Both GDN wire layouts load: legacy fused ssm_in -> in_proj_qkvz; the newer split attn_qkv/attn_gate via a load-time module split (loader swaps in_proj_qkvz for in_proj_qkv/in_proj_z, skipping the runtime de-interleave). V heads HF-grouped => the qwen3.5 tiled-V patch is excluded; +1 norm bake is what mlx-lm expects (passthrough)"), @@ -157,6 +158,10 @@ def has_synth(gguf_arch: str) -> bool: # llama.cpp PR #26185 (Kimi-K3); no upstream mlx-lm class (kimi_linear is # the nearest relative and lacks the five K3-only mechanisms). "kimi_k3": "gmlx.kimi_k3_model", + # llama.cpp LLM_ARCH_MUSE_GLIMMER; no upstream mlx-lm class (afmoe is the + # nearest relative and has neither the NoPE/RoPE inversion nor the second + # norm epsilon). + "muse_glimmer": "gmlx.muse_glimmer_model", } diff --git a/gmlx/config_synth.py b/gmlx/config_synth.py index bd937f9..0c1ad39 100644 --- a/gmlx/config_synth.py +++ b/gmlx/config_synth.py @@ -128,6 +128,13 @@ # per-head qk-norm, plain rope. Native MTP/NextN block past the trunk. # Model class vendored from mlx-lm PR #1485. "hy_v3": "hy_v3", + # Meta Muse Glimmer 30B: dense sandwich-norm decoder with an attention output + # gate, per-head qk-norm absorbing qk_scale_factor, and RoPE on the + # sliding-window layers only (full-attention layers are NoPE). Logit scale + + # tanh softcap on the head. Model class vendored (no mlx-lm class; afmoe is + # the nearest relative). Pairs with the Muse Glimmer mmproj (--mmproj) and + # the DFlash drafter (--draft-gguf). + "muse-glimmer": "muse_glimmer", # Moonshot Kimi-K3 (2.8T-A50B): hybrid KDA (linear, per-channel decay) + # nope-only MLA layers from the per-layer head_count_kv schedule, latent # 896-expert sigmoid MoE behind down/up projections, situ activation, @@ -1639,6 +1646,64 @@ def _synth_hy_v3(meta, shapes, config: dict) -> None: config["rope_parameters"] = rope_parameters +# muse-glimmer (Meta Muse Glimmer 30B) + +def _synth_muse_glimmer(meta, shapes, config: dict) -> None: + """Synthesize a muse_glimmer config from a 'muse-glimmer'-arch GGUF. + + The universal fields cover hidden/layers/heads/kv/ffn/ctx/eps/head_dim/ + rope_theta/tie/vocab. This adds the per-layer sliding/full schedule, the + logit scale + softcap, and the second norm epsilon. + + ``post_norm_eps`` (post-attention and post-FFN norms only) is not in the + GGUF: llama.cpp hardcodes 1e-8 (muse-glimmer.cpp:68) and the HF + text_config carries the same value. Pinned here, in both those terms, so a + future variant that changes it is caught by the parity gate rather than + silently mis-normed. + """ + arch = "muse-glimmer" + n_layers = config["num_hidden_layers"] + + config["sliding_window"] = _require( + _read_int(meta, f"{arch}.attention.sliding_window"), + arch=arch, gguf_field=f"{arch}.attention.sliding_window") + + # Per-layer schedule: 1 = sliding (and rope'd), 0 = full (and NoPE). The KV + # is a per-layer bool array on every known conversion; llama.cpp also + # accepts a scalar period (set_swa_pattern: layer is full when + # (i + 1) % period == 0), so honour that form too. + key = f"{arch}.attention.sliding_window_pattern" + if _array_len(meta, key) == 1: + period = _require(_read_int(meta, key), arch=arch, gguf_field=key) + pattern = [(i + 1) % period != 0 for i in range(n_layers)] + else: + pattern = _read_bool_array(meta, key) + if pattern is None: + # llama.cpp's default period when the key is absent entirely. + pattern = [(i + 1) % 4 != 0 for i in range(n_layers)] + elif len(pattern) != n_layers: + raise ValueError( + f"muse-glimmer synth: {key} has {len(pattern)} entries for " + f"{n_layers} layers") + config["layer_types"] = [ + "sliding_attention" if v else "full_attention" for v in pattern] + + config["output_multiplier"] = _require( + _read_float(meta, f"{arch}.logit_scale"), + arch=arch, gguf_field=f"{arch}.logit_scale") + # llama.cpp reads the softcap as optional and leaves it 0 (disabled) when + # absent; the converter always writes 20.0. + config["final_logit_softcapping"] = ( + _read_float(meta, f"{arch}.final_logit_softcapping") or 0.0) + config["post_norm_eps"] = 1e-8 + + config["rope_parameters"] = { + "rope_theta": _require(config.get("rope_theta"), + arch=arch, gguf_field=f"{arch}.rope.freq_base"), + "rope_type": "default", + } + + # granitehybrid (IBM Granite 4.x hybrid: H-Micro / H-Tiny / H-Small) def _synth_granite_hybrid(meta, shapes, config: dict) -> None: @@ -2892,6 +2957,7 @@ def _print_summary(config: dict, arch: str) -> None: "minimax-m3": _synth_minimax_m3, "hunyuan-moe": _synth_hunyuan, "hy_v3": _synth_hy_v3, + "muse-glimmer": _synth_muse_glimmer, "granitehybrid": _synth_granite_hybrid, "falcon-h1": _synth_falcon_h1, "qwen3next": _synth_qwen3next, diff --git a/gmlx/loader.py b/gmlx/loader.py index bec74d9..ad39cbe 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -907,6 +907,14 @@ def build_model(config_dict: dict, *, mtp: bool = False): from . import kimi_k3_model kimi_k3_model.ensure_registered() + if mt == "muse_glimmer": + # mlx-lm ships no muse_glimmer module (afmoe is the nearest relative); + # same vendored-registration pattern as kimi_k3. The tool parser + # registers with the model so a later serve template-inference + # resolves it. + from . import muse_glimmer_model + + muse_glimmer_model.ensure_registered() Model, ModelArgs = _get_classes(config) model_args = ModelArgs.from_dict(config) model = Model(model_args) diff --git a/gmlx/muse_glimmer_model.py b/gmlx/muse_glimmer_model.py new file mode 100644 index 0000000..7ec149e --- /dev/null +++ b/gmlx/muse_glimmer_model.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Asher Feldman +# Portions copyright (c) 2024 Apple Inc. (mlx-lm afmoe skeleton, MIT) +"""Vendored mlx-lm-style model for Meta Muse Glimmer (GGUF arch ``muse-glimmer``). + +mlx-lm has no muse_glimmer class; this module is the runtime for llama.cpp's +``LLM_ARCH_MUSE_GLIMMER`` conversions, built from the pinned mlx-lm 0.31.3 +``afmoe`` skeleton - which already has the attention output gate, per-head +QK-norm, sandwich norms and the sliding/full ``layer_types`` split - with the +MoE stripped and the Glimmer-only mechanics added: + + 1. RoPE rides the sliding-window layers only; full-attention layers are NoPE. + That is the inverse of the usual arrangement, and the reason the declared + 131072 context has no extrapolation ceiling: the largest positional offset + ever resolved is the 2048 window. + 2. an unweighted RMSNorm on the token embeddings, before layer 0. + 3. two norm epsilons: ``rms_norm_eps`` (1e-5) on the pre-norms and the final + norm, ``post_norm_eps`` (1e-8) on the two post-norms. + 4. ``output_multiplier`` on the logits, then a gemma-style tanh softcap. + 5. Q/K arrive interleaved (llama.cpp tags the arch LLAMA_ROPE_TYPE_NORM and + the converter un-permutes HF's rotate_half layout), so rope runs + ``traditional=True`` rather than permuting the wire bytes on load. + +The four per-layer norm weights arrive with the +1 already folded in at +conversion, so they load as plain ``nn.RMSNorm`` weights with no gemma-style +unbake; the final norm is not baked. The QK-norm weights are synthesized at +conversion to absorb ``qk_scale_factor`` (q_norm is a uniform 3.87, k_norm is +ones), which is also why the interleaved Q/K layout is safe: a uniform per-head +norm is invariant under the rope permutation. +""" + +import importlib +import sys +from dataclasses import dataclass +from typing import Any, List, Optional + +import mlx.core as mx +import mlx.nn as nn + +from mlx_lm.models.activations import swiglu +from mlx_lm.models.base import ( + BaseModelArgs, + create_attention_mask, + scaled_dot_product_attention, +) +from mlx_lm.models.cache import KVCache, RotatingKVCache +from mlx_lm.models.rope_utils import initialize_rope + + +def ensure_registered() -> None: + """Make ``import mlx_lm.models.muse_glimmer`` resolve, preferring upstream.""" + if "mlx_lm.models.muse_glimmer" not in sys.modules: + try: + importlib.import_module("mlx_lm.models.muse_glimmer") # upstream wins + except ImportError: + sys.modules["mlx_lm.models.muse_glimmer"] = sys.modules[__name__] + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + head_dim: int + vocab_size: int + layer_types: List[str] + sliding_window: int = 2048 + rms_norm_eps: float = 1e-5 + # Post-attention / post-FFN norms only (llama.cpp muse-glimmer.cpp:68, + # HF text_config.post_norm_eps). Not carried in the GGUF. + post_norm_eps: float = 1e-8 + rope_theta: float = 500000.0 + rope_parameters: Optional[dict] = None + max_position_embeddings: int = 131072 + output_multiplier: float = 1.0 + final_logit_softcapping: float = 0.0 + tie_word_embeddings: bool = False + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs, use_sliding: bool): + super().__init__() + dim = args.hidden_size + self.n_heads = args.num_attention_heads + self.n_kv_heads = args.num_key_value_heads + self.head_dim = args.head_dim + self.use_sliding = use_sliding + self.scale = self.head_dim**-0.5 + + self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False) + self.gate_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False) + + self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + + # Interleaved Q/K wire layout => traditional rope. Full-attention + # layers carry no rope at all. + self.rope = ( + initialize_rope( + self.head_dim, + args.rope_theta, + True, + args.rope_parameters, + args.max_position_embeddings, + ) + if use_sliding + else None + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + + queries = self.q_proj(x).reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3) + keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + values = self.v_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + + queries = self.q_norm(queries) + keys = self.k_norm(keys) + + if self.rope is not None: + offset = cache.offset if cache is not None else 0 + queries = self.rope(queries, offset=offset) + keys = self.rope(keys, offset=offset) + + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + output = output * mx.sigmoid(self.gate_proj(x)) + return self.o_proj(output) + + +class MLP(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + dim, hidden = args.hidden_size, args.intermediate_size + self.gate_proj = nn.Linear(dim, hidden, bias=False) + self.up_proj = nn.Linear(dim, hidden, bias=False) + self.down_proj = nn.Linear(hidden, dim, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class DecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, use_sliding: bool): + super().__init__() + self.use_sliding = use_sliding + self.self_attn = Attention(args, use_sliding) + self.mlp = MLP(args) + + dim, eps, post_eps = args.hidden_size, args.rms_norm_eps, args.post_norm_eps + self.input_layernorm = nn.RMSNorm(dim, eps=eps) + self.post_attention_layernorm = nn.RMSNorm(dim, eps=post_eps) + self.pre_feedforward_layernorm = nn.RMSNorm(dim, eps=eps) + self.post_feedforward_layernorm = nn.RMSNorm(dim, eps=post_eps) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + self.post_attention_layernorm(r) + r = self.mlp(self.pre_feedforward_layernorm(h)) + return h + self.post_feedforward_layernorm(r) + + +class MuseGlimmerModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.sliding_window = args.sliding_window + self.layer_types = args.layer_types + + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + DecoderLayer(args, layer_type == "sliding_attention") + for layer_type in args.layer_types + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + self.fa_idx = next( + (i for i, t in enumerate(args.layer_types) if t != "sliding_attention"), + None, + ) + self.swa_idx = next( + (i for i, t in enumerate(args.layer_types) if t == "sliding_attention"), + None, + ) + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + capture_layers: Optional[tuple] = None, + ): + h = self.embed_tokens(inputs) + h = mx.fast.rms_norm(h, None, self.args.rms_norm_eps) + + if cache is None: + cache = [None] * len(self.layers) + + fa_mask = swa_mask = None + if self.fa_idx is not None: + fa_mask = create_attention_mask(h, cache[self.fa_idx]) + if self.swa_idx is not None: + swa_mask = create_attention_mask( + h, cache[self.swa_idx], window_size=self.sliding_window + ) + + captures = [] + cap_set = capture_layers or () + for idx, (layer, c) in enumerate(zip(self.layers, cache)): + h = layer(h, swa_mask if layer.use_sliding else fa_mask, cache=c) + if idx in cap_set: + captures.append(h) + + if capture_layers is not None: + return self.norm(h), captures + 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 = MuseGlimmerModel(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def head(self, h: mx.array) -> mx.array: + """Logits from a final-normed hidden state: lm_head, output multiplier, + tanh softcap. Computed in fp32 - llama.cpp's parity oracle scales and + softcaps an fp32 ``result_output``, and the softcap is nonlinear enough + that bf16 rounding moves argmax at depth.""" + if self.args.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(h) + else: + out = self.lm_head(h) + out = out.astype(mx.float32) * self.args.output_multiplier + cap = self.args.final_logit_softcapping + if cap: + out = mx.tanh(out / cap) * cap + return out + + def __call__(self, inputs: mx.array, cache: Optional[Any] = None): + return self.head(self.model(inputs, cache)) + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + return [ + RotatingKVCache(max_size=self.model.sliding_window) + if layer.use_sliding + else KVCache() + for layer in self.layers + ] diff --git a/gmlx/profiles.py b/gmlx/profiles.py index 4726fac..c1f62fb 100644 --- a/gmlx/profiles.py +++ b/gmlx/profiles.py @@ -185,6 +185,30 @@ "reasoning-max": {"chat_template_kwargs": {"thinking_effort": "max"}}, }, }, + # https://huggingface.co/meta-models/Muse-Glimmer-30B "Best Practices", + # 2026-08 (generation_config carries no sampling): t=1.0/top_p=0.95/top_k=64. + # The chat template's reasoning_strength takes low/medium/high/xhigh and + # defaults to high; it interpolates the value into the system prompt without + # validating it. The reasoning channel's markers are set so the server's + # open-think detection, budget criteria, and stream splitter see the real + # header rather than the '' default. + "muse": { + "label": "Muse Glimmer", + "arches": ("muse-glimmer",), + "base": {"sampling": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 64, + "thinking_start_token": "<|start|>assistant to=self<|message|>", + "thinking_end_token": "<|eom|>", + }}, + "intents": { + "reasoning-low": {"chat_template_kwargs": {"reasoning_strength": "low"}}, + "reasoning-medium": {"chat_template_kwargs": {"reasoning_strength": "medium"}}, + "reasoning-high": {"chat_template_kwargs": {"reasoning_strength": "high"}}, + "reasoning-xhigh": {"chat_template_kwargs": {"reasoning_strength": "xhigh"}}, + }, + }, # Llama 3.x generation_config (t=0.6/top_p=0.9); SmolLM3 card matches # closely enough to share. "llama": { diff --git a/gmlx/remap.py b/gmlx/remap.py index 0dcf500..0897bdb 100644 --- a/gmlx/remap.py +++ b/gmlx/remap.py @@ -164,6 +164,18 @@ # HF tensor; the split converter's de-interleave preserves group order), so # the qwen3.5 tiled-V patch must not fire - see _needs_tiled_v_patch. "qwen3next": "QWEN3NEXT", + # Meta Muse Glimmer (llama.cpp 'muse-glimmer'): a dense sandwich-norm decoder + # with an attention output gate and per-head qk-norm. Almost everything + # resolves canonically; see the MUSE_GLIMMER block for the two names that + # can't (the ffn_norm collision and attn_gate, which the canonical map homes + # on qwen3.5's linear_attn). llama.cpp tags the arch LLAMA_ROPE_TYPE_NORM and + # the converter un-permutes HF's rotate_half Q/K into the interleaved layout, + # so Q/K pass through un-permuted and the model runs rope traditional=True - + # equivalent to qk_permute + traditional=False, without copying every + # attn_q/attn_k off the wire. The four per-layer norms carry a baked +1 that + # a plain nn.RMSNorm consumes directly, so the arch stays out of + # _GEMMA_NORM_BAKED_ARCHS. + "muse-glimmer": "MUSE_GLIMMER", # DiffusionGemma (llama.cpp 'diffusion-gemma'): an encoder-decoder block- # diffusion model on the Gemma-4 MoE backbone. The decoder backbone uses the # exact Gemma-4 GGUF tensor names, but the mlx-vlm Model nests them under @@ -470,6 +482,18 @@ def _gemma_norm_transform(arch_string: str, hf_name: str | None, (re.compile(r"^blk\.(\d+)\.attn_v\.bias$"), "model.layers.{bid}.self_attn.v_proj.bias", "passthrough"), ], + "MUSE_GLIMMER": [ + # Muse Glimmer keeps a separate pre-FFN norm, so ffn_norm is the + # pre_feedforward_layernorm. Pin past the FFN_NORM/FFN_PRE_NORM + # collision; post_attention_norm and post_ffw_norm resolve canonically. + (re.compile(r"^blk\.(\d+)\.ffn_norm\.weight$"), + "model.layers.{bid}.pre_feedforward_layernorm.weight", "passthrough"), + # Attention output gate (sigmoid(x_norm @ W_gate) * attn_out, before + # o_proj). CANONICAL_HF homes ATTN_GATE on qwen3.5's linear_attn + # in_proj_z, so claim it here for the afmoe-shaped self_attn.gate_proj. + (re.compile(r"^blk\.(\d+)\.attn_gate\.weight$"), + "model.layers.{bid}.self_attn.gate_proj.weight", "passthrough"), + ], "ERNIE4_5_MOE": [ # Baidu ERNIE-4.5-MoE: shared-expert fine-grained MoE with leading dense # layers. mlx-lm's ernie4_5_moe uses traditional=True rope, which consumes From d1d7fa44721aa83710011faba67be05a0472fcf7 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:57:43 -0700 Subject: [PATCH 02/17] feat(chat): muse-glimmer ATEM reasoning channel, atem tool parser, reasoning_strength controls --- docs/server-config.md | 1 + gmlx/chat.py | 32 ++++++- gmlx/generation.py | 5 +- gmlx/loader.py | 3 +- gmlx/muse_glimmer_tools.py | 135 +++++++++++++++++++++++++++ gmlx/reasoning.py | 119 +++++++++++++++++++---- gmlx/server_patches/chat_behavior.py | 32 ++++--- gmlx/server_patches/render.py | 36 +++++++ gmlx/sessions.py | 14 ++- gmlx/thinking_budget.py | 8 +- gmlx/upstream_seams.py | 7 +- 11 files changed, 345 insertions(+), 47 deletions(-) create mode 100644 gmlx/muse_glimmer_tools.py diff --git a/docs/server-config.md b/docs/server-config.md index 5f692ac..edca07b 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -379,6 +379,7 @@ model fully resolved). Values are cited to the primary model cards in | `hunyuan` | `hunyuan-moe` | temperature=0.7 top_p=0.8 top_k=20 repetition_penalty=1.05 | - | | `hy3` | `hy_v3` | temperature=0.9 thinking_start_token= thinking_end_token= | `@reasoning-high`: temperature=0.9 thinking_start_token= thinking_end_token= reasoning_effort=high; `@reasoning-low`: temperature=0.9 thinking_start_token= thinking_end_token= reasoning_effort=low | | `kimi` | `kimi-k3` | temperature=1.0 top_p=0.95 thinking_start_token=<|open|>think<|sep|> thinking_end_token=<|close|>think<|sep|> | `@reasoning-high`: temperature=1.0 top_p=0.95 thinking_start_token=<|open|>think<|sep|> thinking_end_token=<|close|>think<|sep|> thinking_effort=high; `@reasoning-low`: temperature=1.0 top_p=0.95 thinking_start_token=<|open|>think<|sep|> thinking_end_token=<|close|>think<|sep|> thinking_effort=low; `@reasoning-max`: temperature=1.0 top_p=0.95 thinking_start_token=<|open|>think<|sep|> thinking_end_token=<|close|>think<|sep|> thinking_effort=max | +| `muse` | `muse-glimmer` | temperature=1.0 top_p=0.95 top_k=64 thinking_start_token=<|start|>assistant to=self<|message|> thinking_end_token=<|eom|> | `@reasoning-high`: temperature=1.0 top_p=0.95 top_k=64 thinking_start_token=<|start|>assistant to=self<|message|> thinking_end_token=<|eom|> reasoning_strength=high; `@reasoning-low`: temperature=1.0 top_p=0.95 top_k=64 thinking_start_token=<|start|>assistant to=self<|message|> thinking_end_token=<|eom|> reasoning_strength=low; `@reasoning-medium`: temperature=1.0 top_p=0.95 top_k=64 thinking_start_token=<|start|>assistant to=self<|message|> thinking_end_token=<|eom|> reasoning_strength=medium; `@reasoning-xhigh`: temperature=1.0 top_p=0.95 top_k=64 thinking_start_token=<|start|>assistant to=self<|message|> thinking_end_token=<|eom|> reasoning_strength=xhigh | | `llama` | `llama`, `smollm3` | temperature=0.6 top_p=0.9 | - | | `mistral` | `mistral3` | temperature=0.15 | - | | `default` | *(anything else)* | temperature=0.7 top_p=0.95 | `@coding`: temperature=0.3 top_p=0.95; `@creative`: temperature=1.0 top_p=0.95 min_p=0.05; `@instruct`: temperature=0.7 top_p=0.95 | diff --git a/gmlx/chat.py b/gmlx/chat.py index 10d7f3f..7ee5a38 100644 --- a/gmlx/chat.py +++ b/gmlx/chat.py @@ -426,6 +426,7 @@ class ChatState: last_stats: dict | None = None last_tps: float | None = None last_think_open: bool = False + last_header_open: bool = False # session bookkeeping session_stats: dict | None = None @@ -875,6 +876,7 @@ def _end_turn(state: ChatState, reply: str, canceled: bool, cache=None) -> None: "ts": now, "canceled": bool(canceled), "think_open": state.last_think_open, + "header_open": state.last_header_open, "stats": stats, }, "cache_before": cp.get("cache_before", 0), @@ -1007,11 +1009,12 @@ def _print_model_info(state: ChatState) -> None: print(f" adapter {info['adapter']}") -def _strip_thinking(text: str, start_in_thinking: bool = False) -> str: +def _strip_thinking(text: str, start_in_thinking: bool = False, + start_in_header: bool = False) -> str: """The answer portion of a raw reply (reasoning spans + markers removed).""" from .sessions import split_thinking - return split_thinking(text, start_in_thinking)[1] + return split_thinking(text, start_in_thinking, start_in_header)[1] def _session_doc(state: ChatState) -> dict: @@ -1069,7 +1072,8 @@ def _copy_last_answer(state: ChatState) -> None: print("[chat] nothing to copy yet") return a = transcript[-1]["assistant"] - text = _strip_thinking(a.get("content", ""), a.get("think_open", False)) + text = _strip_thinking(a.get("content", ""), a.get("think_open", False), + a.get("header_open", False)) if not text: print("[chat] the last reply has no answer text to copy") return @@ -1766,13 +1770,16 @@ def _stream_reply( state: ChatState, stops: list | None = None, start_in_thinking: bool = False, + start_in_header: bool = False, drafter=None, ) -> tuple[str, bool]: """Print a streaming reply; Esc or Ctrl-C cancels it (the session keeps running) and a ``stops`` sequence ends it cleanly (trimmed). Reasoning ("thinking") spans are stripped of their control markers and dimmed (or hidden) per ``state.reasoning`` - ``start_in_thinking`` seeds the case - where the chat template pre-opens ```` so only the close is streamed. + where the chat template pre-opens ```` so only the close is streamed, + and ``start_in_header`` the harmony/ATEM case where the prompt stops + mid-header at ``<|start|>assistant``. Returns ``(text_so_far, canceled)`` - the *raw* text (markers intact) so multi-turn history stays faithful - and records the reply's tok/s for the stat line + toolbar when it completes.""" @@ -1781,8 +1788,10 @@ def _stream_reply( scanner = StopScanner(stops) if stops else None state.last_think_open = bool(start_in_thinking) + state.last_header_open = bool(start_in_header) display = state.reasoning - rf = None if display == "raw" else ReasoningFilter(start_in_thinking=start_in_thinking) + rf = None if display == "raw" else ReasoningFilter( + start_in_thinking=start_in_thinking, start_in_header=start_in_header) theme = state.theme renderer = None if state.render in ("lite", "rich") and display != "raw": @@ -1895,6 +1904,14 @@ def _opens_thinking(prompt) -> bool: return prompt_opens_thinking(prompt) +def _opens_header(prompt) -> bool: + """Whether a rendered ``prompt`` stops inside a harmony/ATEM message header + (see ``reasoning.prompt_opens_header``).""" + from .reasoning import prompt_opens_header + + return prompt_opens_header(prompt) + + def _vlm_message( model_type: str, content: str, @@ -3460,6 +3477,7 @@ def _apply_session(doc: dict, name: str) -> None: state, stops=args.stop, start_in_thinking=_opens_thinking(prompt_text), + start_in_header=_opens_header(prompt_text), drafter=drafter, ) if reply: @@ -3537,6 +3555,7 @@ def _apply_session(doc: dict, name: str) -> None: state, stops=args.stop, start_in_thinking=_opens_thinking(prompt), + start_in_header=_opens_header(prompt), ) finally: clear_finish_key_target() @@ -3581,6 +3600,7 @@ def _apply_session(doc: dict, name: str) -> None: state, stops=args.stop, start_in_thinking=_opens_thinking(prompt_text), + start_in_header=_opens_header(prompt_text), drafter=drafter, ) _end_turn(state, reply, canceled, cache=cache) @@ -3615,6 +3635,7 @@ def _apply_session(doc: dict, name: str) -> None: state, stops=args.stop, start_in_thinking=_opens_thinking(prompt_text), + start_in_header=_opens_header(prompt_text), ) _end_turn(state, reply, canceled) continue @@ -3703,6 +3724,7 @@ def _apply_session(doc: dict, name: str) -> None: state, stops=args.stop, start_in_thinking=_opens_thinking(prompt_text), + start_in_header=_opens_header(prompt_text), ) finally: clear_finish_key_target() diff --git a/gmlx/generation.py b/gmlx/generation.py index 3767e01..9ed14d4 100644 --- a/gmlx/generation.py +++ b/gmlx/generation.py @@ -186,9 +186,10 @@ def _verbose_emitter(prompt, tokenizer, reasoning): echoing the prompt-opened think tag so the bare close marker still reads.""" open_tag = _echo_think_tag(prompt, tokenizer) if reasoning in ("show", "hide"): - from .reasoning import StreamRenderer + from .reasoning import StreamRenderer, prompt_opens_header - r = StreamRenderer(reasoning, start_in_thinking=open_tag is not None) + r = StreamRenderer(reasoning, start_in_thinking=open_tag is not None, + start_in_header=prompt_opens_header(prompt)) return r.write, r.close if open_tag is not None: print(open_tag, flush=True) diff --git a/gmlx/loader.py b/gmlx/loader.py index ad39cbe..98fa870 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -912,9 +912,10 @@ def build_model(config_dict: dict, *, mtp: bool = False): # same vendored-registration pattern as kimi_k3. The tool parser # registers with the model so a later serve template-inference # resolves it. - from . import muse_glimmer_model + from . import muse_glimmer_model, muse_glimmer_tools muse_glimmer_model.ensure_registered() + muse_glimmer_tools.ensure_registered() Model, ModelArgs = _get_classes(config) model_args = ModelArgs.from_dict(config) model = Model(model_args) diff --git a/gmlx/muse_glimmer_tools.py b/gmlx/muse_glimmer_tools.py new file mode 100644 index 0000000..ae97c0c --- /dev/null +++ b/gmlx/muse_glimmer_tools.py @@ -0,0 +1,135 @@ +"""Muse Glimmer tool-call parser (the Onyx ATEM XML format). + +Muse Glimmer's chat template emits Claude-shaped tool calls inside a wrapper +block, one ```` segment per call: + + + + {value} + ... + + + +mlx-vlm's parser registry knows none of this, and its template-marker +inference (``mlx_vlm.tool_parsers._TEMPLATE_MARKERS``) matches no ATEM tag. +``ensure_registered()`` grafts this module in as +``mlx_vlm.tool_parsers.muse_glimmer`` (upstream-first, same pattern as +``hy_v3_tools``) and prepends the ATEM markers so +``_infer_tool_parser_from_processor`` resolves it from the template. + +Parser contract (see mlx-vlm's ``tool_parsers/__init__.py`` consumers): module +attributes ``tool_call_start`` / ``tool_call_end`` delimit the call block in +the generated text; ``parse_tool_call(text, tools)`` receives the inner text +and returns ``{"name", "arguments"}`` or a list of them. The delimiters are +the *wrapper* tags: the server's streaming suppressor holds deltas only while +the tail is a prefix of ``tool_call_start``, so keying on the inner +``(.*?)', re.DOTALL +) +_PARAM_RE = re.compile( + r'(.*?)', re.DOTALL +) + +# The raw Jinja carries these literals verbatim (the tags are not built by +# .format(), unlike Hy3's). +_TEMPLATE_MARKERS = [ + ((" dict: + """The declared ``properties`` schema for ``tool_name``, or empty.""" + for tool in tools or (): + func = tool.get("function") if isinstance(tool, dict) else None + func = func or tool + if not isinstance(func, dict) or func.get("name") != tool_name: + continue + return (func.get("parameters") or {}).get("properties") or {} + return {} + + +def _coerce(value: str, spec: dict | None) -> Any: + """Invert the template's per-type serialization for one parameter. With no + schema entry the value stays literal text - guessing would turn a string + argument that happens to look numeric into a number.""" + declared = (spec or {}).get("type") + if declared == "string": + return value + stripped = value.strip() + if declared == "boolean": + return stripped == "true" + if declared in ("number", "integer"): + try: + return json.loads(stripped) + except ValueError: + return value + if declared in ("object", "array"): + try: + return json.loads(stripped) + except ValueError: + return value + if declared == "null" or stripped == "null": + return None + return value + + +def _parse_single(name: str, body: str, tools: list[Any] | None) -> dict: + """One ```` body -> ``{"name", "arguments"}``.""" + properties = _schema(name, tools) + arguments: dict[str, Any] = {} + for m in _PARAM_RE.finditer(body): + key = m.group(1) + arguments[key] = _coerce(m.group(2), properties.get(key)) + return {"name": name, "arguments": arguments} + + +def parse_tool_call(text: str, tools: list[Any] | None = None): + """Parse a wrapper block's inner text: a list of parsed calls when + ```` segments are present, else an unknown-call envelope + carrying the raw text (never a guess at freeform prose).""" + calls = _INVOKE_RE.findall(text) + if calls: + return [_parse_single(name, body, tools) for name, body in calls] + return {"name": "unknown", "arguments": {"raw": text.strip()}} + + +def ensure_registered() -> None: + """Make ``mlx_vlm.tool_parsers.muse_glimmer`` resolve (upstream wins) and + teach the template-marker inference the ATEM spellings. Idempotent.""" + import importlib + + if "mlx_vlm.tool_parsers.muse_glimmer" not in sys.modules: + try: + importlib.import_module("mlx_vlm.tool_parsers.muse_glimmer") + except ImportError: + sys.modules["mlx_vlm.tool_parsers.muse_glimmer"] = sys.modules[__name__] + try: + registry = importlib.import_module("mlx_vlm.tool_parsers") + except ImportError: + return + markers = getattr(registry, "_TEMPLATE_MARKERS", None) + if isinstance(markers, list): + for entry in _TEMPLATE_MARKERS: + if entry not in markers: + markers.insert(0, entry) diff --git a/gmlx/reasoning.py b/gmlx/reasoning.py index 6c72277..044aa3b 100644 --- a/gmlx/reasoning.py +++ b/gmlx/reasoning.py @@ -12,6 +12,11 @@ * harmony models (gpt-oss): ``<|channel|>analysis<|message|>`` ... ``<|end|>`` for reasoning, ``<|start|>assistant<|channel|>final<|message|>`` ... for the answer (``commentary`` channels carry tool preludes - treated as reasoning). + * Onyx ATEM (Muse Glimmer): the routing is a recipient in the message header, + ``<|start|>assistant to=self<|message|>`` ... ``<|eom|>`` for reasoning and + ``to=user`` (or a tool name) for everything else. The prompt ends at + ``<|start|>assistant``, so the first header arrives without its opener - + seed with ``start_in_header=True``. * Gemma-style ``<|channel>thought`` ... ```` (as detokenized). ``ReasoningFilter`` is a streaming state machine that strips the markers and @@ -57,6 +62,13 @@ ("<|start|>", _DROP), ("<|call|>", _DROP), ("<|end|>", _DROP), + # Onyx ATEM (Muse Glimmer): no channel marker - the recipient lives in the + # header that "<|start|>assistant" opens and "<|message|>" closes, so the + # routing is decided in _close_header. "<|eom|>" ends one message of a + # multi-message turn; the next header re-decides, and until it arrives the + # safe assumption is answer. + ("<|eom|>", _ANSWER), + ("<|eot|>", _DROP), # gemma-style channel (as detokenized - note the lopsided pipes). ("<|channel>thought", _REASON), ("", _ANSWER), @@ -85,6 +97,15 @@ ("<|end_of_msg|>", _DROP), ) +# Markers that open a message header - text from here to the closing +# "<|message|>" is routing metadata, never a display span. Harmony puts the +# channel in the marker itself; ATEM puts the recipient in the header body. +_HEADER_PREFIXES = ("<|channel|>", "<|start|>assistant") + +# ATEM's reasoning recipient. "self" is the only recipient that is not +# user-visible; "user" and tool namespaces are answer-side. +_SELF_RECIPIENT = "to=self" + class ReasoningFilter: """Strip reasoning control markers from a token stream and tag the rest. @@ -96,12 +117,16 @@ class ReasoningFilter: remains at end-of-stream (a partial marker there is just literal text). """ - def __init__(self, *, start_in_thinking: bool = False): + def __init__(self, *, start_in_thinking: bool = False, + start_in_header: bool = False): self._markers = sorted(_MARKERS, key=lambda m: len(m[0]), reverse=True) self.mode = _REASON if start_in_thinking else _ANSWER self.buf = "" - self._swallow = False # inside a harmony channel header (drop text) + self._swallow = False # inside a message header (drop text) self._swallow_budget = 0 + self._header = "" + if start_in_header: + self._open_header() def feed(self, text: str) -> list[tuple[str, str]]: self.buf += text @@ -137,26 +162,43 @@ def _consume(self, *, final: bool) -> list[tuple[str, str]]: self.buf = self.buf[len(marker):] if action != _DROP: self.mode = action - # Harmony headers run "<|channel|>NAME [annotations]<|message|>": - # the routing marker opens the header, "<|message|>" closes it, - # and annotation text in between is never a display span. - if marker.startswith("<|channel|>"): - self._swallow = True - self._swallow_budget = 256 + # Message headers run "<|channel|>NAME [annotations]<|message|>" + # (harmony) or "<|start|>assistant to=RECIPIENT<|message|>" (ATEM): + # an opener marker starts the header, "<|message|>" closes it, and + # the text in between is routing, never a display span. + if marker.startswith(_HEADER_PREFIXES): + self._open_header() elif marker == "<|message|>": - self._swallow = False + self._close_header() return [s for s in spans if s[0]] + def _open_header(self) -> None: + self._swallow = True + self._swallow_budget = 256 + self._header = "" + + def _close_header(self) -> None: + """End the header and apply its routing. Only an explicit ``to=self`` + recipient moves the mode: a harmony header is empty (its channel marker + already routed), and an ATEM answer/tool header must leave a mode that + "<|eom|>" or the initial state already set.""" + self._swallow = False + if _SELF_RECIPIENT in self._header: + self.mode = _REASON + self._header = "" + def _emit(self, text: str, spans: list[tuple[str, str]]) -> None: - """Append a display span, unless a channel header is being swallowed. + """Append a display span, unless a message header is being swallowed. The budget bounds the swallow: a literal "<|channel|>" in a non-harmony reply (no "<|message|>" ever follows) must not eat the rest of the message.""" if self._swallow: self._swallow_budget -= len(text) if self._swallow_budget >= 0: + self._header += text return self._swallow = False + self._header = "" spans.append((text, self.mode)) def _next_marker_pos(self, final: bool) -> int | None: @@ -202,8 +244,19 @@ def _could_extend(self) -> bool: return False -def split_harmony_reply(text: str) -> tuple[str | None, str]: - """Split a complete harmony (gpt-oss) reply into ``(reasoning, content)``. +def prompt_opens_header(prompt) -> bool: + """Whether a rendered ``prompt`` stops inside a message header, so the + filter must start mid-header. Both harmony and ATEM generation prompts end + at ``<|start|>assistant`` and leave the channel or recipient to the model. + Tolerant of token-id prompts (False for non-strings).""" + return (isinstance(prompt, str) + and prompt.rstrip().endswith("<|start|>assistant")) + + +def split_harmony_reply(text: str, *, + start_in_header: bool = False) -> tuple[str | None, str]: + """Split a complete harmony (gpt-oss) or ATEM reply into + ``(reasoning, content)``. The serve path's stock splitter knows none of the harmony markers, so it returns the raw channel markup as content - which the model's own chat @@ -213,8 +266,12 @@ def split_harmony_reply(text: str) -> tuple[str | None, str]: commentary channels (tool preludes included) become reasoning, the final channel becomes content, and a reply truncated inside analysis returns all-reasoning with empty content (the convention the truncated-thinking - handling already uses for think-tag models).""" - filt = ReasoningFilter() + handling already uses for think-tag models). + + ``start_in_header`` seeds the ATEM case, where the generation prompt ends + mid-header at ``<|start|>assistant`` and the reply opens with `` + to=self<|message|>``.""" + filt = ReasoningFilter(start_in_header=start_in_header) spans = filt.feed(text) spans += filt.flush() reasoning = "".join(t for t, m in spans if m == _REASON).strip() @@ -420,11 +477,14 @@ class StreamRenderer: generate paths stream a thinking model through the same show/hide styling the chat REPL uses. ``start_in_thinking`` seeds the pre-fill-template case (the prompt opened the think block, so the stream carries only the close - marker).""" + marker); ``start_in_header`` the harmony/ATEM case (the prompt stopped + mid-header).""" def __init__(self, display: str = "show", *, - start_in_thinking: bool = False, color: bool | None = None): - self._filter = ReasoningFilter(start_in_thinking=start_in_thinking) + start_in_thinking: bool = False, + start_in_header: bool = False, color: bool | None = None): + self._filter = ReasoningFilter(start_in_thinking=start_in_thinking, + start_in_header=start_in_header) self._printer = ReasoningPrinter( display=display, color=want_color() if color is None else color) @@ -480,6 +540,23 @@ def normalize_template_kwargs(kwargs: dict) -> dict: True: "on", False: "off"} +# Templates that grade reasoning depth under a name of their own. The control +# is spelled reasoning_effort throughout gmlx, so map it onto the template's +# spelling rather than emitting a variable the template ignores (Muse Glimmer's +# ATEM template reads reasoning_strength). +_EFFORT_ALIASES = ("reasoning_strength",) + + +def _effort_variable(template: str) -> str: + """The reasoning-depth variable ``template`` reads, defaulting to the + canonical ``reasoning_effort``.""" + if template and "reasoning_effort" not in template: + for alias in _EFFORT_ALIASES: + if alias in template: + return alias + return "reasoning_effort" + + def map_thinking_controls(base: dict, thinking=None, reasoning_effort=None, template: str = "", warn=None) -> dict: """Overlay the dedicated thinking controls onto ``base`` template kwargs, @@ -503,8 +580,9 @@ def map_thinking_controls(base: dict, thinking=None, reasoning_effort=None, _warn = warn or (lambda msg: None) out = dict(base) if reasoning_effort is not None: - out["reasoning_effort"] = reasoning_effort - if template and "reasoning_effort" not in template: + name = _effort_variable(template) + out[name] = reasoning_effort + if template and name not in template: _warn("this model's chat template has no reasoning_effort " "variable; reasoning_effort is likely a no-op") if thinking is None: @@ -532,6 +610,9 @@ def map_thinking_controls(base: dict, thinking=None, reasoning_effort=None, elif "reasoning_effort" in template: _warn("this model has no thinking switch (reasoning always runs); " "use reasoning_effort low|medium|high to size it") + elif any(alias in template for alias in _EFFORT_ALIASES): + _warn("this model has no thinking switch (reasoning always runs); " + "use reasoning_effort to size it") else: _warn("this model's chat template has no thinking switch; the " "thinking control is likely a no-op") diff --git a/gmlx/server_patches/chat_behavior.py b/gmlx/server_patches/chat_behavior.py index dbd1780..3f792e6 100644 --- a/gmlx/server_patches/chat_behavior.py +++ b/gmlx/server_patches/chat_behavior.py @@ -296,11 +296,13 @@ def __init__(self, *args, **kw): self.in_thinking = _prompt_tail_opens_thinking( prompt, self.open_close_markers) if prompt.rstrip().endswith("<|start|>assistant"): - # harmony (gpt-oss): the state machine's open/close - # pairs cannot express channel routing, so the stream - # splits through the REPL's marker filter instead. + # harmony (gpt-oss) and ATEM (Muse Glimmer): the state + # machine's open/close pairs cannot express channel or + # recipient routing, so the stream splits through the + # REPL's marker filter instead. Both prompts stop + # mid-header, so the filter starts inside one. from ..reasoning import ReasoningFilter - self._kq_harmony = ReasoningFilter() + self._kq_harmony = ReasoningFilter(start_in_header=True) self._kq_harmony_closed = False __init__.__dict__[_STREAM_SEED_FLAG] = True @@ -353,15 +355,21 @@ def apply_chat_template(*a, **kw): if split is not None and not getattr(split, _STREAM_SEED_FLAG, False): def _split_thinking_text(text, thinking_start_token=None, thinking_end_token=None): - if text and "<|channel|>" in text: - # harmony (gpt-oss): the stock splitter knows none of these - # markers and returns the raw markup as content, which the - # model's own chat template rejects with a 500 once a client - # sends the reply back as history. (Gemma's lopsided - # "<|channel>thought" lacks the closing pipe, so this gate - # cannot misfire on it.) + rendered = _LAST_RENDERED_PROMPT.get() + in_header = bool(rendered) and rendered.rstrip().endswith( + "<|start|>assistant") + if text and ("<|channel|>" in text or (in_header + and "<|message|>" in text)): + # harmony (gpt-oss) and ATEM (Muse Glimmer): the stock splitter + # knows none of these markers and returns the raw markup as + # content, which the model's own chat template rejects with a + # 500 once a client sends the reply back as history. ATEM emits + # no channel marker at all, so it is recognised by the + # mid-header prompt tail plus a header close in the reply. + # (Gemma's lopsided "<|channel>thought" lacks the closing pipe, + # so this gate cannot misfire on it.) from ..reasoning import split_harmony_reply - return split_harmony_reply(text) + return split_harmony_reply(text, start_in_header=in_header) reasoning, content = split( text, thinking_start_token, thinking_end_token) if reasoning is None and content and retire_key.truncated_thinking( diff --git a/gmlx/server_patches/render.py b/gmlx/server_patches/render.py index a343bb0..35b9063 100644 --- a/gmlx/server_patches/render.py +++ b/gmlx/server_patches/render.py @@ -26,6 +26,7 @@ from __future__ import annotations import importlib +import json import os # Stock tail: these model types return only the last message when a @@ -33,6 +34,39 @@ _LAST_MESSAGE_ONLY = ("paligemma", "molmo", "florence2", "falcon_ocr") +def _decode_tool_arguments(msgs: list) -> None: + """Decode ``tool_calls[*].function.arguments`` from the OpenAI wire form + (a JSON string) into a mapping, in place. + + Muse Glimmer's ATEM template raises outright on a string - "a JSON string + cannot be parsed in the HF jinja sandbox" - so a client replaying its own + tool turn would 500 on render. Undecodable arguments become an empty + mapping, the same fallback the retirement mirror uses + (``retire_key.assistant_message_from_reply``). + """ + for i, m in enumerate(msgs): + calls = m.get("tool_calls") if isinstance(m, dict) else None + if not isinstance(calls, list): + continue + rebuilt, changed = [], False + for call in calls: + fn = call.get("function") if isinstance(call, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if not isinstance(args, str): + rebuilt.append(call) + continue + try: + decoded = json.loads(args) + except (json.JSONDecodeError, TypeError): + decoded = {} + if not isinstance(decoded, dict): + decoded = {} + rebuilt.append({**call, "function": {**fn, "arguments": decoded}}) + changed = True + if changed: + msgs[i] = {**m, "tool_calls": rebuilt} + + def install_faithful_history() -> None: """Wrap every captured ``apply_chat_template`` binding with the key merge: the openai and anthropic protocol modules plus @@ -80,6 +114,8 @@ def apply_chat_template(processor, config, prompt, if (isinstance(m, dict) and "thinking" not in m and m.get("reasoning_content")): msgs[i] = {**m, "thinking": m["reasoning_content"]} + elif cfg.get("model_type") == "muse_glimmer": + _decode_tool_arguments(msgs) if return_messages: return msgs if cfg.get("model_type") in _LAST_MESSAGE_ONLY: diff --git a/gmlx/sessions.py b/gmlx/sessions.py index 9d4a932..6a29586 100644 --- a/gmlx/sessions.py +++ b/gmlx/sessions.py @@ -125,11 +125,16 @@ def latest_for_model(model_path: str) -> str | None: return None -def split_thinking(text: str, think_open: bool = False) -> tuple[str, str]: - """``(reasoning, answer)`` portions of a raw reply (markers dropped).""" +def split_thinking(text: str, think_open: bool = False, + header_open: bool = False) -> tuple[str, str]: + """``(reasoning, answer)`` portions of a raw reply (markers dropped). + + ``think_open`` seeds a prompt-opened ```` block, ``header_open`` a + harmony/ATEM reply whose prompt stopped mid-header.""" from .reasoning import ReasoningFilter - rf = ReasoningFilter(start_in_thinking=think_open) + rf = ReasoningFilter(start_in_thinking=think_open, + start_in_header=header_open) spans = rf.feed(text) + rf.flush() reason = "".join(t for t, m in spans if m == "reason").strip() answer = "".join(t for t, m in spans if m == "answer").strip() @@ -165,7 +170,8 @@ def export_markdown(doc: dict, path: str) -> Path: canceled = " *(canceled)*" if m.get("canceled") else "" lines += [f"## Assistant{ts}{canceled}", ""] reason, answer = split_thinking( - m.get("content", ""), m.get("think_open", False) + m.get("content", ""), m.get("think_open", False), + m.get("header_open", False) ) if reason: lines += [ diff --git a/gmlx/thinking_budget.py b/gmlx/thinking_budget.py index e634c57..90797c3 100644 --- a/gmlx/thinking_budget.py +++ b/gmlx/thinking_budget.py @@ -41,7 +41,13 @@ ("", ""), # Kimi-K3 XTML sections (multi-token markers; the prompt # pre-opens the think section). - ("<|open|>think<|sep|>", "<|close|>think<|sep|>")) + ("<|open|>think<|sep|>", "<|close|>think<|sep|>"), + # Onyx ATEM (Muse Glimmer): the reasoning message is addressed + # to self and ends at <|eom|>. Both literals appear verbatim in + # the template source. The generation prompt stops one marker + # short (at "<|start|>assistant"), so the block is not + # prompt-opened and the model emits the whole header itself. + ("<|start|>assistant to=self<|message|>", "<|eom|>")) # Forced into the thinking block ahead of the first budget-triggered close. # The model must see itself DECIDE to answer: a bare close tag cuts the diff --git a/gmlx/upstream_seams.py b/gmlx/upstream_seams.py index c0cf45c..b2d978d 100644 --- a/gmlx/upstream_seams.py +++ b/gmlx/upstream_seams.py @@ -241,11 +241,12 @@ class Seam: "server_patches (HF download gate)", critical=True), Seam("mlx_vlm.utils", "StoppingCriteria.__call__", "server_patches (ignore-EOS)"), - # --- tool-parser registry (hy_v3_tools.ensure_registered) --- + # --- tool-parser registry (hy_v3_tools / muse_glimmer_tools) --- Seam("mlx_vlm.tool_parsers", "_TEMPLATE_MARKERS", - "hy_v3_tools.ensure_registered (Hy3 marker prepend)"), + "hy_v3_tools / muse_glimmer_tools ensure_registered (marker prepend)"), Seam("mlx_vlm.tool_parsers", "load_tool_module", - "hy_v3_tools (sys.modules graft resolves through it)"), + "hy_v3_tools / muse_glimmer_tools (sys.modules graft resolves " + "through it)"), # --- APC internals (lone-harvest patch, gmlx manager subclass, apc_pooling) --- Seam("mlx_vlm.apc", "harvest_blocks_from_batch_cache", "server_patches.install_apc_lone_harvest", critical=True), From f02da125a6f5b18390b9f1f4babd5c0a7be03f1f Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:24:33 -0700 Subject: [PATCH 03/17] fix(attn): chunked prefill must not causal-mask an unmasked bidirectional call --- gmlx/attn_hd512.py | 20 +++++++++++-------- tests/test_attn_hd512_verify.py | 35 +++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 8 deletions(-) diff --git a/gmlx/attn_hd512.py b/gmlx/attn_hd512.py index e1a63c6..8311901 100644 --- a/gmlx/attn_hd512.py +++ b/gmlx/attn_hd512.py @@ -205,6 +205,9 @@ def _prefill_eligible(q, k, v, mask): qL = q.shape[2] if qL <= 8 or qL > k.shape[2]: return False + # None (a bidirectional encoder's full attention) and "causal" both tile by + # query rows; _chunked_prefill keeps them apart, since only causal may also + # slice keys to the tile horizon. if _causal_str(mask): return True # an array mask we can slice along the query axis (sliding-window layers) @@ -217,17 +220,18 @@ def _chunked_prefill(q, k, v, scale, mask, tile, sinks=None): under attention sinks, a per-head additive logit in every row's softmax denominator regardless of key slicing. Causal: also slice keys to the tile horizon (mask='causal' keeps it causal). Array mask (sliding): slice the - mask's query rows, keep all keys. + mask's query rows, keep all keys. ``mask=None`` is *unmasked*, not causal: + a bidirectional encoder (a vision tower) attends to every key from every + query row, so its keys are never sliced. Each tile is eval'd before the next so its [Hq, tile, kL] score is freed instead of accumulating across tiles (and layers) in one lazy graph -- without this, peak memory stays as high as the full materialization and still swaps.""" skw = {} if sinks is None else {"sinks": sinks} qL = q.shape[2] - arr = isinstance(mask, mx.array) + causal = isinstance(mask, str) and mask == "causal" if qL <= tile: - return _orig_sdpa(q, k, v, scale=scale, - mask=(mask if arr else "causal"), **skw) + return _orig_sdpa(q, k, v, scale=scale, mask=mask, **skw) # With a cached prefix (kL > qL, chunk 2+ of a chunked prefill) the causal # horizon of query row t is offset + t, not t: slicing keys to t1 would # select only the head of the cached prefix and drop the chunk's own keys. @@ -236,13 +240,13 @@ def _chunked_prefill(q, k, v, scale, mask, tile, sinks=None): for t0 in range(0, qL, tile): t1 = min(t0 + tile, qL) qt = q[:, :, t0:t1, :] - if arr: - ot = _orig_sdpa(qt, k, v, scale=scale, mask=mask[..., t0:t1, :], - **skw) - else: + if causal: ot = _orig_sdpa(qt, k[:, :, :offset + t1, :], v[:, :, :offset + t1, :], scale=scale, mask="causal", **skw) + else: + sliced = mask[..., t0:t1, :] if isinstance(mask, mx.array) else None + ot = _orig_sdpa(qt, k, v, scale=scale, mask=sliced, **skw) mx.eval(ot) outs.append(ot) return mx.concatenate(outs, axis=2) diff --git a/tests/test_attn_hd512_verify.py b/tests/test_attn_hd512_verify.py index 1fea9ae..5f026e0 100644 --- a/tests/test_attn_hd512_verify.py +++ b/tests/test_attn_hd512_verify.py @@ -47,6 +47,41 @@ def test_chunked_prefill_causal_with_cached_prefix(qL, kL, monkeypatch): assert err < 2e-2, f"qL={qL} kL={kL} err={err}" +@pytest.mark.parametrize("d", [96, 512]) +@pytest.mark.parametrize("qL", [32, 96]) +def test_chunked_prefill_unmasked_stays_unmasked(d, qL, monkeypatch): + # mask=None is *unmasked*, not causal. A bidirectional encoder (the + # muse-glimmer ViT, hd 96) attends to every key from every query row; + # treating None as "causal" silently halved its receptive field. + monkeypatch.setattr( + attn_hd512, "_orig_sdpa", mx.fast.scaled_dot_product_attention) + scale = d**-0.5 + q, k, v = _rand(qL, kL=qL, hq=16, hkv=16, d=d) + out = attn_hd512._chunked_prefill(q, k, v, scale, None, tile=32) + err = mx.abs(out.astype(mx.float32) + - _ref(q, k, v, False, scale=scale)).max().item() + assert err < 2e-2, f"d={d} qL={qL} err={err}" + if qL > 32: + causal_err = mx.abs(out.astype(mx.float32) + - _ref(q, k, v, True, scale=scale)).max().item() + assert causal_err > 1e-2, "unmasked output collapsed onto the causal one" + + +def test_chunked_prefill_block_diagonal_mask(monkeypatch): + # the ViT's window attention: a non-causal array mask, sliced per tile + monkeypatch.setattr( + attn_hd512, "_orig_sdpa", mx.fast.scaled_dot_product_attention) + qL, d = 96, 96 + q, k, v = _rand(qL, kL=qL, hq=16, hkv=16, d=d) + seg = mx.arange(qL) // 32 + mask = (seg[:, None] == seg[None, :])[None, None] + out = attn_hd512._chunked_prefill(q, k, v, d**-0.5, mask, tile=32) + ref = mx.fast.scaled_dot_product_attention( + q, k, v, scale=d**-0.5, mask=mask) + err = mx.abs(out.astype(mx.float32) - ref.astype(mx.float32)).max().item() + assert err < 2e-2, f"err={err}" + + @pytest.mark.parametrize("qL", [3, 4, 6]) @pytest.mark.parametrize("causal", [True, False]) def test_verify_gemm_matches_reference(qL, causal): From 89b1f725530e5d14d91831492bb38587fdab305e Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:24:33 -0700 Subject: [PATCH 04/17] test(long-context): add muse-glimmer to the parity and decode-integrity gates --- tests/test_long_context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_long_context.py b/tests/test_long_context.py index f4aa506..ef1356d 100644 --- a/tests/test_long_context.py +++ b/tests/test_long_context.py @@ -43,7 +43,7 @@ "nemotron_h_moe", "deepseek2", "mixtral", "glm4moe", "gpt-oss", "seed_oss", "smollm3", "granite", "ernie4_5-moe", "minimax-m2", "minimax-m3", "hunyuan-moe", "granitehybrid", "falcon-h1", "qwen3next", "hy_v3", - "kimi-k3", + "kimi-k3", "muse-glimmer", ] TARGET = int(os.environ.get("KQUANT_LONGCTX_TOKENS", "16384")) From 0d56cef63c960ced19fe3cff725afa950e9f129e Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:24:33 -0700 Subject: [PATCH 05/17] feat(vlm): muse-glimmer vision tower, remap, config synth, and image processor --- LICENSE | 12 +- gmlx/loader.py | 6 + gmlx/muse_glimmer_model.py | 27 +- gmlx/muse_glimmer_vlm_model.py | 449 +++++++++++++++++++++++++++++++++ gmlx/upstream_seams.py | 32 +++ gmlx/vlm.py | 351 +++++++++++++++++++++++++- pyproject.toml | 2 +- 7 files changed, 862 insertions(+), 17 deletions(-) create mode 100644 gmlx/muse_glimmer_vlm_model.py diff --git a/LICENSE b/LICENSE index a37ee69..53927ba 100644 --- a/LICENSE +++ b/LICENSE @@ -27,9 +27,9 @@ please contact asher_feldman@icloud.com. File-level exception: files that carry an "SPDX-License-Identifier: MIT" header (the DSpark draft-model module and the DSpark sidecar converter, the -mlx-lm-style model modules for kimi-k3, minimax-m3, and hy-v3, and the -kimi-k3 tests) are licensed under the MIT license (see LICENSE-MIT), not -under this License. +mlx-lm-style model modules for kimi-k3, minimax-m3, hy-v3, and muse-glimmer, +the mlx-vlm-style vision module for muse-glimmer, and the kimi-k3 tests) are +licensed under the MIT license (see LICENSE-MIT), not under this License. Notice @@ -110,8 +110,10 @@ was taken: - gmlx/server_patches/apc.py and gmlx/thinking_budget.py carry code copied or ported from mlx-vlm (MIT, Copyright © Prince Canuma; licenses/mlx-vlm-LICENSE). -- The pre-tokenizer split regex patterns in gmlx/tokenizer.py come from - llama.cpp (MIT; licenses/llama.cpp-LICENSE). +- The pre-tokenizer split regex patterns in gmlx/tokenizer.py, and the + Muse Glimmer vision graph, grid index math, and image preprocessing in + gmlx/muse_glimmer_vlm_model.py and gmlx/vlm.py, come from llama.cpp + (MIT; licenses/llama.cpp-LICENSE). The test asset tests/e2e/assets/cats.jpg is reproduced from the mlx-vlm repository (MIT License, Copyright © Prince Canuma); see diff --git a/gmlx/loader.py b/gmlx/loader.py index 98fa870..dc94bad 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -2697,6 +2697,12 @@ def _warm_touch_pass( # computed fp32 (the vendored cast_predicate pins the same set). "kimi_k3": (".mlp.gate.weight", ".e_score_correction_bias", ".a_folded", ".dt_bias", "_res_score"), + # muse_glimmer's mmproj is native F16 and llama.cpp runs the tower with f32 + # activations. 50 residual layers with large outliers (features span +-76) + # compound bf16 rounding into ~10% relative RMS on the projected embeddings + # against an f32 run; f16->f32 is lossless, so this reproduces the oracle. + # Vision only - the text tower's bf16 holds 16k parity. + "muse_glimmer": ("vision_tower.", "vision_adapter.", "vision_projection."), } diff --git a/gmlx/muse_glimmer_model.py b/gmlx/muse_glimmer_model.py index 7ec149e..f6fb3c2 100644 --- a/gmlx/muse_glimmer_model.py +++ b/gmlx/muse_glimmer_model.py @@ -209,8 +209,11 @@ def __call__( inputs: mx.array, cache: Optional[Any] = None, capture_layers: Optional[tuple] = None, + inputs_embeds: Optional[mx.array] = None, ): - h = self.embed_tokens(inputs) + # The embedding norm sits after llama.cpp's build_inp_embd, so injected + # multimodal embeddings are normed alongside token embeddings. + h = self.embed_tokens(inputs) if inputs_embeds is None else inputs_embeds h = mx.fast.rms_norm(h, None, self.args.rms_norm_eps) if cache is None: @@ -236,6 +239,17 @@ def __call__( return self.norm(h) +def scale_and_softcap(out: mx.array, multiplier: float, cap: float) -> mx.array: + """Logit tail shared with the vision-language wrapper: the output multiplier + then the gemma-style tanh softcap. Computed in fp32 - llama.cpp's parity + oracle scales and softcaps an fp32 ``result_output``, and the softcap is + nonlinear enough that bf16 rounding moves argmax at depth.""" + out = out.astype(mx.float32) * multiplier + if cap: + out = mx.tanh(out / cap) * cap + return out + + class Model(nn.Module): def __init__(self, args: ModelArgs): super().__init__() @@ -247,18 +261,13 @@ def __init__(self, args: ModelArgs): def head(self, h: mx.array) -> mx.array: """Logits from a final-normed hidden state: lm_head, output multiplier, - tanh softcap. Computed in fp32 - llama.cpp's parity oracle scales and - softcaps an fp32 ``result_output``, and the softcap is nonlinear enough - that bf16 rounding moves argmax at depth.""" + tanh softcap.""" if self.args.tie_word_embeddings: out = self.model.embed_tokens.as_linear(h) else: out = self.lm_head(h) - out = out.astype(mx.float32) * self.args.output_multiplier - cap = self.args.final_logit_softcapping - if cap: - out = mx.tanh(out / cap) * cap - return out + return scale_and_softcap( + out, self.args.output_multiplier, self.args.final_logit_softcapping) def __call__(self, inputs: mx.array, cache: Optional[Any] = None): return self.head(self.model(inputs, cache)) diff --git a/gmlx/muse_glimmer_vlm_model.py b/gmlx/muse_glimmer_vlm_model.py new file mode 100644 index 0000000..4b68a00 --- /dev/null +++ b/gmlx/muse_glimmer_vlm_model.py @@ -0,0 +1,449 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Asher Feldman +"""Vendored mlx-vlm model for Meta Muse Glimmer (mmproj projector ``muse-glimmer``). + +mlx-vlm has no muse_glimmer package, so this module supplies the vision half of +the pair: a 50-layer ViT, the pixel-shuffle downsample, and the adapter MLP that +lands in the text tower's residual width. The text half is the same vendored +class the text-only path uses (:mod:`gmlx.muse_glimmer_model`), wrapped here in +the ``language_model`` shape mlx-vlm's generate stack expects. + +The tower is ported from llama.cpp's ``clip_graph_muse_glimmer::build`` plus the +host-side index math in ``clip.cpp`` (``PROJECTOR_TYPE_MUSE_GLIMMER`` set_input). +Four mechanics are specific to this family: + + 1. Patches are reordered into 32x32 windows and 3 of every 4 layers attend + only within a window; every 4th layer and the last one are global. + 2. 2-D RoPE: the first half of each head's dimensions is rotated by the patch's + 1-indexed column, the second half by its row, both on the same frequency + ladder (llama.cpp ``build_rope_2d`` with ``interleave_freq`` false). + 3. The learned 32x32 position grid is bilinearly resampled to the image's patch + grid, matching ggml's non-antialiased ``GGML_SCALE_MODE_BILINEAR``. + 4. Output tokens are pixel-shuffled 2x2 channel-outer (1536 -> 6144) before the + adapter, so one soft token covers a 28x28 pixel cell. + +Q/K pass through un-permuted for the same reason the text tower's do: the +converter emits the interleaved layout llama.cpp's rope mode 0 consumes. +""" + +from __future__ import annotations + +import importlib +import math +import sys +from dataclasses import dataclass, field +from typing import List, Optional + +import mlx.core as mx +import mlx.nn as nn + +from mlx_vlm.models.base import BaseModelConfig, InputEmbeddingsFeatures, LanguageModelOutput +from mlx_vlm.models.cache import KVCache, RotatingKVCache +from mlx_vlm.models.interpolate import bilinear_interpolate + +from .muse_glimmer_model import MuseGlimmerModel, scale_and_softcap + + +def ensure_registered() -> None: + """Make ``mlx_vlm.models.muse_glimmer`` resolve, preferring upstream.""" + if "mlx_vlm.models.muse_glimmer" not in sys.modules: + try: + importlib.import_module("mlx_vlm.models.muse_glimmer") # upstream wins + except ImportError: + sys.modules["mlx_vlm.models.muse_glimmer"] = sys.modules[__name__] + + +@dataclass +class TextConfig(BaseModelConfig): + model_type: str = "muse_glimmer" + hidden_size: int = 6656 + intermediate_size: int = 19968 + num_hidden_layers: int = 52 + num_attention_heads: int = 32 + num_key_value_heads: int = 2 + head_dim: int = 128 + vocab_size: int = 202048 + layer_types: List[str] = field(default_factory=list) + sliding_window: int = 2048 + rms_norm_eps: float = 1e-5 + post_norm_eps: float = 1e-8 + rope_theta: float = 500000.0 + rope_parameters: Optional[dict] = None + max_position_embeddings: int = 131072 + output_multiplier: float = 1.0 + final_logit_softcapping: float = 0.0 + tie_word_embeddings: bool = False + + +@dataclass +class VisionConfig(BaseModelConfig): + model_type: str = "muse_glimmer" + num_hidden_layers: int = 50 + hidden_size: int = 1536 + intermediate_size: int = 8960 + num_attention_heads: int = 16 + image_size: int = 896 + patch_size: int = 14 + num_channels: int = 3 + projection_dim: int = 6656 + adapter_hidden_size: int = 4096 + layer_norm_eps: float = 1e-5 + rope_theta: float = 10000.0 + spatial_merge_size: int = 2 + # 3 window layers then 1 global, repeating; the last layer is always global. + sparse_factor: int = 4 + num_position_embeddings: int = 1024 + + +@dataclass +class ModelConfig(BaseModelConfig): + text_config: TextConfig = field(default_factory=TextConfig) + vision_config: VisionConfig = field(default_factory=VisionConfig) + model_type: str = "muse_glimmer" + image_token_id: int = 200092 + image_token_index: Optional[int] = None + vocab_size: int = 202048 + eos_token_id: Optional[List[int]] = None + + def __post_init__(self): + if self.image_token_index is None: + self.image_token_index = self.image_token_id + + +# Grid index math (pure functions of the patch grid; unit-tested) + +def window_order(grid_w: int, grid_h: int, window: int) -> tuple[list[int], list[int]]: + """Patch order that makes window attention block-diagonal. + + Returns ``(perm, segment)``: ``perm[i]`` is the row-major patch index sitting + at permuted position ``i``, and ``segment[i]`` is its window id, so the + attention mask is ``segment[:, None] == segment[None, :]``. Windows on the + right and bottom edges are partial, exactly as llama.cpp builds them. + """ + perm: list[int] = [] + segment: list[int] = [] + win_id = 0 + for wy in range(0, grid_h, window): + for wx in range(0, grid_w, window): + count = 0 + for gy in range(wy, min(wy + window, grid_h)): + for gx in range(wx, min(wx + window, grid_w)): + perm.append(gy * grid_w + gx) + segment.append(win_id) + count += 1 + if count: + win_id += 1 + return perm, segment + + +def pixel_shuffle_order(grid_w: int, grid_h: int, merge: int) -> list[int]: + """Gather order that groups each ``merge`` x ``merge`` cell contiguously, in + row-major cell order (llama.cpp's ``ds_perm``).""" + order: list[int] = [] + for oy in range(grid_h // merge): + for ox in range(grid_w // merge): + for ry in range(merge): + for rx in range(merge): + order.append((oy * merge + ry) * grid_w + (ox * merge + rx)) + return order + + +def _rope_tables(pos: mx.array, half_dim: int, base: float): + """cos/sin for an interleaved rope over ``half_dim`` dims at the given + integer positions: pair ``j`` turns at ``base ** (-2j / half_dim)``.""" + n_pair = half_dim // 2 + inv = mx.exp( + -mx.arange(n_pair, dtype=mx.float32) * (math.log(base) * 2.0 / half_dim)) + theta = pos.astype(mx.float32)[:, None] * inv[None, :] + return mx.cos(theta), mx.sin(theta) + + +def _rope_half(v: mx.array, cos: mx.array, sin: mx.array) -> mx.array: + """Interleaved (pairwise) rotation of ``[B, H, L, D]`` by per-position + tables of shape ``[L, D // 2]``.""" + B, H, L, D = v.shape + v = v.reshape(B, H, L, D // 2, 2) + x0, x1 = v[..., 0], v[..., 1] + c, s = cos[None, None].astype(v.dtype), sin[None, None].astype(v.dtype) + return mx.stack([x0 * c - x1 * s, x0 * s + x1 * c], axis=-1).reshape(B, H, L, D) + + +def _rope_2d(x: mx.array, tables_w, tables_h) -> mx.array: + half = x.shape[-1] // 2 + return mx.concatenate( + [_rope_half(x[..., :half], *tables_w), _rope_half(x[..., half:], *tables_h)], + axis=-1, + ) + + +# Vision tower + +class VisionAttention(nn.Module): + def __init__(self, config: VisionConfig): + super().__init__() + dim = config.hidden_size + self.n_heads = config.num_attention_heads + self.head_dim = dim // self.n_heads + self.scale = self.head_dim**-0.5 + self.q_proj = nn.Linear(dim, dim, bias=True) + self.k_proj = nn.Linear(dim, dim, bias=True) + self.v_proj = nn.Linear(dim, dim, bias=True) + self.o_proj = nn.Linear(dim, dim, bias=True) + + def __call__(self, x: mx.array, tables_w, tables_h, mask) -> mx.array: + B, L, _ = x.shape + shape = (B, L, self.n_heads, self.head_dim) + q = self.q_proj(x).reshape(shape).transpose(0, 2, 1, 3) + k = self.k_proj(x).reshape(shape).transpose(0, 2, 1, 3) + v = self.v_proj(x).reshape(shape).transpose(0, 2, 1, 3) + q = _rope_2d(q, tables_w, tables_h) + k = _rope_2d(k, tables_w, tables_h) + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask) + return self.o_proj(out.transpose(0, 2, 1, 3).reshape(B, L, -1)) + + +class VisionMLP(nn.Module): + def __init__(self, config: VisionConfig): + super().__init__() + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=True) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=True) + + def __call__(self, x: mx.array) -> mx.array: + return self.fc2(nn.gelu(self.fc1(x))) + + +class VisionLayer(nn.Module): + def __init__(self, config: VisionConfig): + super().__init__() + eps = config.layer_norm_eps + self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=eps) + self.self_attn = VisionAttention(config) + self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=eps) + self.mlp = VisionMLP(config) + + def __call__(self, x: mx.array, tables_w, tables_h, mask) -> mx.array: + x = x + self.self_attn(self.layer_norm1(x), tables_w, tables_h, mask) + return x + self.mlp(self.layer_norm2(x)) + + +class VisionModel(nn.Module): + """The ViT alone: pixels in, post-normed patch features in row-major grid + order out. The window permutation is applied and undone internally, so the + caller never sees the sparse-attention ordering.""" + + def __init__(self, config: VisionConfig): + super().__init__() + self.config = config + self.model_type = config.model_type + patch = config.patch_size + self.patch_embed = nn.Conv2d( + config.num_channels, config.hidden_size, kernel_size=patch, stride=patch, + bias=False) + self.position_embedding = mx.zeros( + (config.num_position_embeddings, config.hidden_size)) + self.pre_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.layers = [VisionLayer(config) for _ in range(config.num_hidden_layers)] + self.post_layernorm = nn.LayerNorm( + config.hidden_size, eps=config.layer_norm_eps) + # The window side is the learned position grid's side (32), not a + # separate hyperparameter (clip.cpp derives it the same way). + self.window = int(round(math.sqrt(config.num_position_embeddings))) + + def _position_embedding(self, grid_w: int, grid_h: int) -> mx.array: + side = self.window + if grid_w == side and grid_h == side: + return self.position_embedding + grid = self.position_embedding.reshape(side, side, -1) + resized = bilinear_interpolate(grid.astype(mx.float32), grid_h, grid_w) + return resized.reshape(grid_h * grid_w, -1).astype(self.position_embedding.dtype) + + def __call__(self, pixel_values: mx.array) -> mx.array: + """``pixel_values`` is a single image as ``[1, H, W, C]``.""" + patch = self.config.patch_size + grid_h = pixel_values.shape[1] // patch + grid_w = pixel_values.shape[2] // patch + + x = self.patch_embed(pixel_values).reshape(1, grid_h * grid_w, -1) + x = x + self._position_embedding(grid_w, grid_h)[None] + + perm, segment = window_order(grid_w, grid_h, self.window) + perm = mx.array(perm) + seg = mx.array(segment) + window_mask = (seg[:, None] == seg[None, :])[None, None] + + x = self.pre_layernorm(x) + x = mx.take(x, perm, axis=1) + + # 1-indexed column/row of each patch, in the permuted order. + pos_w = perm % grid_w + 1 + pos_h = perm // grid_w + 1 + half = (self.config.hidden_size // self.config.num_attention_heads) // 2 + tables_w = _rope_tables(pos_w, half, self.config.rope_theta) + tables_h = _rope_tables(pos_h, half, self.config.rope_theta) + + n_layer = len(self.layers) + sf = self.config.sparse_factor + for idx, layer in enumerate(self.layers): + is_global = idx == n_layer - 1 or (idx + 1) % sf == 0 + x = layer(x, tables_w, tables_h, None if is_global else window_mask) + + x = self.post_layernorm(x) + inverse = mx.zeros(perm.shape, dtype=mx.int32) + inverse[perm] = mx.arange(perm.size, dtype=mx.int32) + return mx.take(x, inverse, axis=1)[0] + + +class VisionAdapter(nn.Module): + """The mmproj's two-layer adapter; the third linear lives in the LLM as + ``vision_projection``, matching where the HF checkpoint keeps it.""" + + def __init__(self, config: VisionConfig): + super().__init__() + merged = config.hidden_size * config.spatial_merge_size**2 + self.fc1 = nn.Linear(merged, config.adapter_hidden_size, bias=False) + self.fc2 = nn.Linear( + config.adapter_hidden_size, config.adapter_hidden_size, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return nn.gelu(self.fc2(nn.gelu(self.fc1(x)))) + + +# Text tower, in the shape mlx-vlm's generate stack expects + +class LanguageModel(nn.Module): + def __init__(self, config: TextConfig): + super().__init__() + self.config = config + self.model_type = config.model_type + self.model = MuseGlimmerModel(config) + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + inputs_embeds: Optional[mx.array] = None, + mask: Optional[mx.array] = None, + cache=None, + **kwargs, + ): + h = self.model(inputs, cache=cache, inputs_embeds=inputs_embeds) + if self.config.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(h) + else: + out = self.lm_head(h) + return LanguageModelOutput(logits=scale_and_softcap( + out, self.config.output_multiplier, self.config.final_logit_softcapping)) + + @property + def layers(self): + return self.model.layers + + @property + def head_dim(self): + return self.config.head_dim + + @property + def n_kv_heads(self): + return self.config.num_key_value_heads + + def make_cache(self): + return [ + RotatingKVCache(max_size=self.config.sliding_window, keep=0) + if layer.use_sliding + else KVCache() + for layer in self.model.layers + ] + + +class Model(nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.config = config + self.model_type = config.model_type + self.vision_tower = VisionModel(config.vision_config) + self.vision_adapter = VisionAdapter(config.vision_config) + self.vision_projection = nn.Linear( + config.vision_config.adapter_hidden_size, + config.vision_config.projection_dim, bias=False) + self.language_model = LanguageModel(config.text_config) + + def _image_features(self, pixel_values: mx.array, image_sizes) -> mx.array: + """One padded ``[N, C, H, W]`` batch plus its true ``(h, w)`` sizes -> + ``[total_soft_tokens, text_hidden]``. Images are run one at a time: the + patch grid sets the window layout and the rope positions, so a padded + batch would attend over padding.""" + merge = self.config.vision_config.spatial_merge_size + patch = self.config.vision_config.patch_size + feats = [] + for i, (h, w) in enumerate(image_sizes): + image = pixel_values[i, :, :h, :w].transpose(1, 2, 0)[None] + x = self.vision_tower(image.astype(self.vision_projection.weight.dtype)) + grid_h, grid_w = h // patch, w // patch + order = mx.array(pixel_shuffle_order(grid_w, grid_h, merge)) + n_out = (grid_h // merge) * (grid_w // merge) + x = mx.take(x, order, axis=0).reshape(n_out, merge * merge, -1) + x = x.transpose(0, 2, 1).reshape(n_out, -1) + feats.append(self.vision_projection(self.vision_adapter(x))) + return mx.concatenate(feats, axis=0) + + def get_input_embeddings( + self, + input_ids: Optional[mx.array] = None, + pixel_values: Optional[mx.array] = None, + **kwargs, + ): + embeds = self.language_model.model.embed_tokens(input_ids) + if pixel_values is None: + return InputEmbeddingsFeatures(inputs_embeds=embeds) + + features = kwargs.get("cached_image_features") + if features is None: + sizes = kwargs.get("image_sizes") + if sizes is None: + sizes = [pixel_values.shape[-2:]] * pixel_values.shape[0] + features = self._image_features( + pixel_values, [(int(h), int(w)) for h, w in sizes]) + return InputEmbeddingsFeatures( + inputs_embeds=self.merge_input_ids_with_image_features( + self.config.image_token_index, features, embeds, input_ids)) + + @staticmethod + def merge_input_ids_with_image_features( + image_token_index, image_features, inputs_embeds, input_ids + ): + """Scatter ``image_features`` onto the placeholder positions, in order.""" + if image_features.ndim == 3 and image_features.shape[0] == 1: + image_features = image_features.squeeze(0) + positions = input_ids == image_token_index + n_slots = int(mx.sum(positions).item()) + if n_slots != image_features.shape[0]: + raise ValueError( + f"{n_slots} image placeholder tokens but " + f"{image_features.shape[0]} image features") + if n_slots == 0: + return inputs_embeds + features = image_features.astype(inputs_embeds.dtype) + rank = mx.cumsum(positions.astype(mx.int32).reshape(-1)) - 1 + gathered = mx.take(features, mx.maximum(rank, 0), axis=0) + gathered = gathered.reshape(inputs_embeds.shape) + return mx.where(positions[..., None], gathered, inputs_embeds) + + @property + def layers(self): + return self.language_model.model.layers + + def make_cache(self): + return self.language_model.make_cache() + + def __call__( + self, + input_ids: mx.array, + pixel_values: Optional[mx.array] = None, + mask: Optional[mx.array] = None, + cache=None, + **kwargs, + ): + features = self.get_input_embeddings(input_ids, pixel_values, **kwargs) + return self.language_model( + input_ids, cache=cache, inputs_embeds=features.inputs_embeds) diff --git a/gmlx/upstream_seams.py b/gmlx/upstream_seams.py index b2d978d..1d55f31 100644 --- a/gmlx/upstream_seams.py +++ b/gmlx/upstream_seams.py @@ -494,6 +494,38 @@ def vendored_upstream_collisions() -> list[str]: f"{mod_name}: upstream mlx-lm now ships this module; the " f"vendored copy shadows it - reconcile and drop the vendor " f"entry (arch_table._VENDORED_MLX_LM_MODULES)") + hits += _vendored_vlm_collisions() + return hits + + +# gmlx module -> the mlx-vlm namespace its ensure_registered() grafts into. +# Each is a package directory upstream, so a native arrival shows up as either +# a .py module or a / package. +VENDORED_MLX_VLM_MODULES = { + "gmlx.muse_glimmer_vlm_model": "mlx_vlm.models.muse_glimmer", + "gmlx.hy_v3_tools": "mlx_vlm.tool_parsers.hy_v3", + "gmlx.muse_glimmer_tools": "mlx_vlm.tool_parsers.muse_glimmer", +} + + +def _vendored_vlm_collisions() -> list[str]: + """Same check on the mlx-vlm side: our grafts are upstream-first at import + time, but a native module arriving under a name we also register is the + signal to drop the vendored copy rather than keep shadowing it.""" + hits = [] + for mod_name, target in VENDORED_MLX_VLM_MODULES.items(): + pkg, _, leaf = target.rpartition(".") + try: + parent = importlib.import_module(pkg) + except ImportError: + continue + root = os.path.dirname(parent.__file__) + if (os.path.exists(os.path.join(root, f"{leaf}.py")) + or os.path.isdir(os.path.join(root, leaf))): + hits.append( + f"{mod_name}: upstream mlx-vlm now ships {target}; the " + f"vendored copy is only a fallback - reconcile and drop the " + f"vendor entry (upstream_seams.VENDORED_MLX_VLM_MODULES)") return hits diff --git a/gmlx/vlm.py b/gmlx/vlm.py index 41338f1..3dbfd92 100644 --- a/gmlx/vlm.py +++ b/gmlx/vlm.py @@ -34,6 +34,7 @@ ) from .gguf_meta import first_nonzero_int, read_int from .loader import ( + _FP32_KEEP_BY_MODEL_TYPE, _active_now, _install_and_load, load_gguf_wire_bytes, @@ -66,6 +67,11 @@ def resolve_vlm_model_type(llm_arch: str, mm_meta: dict) -> str: # Mistral Pixtral: a plain-float Pixtral ViT (2-D RoPE, RMSNorm, SiLU MLP) # + a 2-layer GELU projector onto a Mistral-Nemo (llama-arch) text tower. return "pixtral" + if proj == "muse-glimmer": + # Meta Muse Glimmer: a 50-layer window-attention ViT + a 2-layer GELU + # adapter onto the muse-glimmer text tower. Both halves are vendored + # (gmlx.muse_glimmer_vlm_model); mlx-vlm ships no class for either. + return "muse_glimmer" if proj == "qwen2vl_merger": # Resolvable in principle (mlx-vlm has qwen2_vl), but none of the # vision remap / config synth / processor synth paths exist for the @@ -217,6 +223,55 @@ def _pixtral_vision_name(name: str): return f"{_PVM}.transformer.layers.{bid}.{tgt}.{leaf}", False +# Muse Glimmer: a LayerNorm/GELU ViT with 2-D RoPE and window attention, onto +# the vendored gmlx.muse_glimmer_vlm_model tower. Every block tensor carries a +# bias. Q/K stay un-permuted - the converter already emits the interleaved +# layout llama.cpp's rope mode 0 (and this port's rope) consumes, the same +# decision the text tower records in remap.ARCH_ALIAS. +_MUSE_GLIMMER_BLK_SUBMAP = { + "attn_q": "self_attn.q_proj", + "attn_k": "self_attn.k_proj", + "attn_v": "self_attn.v_proj", + "attn_out": "self_attn.o_proj", + "ln1": "layer_norm1", + "ln2": "layer_norm2", + "ffn_up": "mlp.fc1", # hidden -> intermediate + "ffn_down": "mlp.fc2", # intermediate -> hidden +} +_MUSE_GLIMMER_TOP_MAP = { + "v.position_embd.weight": "vision_tower.position_embedding", + "v.pre_ln.weight": "vision_tower.pre_layernorm.weight", + "v.pre_ln.bias": "vision_tower.pre_layernorm.bias", + "v.post_ln.weight": "vision_tower.post_layernorm.weight", + "v.post_ln.bias": "vision_tower.post_layernorm.bias", + "mm.0.weight": "vision_adapter.fc1.weight", + "mm.1.weight": "vision_adapter.fc2.weight", + # mm.2 is the LLM-side projection into the text residual width; the HF + # checkpoint keeps it outside the adapter, and so does this tree. + "mm.2.weight": "vision_projection.weight", +} + + +def _muse_glimmer_vision_name(name: str): + """Map an mmproj clip tensor to its vendored muse_glimmer path. + + Returns ``(target_name, is_patch_conv)`` or ``None`` to skip.""" + hit = _MUSE_GLIMMER_TOP_MAP.get(name) + if hit is not None: + return hit, False + if name == "v.patch_embd.weight": + return "vision_tower.patch_embed.weight", True + m = _VISION_BLK_RE.match(name) + if m is None: + return None + bid, rest = m.group(1), m.group(2) + sub, _, leaf = rest.rpartition(".") # leaf = weight | bias + tgt = _MUSE_GLIMMER_BLK_SUBMAP.get(sub) + if tgt is None: + return None + return f"vision_tower.layers.{bid}.{tgt}.{leaf}", False + + # gemma-4 (E2B/E4B-it) omni vision tower onto mlx_vlm.models.gemma4. The vision # attn/MLP linears nest under a ClippableLinear (``...q_proj.linear.weight``); # use_clipped_linears is False, so the GGUF's per-tensor input/output min/max @@ -609,6 +664,19 @@ def remap_vision_arrays( out[hf] = mx.transpose(arr, (0, 2, 3, 1)) if is_patch else arr return out, skipped, vis_kqmeta + if model_type == "muse_glimmer": + for name, arr in arrays.items(): + if name.endswith(".scales") or name.endswith(".biases"): + continue + res = _muse_glimmer_vision_name(name) + if res is None: + skipped.append(name) + continue + hf, is_patch = res + # patch conv: GGUF [out, in, kH, kW] (NCHW) -> nn.Conv2d [out, kH, kW, in]. + out[hf] = mx.transpose(arr, (0, 2, 3, 1)) if is_patch else arr + return out, skipped, vis_kqmeta + if model_type == "gemma4": for name, arr in arrays.items(): if name.endswith(".scales") or name.endswith(".biases"): @@ -991,9 +1059,77 @@ def _synthesize_pixtral_vlm_config( return config +def _other_dim(shape, known: int) -> int: + """The size of a 2-D tensor's other axis, given one axis' length. Reading it + this way rather than by position keeps the caller independent of whether the + shape arrived in GGUF ne order or numpy order.""" + total = 1 + for d in shape: + total *= int(d) + if known <= 0 or total % known: + raise ValueError(f"shape {tuple(shape)} has no axis of size {known}") + return total // known + + +def _synthesize_muse_glimmer_vlm_config( + text_config: dict, mm_meta: dict, llm_meta: dict, mm_shapes: dict | None = None, +) -> dict: + """Muse Glimmer VLM config: the muse-glimmer text synth plus a vision tower + read from ``clip.vision.*``. + + Two values llama.cpp keeps as arch constants rather than GGUF fields are + pinned here with their source (clip.cpp ``PROJECTOR_TYPE_MUSE_GLIMMER``): + the vision rope base 10000 and the 3-sparse-then-1-global layer period. Two + more are read off the mmproj's own tensor shapes, because no metadata key + carries them: the learned position grid's length (llama.cpp likewise takes + ``sqrt(position_embeddings->ne[1])``) and the adapter's hidden width. + """ + vision_config: dict = { + "model_type": "muse_glimmer", + "num_hidden_layers": _mm_int(mm_meta, "clip.vision.block_count"), + "hidden_size": _mm_int(mm_meta, "clip.vision.embedding_length"), + "intermediate_size": _mm_int(mm_meta, "clip.vision.feed_forward_length"), + "num_attention_heads": _mm_int(mm_meta, "clip.vision.attention.head_count"), + "image_size": _mm_int(mm_meta, "clip.vision.image_size"), + "patch_size": _mm_int(mm_meta, "clip.vision.patch_size"), + "num_channels": 3, + "projection_dim": _mm_int(mm_meta, "clip.vision.projection_dim"), + "rope_theta": 10000.0, + "sparse_factor": 4, + } + merge = _mm(mm_meta, "clip.vision.spatial_merge_size") + vision_config["spatial_merge_size"] = int(merge) if merge is not None else 2 + eps = _mm(mm_meta, "clip.vision.attention.layer_norm_epsilon") + if eps is not None: + vision_config["layer_norm_eps"] = float(eps) + # Both dims are read as "the other side" of a known axis rather than by + # shape order, since the mmproj's shapes arrive in GGUF ne order. + hidden = vision_config["hidden_size"] + merged = hidden * vision_config["spatial_merge_size"] ** 2 + pos = (mm_shapes or {}).get("v.position_embd.weight") + if pos: + vision_config["num_position_embeddings"] = _other_dim(pos, hidden) + adapter = (mm_shapes or {}).get("mm.0.weight") + if adapter: + vision_config["adapter_hidden_size"] = _other_dim(adapter, merged) + + config: dict = { + "model_type": "muse_glimmer", + "text_config": text_config, + "vision_config": vision_config, + "vocab_size": int(text_config.get("vocab_size", 202048)), + } + img_id = _gguf_token_id(llm_meta, "<|patch|>") + if img_id is not None: + config["image_token_index"] = img_id + config["image_token_id"] = img_id + return config + + def synthesize_vlm_config( model_type: str, llm_meta: dict, llm_shapes: dict, mm_meta: dict, *, mm_tensor_names: set[str] | None = None, + mm_shapes: dict | None = None, ) -> dict: """Assemble an mlx-vlm config dict from the two GGUFs. @@ -1002,10 +1138,16 @@ def synthesize_vlm_config( ``mm_tensor_names`` (the mmproj's tensor key set) lets the config reflect optional tensors that carry no metadata flag - e.g. gemma-4 vision standardization, present only on the larger (31B) SigLIP encoder. + ``mm_shapes`` (the mmproj's tensor->shape map) supplies the dims some + families record only in their tensors - e.g. Muse Glimmer's position grid. """ text_config = synthesize_config(llm_meta, llm_shapes) names = mm_tensor_names or set() + if model_type == "muse_glimmer": + return _synthesize_muse_glimmer_vlm_config( + text_config, mm_meta, llm_meta, mm_shapes) + if model_type == "gemma4": standardize = "v.std_scale" in names and "v.std_bias" in names return _synthesize_gemma4_vlm_config( @@ -1211,6 +1353,8 @@ def _synthesize_vlm_processor(model_type: str, tokenizer, mm_meta: dict): return _synthesize_qwen3_omni_processor(tokenizer, mm_meta) if model_type == "pixtral": return _synthesize_pixtral_processor(tokenizer, mm_meta) + if model_type == "muse_glimmer": + return _synthesize_muse_glimmer_processor(tokenizer, mm_meta) if model_type != "gemma4": raise UnsupportedVLMError( f"processor synth not implemented for model_type {model_type!r}") @@ -1619,6 +1763,202 @@ def _synthesize_pixtral_processor(tokenizer, mm_meta: dict): return _attach_streaming_helpers(processor, tokenizer) +class _MuseGlimmerGgufImageProcessor(ImageProcessingMixin): + """Torch-free Muse Glimmer image preprocessing (numpy + PIL only). + + Ports ``mtmd_image_preprocessor_muse_glimmer`` / ``muse_glimmer_grid_size`` + (llama.cpp ``tools/mtmd/mtmd-image.cpp``), itself a replica of transformers' + ``get_aspect_ratio_preserving_size``: pick the soft-token grid whose aspect + ratio is closest to the image's (ties going to the larger grid) under a + ``max_image_tokens`` cap, then resize straight to ``grid * cell`` pixels. + + The resize is a plain stretch with no padding, and Lanczos-3 - which PIL's + ``Image.LANCZOS`` matches exactly (llama.cpp says so at mtmd-image.cpp:353), + unlike the BICUBIC-vs-BILINEAR mismatch that bites the gemma-4 path above. + + Subclasses ``ImageProcessingMixin`` for the same reason the Pixtral one does: + ``ProcessorMixin``'s type-check accepts it while mlx-vlm ``prepare_inputs`` + keeps it out of the single-soft-token branch. + """ + + model_input_names = ["pixel_values", "image_sizes"] + + def __init__(self, image_mean, image_std, patch_size=14, + spatial_merge_size=2, max_image_tokens=4096): + super().__init__() + self.image_mean = list(image_mean) + self.image_std = list(image_std) + self.patch_size = int(patch_size) + self.spatial_merge_size = int(spatial_merge_size) + self.max_image_tokens = int(max_image_tokens) + self.cell = self.patch_size * self.spatial_merge_size + self.size = {"height": self.cell, "width": self.cell} + + def soft_tokens(self, height: int, width: int) -> int: + """Soft tokens a preprocessed ``height`` x ``width`` image occupies.""" + return (height // self.cell) * (width // self.cell) + + def _target_hw(self, h: int, w: int) -> tuple[int, int]: + import math + cell = self.cell + cap = self.max_image_tokens + i_nph = h / cell + i_npw = w / cell + ratio = (i_npw / i_nph) if i_nph > 0 else 1.0 + if i_nph * i_npw > cap: + i_nph = math.sqrt(cap / ratio) + i_npw = i_nph * ratio + target_ar = h / w + best = None + for nph in (math.floor(i_nph), math.ceil(i_nph)): + for npw in (math.floor(i_npw), math.ceil(i_npw)): + if nph < 1 or npw < 1 or nph * npw > cap: + continue + d = abs(nph / npw - target_ar) + if best is None or d < best[0] or ( + d == best[0] and nph * npw > best[1] * best[2]): + best = (d, nph, npw) + if best is None: # nothing fit under the cap: round and clamp + nph = max(1, math.floor(i_nph + 0.5)) + npw = max(1, math.floor(i_npw + 0.5)) + else: + _, nph, npw = best + return nph * cell, npw * cell + + def _one(self, img): + import numpy as np + from PIL import Image + if not isinstance(img, Image.Image): + img = Image.fromarray(np.asarray(img)) + if img.mode != "RGB": + img = img.convert("RGB") + h_out, w_out = self._target_hw(img.height, img.width) + img = img.resize((w_out, h_out), Image.Resampling.LANCZOS) # PIL: (W, H) + arr = np.asarray(img, dtype=np.float32) / 255.0 # [H, W, C] + mean = np.array(self.image_mean, dtype=np.float32) + std = np.array(self.image_std, dtype=np.float32) + arr = (arr - mean) / std + return np.transpose(arr, (2, 0, 1)), (h_out, w_out) # [C, H, W] + + def __call__(self, images, **kwargs): + import numpy as np + flat = _PixtralGgufImageProcessor._flatten(images) + processed, sizes = [], [] + for img in flat: + chw, hw = self._one(img) + processed.append(chw) + sizes.append(hw) + max_h = max(s[0] for s in sizes) + max_w = max(s[1] for s in sizes) + padded = np.zeros((len(processed), 3, max_h, max_w), dtype=np.float32) + for i, (chw, (h, w)) in enumerate(zip(processed, sizes)): + padded[i, :, :h, :w] = chw + return {"pixel_values": padded, "image_sizes": sizes} + + +def _synthesize_muse_glimmer_processor(tokenizer, mm_meta: dict): + """Build the Muse Glimmer processor from the GGUFs alone - no HF download. + + mlx-vlm has no processor for this family either, so the marker expansion + lives here: the chat template emits one ``<|patch|>`` per image, and the + model wants that placeholder repeated once per soft token and wrapped in + ``<|image_start|>`` / ``<|image_end|>`` - the same bracketing llama.cpp's + mtmd adds around the image embeddings (mtmd.cpp, MUSE_GLIMMER case). + + The 4096-soft-token cap is a clip.cpp arch constant + (``set_limit_image_tokens(1, 4096)``), not GGUF metadata. + """ + from transformers.feature_extraction_utils import BatchFeature + from transformers.processing_utils import ProcessorMixin + + from mlx_vlm.models.base import to_mlx + + patch_size = _mm_int(mm_meta, "clip.vision.patch_size") + merge = _mm(mm_meta, "clip.vision.spatial_merge_size") + image_mean = _mm_floats(mm_meta, "clip.vision.image_mean") or [0.5, 0.5, 0.5] + image_std = _mm_floats(mm_meta, "clip.vision.image_std") or [0.5, 0.5, 0.5] + + image_processor = _MuseGlimmerGgufImageProcessor( + image_mean=image_mean, image_std=image_std, patch_size=patch_size, + spatial_merge_size=int(merge) if merge is not None else 2, + max_image_tokens=4096) + + class MuseGlimmerProcessor(ProcessorMixin): + attributes = ["image_processor", "tokenizer"] + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + image_token = "<|patch|>" + image_start_token = "<|image_start|>" + image_end_token = "<|image_end|>" + + def __call__(self, images=None, text=None, **kwargs): + if text is None and images is None: + raise ValueError("You must provide either text or images.") + if isinstance(text, str): + text = [text] + + image_inputs = {} + if images is not None: + if not isinstance(images, (list, tuple)): + images = [images] + image_inputs = self.image_processor(images) + if text is not None: + text = self._expand(text, image_inputs["image_sizes"]) + + kwargs.pop("return_tensors", None) + data = dict(image_inputs) + if text is not None: + data = {**self.tokenizer(text, **kwargs), **data} + return BatchFeature(data=to_mlx(data)) + + def _expand(self, texts, sizes): + """Replace each ``<|patch|>`` with its image's full token block. Sizes + are consumed in order across the batch, matching how the images were + flattened for preprocessing.""" + out, index = [], 0 + for sample in texts: + parts = sample.split(self.image_token) + rebuilt = parts[0] + for part in parts[1:]: + if index < len(sizes): + h, w = sizes[index] + n = self.image_processor.soft_tokens(h, w) + rebuilt += (self.image_start_token + self.image_token * n + + self.image_end_token) + index += 1 + else: + rebuilt += self.image_token + rebuilt += part + out.append(rebuilt) + return out + + def batch_decode(self, *args, **kwargs): + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + return list(dict.fromkeys( + list(self.tokenizer.model_input_names) + + list(self.image_processor.model_input_names))) + + processor = MuseGlimmerProcessor( + image_processor=image_processor, tokenizer=tokenizer, + chat_template=getattr(tokenizer, "chat_template", None)) + + # mlx-vlm's message formatter is a closed table; without a row + # ``chat._vlm_message`` degrades to a plain text dict and the image part + # never reaches the template. + import mlx_vlm.prompt_utils as _prompt_utils + _prompt_utils.MODEL_CONFIG.setdefault( + "muse_glimmer", _prompt_utils.MessageFormat.LIST_WITH_IMAGE_FIRST) + + return _attach_streaming_helpers(processor, tokenizer) + + # Public entry point @loadlog.seeds @@ -1672,9 +2012,15 @@ def load_vlm_model( # model_type, which fixes where the text tower nests. loadlog.stage("reading mmproj") loadlog.fact("mmproj", True) - mm_arrays, mm_codecs, _mm_arch, mm_meta, _mm_shapes = load_gguf_wire_bytes( + mm_arrays, mm_codecs, _mm_arch, mm_meta, mm_shapes = load_gguf_wire_bytes( mmproj_path, zero_copy=zero_copy, expect_quant=False) model_type = resolve_vlm_model_type(llm_arch, mm_meta) + if model_type == "muse_glimmer": + # mlx-vlm ships no muse_glimmer package; graft the vendored model + + # tool parser in before get_model_and_args resolves the model_type. + from . import muse_glimmer_tools, muse_glimmer_vlm_model + muse_glimmer_vlm_model.ensure_registered() + muse_glimmer_tools.ensure_registered() with_audio = bool(mm_meta.get("clip.has_audio_encoder")) _log(f"[vlm] model_type={model_type} audio={with_audio}") @@ -1705,7 +2051,7 @@ def load_vlm_model( loadlog.stage("building model") config = synthesize_vlm_config( model_type, llm_meta, llm_shapes, mm_meta, - mm_tensor_names=set(mm_arrays)) + mm_tensor_names=set(mm_arrays), mm_shapes=mm_shapes) text_config = config.get("text_config", {}) model, config = build_vlm_model(config) loadlog.fact("model_type", config.get("model_type")) @@ -1739,6 +2085,7 @@ def load_vlm_model( # (text under [thinker.]language_model.model.*, vision/audio under their # towers), so model.sanitize must not run - it would re-prefix text keys. _install_and_load(model, hf_weights, hf_kquant_meta, log=_log, sanitize=False, + fp32_keep=_FP32_KEEP_BY_MODEL_TYPE.get(model_type, ()), source_key=weights_source_key(*pf.shards, mmproj_path), active_before=active_before) materialize_module_arrays(model) diff --git a/pyproject.toml b/pyproject.toml index 54ad3a4..1d4fe68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ version = "0.2.2" description = "A local inference platform for Apple Silicon: run, chat with, serve, and fine-tune the GGUF ecosystem's quantized models natively on MLX, straight off the file." readme = "README.md" requires-python = ">=3.11" # 3.10 EOLs 2026-10; mlx-kquant's lower floor is a library floor -license = "BUSL-1.1 AND MIT" # MIT: files with an SPDX MIT header (model modules, dspark tooling, kimi-k3 tests) +license = "BUSL-1.1 AND MIT" # MIT: files with an SPDX MIT header (model and vision modules, dspark tooling, kimi-k3 tests) license-files = ["LICENSE", "LICENSE-MIT", "licenses/*"] authors = [{ name = "Asher Feldman" }] classifiers = [ From d8e46fe26d3158a4ef3f5422b980137a6ee70470 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:22:51 -0700 Subject: [PATCH 06/17] feat(mtp): muse-glimmer dflash drafter, speculative target hooks, and companion autodetect --- gmlx/arch_table.py | 1 + gmlx/discovery.py | 2 +- gmlx/loader.py | 18 +++ gmlx/mtp_load.py | 256 ++++++++++++++++++++++++++++++++++++ gmlx/muse_glimmer_dflash.py | 228 ++++++++++++++++++++++++++++++++ gmlx/muse_glimmer_mtp.py | 168 +++++++++++++++++++++++ 6 files changed, 672 insertions(+), 1 deletion(-) create mode 100644 gmlx/muse_glimmer_dflash.py create mode 100644 gmlx/muse_glimmer_mtp.py diff --git a/gmlx/arch_table.py b/gmlx/arch_table.py index 5349b06..981e38d 100644 --- a/gmlx/arch_table.py +++ b/gmlx/arch_table.py @@ -128,6 +128,7 @@ class UnsupportedArchError(Exception): # fails at build. Extend together with the loader dispatch. MTP_WIRED_MODEL_TYPES = frozenset({ "qwen3_5", "qwen3_5_moe", "gemma4_text", "deepseek_v4", "hy_v3", + "muse_glimmer", }) diff --git a/gmlx/discovery.py b/gmlx/discovery.py index fc308d2..94fe93a 100644 --- a/gmlx/discovery.py +++ b/gmlx/discovery.py @@ -66,7 +66,7 @@ # (mmproj/assistant/...) plus imatrix provenance (mradermacher `i1`, `imatrix`) - # that's quant provenance, not part of the model name, and leaving it in splits # one model's quants across two id prefixes (`...instruct.i1-*` vs `...instruct-*`). -_ID_MARKERS = ("mmproj", "assistant", "draft", "mtp", "gguf", +_ID_MARKERS = ("mmproj", "assistant", "draft", "mtp", "dflash", "gguf", "imatrix", "imat", "i1") diff --git a/gmlx/loader.py b/gmlx/loader.py index dc94bad..7c86a62 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -616,6 +616,13 @@ def remap_gemma4_assistant_arrays(arrays: dict, kquant_meta: dict): "speculative_argmax_from_hidden", "speculative_verify_hidden", ), + # MuseGlimmerSpecLM (vendored mlx-lm class): same lean set as deepseek_v4. + "muse_glimmer": ( + "rollback_speculative_cache", + "speculative_logits_from_hidden", + "speculative_argmax_from_hidden", + "speculative_verify_hidden", + ), } @@ -751,6 +758,17 @@ def build(config): return hy_v3_mtp.HyV3SpecLM(ModelArgs.from_dict(config)) return hy_v3_mtp.HyV3SpecLM, build + if model_type == "muse_glimmer": + from . import muse_glimmer_mtp, muse_glimmer_tools + from .muse_glimmer_model import ModelArgs, ensure_registered + + ensure_registered() + muse_glimmer_tools.ensure_registered() + + def build(config): + return muse_glimmer_mtp.MuseGlimmerSpecLM(ModelArgs.from_dict(config)) + + return muse_glimmer_mtp.MuseGlimmerSpecLM, build from .arch_table import MTP_WIRED_MODEL_TYPES raise NotImplementedError( diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index 4e4ece9..e943f0c 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -77,6 +77,7 @@ "hy_v3": 1, "deepseek_v4": 1, "deepseek4": 1, + "muse_glimmer": 1, } # Unknown arch: cap conservatively rather than opting a new family into the # losing regime. Uncapped is earned by measurement, not inherited by default. @@ -486,6 +487,12 @@ def _load_deepseek4_mtp_drafter( draft_gguf_path, zero_copy=zero_copy ) if d_arch == "dflash": + container = dflash_container(arrays) + if container != "dspark": + raise ValueError( + f"{draft_gguf_path}: this dflash GGUF holds the {container} " + f"drafter, which a deepseek_v4 target cannot drive" + ) arrays, kquant_meta, _meta = normalize_dflash_arrays( arrays, kquant_meta, _meta ) @@ -642,6 +649,236 @@ def _dflash_rename(name: str, last_stage: int) -> str: return mapped.replace("{L}", str(last_stage)) + suffix +def dflash_container(arrays: dict) -> str: + """Which drafter a llama.cpp ``dflash`` GGUF actually holds. + + The arch tag is shared: llama.cpp packages both the DeepSeek-V4 DSpark + drafter and the Muse Glimmer one under ``dflash``, and picks its graph on + ``dsv4_hc_mult``. Tensor presence is the equivalent split here - DSpark + carries the markov/confidence heads and MLA's ``attn_q_a``, Muse Glimmer + carries plain ``attn_q`` with per-head QK-norms and no hyper-connections. + """ + if any(n.startswith(("markov_w1", "markov_w2", "conf_proj", "output_hc_")) + or ".attn_q_a" in n for n in arrays): + return "dspark" + if any(".attn_q_norm" in n for n in arrays): + return "muse_glimmer" + raise RuntimeError( + "dflash GGUF matches no known drafter container (expected DSpark's " + "markov/confidence heads or Muse Glimmer's attn_q_norm)" + ) + + +# The closed tensor set of a Muse Glimmer dflash drafter, onto mlx-vlm's +# DFlashDraftModel tree. Per-block leaves (blk.{i}. -> layers.{i}.): +_MUSE_GLIMMER_DFLASH_BLK = { + "attn_norm": "input_layernorm.weight", + "attn_q": "self_attn.q_proj.weight", + "attn_k": "self_attn.k_proj.weight", + "attn_v": "self_attn.v_proj.weight", + "attn_output": "self_attn.o_proj.weight", + "attn_q_norm": "self_attn.q_norm.weight", + "attn_k_norm": "self_attn.k_norm.weight", + "ffn_norm": "post_attention_layernorm.weight", + "ffn_gate": "mlp.gate_proj.weight", + "ffn_up": "mlp.up_proj.weight", + "ffn_down": "mlp.down_proj.weight", +} +# Drafter-level leaves. ``enc.output_norm`` closes the encoder that fuses the +# target captures (llama.cpp's dflash graph); ``output_norm`` is the +# decoder's final norm before the borrowed LM head. +_MUSE_GLIMMER_DFLASH_ROOT = { + "fc": "fc.weight", + "enc.output_norm": "hidden_norm.weight", + "output_norm": "norm.weight", +} + + +def remap_muse_glimmer_dflash_arrays(arrays: dict, kquant_meta: dict): + """Remap a Muse Glimmer ``dflash`` GGUF onto the drafter param tree. + Closed tensor set: unknown names are hard errors (converter drift must + surface at load, not as an unfilled param).""" + hf_weights: dict[str, mx.array] = {} + hf_kquant_meta: dict[str, str] = {} + stats = {"mapped": 0} + for name, arr in arrays.items(): + if name.endswith((".scales", ".biases")): + continue + base = name[: -len(".weight")] if name.endswith(".weight") else name + if base.startswith("blk."): + _, idx, leaf = base.split(".", 2) + target = _MUSE_GLIMMER_DFLASH_BLK.get(leaf) + if target is not None: + target = f"layers.{idx}.{target}" + else: + target = _MUSE_GLIMMER_DFLASH_ROOT.get(base) + if target is None: + raise RuntimeError( + f"muse-glimmer dflash remap: unknown tensor {name!r} " + f"(the drafter tensor set is closed)" + ) + hf_weights[target] = arr + codec = kquant_meta.get(name) + if codec is not None: + hf_weights[_strip_weight(target) + ".scales"] = arrays.get( + _strip_weight(name) + ".scales") + hf_kquant_meta[target] = codec + stats["mapped"] += 1 + return hf_weights, hf_kquant_meta, stats + + +def _load_muse_glimmer_dflash_drafter( + draft_gguf_path: str, + target, + target_config_dict: dict, + *, + arrays: dict, + kquant_meta: dict, + meta: dict, + active_before: float | None = None, + log=loadlog.verbose_print, +): + """Build + load + bind the Muse Glimmer DFlash drafter, and wire the + target's ``_dflash_capture`` so every engine-facing hidden carries the + five captured residuals.""" + from .muse_glimmer_dflash import ( + MuseGlimmerDFlashConfig, + MuseGlimmerDFlashDrafter, + ) + + layers = meta.get("dflash.target_layers") + block_size = meta.get("dflash.block_size") + mask_token_id = meta.get("tokenizer.ggml.mask_token_id") + if not layers or block_size is None or mask_token_id is None: + raise ValueError( + f"{draft_gguf_path}: dflash.target_layers / dflash.block_size / " + f"tokenizer.ggml.mask_token_id missing - re-run the converter" + ) + # llama.cpp indexes the residual ENTERING a layer, so the converter writes + # the HF ids (layer outputs) one higher. Undo that: the capture seam takes + # layer-output indices. + layer_ids = tuple(int(i) - 1 for i in layers) + n_target_layers = int(target_config_dict["num_hidden_layers"]) + if list(layer_ids) != sorted(set(layer_ids)) or not ( + 0 <= layer_ids[0] and layer_ids[-1] < n_target_layers + ): + raise ValueError( + f"{draft_gguf_path}: dflash.target_layers {layers} must be " + f"strictly increasing and within [1, {n_target_layers}]" + ) + n_layers = 1 + max( + int(n.split(".")[1]) for n in arrays if n.startswith("blk.")) + pattern = meta.get("dflash.attention.sliding_window_pattern") or () + layer_types = [ + "sliding_attention" if bool(t) else "full_attention" for t in pattern + ] or ["full_attention"] * n_layers + window = int(meta.get("dflash.attention.sliding_window") or 0) or None + native_total = int(block_size) + block_total = max( + 2, min(env_int("GMLX_MUSE_DFLASH_BLOCK", native_total), native_total)) + + config = MuseGlimmerDFlashConfig( + hidden_size=int(target_config_dict["hidden_size"]), + intermediate_size=int(meta["dflash.feed_forward_length"]), + num_hidden_layers=n_layers, + num_attention_heads=int(meta["dflash.attention.head_count"]), + num_key_value_heads=int(meta["dflash.attention.head_count_kv"]), + head_dim=int(meta["dflash.attention.key_length"]), + rms_norm_eps=float(meta["dflash.attention.layer_norm_rms_epsilon"]), + vocab_size=int(target_config_dict["vocab_size"]), + max_position_embeddings=int(meta.get("dflash.context_length") + or target_config_dict["max_position_embeddings"]), + rope_theta=float(meta["dflash.rope.freq_base"]), + tie_word_embeddings=False, + block_size=block_total, + mask_token_id=int(mask_token_id), + target_layer_ids=list(layer_ids), + num_target_layers=n_target_layers, + layer_types=layer_types, + sliding_window=window, + # The ring is temporal-ordered and rollback-slack backed, which the + # split inject/draft forward needs (it reads cache.state directly). + draft_window_size=window, + final_logit_softcapping=target_config_dict.get( + "final_logit_softcapping") or None, + output_multiplier=float(target_config_dict.get("output_multiplier", 1.0)), + ) + drafter = MuseGlimmerDFlashDrafter(config) + log( + f"[mtp] drafter: muse-glimmer dflash layers={n_layers} " + f"targets={layer_ids} block_total={block_total} window={window}" + ) + + d_weights, d_meta, d_stats = remap_muse_glimmer_dflash_arrays( + arrays, kquant_meta) + log(f"[mtp] drafter remap: {d_stats}") + _install_and_load( + drafter, + d_weights, + d_meta, + log=log, + sanitize=False, + source_key=weights_source_key(draft_gguf_path), + active_before=active_before, + ) + drafter.bind(target) + + lm = getattr(target, "language_model", target) + if not callable(getattr(lm, "set_dflash_capture", None)): + raise RuntimeError( + "DFlash drafter needs a muse_glimmer target carrying the " + f"_dflash_capture seam; got {type(lm).__name__}" + ) + lm.set_dflash_capture(layer_ids) + + from .drafter_protocol import validate_drafter + + validate_drafter(drafter) + # No draft-side head quantization: Muse Glimmer GGUFs ship a quantized + # output.weight, which _patch_draft_head_quantized leaves alone anyway. + log("[mtp] dflash drafter bound; target capture layers wired") + _stamp_mtp_width_cap(drafter, "muse_glimmer", target=target, log=log) + return drafter + + +def _load_dflash_drafter( + draft_gguf_path: str, + target, + target_config_dict: dict, + *, + zero_copy: bool = True, + log=loadlog.verbose_print, +): + """Load a Muse Glimmer ``dflash`` companion drafter.""" + active_before = _active_now() + arrays, kquant_meta, d_arch, meta, _shapes = load_gguf_wire_bytes( + draft_gguf_path, zero_copy=zero_copy + ) + if d_arch != "dflash": + raise ValueError( + f"{draft_gguf_path}: expected a dflash drafter GGUF for a " + f"muse_glimmer target, got arch {d_arch!r}" + ) + container = dflash_container(arrays) + if container != "muse_glimmer": + raise ValueError( + f"{draft_gguf_path}: this dflash GGUF holds the {container} " + f"drafter, which a muse_glimmer target cannot drive" + ) + log(f"[mtp] drafter gguf ({d_arch}/{container}): {len(arrays)} arrays, " + f"{len(kquant_meta)} kquant") + return _load_muse_glimmer_dflash_drafter( + draft_gguf_path, + target, + target_config_dict, + arrays=arrays, + kquant_meta=kquant_meta, + meta=meta, + active_before=active_before, + log=log, + ) + + def normalize_dflash_arrays(arrays: dict, kquant_meta: dict, meta: dict): """Translate a llama.cpp ``dflash`` GGUF (tensor names and metadata) to the ``deepseek4-dspark`` namespace. Returns ``(arrays, kquant_meta, @@ -956,6 +1193,21 @@ def load_mtp_model( assistant = True loadlog.fact("mtp_companion", os.path.basename(draft_gguf_path)) _log(f"[mtp] companion drafter autodetected: {draft_gguf_path}") + if not assistant and config_dict.get("model_type") == "muse_glimmer": + # Muse Glimmer's drafter is likewise a companion GGUF (arch dflash), + # never an in-file nextn block. + from .discovery import find_mtp_companion + + draft_gguf_path = find_mtp_companion(gguf_path, ("dflash",)) + if draft_gguf_path is None: + raise ValueError( + "muse_glimmer MTP needs its companion DFlash drafter GGUF " + f"(arch dflash); none found next to {gguf_path} - pass " + "--draft-gguf ." + ) + assistant = True + loadlog.fact("mtp_companion", os.path.basename(draft_gguf_path)) + _log(f"[mtp] companion drafter autodetected: {draft_gguf_path}") if not assistant and int(config_dict.get("mtp_num_hidden_layers", 0)) < 1: raise ValueError( f"{gguf_path}: no native MTP head " @@ -1060,6 +1312,10 @@ def load_mtp_model( drafter = _load_deepseek4_mtp_drafter( draft_gguf_path, model, config_dict, zero_copy=zero_copy, log=_log ) + elif _mt == "muse_glimmer": + drafter = _load_dflash_drafter( + draft_gguf_path, model, config_dict, zero_copy=zero_copy, log=_log + ) else: drafter = _load_gemma4_assistant_drafter( draft_gguf_path, model, zero_copy=zero_copy, log=_log diff --git a/gmlx/muse_glimmer_dflash.py b/gmlx/muse_glimmer_dflash.py new file mode 100644 index 0000000..c506771 --- /dev/null +++ b/gmlx/muse_glimmer_dflash.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Asher Feldman +# Portions copyright (c) 2026 Prince Canuma (mlx-vlm DFlashDraftModel, MIT) +"""Muse Glimmer DFlash drafter: mlx-vlm's ``DFlashDraftModel`` weights and +module tree, driven by gmlx's owned speculative engine. + +DFlash denoises a whole block at once. Its decoder runs in two modes +(llama.cpp ``src/models/dflash.cpp``): + +- **inject** (embd batch): the target's captured residuals go through the + encoder ``hidden_norm(fc(captures))``, and each layer projects that straight + to K/V (k_norm + RoPE, no ``input_layernorm``, no query, no FFN) and writes + it into the draft KV ring at the committed tokens' positions. +- **draft** (token batch): ``[last_bonus, MASK x (block-1)]`` is embedded with + the target's ``tok_embd``, positioned at ``n_past + i``, and run through the + layers with **non-causal** attention over the ring plus the block's own + transient K/V. Drafts are read from rows 1..block-1, and the target's + ``lm_head`` produces the logits. + +mlx-vlm folds both into one ``draft_block(last_bonus, hidden, ...)`` call +because its engine hands the drafter the newly committed hidden each round. +gmlx's engine splits the same information across ``prefill_from_target_hidden`` +and ``accept_verified_tokens`` (and passes only the last hidden row to +``draft_block``), so the two modes are split here to match - same math, same +weights, different call boundary. + +Positions are relative: seeding only the last ``sliding_window`` prompt rows +shifts context and block queries by the same amount, which RoPE is invariant +to. Correctness never rests on the drafter regardless: the verify walk emits +the target's own tokens, so the drafter moves acceptance, never output. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, List, Optional + +import mlx.core as mx + +from mlx_vlm.speculative.drafters.qwen3_dflash.config import DFlashConfig +from mlx_vlm.speculative.drafters.qwen3_dflash.dflash import DFlashDraftModel + +from . import muse_glimmer_model as mg + + +@dataclass +class MuseGlimmerDFlashConfig(DFlashConfig): + """``DFlashConfig`` plus the Glimmer logit tail. The drafter borrows the + target's LM head, so it must reproduce the target's scale and softcap.""" + + output_multiplier: float = 1.0 + + +class MuseGlimmerDFlashDrafter(DFlashDraftModel): + """DFlash drafter following gmlx's BatchDrafterProtocol (owned engine).""" + + supports_greedy_draft_argmax = True + prefer_requested_block_size = False + cap_at_configured_depth = True + uses_shared_kv = False + supports_kv_sidecar = False + # CLI entry points must route to the owned engine: mlx-vlm's stock MTP + # round doesn't know the muse_glimmer target hooks (packed hidden). + requires_owned_engine = True + + def __init__(self, config: MuseGlimmerDFlashConfig): + super().__init__(config) + self._native_block_size = int(config.block_size) + self._hidden = int(config.hidden_size) + self._n_targets = len(config.target_layer_ids) + # Only the trailing window of the prompt capture is usable; the engine + # trims the retained prefill hidden to this many positions. + self.hidden_capture_limit = int(config.sliding_window or 0) or None + self._cache: List[Any] = [] + + # --- lifecycle ---------------------------------------------------------- + + def reset(self, target_model, left_padding: Optional[List[int]] = None) -> list: + if left_padding is not None: + raise NotImplementedError( + "MuseGlimmerDFlashDrafter is B=1 only (v1): batched rounds " + "need per-row cache offsets in the inject path" + ) + self.bind(target_model) + self.accept_lens = [] + self.draft_lens = [] + self._cache = self.make_cache() + return self._cache + + def draft_eval_state(self) -> List[Any]: + return [c.state for c in self._cache] + + def set_shared_kv(self, *args, **kwargs) -> None: + return None + + def make_cache(self, left_padding: Optional[List[int]] = None) -> list: + if left_padding is not None: + raise NotImplementedError("MuseGlimmerDFlashDrafter is B=1 only (v1)") + return super().make_cache() + + def _logits(self, hidden: mx.array) -> mx.array: + # The borrowed LM head is the target's bare Linear; the scale and + # softcap live in the target's own head, so reapply both here. + return mg.scale_and_softcap( + self.lm_head(hidden), self.config.output_multiplier, + self.config.final_logit_softcapping or 0.0) + + # --- inject mode -------------------------------------------------------- + + def _captures(self, packed: mx.array) -> mx.array: + """Trailing ``n_targets*hidden`` of the packed target hidden.""" + want = self._hidden * (1 + self._n_targets) + if int(packed.shape[-1]) != want: + raise ValueError( + f"packed hidden width {packed.shape[-1]} != {want}; the " + "target's _dflash_capture wiring is missing" + ) + # Materialize: this feeds the quantized fc, whose kernel reads the + # buffer directly and would otherwise see the packed strides. + return mx.contiguous(packed[..., self._hidden:]) + + def append_context(self, captures: mx.array) -> None: + """Encode committed positions and inject their K/V into the ring.""" + h_ctx = self.hidden_norm(self.fc(captures)) + B, S, _ = h_ctx.shape + for layer, c in zip(self.layers, self._cache): + attn = layer.self_attn + keys = attn.k_norm( + attn.k_proj(h_ctx).reshape(B, S, attn.n_kv_heads, -1) + ).transpose(0, 2, 1, 3) + values = attn.v_proj(h_ctx).reshape( + B, S, attn.n_kv_heads, -1).transpose(0, 2, 1, 3) + c.update_and_fetch(self.rope(keys, offset=c.offset), values) + + # --- draft mode --------------------------------------------------------- + + def _draft_hidden(self, tokens: mx.array) -> mx.array: + h = self._embed_input_tokens(tokens) + B, L, _ = h.shape + for layer, c in zip(self.layers, self._cache): + attn = layer.self_attn + x = layer.input_layernorm(h) + q = attn.q_norm( + attn.q_proj(x).reshape(B, L, attn.n_heads, -1) + ).transpose(0, 2, 1, 3) + k = attn.k_norm( + attn.k_proj(x).reshape(B, L, attn.n_kv_heads, -1) + ).transpose(0, 2, 1, 3) + v = attn.v_proj(x).reshape( + B, L, attn.n_kv_heads, -1).transpose(0, 2, 1, 3) + q = self.rope(q, offset=c.offset) + k = self.rope(k, offset=c.offset) + ck, cv = c.state + if ck is not None: + k = mx.concatenate([ck, k], axis=2) + v = mx.concatenate([cv, v], axis=2) + # The block denoises as a whole: every row sees the ring and all + # its siblings. The block's own K/V never enters the ring. + o = mx.fast.scaled_dot_product_attention( + q, k, v, scale=attn.scale, mask=None) + h = h + attn.o_proj(o.transpose(0, 2, 1, 3).reshape(B, L, -1)) + h = h + layer.mlp(layer.post_attention_layernorm(h)) + return self.norm(h) + + def draft_block( + self, + last_bonus, + hidden: mx.array, + cache, + block_size: int, + sampler, + token_dtype: mx.Dtype = mx.int32, + greedy: bool = False, + ) -> mx.array: + """One DFlash round: ``[bonus, MASK x (block_size-1)]`` denoised in a + single forward; drafts are rows 1..block_size-1.""" + del hidden, cache + if block_size > self._native_block_size: + raise RuntimeError( + f"MuseGlimmerDFlashDrafter drafts at most " + f"{self._native_block_size - 1} token(s)/round; got " + f"block_size={block_size} - cap_at_configured_depth should " + f"have clamped it" + ) + if not self._cache: + raise RuntimeError("reset(target_model) must run before draft_block()") + mask_id = int(self.config.mask_token_id) + bonus = (int(last_bonus) if isinstance(last_bonus, int) + else int(last_bonus.reshape(-1)[0].item())) + block = mx.array([[bonus] + [mask_id] * (block_size - 1)], dtype=token_dtype) + logits = self._logits(self._draft_hidden(block)[:, 1:]) + return mx.argmax(logits, axis=-1) if greedy else sampler(logits) + + # --- commit ------------------------------------------------------------- + + def prefill_from_target_hidden( + self, + input_ids: mx.array, + hidden: mx.array, + bonus_token, + sampler, + token_dtype: mx.Dtype = mx.int32, + greedy: bool = False, + ) -> None: + """Seed the ring from the trailing prompt hiddens. DFlash needs no + draft seed: rounds start from the engine-passed bonus token.""" + del input_ids, bonus_token, sampler, token_dtype, greedy + if int(hidden.shape[1]) == 0: + return + limit = self.hidden_capture_limit + self.append_context(self._captures(hidden[:, -limit:] if limit else hidden)) + + def accept_verified_tokens( + self, + verify_hidden: mx.array, + draft_tokens: mx.array, + accepted: int, + new_tokens: List[int], + sampler, + token_dtype: mx.Dtype = mx.int32, + greedy: bool = False, + ) -> None: + """Inject the committed positions' captures. ``verify_hidden[:, p]`` is + the target hidden at verify position ``p``; 0..accepted were committed + (the accepted drafts plus the row the new bonus was sampled from), + matching the rolled-back target 1:1.""" + del draft_tokens, new_tokens, sampler, token_dtype, greedy + self.append_context(self._captures(verify_hidden[:, : int(accepted) + 1])) diff --git a/gmlx/muse_glimmer_mtp.py b/gmlx/muse_glimmer_mtp.py new file mode 100644 index 0000000..8b2ccbb --- /dev/null +++ b/gmlx/muse_glimmer_mtp.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Asher Feldman +"""Muse Glimmer speculative target: the ``speculative_*`` hooks plus the +packed-hidden capture the DFlash drafter reads. + +The drafter (``muse_glimmer_dflash``) consumes the target's residual stream at +five layers. llama.cpp exposes those as ``t_layer_inp[il]`` - the residual +*entering* layer ``il`` - for ``il`` in the GGUF's ``dflash.target_layers`` +([2, 14, 26, 38, 50]); the converter writes those one higher than the HF +``target_layer_ids`` because HF records layer *outputs*. Entering layer 14 is +leaving layer 13, so the capture set here is the GGUF list minus one, taken as +layer outputs (:func:`MuseGlimmerModel.__call__`'s ``capture_layers``). + +Capture rides the DeepSeek-V4 packed-hidden seam rather than mlx-vlm's +``capture_layer_ids``/``hidden_states`` route: every engine-facing hidden is +widened to ``[trunk | cap_1 | cap_13 | ... ]`` so the existing slicing and +capture-trim seams work untouched and no engine change is needed. The drafter +unpacks the trailing ``n_targets*hidden``; the logits hooks slice the lead. + +Rollback needs no undo log here. The sliding layers hold ``keep=0`` +``RotatingKVCache`` leaves, which ``_buffer_mtp_target_cache`` swaps for +``BufferedRotatingKVCache`` before the decode loop - that cache keeps +rollback slack past the window edge, so ``is_trimmable()`` holds however deep +the context runs (a rotated stock ring would refuse, since the evicted slot is +gone). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, List, Optional + +import mlx.core as mx + +from . import muse_glimmer_model as mg + + +@dataclass +class _SpecOutput: + """Duck-typed output for the owned engine's ``return_hidden`` calls.""" + + logits: mx.array + hidden_states: List[mx.array] + shared_kv_states: dict = field(default_factory=dict) + gdn_states: Optional[list] = None + + +class SpecHooks: + """The ``speculative_*`` hook set, shared by the text MTP target and the + VLM ``language_model`` so ``--mmproj`` and ``--speculative`` compose. + + Requires ``self.model`` (a :class:`MuseGlimmerModel`) and a logits tail; + both hosts carry the model's ``args`` on ``self.model.args``. + """ + + # Set by the drafter loader: the trunk layer ids whose outputs the drafter + # consumes. While set, every engine-facing hidden is packed. + _dflash_capture: Optional[tuple] = None + + def set_dflash_capture(self, layer_ids) -> None: + """Arm packed-hidden capture on the given trunk layer outputs. + + Assigned around ``nn.Module.__setattr__``, which routes tuples into + the parameter dict - where this class default would shadow the entry + on read, and tree walks would meet non-array leaves.""" + object.__setattr__(self, "_dflash_capture", tuple(layer_ids)) + + def _spec_logits(self, h: mx.array) -> mx.array: + args = self.model.args + out = (self.model.embed_tokens.as_linear(h) if args.tie_word_embeddings + else self.lm_head(h)) + return mg.scale_and_softcap( + out, args.output_multiplier, args.final_logit_softcapping) + + def _dflash_pack(self, h: mx.array, captures) -> mx.array: + return mx.concatenate([h, *captures], axis=-1) + + def _dflash_trunk(self, hidden: mx.array) -> mx.array: + if self._dflash_capture is None: + return hidden + # The trunk lead is a strided view of the packed hidden, and the logit + # tail is a quantized matmul, whose kernel reads the buffer directly. + # Slicing lazily hands it the packed strides and it reads the wrong + # rows, so materialize before the head sees it. + return mx.contiguous(hidden[..., : self.model.args.hidden_size]) + + def chunked_prefill_policy(self, **kwargs): + # Stock mlx-vlm disables chunked prefill whenever a drafter is + # attached. The DFlash drafter is window-limited + # (hidden_capture_limit trailing positions), so last-chunk capture + # suffices and chunking stays safe. + return True + + def speculative_logits_from_hidden(self, hidden: mx.array) -> mx.array: + return self._spec_logits(self._dflash_trunk(hidden)) + + def speculative_argmax_from_hidden(self, hidden: mx.array) -> mx.array: + return mx.argmax(self.speculative_logits_from_hidden(hidden), axis=-1) + + def speculative_verify_hidden(self, verify_input: mx.array, prompt_cache): + """The single verify forward (qL = drafts + 1): trunk only, no head - + the walk computes logits/argmax from the returned hidden.""" + if self._dflash_capture is not None: + h, caps = self.model( + verify_input, prompt_cache, capture_layers=self._dflash_capture) + return self._dflash_pack(h, caps), {} + return self.model(verify_input, prompt_cache), {} + + def rollback_speculative_cache( + self, prompt_cache, gdn_states, accepted: int, block_size: int + ) -> None: + """Trim the rejected verify tail from every layer cache, two-phase: + verify ALL are trimmable before mutating ANY (the shared attention + mask is built from one layer's offset, so a partial rollback would + desync layers and corrupt decode).""" + del gdn_states + rejected = int(block_size) - int(accepted) - 1 + if rejected <= 0: + return + refused = [ + type(c).__name__ for c in prompt_cache if not c.is_trimmable() + ] + if refused: + raise RuntimeError( + f"Muse Glimmer MTP rollback: untrimmable cache leaves " + f"{refused} (rejected={rejected}); the sliding leaves must be " + f"BufferedRotatingKVCache by now" + ) + for c in prompt_cache: + if c.trim(rejected) != rejected: + raise RuntimeError( + f"Muse Glimmer MTP rollback: {type(c).__name__}.trim" + f"({rejected}) refused after is_trimmable() - cache state " + f"is now inconsistent" + ) + + +class MuseGlimmerSpecLM(SpecHooks, mg.Model): + """Vendored Muse Glimmer ``Model`` + the speculative hooks, in the shape + the owned MTP engine drives (``model.language_model``).""" + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + inputs_embeds: Optional[mx.array] = None, + n_to_process: Optional[int] = None, + return_hidden: bool = False, + return_shared_kv: bool = False, + **kwargs, + ): + # mlx-vlm's chunked prefill calls language_model(inputs=ids, ...) by + # keyword. shared_kv is never used (the drafter owns its KV). + del n_to_process, kwargs + want_hidden = return_hidden or return_shared_kv + if want_hidden and self._dflash_capture is not None: + h, caps = self.model( + inputs, cache, capture_layers=self._dflash_capture, + inputs_embeds=inputs_embeds) + return _SpecOutput(logits=self._spec_logits(h), + hidden_states=[self._dflash_pack(h, caps)]) + h = self.model(inputs, cache, inputs_embeds=inputs_embeds) + logits = self._spec_logits(h) + if not want_hidden: + from mlx_vlm.models.base import LanguageModelOutput + + return LanguageModelOutput(logits=logits) + return _SpecOutput(logits=logits, hidden_states=[h]) From c4ba31aafd000317ac88d5d8f47d98624eb20e69 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:28:38 -0700 Subject: [PATCH 07/17] feat(mtp): compose muse-glimmer vision and dflash speculative decoding on one model --- gmlx/mtp_load.py | 17 ++++++++++++----- gmlx/muse_glimmer_vlm_model.py | 29 +++++++++++++++++++++-------- 2 files changed, 33 insertions(+), 13 deletions(-) diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index e943f0c..b7ecbc0 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -1448,14 +1448,21 @@ def load_vlm_mtp_model( "this VLM arch can't run text-only MTP" ) - # 3. drafter - assistant (a --draft-gguf companion; gemma4) or native-head - # (nextn block inside the LLM GGUF; qwen3.5/3.6). + # 3. drafter - assistant (a --draft-gguf companion; gemma4, or a + # muse-glimmer dflash) or native-head (nextn block inside the LLM GGUF; + # qwen3.5/3.6). loadlog.stage("loading drafter") loadlog.fact("drafter", "assistant" if draft_gguf_path else "native-head") if draft_gguf_path: - drafter = _load_gemma4_assistant_drafter( - draft_gguf_path, model, zero_copy=zero_copy, log=_log - ) + if config.get("model_type") == "muse_glimmer": + drafter = _load_dflash_drafter( + draft_gguf_path, model, config["text_config"], + zero_copy=zero_copy, log=_log + ) + else: + drafter = _load_gemma4_assistant_drafter( + draft_gguf_path, model, zero_copy=zero_copy, log=_log + ) else: # Native head: load_vlm_model already loaded the target and applied the # mlx-lm tiled-V patch, but it discards the raw GGUF arrays the drafter's diff --git a/gmlx/muse_glimmer_vlm_model.py b/gmlx/muse_glimmer_vlm_model.py index 4b68a00..b8a9ff7 100644 --- a/gmlx/muse_glimmer_vlm_model.py +++ b/gmlx/muse_glimmer_vlm_model.py @@ -41,7 +41,8 @@ class the text-only path uses (:mod:`gmlx.muse_glimmer_model`), wrapped here in from mlx_vlm.models.cache import KVCache, RotatingKVCache from mlx_vlm.models.interpolate import bilinear_interpolate -from .muse_glimmer_model import MuseGlimmerModel, scale_and_softcap +from .muse_glimmer_model import MuseGlimmerModel +from .muse_glimmer_mtp import SpecHooks, _SpecOutput def ensure_registered() -> None: @@ -310,7 +311,10 @@ def __call__(self, x: mx.array) -> mx.array: # Text tower, in the shape mlx-vlm's generate stack expects -class LanguageModel(nn.Module): +class LanguageModel(SpecHooks, nn.Module): + """The text tower, carrying the same speculative hooks as the text-only + target so ``--mmproj`` and ``--speculative`` compose on one model.""" + def __init__(self, config: TextConfig): super().__init__() self.config = config @@ -325,15 +329,24 @@ def __call__( inputs_embeds: Optional[mx.array] = None, mask: Optional[mx.array] = None, cache=None, + return_hidden: bool = False, + return_shared_kv: bool = False, **kwargs, ): + # The backbone builds its own sliding/full masks from the caches. + del mask, kwargs + want_hidden = return_hidden or return_shared_kv + if want_hidden and self._dflash_capture is not None: + h, caps = self.model( + inputs, cache, capture_layers=self._dflash_capture, + inputs_embeds=inputs_embeds) + return _SpecOutput(logits=self._spec_logits(h), + hidden_states=[self._dflash_pack(h, caps)]) h = self.model(inputs, cache=cache, inputs_embeds=inputs_embeds) - if self.config.tie_word_embeddings: - out = self.model.embed_tokens.as_linear(h) - else: - out = self.lm_head(h) - return LanguageModelOutput(logits=scale_and_softcap( - out, self.config.output_multiplier, self.config.final_logit_softcapping)) + logits = self._spec_logits(h) + if not want_hidden: + return LanguageModelOutput(logits=logits) + return _SpecOutput(logits=logits, hidden_states=[h]) @property def layers(self): From ecc590a53149aba628ea90846cfe739fe408f1c2 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:33:35 -0700 Subject: [PATCH 08/17] test(muse-glimmer): config synth, dflash container split, and the packed-hidden capture seam --- tests/test_config_synth.py | 90 ++++++++++++ tests/test_dflash_normalize.py | 96 +++++++++++- tests/test_muse_glimmer_mtp.py | 258 +++++++++++++++++++++++++++++++++ 3 files changed, 442 insertions(+), 2 deletions(-) create mode 100644 tests/test_muse_glimmer_mtp.py diff --git a/tests/test_config_synth.py b/tests/test_config_synth.py index 6b415c2..ffa1f65 100644 --- a/tests/test_config_synth.py +++ b/tests/test_config_synth.py @@ -1664,3 +1664,93 @@ def test_qwen3_rope_scaling_none_omitted(): m["qwen3.rope.scaling.type"] = "none" c = synthesize_config(m, tensor_shapes={}) assert "rope_scaling" not in c + + +def _muse_glimmer_meta() -> dict: + arch = "muse-glimmer" + m = _base_meta(arch) + m[f"{arch}.block_count"] = 4 + m[f"{arch}.attention.sliding_window"] = 512 + # 3-of-4 sliding, matching the shipped [T,T,T,F] schedule. + m[f"{arch}.attention.sliding_window_pattern"] = [True, True, True, False] + m[f"{arch}.logit_scale"] = 0.19611613 + m[f"{arch}.final_logit_softcapping"] = 20.0 + return m + + +# Untied head, per-head qk-norms, and the attention output gate. +_MUSE_GLIMMER_SHAPES = { + "output.weight": [64, VOCAB], + "blk.0.attn_q_norm.weight": [16], + "blk.0.attn_k_norm.weight": [16], + "blk.0.attn_gate.weight": [64, 64], +} + + +def test_muse_glimmer_synth_instantiates(): + from gmlx import muse_glimmer_model + muse_glimmer_model.ensure_registered() + + c = synthesize_config(_muse_glimmer_meta(), tensor_shapes=_MUSE_GLIMMER_SHAPES) + assert c["model_type"] == "muse_glimmer" + assert c["num_hidden_layers"] == 4 + assert c["sliding_window"] == 512 + assert c["layer_types"] == [ + "sliding_attention", "sliding_attention", "sliding_attention", + "full_attention"] + assert c["output_multiplier"] == pytest.approx(0.19611613) + assert c["final_logit_softcapping"] == 20.0 + # Not a GGUF field: pinned from llama.cpp + the HF text_config. + assert c["post_norm_eps"] == 1e-8 + assert c["rms_norm_eps"] == pytest.approx(1e-6) + assert c["rope_parameters"]["rope_type"] == "default" + assert not c["tie_word_embeddings"] + + from gmlx.muse_glimmer_model import Model, ModelArgs + + model = Model(ModelArgs.from_dict(c)) + mx.eval(model.parameters()) + out = model(mx.array([[1, 2, 3]])) + logits = getattr(out, "logits", out) + assert logits.shape == (1, 3, VOCAB) + # The softcap bounds every logit, which is what makes it observable. + assert float(mx.abs(logits).max().item()) <= 20.0 + + +def test_muse_glimmer_synth_scalar_pattern_period(): + """llama.cpp also accepts a scalar period: full every ``period``-th layer.""" + m = _muse_glimmer_meta() + m["muse-glimmer.attention.sliding_window_pattern"] = 4 + c = synthesize_config(m, tensor_shapes=_MUSE_GLIMMER_SHAPES) + assert c["layer_types"] == [ + "sliding_attention", "sliding_attention", "sliding_attention", + "full_attention"] + + +def test_muse_glimmer_synth_absent_pattern_defaults_to_period_4(): + m = _muse_glimmer_meta() + del m["muse-glimmer.attention.sliding_window_pattern"] + c = synthesize_config(m, tensor_shapes=_MUSE_GLIMMER_SHAPES) + assert c["layer_types"][3] == "full_attention" + assert c["layer_types"][:3] == ["sliding_attention"] * 3 + + +def test_muse_glimmer_synth_pattern_length_must_match_layers(): + m = _muse_glimmer_meta() + m["muse-glimmer.attention.sliding_window_pattern"] = [True, False] + with pytest.raises(ValueError, match="entries for"): + synthesize_config(m, tensor_shapes=_MUSE_GLIMMER_SHAPES) + + +def test_muse_glimmer_synth_requires_logit_scale(): + m = _muse_glimmer_meta() + del m["muse-glimmer.logit_scale"] + with pytest.raises(Exception): + synthesize_config(m, tensor_shapes=_MUSE_GLIMMER_SHAPES) + + +def test_muse_glimmer_synth_softcap_defaults_off_when_absent(): + m = _muse_glimmer_meta() + del m["muse-glimmer.final_logit_softcapping"] + c = synthesize_config(m, tensor_shapes=_MUSE_GLIMMER_SHAPES) + assert c["final_logit_softcapping"] == 0.0 diff --git a/tests/test_dflash_normalize.py b/tests/test_dflash_normalize.py index 6362e0d..980b228 100644 --- a/tests/test_dflash_normalize.py +++ b/tests/test_dflash_normalize.py @@ -1,6 +1,8 @@ #!/usr/bin/env python3 -"""dflash -> deepseek4-dspark normalization: llama.cpp's container for the -DSpark drafter (the unsloth release) translated to the gmlx-native namespace. +"""The ``dflash`` GGUF container holds two unrelated drafters: DeepSeek-V4's +DSpark (the unsloth release, normalized into the gmlx-native namespace) and +Muse Glimmer's, which keeps its own Qwen3-shaped remap. Both are covered here, +along with the tensor-presence test that tells them apart. Name/metadata logic only - synthetic arrays, no GGUF files, no model load.""" from __future__ import annotations @@ -9,8 +11,10 @@ mx = pytest.importorskip("mlx.core") from gmlx.mtp_load import ( # noqa: E402 + dflash_container, normalize_dflash_arrays, remap_deepseek4_dspark_arrays, + remap_muse_glimmer_dflash_arrays, ) # Tensor skeleton of the unsloth dspark-DeepSeek-V4-Flash-0731-Q8_0.gguf: @@ -115,3 +119,91 @@ def test_normalized_set_survives_dspark_remap(): assert stats["mapped"] == len( [n for n in n_arrays if not n.endswith(".scales")] ) + + +# --- Muse Glimmer's dflash container ----------------------------------------- + +# 5 Qwen3-shaped layers plus three roots; no markov/confidence/hc head and no +# attn_q_a, which is what separates it from DSpark. +_MUSE_BLK_LEAVES = ( + "attn_norm.weight", "attn_q.weight", "attn_k.weight", "attn_v.weight", + "attn_output.weight", "attn_q_norm.weight", "attn_k_norm.weight", + "ffn_norm.weight", "ffn_gate.weight", "ffn_up.weight", "ffn_down.weight", +) +_MUSE_ROOT_LEAVES = ("fc.weight", "enc.output_norm.weight", "output_norm.weight") + + +def _muse_fixture(n_layers=5): + arrays = {} + for i in range(n_layers): + for leaf in _MUSE_BLK_LEAVES: + arrays[f"blk.{i}.{leaf}"] = mx.zeros((2, 2)) + for leaf in _MUSE_ROOT_LEAVES: + arrays[leaf] = mx.zeros((2, 2)) + kquant = {"blk.0.attn_q.weight": "q4_k", "fc.weight": "q4_k"} + arrays["blk.0.attn_q.scales"] = mx.zeros((1,)) + arrays["fc.scales"] = mx.zeros((1,)) + return arrays, kquant + + +def test_container_tells_the_two_drafters_apart(): + muse, _ = _muse_fixture() + dspark, _, _ = _dflash_fixture() + assert dflash_container(muse) == "muse_glimmer" + assert dflash_container(dspark) == "dspark" + + +def test_container_rejects_an_unrecognized_dflash(): + with pytest.raises(Exception): + dflash_container({"blk.0.mystery.weight": mx.zeros((2, 2))}) + + +def test_muse_remap_produces_every_drafter_param(): + arrays, kquant = _muse_fixture() + hf, hf_kq, stats = remap_muse_glimmer_dflash_arrays(arrays, kquant) + for key in ("fc.weight", "hidden_norm.weight", "norm.weight", + "layers.0.input_layernorm.weight", + "layers.0.self_attn.q_proj.weight", + "layers.0.self_attn.k_proj.weight", + "layers.0.self_attn.v_proj.weight", + "layers.0.self_attn.o_proj.weight", + "layers.0.self_attn.q_norm.weight", + "layers.0.self_attn.k_norm.weight", + "layers.0.post_attention_layernorm.weight", + "layers.4.mlp.gate_proj.weight", + "layers.4.mlp.up_proj.weight", + "layers.4.mlp.down_proj.weight"): + assert key in hf, key + assert hf_kq["layers.0.self_attn.q_proj.weight"] == "q4_k" + assert hf_kq["fc.weight"] == "q4_k" + assert not any(n.startswith("blk.") for n in hf) + assert stats["mapped"] == len( + [n for n in arrays if not n.endswith(".scales")]) + + +def test_muse_remap_keeps_the_two_norms_distinct(): + """``enc.output_norm`` is the post-fc encoder norm and ``output_norm`` the + drafter's final norm; swapping them silently corrupts the borrowed head.""" + arrays, kquant = _muse_fixture() + arrays["enc.output_norm.weight"] = mx.full((2, 2), 3.0) + arrays["output_norm.weight"] = mx.full((2, 2), 7.0) + hf, _, _ = remap_muse_glimmer_dflash_arrays(arrays, kquant) + assert float(hf["hidden_norm.weight"][0, 0]) == 3.0 + assert float(hf["norm.weight"][0, 0]) == 7.0 + + +def test_muse_remap_unknown_tensor_is_hard_error(): + arrays, kquant = _muse_fixture() + arrays["blk.0.mystery.weight"] = mx.zeros((2, 2)) + with pytest.raises(Exception): + remap_muse_glimmer_dflash_arrays(arrays, kquant) + + +def test_dspark_path_is_unchanged_by_the_split(): + """The muse container must not perturb DSpark: its fixture still lands on + the dspark namespace with the same mapped count.""" + arrays, kquant, meta = _dflash_fixture() + n_arrays, n_kquant, n_meta = normalize_dflash_arrays(arrays, kquant, meta) + assert "mtp.0.main_proj.weight" in n_arrays + assert n_meta["dspark.target_layer_ids"] == [40, 41, 42] + assert not any(n.startswith("blk.") for n in n_arrays) diff --git a/tests/test_muse_glimmer_mtp.py b/tests/test_muse_glimmer_mtp.py new file mode 100644 index 0000000..2e67f4d --- /dev/null +++ b/tests/test_muse_glimmer_mtp.py @@ -0,0 +1,258 @@ +"""Muse Glimmer speculative target: the packed-hidden capture seam and the +verify/rollback walk, on a tiny random model (no GGUF, no weights). + +The seam widens every engine-facing hidden to ``[trunk | cap ...]`` so the +drafter can read the target's residuals without an engine change. Two things +have to hold for that to be safe, and both are load-bearing: + +- the logits hooks must slice the trunk back out, and +- the slice must be materialized before it reaches the logit head, which is a + quantized kernel on a real model and reads the buffer directly. A lazy + strided view hands it the packed strides and it reads the wrong rows - + the target emitted token soup and draft acceptance fell to ~3%. Nothing in + a float model reproduces that, so the invariant is pinned directly here and + the numeric end of it rides the integration tier. +""" + +import mlx.core as mx +import pytest + +from gmlx.config_synth import synthesize_config +from gmlx.muse_glimmer_model import ModelArgs, ensure_registered +from gmlx.muse_glimmer_mtp import MuseGlimmerSpecLM + +from test_config_synth import _MUSE_GLIMMER_SHAPES, _muse_glimmer_meta + +CAPTURE = (0, 2) +N_GEN = 16 +BLOCK = 4 +# A tiny random model's argmax is near-tied at many steps, and the verify path +# derives its tokens through a block SDPA (qL = drafts + 1) whose rounding +# differs from the 1-token decode path by ~1e-3 in logit space. Stop the +# identity claim at the first step whose top-2 margin sits under that floor. +GREEDY_TIE_TOL = 1e-2 + + +def _build(): + ensure_registered() + cfg = synthesize_config(_muse_glimmer_meta(), + tensor_shapes=_MUSE_GLIMMER_SHAPES) + lm = MuseGlimmerSpecLM(ModelArgs.from_dict(cfg)) + mx.eval(lm.parameters()) + return lm, cfg + + +def _packed_width(cfg): + return cfg["hidden_size"] * (1 + len(CAPTURE)) + + +def test_capture_arms_the_packed_hidden(): + lm, cfg = _build() + ids = mx.array([[1, 2, 3, 4]]) + plain, _ = lm.speculative_verify_hidden(ids, lm.make_cache()) + assert plain.shape[-1] == cfg["hidden_size"] + + lm.set_dflash_capture(CAPTURE) + packed, _ = lm.speculative_verify_hidden(ids, lm.make_cache()) + assert packed.shape[-1] == _packed_width(cfg) + assert lm._dflash_capture == CAPTURE + + +def test_trunk_slice_recovers_the_unpacked_hidden(): + lm, cfg = _build() + ids = mx.array([[1, 2, 3, 4]]) + bare, _ = lm.speculative_verify_hidden(ids, lm.make_cache()) + lm.set_dflash_capture(CAPTURE) + packed, _ = lm.speculative_verify_hidden(ids, lm.make_cache()) + + trunk = lm._dflash_trunk(packed) + mx.eval(bare, trunk) + assert trunk.shape == bare.shape + assert float(mx.abs(trunk - bare).max().item()) == 0.0 + # and the hooks agree with the unpacked logits + lm.set_dflash_capture(()) + ref = lm.speculative_logits_from_hidden(bare) + lm.set_dflash_capture(CAPTURE) + got = lm.speculative_logits_from_hidden(packed) + mx.eval(ref, got) + assert float(mx.abs(ref - got).max().item()) == 0.0 + + +def test_trunk_is_materialized_before_the_logit_head(monkeypatch): + """White-box on purpose: a float head cannot show the difference, but the + real head is a quantized kernel that reads the buffer directly.""" + lm, cfg = _build() + lm.set_dflash_capture(CAPTURE) + calls = [] + real = mx.contiguous + monkeypatch.setattr( + mx, "contiguous", lambda x, *a, **k: (calls.append(x.shape), real(x, *a, **k))[1]) + packed = mx.zeros((1, 3, _packed_width(cfg))) + lm.speculative_logits_from_hidden(packed) + assert calls, ( + "the packed trunk slice must be materialized before the logit head; " + "a lazy strided view makes the quantized kernel read the wrong rows" + ) + + +def test_argmax_hook_matches_the_logits_hook(): + lm, _ = _build() + lm.set_dflash_capture(CAPTURE) + packed, _ = lm.speculative_verify_hidden( + mx.array([[1, 2, 3, 4]]), lm.make_cache()) + logits = lm.speculative_logits_from_hidden(packed) + am = lm.speculative_argmax_from_hidden(packed) + mx.eval(logits, am) + assert am.tolist() == mx.argmax(logits, axis=-1).tolist() + + +def test_rollback_trims_every_layer_cache(): + lm, _ = _build() + cache = lm.make_cache() + lm.speculative_verify_hidden(mx.array([[1, 2, 3, 4, 5, 6]]), cache) + before = [c.offset for c in cache] + lm.speculative_verify_hidden(mx.array([[7] * BLOCK]), cache) + assert all(c.offset == b + BLOCK for c, b in zip(cache, before)) + # accepted=1 keeps the bonus row plus one draft; the rest is rejected + lm.rollback_speculative_cache(cache, None, 1, BLOCK) + assert all(c.offset == b + 2 for c, b in zip(cache, before)) + + +def test_rollback_is_a_noop_when_the_whole_block_is_accepted(): + lm, _ = _build() + cache = lm.make_cache() + lm.speculative_verify_hidden(mx.array([[1, 2, 3, 4]]), cache) + lm.speculative_verify_hidden(mx.array([[5] * BLOCK]), cache) + offsets = [c.offset for c in cache] + lm.rollback_speculative_cache(cache, None, BLOCK - 1, BLOCK) + assert [c.offset for c in cache] == offsets + + +@pytest.mark.parametrize("armed", [False, True]) +def test_verify_walk_is_token_identical_to_greedy(armed): + """The engine contract: whatever the drafter proposes, the walk emits the + target's own greedy tokens. Driven with deliberately wrong drafts so every + round takes the reject-and-rollback path. + + Parametrized over the capture seam because the bug this guards only + appeared with capture armed - the packed slice reached the head unevaluated. + """ + lm, _ = _build() + prompt = mx.array([[1, 2, 3, 4, 5]]) + + cache = lm.make_cache() + h = lm.model(prompt, cache) + ref_logits = lm._spec_logits(h)[0, -1] + ref, margins = [], [] + for _ in range(N_GEN): + top = mx.sort(ref_logits)[-2:] + margins.append(float((top[1] - top[0]).item())) + t = int(mx.argmax(ref_logits).item()) + ref.append(t) + ref_logits = lm._spec_logits(lm.model(mx.array([[t]]), cache))[0, -1] + + if armed: + lm.set_dflash_capture(CAPTURE) + cache2 = lm.make_cache() + hid, _ = lm.speculative_verify_hidden(prompt, cache2) + tok = int(lm.speculative_argmax_from_hidden(hid)[0, -1].item()) + got = [tok] + while len(got) < N_GEN: + drafts = [900 + i for i in range(BLOCK - 1)] # never the target's pick + hid, _ = lm.speculative_verify_hidden( + mx.array([[tok] + drafts]), cache2) + rows = lm.speculative_argmax_from_hidden(hid)[0].tolist() + accepted = 0 + for i, d in enumerate(drafts): + if int(rows[i]) != d: + break + accepted += 1 + got.extend(drafts[:accepted]) + tok = int(rows[accepted]) + got.append(tok) + lm.rollback_speculative_cache(cache2, None, accepted, BLOCK) + + # compare over the prefix whose greedy pick is unambiguous + limit = next((i for i, m in enumerate(margins) if m < GREEDY_TIE_TOL), N_GEN) + assert limit > 0, "tiny model produced no unambiguous step" + assert got[:limit] == ref[:limit] + + +# --- the drafter side of the same seam --------------------------------------- + +def _build_drafter(cfg, n_layers=2): + from gmlx.muse_glimmer_dflash import ( + MuseGlimmerDFlashConfig, + MuseGlimmerDFlashDrafter, + ) + + return MuseGlimmerDFlashDrafter(MuseGlimmerDFlashConfig( + hidden_size=cfg["hidden_size"], + intermediate_size=64, + num_hidden_layers=n_layers, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + rms_norm_eps=1e-6, + vocab_size=cfg["vocab_size"], + max_position_embeddings=1024, + rope_theta=10000.0, + tie_word_embeddings=False, + block_size=BLOCK, + mask_token_id=7, + target_layer_ids=list(CAPTURE), + num_target_layers=cfg["num_hidden_layers"], + layer_types=["sliding_attention"] * n_layers, + sliding_window=512, + draft_window_size=512, + final_logit_softcapping=cfg["final_logit_softcapping"], + output_multiplier=cfg["output_multiplier"], + )) + + +def test_drafter_captures_are_materialized_for_the_quantized_fc(monkeypatch): + """Same trap as the logit head: ``fc`` is a quantized 5*hidden -> hidden + matmul, so the trailing slice of the packed hidden must not reach it lazily.""" + lm, cfg = _build() + drafter = _build_drafter(cfg) + calls = [] + real = mx.contiguous + monkeypatch.setattr( + mx, "contiguous", lambda x, *a, **k: (calls.append(x.shape), real(x, *a, **k))[1]) + drafter._captures(mx.zeros((1, 3, _packed_width(cfg)))) + assert calls, "the packed capture slice must be materialized before fc" + + +def test_drafter_captures_reject_an_unpacked_hidden(): + """A target whose ``_dflash_capture`` was never armed hands over a bare + trunk; that must fail loudly rather than matmul against garbage.""" + lm, cfg = _build() + drafter = _build_drafter(cfg) + with pytest.raises(ValueError, match="packed hidden width"): + drafter._captures(mx.zeros((1, 3, cfg["hidden_size"]))) + + +def test_drafter_captures_take_the_trailing_block(): + lm, cfg = _build() + drafter = _build_drafter(cfg) + h = cfg["hidden_size"] + packed = mx.concatenate( + [mx.zeros((1, 2, h)), mx.ones((1, 2, h)), mx.full((1, 2, h), 2.0)], + axis=-1) + caps = drafter._captures(packed) + mx.eval(caps) + assert caps.shape[-1] == h * len(CAPTURE) + assert float(caps[0, 0, 0].item()) == 1.0 # first capture + assert float(caps[0, 0, h].item()) == 2.0 # second, in order + + +def test_drafter_satisfies_the_protocol(): + from gmlx.drafter_protocol import validate_drafter + + lm, cfg = _build() + drafter = _build_drafter(cfg) + mx.eval(drafter.parameters()) + drafter.bind(lm) + validate_drafter(drafter) + assert drafter.uses_shared_kv is False + assert drafter.requires_owned_engine is True From 188cc5435a050204bd25a127b31382441a3e57cd Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:34:47 -0700 Subject: [PATCH 09/17] test(muse-glimmer): tensor remap pins and the atem tool parser --- tests/test_muse_glimmer_tools.py | 123 +++++++++++++++++++++++++++++++ tests/test_remap.py | 64 ++++++++++++++++ 2 files changed, 187 insertions(+) create mode 100644 tests/test_muse_glimmer_tools.py diff --git a/tests/test_muse_glimmer_tools.py b/tests/test_muse_glimmer_tools.py new file mode 100644 index 0000000..1f929b4 --- /dev/null +++ b/tests/test_muse_glimmer_tools.py @@ -0,0 +1,123 @@ +"""Muse Glimmer ATEM tool-call parser (``gmlx.muse_glimmer_tools``). +CPU-only, no model load. + +Exercises ``parse_tool_call`` on wrapper-inner text as mlx-vlm's +``process_tool_calls`` hands it over (the caller strips the module's +``tool_call_start``/``tool_call_end``), the schema-driven value coercion, and +the ``ensure_registered`` graft into mlx-vlm's parser registry. +""" + +from __future__ import annotations + +import sys + +import pytest + +from gmlx import muse_glimmer_tools +from gmlx.muse_glimmer_tools import ( + parse_tool_call, + tool_call_end, + tool_call_start, +) + +_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + "precise": {"type": "boolean"}, + "ratio": {"type": "number"}, + "opts": {"type": "object"}, + "tags": {"type": "array"}, + }, + }, + }, + } +] + + +def _invoke(name, args=()): + parts = [f''] + for k, v in args: + parts.append(f'{v}') + parts.append("") + return "".join(parts) + + +def test_wrapper_tags_are_the_outer_block(): + # The streaming suppressor keys on the wrapper, not the per-call tag, so a + # partial tail is always a prefix of tool_call_start. + assert tool_call_start == "" + assert tool_call_end == "" + + +def test_single_call_parses_name_and_arguments(): + calls = parse_tool_call(_invoke("get_weather", [("city", "Paris")]), _TOOLS) + assert calls == [{"name": "get_weather", "arguments": {"city": "Paris"}}] + + +def test_parallel_calls_parse_in_order(): + text = _invoke("get_weather", [("city", "Paris")]) + _invoke( + "get_weather", [("city", "Rome")]) + calls = parse_tool_call(text, _TOOLS) + assert [c["arguments"]["city"] for c in calls] == ["Paris", "Rome"] + + +def test_values_are_coerced_by_the_tool_schema(): + calls = parse_tool_call(_invoke("get_weather", [ + ("city", "Paris"), + ("days", "3"), + ("precise", "true"), + ("ratio", "0.5"), + ("opts", '{"a": 1}'), + ("tags", '["x", "y"]'), + ]), _TOOLS) + args = calls[0]["arguments"] + assert args["city"] == "Paris" + assert args["days"] == 3 and isinstance(args["days"], int) + assert args["precise"] is True + assert args["ratio"] == pytest.approx(0.5) + assert args["opts"] == {"a": 1} + assert args["tags"] == ["x", "y"] + + +def test_unknown_parameter_stays_a_string(): + # No schema entry means no guess at the type. + calls = parse_tool_call( + _invoke("get_weather", [("mystery", "3")]), _TOOLS) + assert calls[0]["arguments"]["mystery"] == "3" + + +def test_unschemad_tool_leaves_every_value_a_string(): + calls = parse_tool_call(_invoke("other", [("days", "3")]), _TOOLS) + assert calls[0]["arguments"]["days"] == "3" + + +def test_no_tools_argument_is_tolerated(): + calls = parse_tool_call(_invoke("get_weather", [("city", "Paris")])) + assert calls[0]["name"] == "get_weather" + + +def test_freeform_text_becomes_an_unknown_envelope(): + # Never guess a call out of prose - hand the raw text back instead. + out = parse_tool_call("just some prose", _TOOLS) + assert out == {"name": "unknown", "arguments": {"raw": "just some prose"}} + + +def test_multiline_parameter_value_is_preserved(): + body = _invoke("get_weather", [("city", "Paris\nFrance")]) + assert parse_tool_call(body, _TOOLS)[0]["arguments"]["city"] == "Paris\nFrance" + + +def test_ensure_registered_grafts_into_mlx_vlm(): + pytest.importorskip("mlx_vlm.tool_parsers") + muse_glimmer_tools.ensure_registered() + assert "mlx_vlm.tool_parsers.muse_glimmer" in sys.modules + mod = sys.modules["mlx_vlm.tool_parsers.muse_glimmer"] + assert hasattr(mod, "parse_tool_call") + muse_glimmer_tools.ensure_registered() # idempotent diff --git a/tests/test_remap.py b/tests/test_remap.py index 411c83a..ef613e6 100644 --- a/tests/test_remap.py +++ b/tests/test_remap.py @@ -1231,3 +1231,67 @@ def test_unknown_arch_skips_not_fails(): def test_unknown_tensor_on_known_arch_hard_fails(): r = d("llama", "blk.0.this_is_not_a_real_tensor.weight") assert r.kind == FAIL + + +# muse-glimmer: sandwich norms with the +1 already baked (consumed as-is), +# an attention output gate, and interleaved Q/K consumed without a permute. +def test_muse_glimmer_ffn_norm_pins_to_pre_feedforward(): + # FFN_NORM/FFN_PRE_NORM collide in TENSOR_NAMES; muse-glimmer's ffn_norm + # is the PRE-FFN norm, same as gemma2/gemma3. + r = d("muse-glimmer", "blk.1.ffn_norm.weight") + assert r.kind == MAP + assert r.hf_name == "model.layers.1.pre_feedforward_layernorm.weight" + + +def test_muse_glimmer_norms_are_not_unbaked(): + # The converter folds the +1 into the four per-layer norms and a plain + # nn.RMSNorm consumes them directly - no gemma-style unbake. + for t in ("attn_norm", "ffn_norm", "post_attention_norm", "post_ffw_norm"): + assert d("muse-glimmer", f"blk.0.{t}.weight").transform == "passthrough" + assert d("muse-glimmer", "output_norm.weight").transform == "passthrough" + + +def test_muse_glimmer_post_norms_land_on_the_sandwich_slots(): + assert (d("muse-glimmer", "blk.2.post_attention_norm.weight").hf_name + == "model.layers.2.post_attention_layernorm.weight") + assert (d("muse-glimmer", "blk.2.post_ffw_norm.weight").hf_name + == "model.layers.2.post_feedforward_layernorm.weight") + + +def test_muse_glimmer_attn_gate_claims_the_gate_projection(): + # CANONICAL_HF["ATTN_GATE"] points at qwen3.5's linear_attn.in_proj_z, so + # muse-glimmer has to claim this row first. + r = d("muse-glimmer", "blk.0.attn_gate.weight") + assert r.kind == MAP + assert r.hf_name == "model.layers.0.self_attn.gate_proj.weight" + assert "linear_attn" not in r.hf_name + + +def test_muse_glimmer_qk_are_not_permuted(): + # LLAMA_ROPE_TYPE_NORM: the converter already un-permuted HF's rotate_half + # into interleaved, and the model class ropes with traditional=True. A + # qk_permute here would double-apply it (and force _own() copies). + for t in ("attn_q", "attn_k"): + r = d("muse-glimmer", f"blk.3.{t}.weight") + assert r.kind == MAP and r.bid == 3 + assert r.transform == "passthrough" + + +def test_muse_glimmer_qk_norms_and_projections_resolve_canonically(): + for gguf, hf in ( + ("blk.0.attn_q_norm.weight", "model.layers.0.self_attn.q_norm.weight"), + ("blk.0.attn_k_norm.weight", "model.layers.0.self_attn.k_norm.weight"), + ("blk.0.attn_v.weight", "model.layers.0.self_attn.v_proj.weight"), + ("blk.0.attn_output.weight", "model.layers.0.self_attn.o_proj.weight"), + ("blk.0.ffn_gate.weight", "model.layers.0.mlp.gate_proj.weight"), + ("blk.0.ffn_up.weight", "model.layers.0.mlp.up_proj.weight"), + ("blk.0.ffn_down.weight", "model.layers.0.mlp.down_proj.weight"), + ): + r = d("muse-glimmer", gguf) + assert r.kind == MAP and r.hf_name == hf, gguf + + +def test_muse_glimmer_globals_resolve(): + assert d("muse-glimmer", "token_embd.weight").hf_name == "model.embed_tokens.weight" + assert d("muse-glimmer", "output_norm.weight").hf_name == "model.norm.weight" + assert d("muse-glimmer", "output.weight").hf_name == "lm_head.weight" From 49b8e1376301d49d6ffcb4e4e946ff13247a378d Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:35:30 -0700 Subject: [PATCH 10/17] test(muse-glimmer): atem reasoning channel transcripts across chunk boundaries --- tests/test_reasoning.py | 80 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/tests/test_reasoning.py b/tests/test_reasoning.py index f167ea2..5a0397e 100644 --- a/tests/test_reasoning.py +++ b/tests/test_reasoning.py @@ -537,3 +537,83 @@ def test_split_harmony_reply_shapes(): assert split_harmony_reply(capped) == ("Count the", "") plain = "Six." assert split_harmony_reply(plain) == (None, "Six.") + + +# --- muse-glimmer ATEM channel ------------------------------------------------ +# +# Routing is on the message HEADER, not on a "to=self" marker: any marker whose +# text can occur in prose would misclassify other models, and one starting with +# a space or letter would make _partial_tail_start hold that character back at +# every chunk boundary. The header is swallowed whole and only an explicit +# "to=self" recipient routes to reasoning. + +_MUSE = ( + "<|start|>assistant to=self<|message|>Work out the capital.<|eom|>" + "<|start|>assistant<|message|>The capital of France is Paris.<|eot|>" +) +# The generation prompt already ends with "<|start|>assistant", so the FIRST +# generated turn streams the header tail only - hence the start_in_header seed. +_MUSE_FIRST = ( + " to=self<|message|>Work out the capital.<|eom|>" + "<|start|>assistant to=user<|message|>The capital of France is Paris.<|eot|>" +) +_MUSE_TOOL = ( + "<|start|>assistant to=self<|message|>Need the weather.<|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '' + 'Paris' + "<|eot|>" +) + + +def _segment_header(text: str, *, chunk: int = 0): + f = ReasoningFilter(start_in_header=True) + spans: list[tuple[str, str]] = [] + if chunk: + for i in range(0, len(text), chunk): + spans += f.feed(text[i : i + chunk]) + else: + spans += f.feed(text) + spans += f.flush() + return ("".join(t for t, m in spans if m == "reason"), + "".join(t for t, m in spans if m == "answer")) + + +@pytest.mark.parametrize("chunk", [0, 1, 3, 7, 13]) +def test_muse_glimmer_self_channel_then_answer(chunk): + reason, answer, _ = _segment(_MUSE, chunk=chunk) + assert reason == "Work out the capital." + assert answer == "The capital of France is Paris." + for frag in ("<|", "|>", "to=self", "message", "eom", "eot"): + assert frag not in reason + answer + + +@pytest.mark.parametrize("chunk", [0, 1, 3, 7, 13]) +def test_muse_glimmer_first_turn_needs_the_header_seed(chunk): + reason, answer = _segment_header(_MUSE_FIRST, chunk=chunk) + assert reason == "Work out the capital." + assert answer == "The capital of France is Paris." + # " to=user" is header too - it must be swallowed, not leaked as answer text + assert "to=user" not in reason + answer + assert not answer.startswith(" to=") + + +@pytest.mark.parametrize("chunk", [0, 1, 3, 7, 13]) +def test_muse_glimmer_tool_recipient_is_swallowed_and_call_is_answer(chunk): + reason, answer, _ = _segment(_MUSE_TOOL, chunk=chunk) + assert reason == "Need the weather." + assert "to=weather.get" not in reason + answer + # the tool block itself is answer-side, so the tool parser still sees it + assert answer.startswith("") + assert 'name="get_weather"' in answer + + +@pytest.mark.parametrize("chunk", [0, 1, 3, 7, 13]) +def test_muse_glimmer_eom_returns_to_answer_without_a_self_header(chunk): + """A non-final message that is not addressed to self stays answer-side.""" + text = ("<|start|>assistant to=self<|message|>think<|eom|>" + "<|start|>assistant<|message|>first<|eom|>" + "<|start|>assistant<|message|>second<|eot|>") + reason, answer, _ = _segment(text, chunk=chunk) + assert reason == "think" + assert answer == "firstsecond" From 863fcd2dd02274168772c7c509872c8672d57564 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:36:45 -0700 Subject: [PATCH 11/17] test(muse-glimmer): profiles, thinking markers, projector dispatch, and drafter discovery --- tests/test_discovery.py | 33 +++++++++++++++++++++++++++++ tests/test_profiles.py | 19 +++++++++++++++++ tests/test_thinking_budget.py | 39 +++++++++++++++++++++++++++++++++++ tests/test_vlm_resolve.py | 3 +++ 4 files changed, 94 insertions(+) diff --git a/tests/test_discovery.py b/tests/test_discovery.py index bda4919..6099db6 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -848,3 +848,36 @@ def test_read_sampling_normalizes_disabled_sentinels(): "general.sampling.top_k": -1, "general.sampling.repeat_penalty": 1.0} assert disc._read_sampling(kv) == {"temperature": 0.9, "top_p": 0.0, "top_k": 0} + + +# muse-glimmer: the per-family arch tuple, and the "dflash" id marker + + +def test_muse_glimmer_companion_is_found_by_arch_tuple(tmp_path, monkeypatch): + """A muse-glimmer target asks for ``dflash`` only - the deepseek4 arches + are not in its tuple, so a dspark sidecar next door is not picked up.""" + target = tmp_path / "Muse-Glimmer-30B-Q6_K_L.gguf" + for name in ("dflash-kquant.gguf", target.name): + (tmp_path / name).write_bytes(b"GGUF") + metas = {str(tmp_path / "dflash-kquant.gguf"): {"arch": "dflash"}, + str(target): {"arch": "muse-glimmer"}} + monkeypatch.setattr(disc, "header_meta", lambda p: metas.get(str(p))) + assert disc.find_mtp_companion(str(target), ("dflash",)) == str( + tmp_path / "dflash-kquant.gguf") + + +def test_muse_glimmer_companion_ignores_a_native_dspark_sidecar( + tmp_path, monkeypatch): + target = tmp_path / "Muse-Glimmer-30B-Q6_K_L.gguf" + for name in ("dspark-sidecar.gguf", target.name): + (tmp_path / name).write_bytes(b"GGUF") + metas = {str(tmp_path / "dspark-sidecar.gguf"): {"arch": "deepseek4-dspark"}, + str(target): {"arch": "muse-glimmer"}} + monkeypatch.setattr(disc, "header_meta", lambda p: metas.get(str(p))) + assert disc.find_mtp_companion(str(target), ("dflash",)) is None + + +def test_dflash_is_an_id_marker(): + """Without this the drafter quant splits the model id and a dflash sidecar + is mistaken for a separate model.""" + assert "dflash" in disc._ID_MARKERS diff --git a/tests/test_profiles.py b/tests/test_profiles.py index a87e593..3e2e520 100644 --- a/tests/test_profiles.py +++ b/tests/test_profiles.py @@ -157,3 +157,22 @@ def test_describe_complete(): for r in rows: assert set(r["intents"]) == set(profiles.BUILTIN_INTENTS) assert r["label"] + + +def test_muse_glimmer_family_carries_the_atem_think_markers_and_four_levels(): + """The model card's four reasoning levels map onto the template's own + ``reasoning_strength`` spelling, and the base sampling seeds the open-think + detector with the ATEM header rather than the '' default.""" + from gmlx.profiles import FAMILIES + + fam = FAMILIES["muse"] + assert fam["arches"] == ("muse-glimmer",) + s = fam["base"]["sampling"] + assert s["thinking_start_token"] == "<|start|>assistant to=self<|message|>" + assert s["thinking_end_token"] == "<|eom|>" + levels = { + name.split("-", 1)[1]: spec["chat_template_kwargs"]["reasoning_strength"] + for name, spec in fam["intents"].items() + } + assert levels == {"low": "low", "medium": "medium", + "high": "high", "xhigh": "xhigh"} diff --git a/tests/test_thinking_budget.py b/tests/test_thinking_budget.py index 56ad31f..8ae9815 100644 --- a/tests/test_thinking_budget.py +++ b/tests/test_thinking_budget.py @@ -601,3 +601,42 @@ class _Marker: assert prompt_open_think_tag( "x <|channel>thought reasoning", tokenizer=_Marker() ) == "<|channel>thought" + + +class _MuseTok: + """Just the surface ``_template_think_pair`` reads.""" + chat_template = ( + "{%- if reasoning_strength %}Reasoning strength: {{ reasoning_strength }}" + "{%- endif %}<|start|>assistant to=self<|message|>...<|eom|>" + ) + + +def test_muse_glimmer_think_pair_resolves_from_the_template(): + from gmlx.thinking_budget import _template_think_pair + + assert _template_think_pair(_MuseTok()) == ( + "<|start|>assistant to=self<|message|>", "<|eom|>") + + +def test_muse_glimmer_generation_prompt_is_not_open_thinking(): + """The generation prompt stops at ``<|start|>assistant``; the reasoning + header is only completed by what the model generates, so the budget must + not start counting yet.""" + from gmlx.thinking_budget import prompt_opens_thinking + + prompt = "<|start|>user<|message|>hi<|eot|><|start|>assistant" + assert not prompt_opens_thinking(prompt, tokenizer=_MuseTok()) + + +def test_muse_glimmer_open_reasoning_message_is_detected(): + from gmlx.thinking_budget import prompt_opens_thinking + + prompt = "<|start|>assistant to=self<|message|>half a thought" + assert prompt_opens_thinking(prompt, tokenizer=_MuseTok()) + + +def test_muse_glimmer_closed_reasoning_message_is_not_open(): + from gmlx.thinking_budget import prompt_opens_thinking + + prompt = "<|start|>assistant to=self<|message|>done<|eom|>" + assert not prompt_opens_thinking(prompt, tokenizer=_MuseTok()) diff --git a/tests/test_vlm_resolve.py b/tests/test_vlm_resolve.py index dcabc20..50cbffb 100644 --- a/tests/test_vlm_resolve.py +++ b/tests/test_vlm_resolve.py @@ -18,6 +18,9 @@ ("qwen3vlmoe", {"clip.vision.projector_type": "qwen3vl_merger"}, "qwen3_omni_moe"), ("gemma4", {"clip.vision.projector_type": "gemma4v"}, "gemma4"), ("gemma4", {"clip.vision.projector_type": "gemma4uv"}, "gemma4_unified"), + ("muse-glimmer", {"clip.projector_type": "muse-glimmer"}, "muse_glimmer"), + ("muse-glimmer", {"clip.vision.projector_type": "muse-glimmer"}, + "muse_glimmer"), ]) def test_supported_families_resolve(llm_arch, mm_meta, expected): assert resolve_vlm_model_type(llm_arch, mm_meta) == expected From 8b8adadd6e7b2922f0e900bbbce511b730c9c956 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:44:25 -0700 Subject: [PATCH 12/17] test(vlm): muse-glimmer vision index math, mmproj remap, grid search, and the eot-not-eom stop set --- tests/test_muse_glimmer_vision.py | 240 ++++++++++++++++++++++++++++++ tests/test_tokenizer.py | 59 ++++++++ 2 files changed, 299 insertions(+) create mode 100644 tests/test_muse_glimmer_vision.py diff --git a/tests/test_muse_glimmer_vision.py b/tests/test_muse_glimmer_vision.py new file mode 100644 index 0000000..041e296 --- /dev/null +++ b/tests/test_muse_glimmer_vision.py @@ -0,0 +1,240 @@ +"""Muse Glimmer vision: the pure index math, the mmproj remap and the image +preprocessor's grid search. CPU-only - no GGUF, no weights, no image decode. + +Everything here is a function of the patch grid or of a tensor name, and every +one of them is a silent-failure surface: a wrong permutation, a transposed patch +conv or an off-by-one grid produces a model that runs and describes the wrong +picture. The expectations are derived from llama.cpp (``clip.cpp`` window/ds +permutations and ``muse_glimmer_grid_size`` in ``tools/mtmd/mtmd-image.cpp``), +not read back off this port. +""" + +from __future__ import annotations + +import pytest + +mx = pytest.importorskip("mlx.core") + +from gmlx.muse_glimmer_vlm_model import ( # noqa: E402 + pixel_shuffle_order, + window_order, +) +from gmlx.vlm import ( # noqa: E402 + _MuseGlimmerGgufImageProcessor, + _muse_glimmer_vision_name, + remap_vision_arrays, +) + +# --- window permutation ------------------------------------------------------- + + +def test_window_order_is_a_permutation_with_one_segment_per_patch(): + perm, segment = window_order(7, 5, 3) + assert sorted(perm) == list(range(35)) + assert len(segment) == 35 + + +def test_window_order_is_the_identity_when_the_grid_fits_one_window(): + perm, segment = window_order(4, 4, 32) + assert perm == list(range(16)) + assert set(segment) == {0} + + +def test_window_order_partial_edge_windows(): + # 3x3 grid, window 2: windows of 4, 2, 2 and 1 patches, in row-major window + # order, exactly as llama.cpp builds them. + perm, segment = window_order(3, 3, 2) + assert perm == [0, 1, 3, 4, 2, 5, 6, 7, 8] + assert segment == [0, 0, 0, 0, 1, 1, 2, 2, 3] + + +def test_segments_are_contiguous_so_the_mask_is_block_diagonal(): + _, segment = window_order(9, 6, 4) + assert segment == sorted(segment) + + +def test_two_patches_share_a_segment_iff_they_share_a_window(): + grid_w, grid_h, window = 7, 5, 3 + perm, segment = window_order(grid_w, grid_h, window) + for i, p in enumerate(perm): + for j, q in enumerate(perm): + same_window = ((p % grid_w) // window == (q % grid_w) // window + and (p // grid_w) // window == (q // grid_w) // window) + assert (segment[i] == segment[j]) is same_window + + +def test_inverse_permutation_restores_row_major_order(): + # The tower undoes the permutation with a scatter of arange; if that is + # wrong the patch features land under the wrong soft tokens. + perm = mx.array(window_order(7, 5, 3)[0]) + inverse = mx.zeros(perm.shape, dtype=mx.int32) + inverse[perm] = mx.arange(perm.size, dtype=mx.int32) + x = mx.arange(35).reshape(35, 1) + assert mx.take(mx.take(x, perm, axis=0), inverse, axis=0).reshape(-1).tolist() \ + == list(range(35)) + + +# --- pixel shuffle ------------------------------------------------------------ + + +def test_pixel_shuffle_groups_each_cell_contiguously(): + # 4x2 grid, merge 2: cell (0,0) is patches 0,1,4,5 and cell (0,1) is 2,3,6,7. + assert pixel_shuffle_order(4, 2, 2) == [0, 1, 4, 5, 2, 3, 6, 7] + + +def test_pixel_shuffle_covers_every_patch_once(): + order = pixel_shuffle_order(6, 4, 2) + assert sorted(order) == list(range(24)) + + +def test_pixel_shuffle_merge_one_is_the_identity(): + assert pixel_shuffle_order(5, 3, 1) == list(range(15)) + + +def test_pixel_shuffle_drops_the_ragged_edge(): + # An odd grid side has no partial cell: llama.cpp merges floor(side/merge) + # cells and the preprocessor guarantees even sides anyway. + assert len(pixel_shuffle_order(5, 3, 2)) == (5 // 2) * (3 // 2) * 4 + + +# --- mmproj tensor remap ------------------------------------------------------ + +_ROOTS = ( + "mm.0.weight", "mm.1.weight", "mm.2.weight", + "v.patch_embd.weight", "v.position_embd.weight", + "v.post_ln.bias", "v.post_ln.weight", "v.pre_ln.bias", "v.pre_ln.weight", +) +_BLK_LEAVES = ( + "attn_q", "attn_k", "attn_v", "attn_out", "ln1", "ln2", "ffn_up", "ffn_down", +) + + +def _mmproj_names(n_layers=2): + names = list(_ROOTS) + for i in range(n_layers): + for leaf in _BLK_LEAVES: + names += [f"v.blk.{i}.{leaf}.weight", f"v.blk.{i}.{leaf}.bias"] + return names + + +def test_every_mmproj_tensor_is_claimed(): + assert all(_muse_glimmer_vision_name(n) is not None for n in _mmproj_names()) + + +def test_only_the_patch_conv_is_flagged_for_transpose(): + flagged = [n for n in _mmproj_names() if _muse_glimmer_vision_name(n)[1]] + assert flagged == ["v.patch_embd.weight"] + + +def test_projector_splits_across_the_adapter_and_the_llm(): + # mm.0/mm.1 are the mmproj-side adapter; mm.2 projects into the text + # residual width and lives outside it, where the HF checkpoint keeps it. + assert _muse_glimmer_vision_name("mm.0.weight")[0] == "vision_adapter.fc1.weight" + assert _muse_glimmer_vision_name("mm.1.weight")[0] == "vision_adapter.fc2.weight" + assert _muse_glimmer_vision_name("mm.2.weight")[0] == "vision_projection.weight" + + +def test_ffn_up_and_down_keep_their_direction(): + # fc1 widens and fc2 narrows; swapping them still loads (both linears exist) + # and silently produces garbage features. + assert _muse_glimmer_vision_name("v.blk.3.ffn_up.weight")[0] \ + == "vision_tower.layers.3.mlp.fc1.weight" + assert _muse_glimmer_vision_name("v.blk.3.ffn_down.bias")[0] \ + == "vision_tower.layers.3.mlp.fc2.bias" + + +def test_block_norms_and_attention_land_on_the_vendored_paths(): + got = {leaf: _muse_glimmer_vision_name(f"v.blk.0.{leaf}.weight")[0] + for leaf in _BLK_LEAVES} + assert got["ln1"] == "vision_tower.layers.0.layer_norm1.weight" + assert got["ln2"] == "vision_tower.layers.0.layer_norm2.weight" + assert got["attn_q"] == "vision_tower.layers.0.self_attn.q_proj.weight" + assert got["attn_out"] == "vision_tower.layers.0.self_attn.o_proj.weight" + + +def test_position_embedding_is_a_bare_array_not_a_module_weight(): + # The tower holds it as a plain mx.array attribute, so the target name has + # no ``.weight`` suffix. + assert _muse_glimmer_vision_name("v.position_embd.weight")[0] \ + == "vision_tower.position_embedding" + + +def test_unknown_tensor_is_skipped_not_guessed(): + assert _muse_glimmer_vision_name("v.blk.0.mystery.weight") is None + assert _muse_glimmer_vision_name("a.blk.0.attn_q.weight") is None + + +def test_remap_transposes_the_patch_conv_and_skips_nothing(): + arrays = {n: mx.zeros((2, 2)) for n in _mmproj_names()} + arrays["v.patch_embd.weight"] = mx.zeros((8, 3, 14, 14)) # [out, in, kH, kW] + out, skipped, kq = remap_vision_arrays(arrays, "muse_glimmer") + assert skipped == [] and kq == {} + assert out["vision_tower.patch_embed.weight"].shape == (8, 14, 14, 3) + assert len(out) == len(arrays) + + +# --- image preprocessing grid search ----------------------------------------- + +CELL = 28 # patch_size 14 * spatial_merge_size 2 + + +def _proc(max_image_tokens=4096): + return _MuseGlimmerGgufImageProcessor( + image_mean=[0.5] * 3, image_std=[0.5] * 3, + max_image_tokens=max_image_tokens) + + +@pytest.mark.parametrize("h,w,expect", [ + (280, 280, (280, 280)), # exact square, no search needed + (280, 140, (280, 140)), # exact 2:1 + (300, 200, (308, 196)), # 10.71 x 7.14: the 11x7 tie wins on tokens + (10, 100, (28, 112)), # thinner than one cell: clamps the short side + (10, 10, (28, 28)), # smaller than one cell +]) +def test_target_grid_matches_the_llama_cpp_search(h, w, expect): + assert _proc()._target_hw(h, w) == expect + + +def test_ties_go_to_the_larger_grid(): + # 300x200 has |11/7 - 3/2| == |10/7 - 3/2|; llama.cpp breaks the tie toward + # more tokens, which is 11x7 rather than 10x7. + h_out, w_out = _proc()._target_hw(300, 200) + assert (h_out // CELL, w_out // CELL) == (11, 7) + + +@pytest.mark.parametrize("h,w", [(4000, 4000), (8000, 1000), (1000, 8000), + (1234, 567), (33, 4001)]) +def test_output_is_always_whole_cells_under_the_token_cap(h, w): + proc = _proc() + h_out, w_out = proc._target_hw(h, w) + assert h_out % CELL == 0 and w_out % CELL == 0 + assert h_out >= CELL and w_out >= CELL + assert proc.soft_tokens(h_out, w_out) <= proc.max_image_tokens + + +def test_oversized_image_shrinks_to_the_cap_keeping_its_ratio(): + proc = _proc() + h_out, w_out = proc._target_hw(20000, 10000) + assert (h_out // CELL, w_out // CELL) == (90, 45) + assert proc.soft_tokens(h_out, w_out) == 4050 + + +def test_square_oversized_image_lands_exactly_on_the_cap(): + proc = _proc() + h_out, w_out = proc._target_hw(10000, 10000) + assert (h_out, w_out) == (64 * CELL, 64 * CELL) + assert proc.soft_tokens(h_out, w_out) == 4096 + + +def test_no_candidate_under_the_cap_falls_back_to_round_and_clamp(): + # A 100:1 image against a 1-token budget: every floor/ceil pair either has a + # zero side or exceeds the cap, so the search rounds and clamps instead of + # returning nothing. + h_out, w_out = _proc(max_image_tokens=1)._target_hw(1000, 10) + assert (h_out // CELL, w_out // CELL) == (10, 1) + + +def test_soft_tokens_counts_merged_cells(): + proc = _proc() + assert proc.soft_tokens(280, 140) == 10 * 5 + assert proc.soft_tokens(CELL, CELL) == 1 diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py index adbaded..dfb9dd2 100644 --- a/tests/test_tokenizer.py +++ b/tests/test_tokenizer.py @@ -447,6 +447,65 @@ def test_stop_set_dedups_when_eot_equals_eos(): assert tok._gguf_eos_token_ids.count(1) == 1 # no duplicate +# Muse Glimmer declares eot and deliberately not eom. <|eom|> closes a non-final +# message - the reasoning channel ends on it and the answer follows - so folding +# it into the stop set would cut every reply off at end-of-thinking. Both routes +# into the stop set have to agree on that: the metadata read and the chat +# template heuristic. +_MUSE_SPECIALS = ["", "", "", "<|eom|>", "<|eot|>"] +_MUSE_EOM, _MUSE_EOT = 3, 4 +_MUSE_TEMPLATE = ( + "{%- for m in messages -%}" + "{{- '<|start|>' + m['role'] + '<|message|>' + m['content'] -}}" + "{{- '<|eot|>' if loop.last else '<|eom|>' -}}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%}" +) + + +def _muse_meta() -> dict: + toks = _MUSE_SPECIALS + _ALPHABET + _MERGED + meta = { + "general.architecture": "muse-glimmer", + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.pre": "llama4", + "tokenizer.ggml.tokens": toks, + "tokenizer.ggml.merges": _MERGES, + "tokenizer.ggml.token_type": + [3] * len(_MUSE_SPECIALS) + [1] * (len(toks) - len(_MUSE_SPECIALS)), + "tokenizer.ggml.bos_token_id": 0, + "tokenizer.ggml.eos_token_id": 1, + "tokenizer.ggml.padding_token_id": 2, + "tokenizer.chat_template": _MUSE_TEMPLATE, + } + return meta + + +def test_muse_glimmer_metadata_declares_eot_only(): + from gmlx.tokenizer import _metadata_stop_ids + + # The real ids, as the 30B GGUF carries them: eos <|end_of_text|> 200001, + # eot <|eot|> 200008. <|eom|> is 200007 and appears in no metadata key. + meta = {"tokenizer.ggml.eot_token_id": 200008} + assert _metadata_stop_ids(meta, 202048) == [200008] + + +def test_muse_glimmer_stop_set_keeps_eot_and_not_eom(): + meta = _muse_meta() + meta["tokenizer.ggml.eot_token_id"] = _MUSE_EOT + tok = load_tokenizer_from_gguf(meta, "muse-glimmer") + assert _MUSE_EOT in tok._gguf_eos_token_ids + assert _MUSE_EOM not in tok._gguf_eos_token_ids + + +def test_muse_glimmer_template_heuristic_also_lands_on_eot(): + # Without the metadata key, the template heuristic is the only route; the + # rendered assistant turn ends on <|eot|>, so it must not adopt <|eom|>. + tok = load_tokenizer_from_gguf(_muse_meta(), "muse-glimmer") + assert _MUSE_EOT in tok._gguf_eos_token_ids + assert _MUSE_EOM not in tok._gguf_eos_token_ids + + # add_eos_token: a post-processor appends EOS on the raw path (parity with # llama.cpp), bypassed on the chat path (add_special_tokens=False) so a template # carrying its own EOS never doubles up - symmetric with add_bos_token. From f78e632a14e20dd7ced82cd1b3baca83cf39eed6 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:51:00 -0700 Subject: [PATCH 13/17] docs(arch): muse-glimmer coverage row, vision family, drafter and reasoning mentions, changelog --- CHANGELOG.md | 9 +++++++ README.md | 4 +-- docs/adding-architectures.md | 14 +++++++--- docs/arch-coverage.md | 3 ++- docs/cli.md | 51 ++++++++++++++++++------------------ docs/performance.md | 5 ++-- docs/server-config.md | 5 ++++ docs/vlm.md | 6 +++++ 8 files changed, 63 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index abad8c9..0e25f31 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,15 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Added +- Muse Glimmer support (GGUF arch `muse-glimmer`, Meta Muse Glimmer 30B): + a vendored text decoder with sandwich norms at two epsilons, an + attention output gate, and NoPE on the full-attention layers with RoPE + only on the 2048-window sliding ones. Vision rides the `muse-glimmer` + mmproj through a vendored ViT and a GGUF-only image processor. The ATEM + reasoning channel and its XML tool calls are wired through chat, serve, + thinking budgets, and the muse profile family (reasoning strength + low/medium/high/xhigh). `--draft-gguf` loads the DFlash drafter for + speculative decoding. - serve --speculative: a request arriving while one stream decodes with MTP no longer waits for it to finish; the stream converts to shared batch decode and speculation resumes once the batch drains back under diff --git a/README.md b/README.md index bf17c4c..11e550e 100644 --- a/README.md +++ b/README.md @@ -265,8 +265,8 @@ Coverage runs across the major open-weight families: Llama and Mistral; Qwen 2 through 3.6, dense and MoE, including the gated-DeltaNet hybrids and Qwen3-Next; Gemma 1 through 4 plus DiffusionGemma; DeepSeek V3/R1 and V4-Flash; GLM 4 through 5.2; gpt-oss; Kimi-K3; MiniMax M2 and M3; Hunyuan -A13B and Hy3; Granite, including the 4.x hybrids; Nemotron-H; Falcon-H1; -ERNIE-4.5; Phi-3; Seed-OSS; and SmolLM3. New architectures land regularly, +A13B and Hy3; Muse Glimmer; Granite, including the 4.x hybrids; Nemotron-H; +Falcon-H1; ERNIE-4.5; Phi-3; Seed-OSS; and SmolLM3. New architectures land regularly, and a family is listed in the generated [architecture coverage matrix](https://github.com/asher/gmlx/blob/main/docs/arch-coverage.md) only after token-parity certification against llama.cpp at 16k context; the diff --git a/docs/adding-architectures.md b/docs/adding-architectures.md index a66a48d..798ed34 100644 --- a/docs/adding-architectures.md +++ b/docs/adding-architectures.md @@ -4,10 +4,16 @@ What it takes for a new model family to become a supported architecture, and the acceptance gate every family clears before its row appears in the [coverage matrix](arch-coverage.md). -Precondition: a GGUF arch is only reachable if the installed mlx-lm defines a -`class Model` for the corresponding `model_type`. If mlx-lm doesn't implement -the architecture, support is blocked upstream. gmlx never reimplements -model math. +Precondition: a GGUF arch needs a model class for its `model_type`. That class +normally comes from the installed mlx-lm or mlx-vlm, and gmlx supplies only the +tensor map and the config. + +A few families have no upstream class at all (kimi-k3, muse-glimmer). gmlx +vendors the model math for those, in its own module, grafted into the upstream +namespace so a later upstream implementation wins. Vendoring is the exception. +It is worth the cost only when the family is otherwise unreachable, and it adds +two obligations: numeric parity against llama.cpp, and a collision check that +reports the graft once upstream ships its own class. ## The shape of the work diff --git a/docs/arch-coverage.md b/docs/arch-coverage.md index 65789f3..9916b9b 100644 --- a/docs/arch-coverage.md +++ b/docs/arch-coverage.md @@ -4,7 +4,7 @@ A GGUF loads iff gmlx maps its `general.architecture` to a `model_type` with a model class in one of its installed backend packages (usually mlx-lm; a few archs are backed by mlx-vlm (e.g. DiffusionGemma) or mlx-embeddings (gemma-embedding)) and a config synthesizer exists (or `hf_source` supplies a config.json). -**35 / 36 mapped architectures load end-to-end today** (no `hf_source` needed). +**36 / 37 mapped architectures load end-to-end today** (no `hf_source` needed). | GGUF arch | model_type | model class | remap | synth | status | notes | |-----------|------------|:-----------:|-------|:-----:|--------|-------| @@ -32,6 +32,7 @@ A GGUF loads iff gmlx maps its `general.architecture` to a `model_type` with a m | `minimax-m2` | `minimax` | yes | yes | yes | loadable | MiniMax-M2 (230B-A10B); every-layer fine-grained sigmoid-gated MoE (no dense layers, no shared expert) with full attention, full-width qk-norm (RMSNorm over head_dim*n_heads), and partial rotary (rotary_dim < head_dim). NEOX rope (no qk-permute); router/experts/correction-bias nested under block_sparse_moe.*; head_dim != hidden//heads (from key_length) | | `minimax-m3` | `minimax_m3` | yes | yes | yes | loadable | MiniMax-M3 (428B-A23B); M2's GQA base (partial rotary, head_dim from key_length, NEOX rope) plus gemma-style +1 RMSNorms (unbaked on load), per-head qk-norm, and a DeepSeek-V3-shaped MoE: leading dense layers, sigmoid gating + correction bias, routed weights renormalized x expert_weights_scale, per-layer shared expert (block_sparse_moe.shared_experts). SwiGLU-OAI activation. MSA sparse attention (llama.cpp PR #24908 semantics): a per-GQA-group indexer max-pools scores into 128-token blocks and top-16 blocks (local block forced) bound attention to 2048 KV per query - the form the model is trained with; runs whenever the GGUF carries the blk.N.indexer.* tensors or a `*indexer*.gguf` sidecar sits next to the model (GMLX_INDEXER_SIDECAR overrides discovery, GMLX_MSA_DISABLE=1 forces dense for A/B). Indexless GGUFs fall back to dense with a one-time warning (exact to 2048 tokens, degrades beyond - reasoning loops). Thinking tags are ``/`` (template-detected; the vocab's legacy `` entries are decoys). Model class vendored from mlx-lm PR #1401 (+ gmlx MSA extension) until upstream ships models/minimax_m3.py | | `mistral3` | `ministral3` | yes | yes | yes | loadable | llama.cpp 'mistral3' = Ministral-3 / Mistral-Small-3.1, Llama layout | +| `muse-glimmer` | `muse_glimmer` | yes | yes | yes | loadable | Meta Muse Glimmer 30B (llama.cpp 'muse-glimmer'); dense 52-layer sandwich-norm decoder - pre/post norms around both attention and FFN, at two epsilons (1e-5 pre/final, 1e-8 post) - with an afmoe-shaped attention output gate (sigmoid(x_norm @ W_gate) applied before o_proj), per-head qk-norm whose weights are synthesized at conversion to absorb qk_scale_factor (q_norm a uniform 3.87, k_norm ones), an unweighted RMSNorm on the token embeddings before layer 0, and a logit scale (output_multiplier 0.196) + gemma-style tanh softcap at 20. Attention runs a 3-of-4 sliding window (2048) with RoPE, and full attention with NoPE on every 4th layer - the inverse of the usual arrangement, and the reason the model card says '131072+': the largest positional offset ever resolved is the 2048 window, so there is no rope extrapolation ceiling and past 131k the only limit is KV memory at ~13 KB/token (just the 13 global layers grow; the 39 sliding ones are pinned by their rotating cache). The four per-layer norms carry a baked +1 that a plain RMSNorm consumes as-is (no gemma unbake). NORM rope with Q/K un-permuted at conversion => traditional=True, no qk_permute. Model class vendored (no upstream mlx-lm class; afmoe is the nearest relative). Pairs with the Muse Glimmer mmproj (--mmproj) for vision and with the DFlash drafter GGUF (--draft-gguf) for speculative decoding | | `nemotron_h_moe` | `nemotron_h` | yes | yes | yes | loadable | NVIDIA Nemotron-H MoE hybrid: Mamba2 SSM layers + sparse attention layers + MoE MLPs; layer_norm_epsilon (not rms_norm_eps) | | `phi3` | `phi3` | yes | yes | yes | loadable | Phi-3 mini/small/medium; fused qkv + gate_up; pass hf_source for 128K long-context (su/longrope) variants | | `qwen2` | `qwen2` | yes | yes | yes | loadable | Qwen2 / Qwen2.5 dense; QKV biases, tied embeddings on 0.5B/1.5B | diff --git a/docs/cli.md b/docs/cli.md index f7b610a..5aafe67 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -121,11 +121,12 @@ plumbing at all (`--bench`, `--bench-depths`, `--report-only`, `--stream-cpu`, falling back to plain text generation. A text-only request under `--mmproj` runs through the MTP speculative path -whenever a drafter is available: a `--draft-gguf` assistant (gemma4) or a -native `nextn` head in the LLM GGUF (qwen3.5/3.6, no companion needed). The -verify walk only touches the language model, so a resident VLM gets the decode -speedup on text turns, token-identical to the same model's text-only MTP. An -image or audio request uses the plain VLM path; the drafter is idle that turn. +whenever a drafter is available: a `--draft-gguf` assistant (gemma4, +muse-glimmer) or a native `nextn` head in the LLM GGUF (qwen3.5/3.6, no +companion needed). The verify walk only touches the language model, so a +resident VLM gets the decode speedup on text turns, token-identical to the same +model's text-only MTP. An image or audio request uses the plain VLM path; the +drafter is idle that turn. ### Speculative / MTP @@ -138,9 +139,9 @@ with a warning); `--no-speculative`/`--no-mtp` forces it off. | Flag | Meaning | |------|---------| -| `--speculative` / `--mtp` | Force MTP speculative decoding on. Native-head models (qwen3.5/3.6 `nextn`) need no companion; gemma4 needs `--draft-gguf`. Native heads are auto-enabled without this. Use it to force the path when a sampler flag would otherwise defer. | +| `--speculative` / `--mtp` | Force MTP speculative decoding on. Native-head models (qwen3.5/3.6 `nextn`) need no companion; gemma4 and muse-glimmer need `--draft-gguf`. Native heads are auto-enabled without this. Use it to force the path when a sampler flag would otherwise defer. | | `--no-speculative` / `--no-mtp` | Disable MTP. Overrides the native-head auto-enable and config `speculative: true`. | -| `--draft-gguf PATH` | Separate assistant-drafter GGUF (gemma4 two-GGUF MTP shape, or a deepseek4 DSpark/MTP sidecar - gmlx `deepseek4-dspark`, llama.cpp `dflash`, or legacy `deepseek4_mtp_support`); implies `--speculative` (same as `serve`). A sidecar in the target's directory is autodetected without the flag. | +| `--draft-gguf PATH` | Separate assistant-drafter GGUF (gemma4 two-GGUF MTP shape, a muse-glimmer DFlash drafter, or a deepseek4 DSpark/MTP sidecar - gmlx `deepseek4-dspark`, llama.cpp `dflash`, or legacy `deepseek4_mtp_support`); implies `--speculative` (same as `serve`). A sidecar in the target's directory is autodetected without the flag. | | `--draft-block-size N` | Override the MTP draft block size. | Speculative generation takes only `--temp`/`--top-p`/`--top-k`/`--min-p` plus a @@ -351,28 +352,28 @@ every command. The terminal is upgraded on top: for now. - MTP speculative decoding (auto for native heads; `--no-mtp` to disable): a native-head model (qwen3.5/3.6 `nextn`) drafts and verifies multiple tokens - per step for a decode speedup; gemma4 needs a `--draft-gguf` assistant. The - reply streams the same way and ends with the same `tok/s` stat, and the - persistent KV cache is reused across turns exactly like the text path. Not - combinable with `--adapter` / `--stream-*`. Sampling is + per step for a decode speedup; gemma4 and muse-glimmer need a `--draft-gguf` + assistant. The reply streams the same way and ends with the same `tok/s` + stat, and the persistent KV cache is reused across turns exactly like the + text path. Not combinable with `--adapter` / `--stream-*`. Sampling is temperature/top-p/top-k/min-p only; the MTP verify walk has no penalty/bias hooks, so the other `/` sampling commands don't apply on this path. - VLM + MTP: a `--mmproj` VLM with a drafter (a `--draft-gguf` assistant - for gemma4, or a native `nextn` head for qwen3.5/3.6) keeps MTP on for - text-only turns (the fast path above) while `/image` / `/audio` turns - fall back to the plain VLM stream. The first media turn upgrades the - session to the VLM path for the rest of the conversation, since the text - tokenizer can't render a history that holds image markers. The prior text - turns are carried into that re-prefill so nothing is lost. + for gemma4 or muse-glimmer, or a native `nextn` head for qwen3.5/3.6) + keeps MTP on for text-only turns (the fast path above) while `/image` / + `/audio` turns fall back to the plain VLM stream. The first media turn + upgrades the session to the VLM path for the rest of the conversation, + since the text tokenizer can't render a history that holds image markers. + The prior text turns are carried into that re-prefill so nothing is lost. - Reasoning display: for thinking models (Qwen3/DeepSeek-R1/GLM ``, - gpt-oss harmony channels, Gemma `<|channel>thought`), the chain-of-thought - is stripped of its control markers and streamed in the theme's thinking - style (italic bright blue under the default `dark` theme) inside a - gutter-framed block that closes with a payoff line showing how long the - model thought and how many tokens it spent; the final answer follows in - normal weight. `--reasoning hide` collapses the reasoning to a single live - spinner that resolves to the same payoff, so you see it working without - reading it. Ctrl-O toggles expand and collapse live during a reply (and + gpt-oss harmony channels, Gemma `<|channel>thought`, Muse Glimmer's ATEM + `to=self` channel), the chain-of-thought is stripped of its control markers + and streamed in the theme's thinking style (italic bright blue under the + default `dark` theme) inside a gutter-framed block that closes with a payoff + line showing how long the model thought and how many tokens it spent; the + final answer follows in normal weight. `--reasoning hide` collapses the + reasoning to a single live spinner that resolves to the same payoff, so you + see it working without reading it. Ctrl-O toggles expand and collapse live during a reply (and persists as the default for the next). `--reasoning raw` / `/reasoning raw` passes everything through verbatim (the old behavior, for when a model's markers segment oddly). The stored conversation keeps the raw text in every diff --git a/docs/performance.md b/docs/performance.md index d658dc0..e785b2c 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -130,8 +130,9 @@ Models that ship a native multi-token-prediction head (Qwen3.5 and Qwen3.6) get speculative decoding automatically on `run` and `chat`: the head drafts tokens ahead and the base model verifies them. Output is exactly what the base model would have produced, just faster when drafts are accepted. `--no-mtp` turns it off. -gemma-4 models take the two-file shape instead: a small companion drafter GGUF via -`--draft-gguf`. On the server it is the `speculative:` config key. +gemma-4 and Muse Glimmer take the two-file shape instead: a small companion +drafter GGUF via `--draft-gguf`. On the server it is the `speculative:` config +key. Gains depend on acceptance rate and context depth. In our serve benchmarks (M5 Max, the same server with MTP off as the baseline), speculation roughly diff --git a/docs/server-config.md b/docs/server-config.md index edca07b..ff75aa8 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -407,6 +407,11 @@ Notes on individual families: (`<|open|>think<|sep|>` / `<|close|>think<|sep|>`) are set as the family's thinking tokens so open-think detection, thinking budgets, and the stream splitter track the model's real section tags. +- muse: the `@reasoning-*` intents set `reasoning_strength`, the Muse Glimmer + template's variable name. It takes `low`/`medium`/`high`/`xhigh` and defaults + to `high`. Reasoning is a message channel rather than a tag pair, so the + thinking markers are the channel's own delimiters + (`<|start|>assistant to=self<|message|>` / `<|eom|>`). - qwen3.6 / qwen3: `@instruct` also sets `enable_thinking: false` (the card's non-thinking operating point). - `default`: the fallback for unknown architectures, the historic scaffold diff --git a/docs/vlm.md b/docs/vlm.md index be7a338..b9e1268 100644 --- a/docs/vlm.md +++ b/docs/vlm.md @@ -55,6 +55,12 @@ families that share one. An unsupported pairing fails loudly at load with both n | Qwen3-Omni | `qwen3vl_merger` + `qwen3vlmoe` | Qwen3-Omni (vision + audio) | | gemma-4 omni | `gemma4v`/`gemma4a` | gemma-4-E2B / E4B (vision + audio) | | gemma-4 unified | `gemma4uv` | gemma-4-12B (encoder-free unified embedder) | +| Muse Glimmer | `muse-glimmer` + `muse-glimmer` | Muse-Glimmer-30B | + +Muse Glimmer's vision tower and image processor are implemented in gmlx. Neither +mlx-vlm nor the installed transformers ships the family, so the preprocessing +ports llama.cpp's `mtmd_image_preprocessor_muse_glimmer`. `--hf-source` is not +needed. Qwen2-VL / Qwen2.5-VL mmprojs (`qwen2vl_merger`) are not supported yet. The load fails up front with the family named. LLaVA's image processor isn't From 5de249d2f5a2187b7f4e1d0e6d8642d19b61c729 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 18:01:17 -0700 Subject: [PATCH 14/17] fix(mtp): default muse-glimmer dflash to 3 drafts per round, not the gguf block size --- gmlx/mtp_load.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index b7ecbc0..ef15bb1 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -83,6 +83,10 @@ # losing regime. Uncapped is earned by measurement, not inherited by default. _MTP_WIDTH_CAP_FALLBACK = 2 +# Drafted depth per DFlash round, matching llama.cpp's n_max=3. +# GMLX_MUSE_DFLASH_BLOCK overrides, up to the GGUF's dflash.block_size. +_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT = 4 + def _stamp_mtp_width_cap(drafter, model_type: str, *, target=None, log=loadlog.verbose_print): @@ -774,8 +778,9 @@ def _load_muse_glimmer_dflash_drafter( ] or ["full_attention"] * n_layers window = int(meta.get("dflash.attention.sliding_window") or 0) or None native_total = int(block_size) + default_total = min(_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT, native_total) block_total = max( - 2, min(env_int("GMLX_MUSE_DFLASH_BLOCK", native_total), native_total)) + 2, min(env_int("GMLX_MUSE_DFLASH_BLOCK", default_total), native_total)) config = MuseGlimmerDFlashConfig( hidden_size=int(target_config_dict["hidden_size"]), From 541e1db27c38911a7def11ace013132822ab21bf Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:04:25 -0700 Subject: [PATCH 15/17] perf(vlm): muse-glimmer f16 vision weights and size-batched window attention --- gmlx/loader.py | 15 ++++- gmlx/muse_glimmer_vlm_model.py | 108 ++++++++++++++++++++++++------ gmlx/vlm.py | 2 + tests/test_muse_glimmer_vision.py | 103 ++++++++++++++++++++++++++++ 4 files changed, 207 insertions(+), 21 deletions(-) diff --git a/gmlx/loader.py b/gmlx/loader.py index 7c86a62..3e08447 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -2715,10 +2715,17 @@ def _warm_touch_pass( # computed fp32 (the vendored cast_predicate pins the same set). "kimi_k3": (".mlp.gate.weight", ".e_score_correction_bias", ".a_folded", ".dt_bias", "_res_score"), +} + +# Params kept at their native f16 through the bf16 cast (no upcast). MLX +# promotes an f16-weight matmul against f32 activations to f32, so these read +# half the bytes of an fp32 pin while computing the same values. +_F16_KEEP_BY_MODEL_TYPE: dict[str, tuple[str, ...]] = { # muse_glimmer's mmproj is native F16 and llama.cpp runs the tower with f32 # activations. 50 residual layers with large outliers (features span +-76) # compound bf16 rounding into ~10% relative RMS on the projected embeddings - # against an f32 run; f16->f32 is lossless, so this reproduces the oracle. + # against an f32 run. The tower entry casts its input to f32, so activations + # ride fp32 promotion while the weights stay F16 - the oracle's own layout. # Vision only - the text tower's bf16 holds 16k parity. "muse_glimmer": ("vision_tower.", "vision_adapter.", "vision_projection."), } @@ -2822,6 +2829,7 @@ def _install_and_load( sanitize: bool = True, no_alias: set[str] | None = None, fp32_keep: tuple[str, ...] = (), + f16_keep: tuple[str, ...] = (), source_key: tuple | None = None, active_before: float | None = None, ) -> None: @@ -2842,7 +2850,8 @@ def _install_and_load( the same suffix match used for the kquant meta. ``fp32_keep``: target-name substrings pinned to float32 through the bf16 - cast (see ``_FP32_KEEP_BY_MODEL_TYPE``). + cast (see ``_FP32_KEEP_BY_MODEL_TYPE``). ``f16_keep``: substrings kept at + their native f16 instead (see ``_F16_KEEP_BY_MODEL_TYPE``). ``active_before``: active-memory baseline for the untracked-weights split. Callers that read wire bytes before installing must pass the pre-read @@ -2949,6 +2958,8 @@ def _install_and_load( if v.dtype != mx.float32: # e.g. F16 ape tables loadable[k] = v.astype(mx.float32) continue + if f16_keep and any(s in k for s in f16_keep): + continue if v.dtype == mx.float16: # Same-itemsize f16->bf16 gets buffer-donated into the source # view -- a write through the zero-copy file mapping (dropped diff --git a/gmlx/muse_glimmer_vlm_model.py b/gmlx/muse_glimmer_vlm_model.py index b8a9ff7..92f59da 100644 --- a/gmlx/muse_glimmer_vlm_model.py +++ b/gmlx/muse_glimmer_vlm_model.py @@ -137,6 +137,41 @@ def window_order(grid_w: int, grid_h: int, window: int) -> tuple[list[int], list return perm, segment +def window_partition( + grid_w: int, grid_h: int, window: int +) -> tuple[list[int], list[tuple[int, int, int]]]: + """Patch order that makes window attention batchable without a mask. + + Same window membership as :func:`window_order`, but windows are laid out + grouped by size (largest first) instead of row-major, so each group is a + contiguous run of equal-length windows. Returns ``(perm, groups)`` with + ``groups`` entries ``(start, n_windows, window_len)``: rows + ``perm[start : start + n_windows * window_len]`` reshape to + ``[n_windows, window_len]`` and attend without any mask. A grid has at + most four sizes (interior, right edge, bottom edge, corner), and window + order within a group stays row-major. Attention is permutation-invariant + over its keys, so the layout change cannot alter the math. + """ + windows: dict[int, list[list[int]]] = {} + for wy in range(0, grid_h, window): + for wx in range(0, grid_w, window): + rows = [ + gy * grid_w + gx + for gy in range(wy, min(wy + window, grid_h)) + for gx in range(wx, min(wx + window, grid_w)) + ] + if rows: + windows.setdefault(len(rows), []).append(rows) + perm: list[int] = [] + groups: list[tuple[int, int, int]] = [] + for length in sorted(windows, reverse=True): + group = windows[length] + groups.append((len(perm), len(group), length)) + for rows in group: + perm.extend(rows) + return perm, groups + + def pixel_shuffle_order(grid_w: int, grid_h: int, merge: int) -> list[int]: """Gather order that groups each ``merge`` x ``merge`` cell contiguously, in row-major cell order (llama.cpp's ``ds_perm``).""" @@ -161,11 +196,14 @@ def _rope_tables(pos: mx.array, half_dim: int, base: float): def _rope_half(v: mx.array, cos: mx.array, sin: mx.array) -> mx.array: """Interleaved (pairwise) rotation of ``[B, H, L, D]`` by per-position - tables of shape ``[L, D // 2]``.""" + tables of shape ``[L, D // 2]`` (or already broadcast to 4-D, e.g. + ``[B, 1, L, D // 2]`` for window-batched attention).""" B, H, L, D = v.shape v = v.reshape(B, H, L, D // 2, 2) x0, x1 = v[..., 0], v[..., 1] - c, s = cos[None, None].astype(v.dtype), sin[None, None].astype(v.dtype) + if cos.ndim == 2: + cos, sin = cos[None, None], sin[None, None] + c, s = cos.astype(v.dtype), sin.astype(v.dtype) return mx.stack([x0 * c - x1 * s, x0 * s + x1 * c], axis=-1).reshape(B, H, L, D) @@ -191,16 +229,45 @@ def __init__(self, config: VisionConfig): self.v_proj = nn.Linear(dim, dim, bias=True) self.o_proj = nn.Linear(dim, dim, bias=True) - def __call__(self, x: mx.array, tables_w, tables_h, mask) -> mx.array: + def __call__(self, x: mx.array, tables_w, tables_h, groups) -> mx.array: + """``groups`` is None for a global layer (full attention over all + patches), or the :func:`window_partition` groups for a window layer: + each group's equal-length windows run as one unmasked batched SDPA, + skipping the dense scores a block-diagonal mask would compute.""" B, L, _ = x.shape - shape = (B, L, self.n_heads, self.head_dim) - q = self.q_proj(x).reshape(shape).transpose(0, 2, 1, 3) - k = self.k_proj(x).reshape(shape).transpose(0, 2, 1, 3) - v = self.v_proj(x).reshape(shape).transpose(0, 2, 1, 3) - q = _rope_2d(q, tables_w, tables_h) - k = _rope_2d(k, tables_w, tables_h) - out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask) - return self.o_proj(out.transpose(0, 2, 1, 3).reshape(B, L, -1)) + q = self.q_proj(x) + k = self.k_proj(x) + v = self.v_proj(x) + + if groups is None: + shape = (B, L, self.n_heads, self.head_dim) + q = _rope_2d(q.reshape(shape).transpose(0, 2, 1, 3), tables_w, tables_h) + k = _rope_2d(k.reshape(shape).transpose(0, 2, 1, 3), tables_w, tables_h) + v = v.reshape(shape).transpose(0, 2, 1, 3) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=self.scale, mask=None) + return self.o_proj(out.transpose(0, 2, 1, 3).reshape(B, L, -1)) + + outs = [] + for start, n_win, w_len in groups: + end = start + n_win * w_len + shape = (n_win, w_len, self.n_heads, self.head_dim) + + def _win(t): + return t[:, start:end].reshape(shape).transpose(0, 2, 1, 3) + + def _tabs(tables): + return tuple(t[start:end].reshape(n_win, 1, w_len, -1) + for t in tables) + + tw, th = _tabs(tables_w), _tabs(tables_h) + qg = _rope_2d(_win(q), tw, th) + kg = _rope_2d(_win(k), tw, th) + og = mx.fast.scaled_dot_product_attention( + qg, kg, _win(v), scale=self.scale, mask=None) + outs.append(og.transpose(0, 2, 1, 3).reshape(1, n_win * w_len, -1)) + out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=1) + return self.o_proj(out) class VisionMLP(nn.Module): @@ -222,8 +289,8 @@ def __init__(self, config: VisionConfig): self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=eps) self.mlp = VisionMLP(config) - def __call__(self, x: mx.array, tables_w, tables_h, mask) -> mx.array: - x = x + self.self_attn(self.layer_norm1(x), tables_w, tables_h, mask) + def __call__(self, x: mx.array, tables_w, tables_h, groups) -> mx.array: + x = x + self.self_attn(self.layer_norm1(x), tables_w, tables_h, groups) return x + self.mlp(self.layer_norm2(x)) @@ -255,8 +322,10 @@ def _position_embedding(self, grid_w: int, grid_h: int) -> mx.array: if grid_w == side and grid_h == side: return self.position_embedding grid = self.position_embedding.reshape(side, side, -1) + # Stays f32: the add promotes anyway, and rounding the interpolated + # table back to the f16 weight dtype would only lose precision. resized = bilinear_interpolate(grid.astype(mx.float32), grid_h, grid_w) - return resized.reshape(grid_h * grid_w, -1).astype(self.position_embedding.dtype) + return resized.reshape(grid_h * grid_w, -1) def __call__(self, pixel_values: mx.array) -> mx.array: """``pixel_values`` is a single image as ``[1, H, W, C]``.""" @@ -267,10 +336,8 @@ def __call__(self, pixel_values: mx.array) -> mx.array: x = self.patch_embed(pixel_values).reshape(1, grid_h * grid_w, -1) x = x + self._position_embedding(grid_w, grid_h)[None] - perm, segment = window_order(grid_w, grid_h, self.window) + perm, groups = window_partition(grid_w, grid_h, self.window) perm = mx.array(perm) - seg = mx.array(segment) - window_mask = (seg[:, None] == seg[None, :])[None, None] x = self.pre_layernorm(x) x = mx.take(x, perm, axis=1) @@ -286,7 +353,7 @@ def __call__(self, pixel_values: mx.array) -> mx.array: sf = self.config.sparse_factor for idx, layer in enumerate(self.layers): is_global = idx == n_layer - 1 or (idx + 1) % sf == 0 - x = layer(x, tables_w, tables_h, None if is_global else window_mask) + x = layer(x, tables_w, tables_h, None if is_global else groups) x = self.post_layernorm(x) inverse = mx.zeros(perm.shape, dtype=mx.int32) @@ -391,7 +458,10 @@ def _image_features(self, pixel_values: mx.array, image_sizes) -> mx.array: feats = [] for i, (h, w) in enumerate(image_sizes): image = pixel_values[i, :, :h, :w].transpose(1, 2, 0)[None] - x = self.vision_tower(image.astype(self.vision_projection.weight.dtype)) + # f32 activations against F16 weights (the loader's f16_keep set): + # dtype promotion computes the whole stack in f32, the oracle's own + # layout, at half the resident bytes of an fp32 weight pin. + x = self.vision_tower(image.astype(mx.float32)) grid_h, grid_w = h // patch, w // patch order = mx.array(pixel_shuffle_order(grid_w, grid_h, merge)) n_out = (grid_h // merge) * (grid_w // merge) diff --git a/gmlx/vlm.py b/gmlx/vlm.py index 3dbfd92..2a7ae0e 100644 --- a/gmlx/vlm.py +++ b/gmlx/vlm.py @@ -34,6 +34,7 @@ ) from .gguf_meta import first_nonzero_int, read_int from .loader import ( + _F16_KEEP_BY_MODEL_TYPE, _FP32_KEEP_BY_MODEL_TYPE, _active_now, _install_and_load, @@ -2086,6 +2087,7 @@ def load_vlm_model( # towers), so model.sanitize must not run - it would re-prefix text keys. _install_and_load(model, hf_weights, hf_kquant_meta, log=_log, sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE.get(model_type, ()), + f16_keep=_F16_KEEP_BY_MODEL_TYPE.get(model_type, ()), source_key=weights_source_key(*pf.shards, mmproj_path), active_before=active_before) materialize_module_arrays(model) diff --git a/tests/test_muse_glimmer_vision.py b/tests/test_muse_glimmer_vision.py index 041e296..6e437f1 100644 --- a/tests/test_muse_glimmer_vision.py +++ b/tests/test_muse_glimmer_vision.py @@ -18,6 +18,7 @@ from gmlx.muse_glimmer_vlm_model import ( # noqa: E402 pixel_shuffle_order, window_order, + window_partition, ) from gmlx.vlm import ( # noqa: E402 _MuseGlimmerGgufImageProcessor, @@ -74,6 +75,108 @@ def test_inverse_permutation_restores_row_major_order(): == list(range(35)) +# --- size-grouped partition (the batched-attention layout) -------------------- + + +def test_partition_windows_match_window_order_on_every_grid(): + # window_partition reorders windows for maskless batching; the membership + # of every window must be identical to llama.cpp's row-major layout, or + # the attention pattern silently changes. + for grid_h in range(1, 12): + for grid_w in range(1, 12): + perm_o, segment = window_order(grid_w, grid_h, 4) + perm_p, groups = window_partition(grid_w, grid_h, 4) + assert sorted(perm_p) == list(range(grid_h * grid_w)) + reference: dict[int, list[int]] = {} + for p, s in zip(perm_o, segment): + reference.setdefault(s, []).append(p) + expected = {frozenset(v) for v in reference.values()} + got = set() + for start, n_win, w_len in groups: + for i in range(n_win): + got.add(frozenset( + perm_p[start + i * w_len:start + (i + 1) * w_len])) + assert got == expected, (grid_h, grid_w) + + +def test_partition_groups_tile_the_permutation_exactly(): + perm, groups = window_partition(9, 6, 4) + assert groups[0][0] == 0 + for (s0, n0, w0), (s1, _, _) in zip(groups, groups[1:]): + assert s0 + n0 * w0 == s1 + s, n, w = groups[-1] + assert s + n * w == len(perm) == 54 + + +def test_partition_orders_groups_largest_first(): + # 9 wide x 6 high, window 4: two interior 4x4=16, one right-edge 4x1=4, + # two bottom 2x4=8, one corner 2x1=2. Largest first keeps the dominant + # batch leading. + _, groups = window_partition(9, 6, 4) + assert [(n, w) for _, n, w in groups] == [(2, 16), (2, 8), (1, 4), (1, 2)] + + +def test_partition_single_window_grid_is_one_group(): + perm, groups = window_partition(3, 2, 32) + assert perm == list(range(6)) + assert groups == [(0, 1, 6)] + + +def test_batched_attention_matches_the_masked_reference(): + """The whole tiny tower, new batched path vs the dense block-diagonal + mask it replaced. Any partition or reshape mistake shows up as a large + error here; float noise does not.""" + from gmlx.muse_glimmer_vlm_model import ( + VisionConfig, VisionModel, _rope_2d, _rope_tables) + + cfg = VisionConfig(num_hidden_layers=4, hidden_size=64, + intermediate_size=128, num_attention_heads=4, + num_position_embeddings=16) # window side 4 + model = VisionModel(cfg) + mx.eval(model.parameters()) + + def masked_reference(pixel_values): + patch = cfg.patch_size + grid_h = pixel_values.shape[1] // patch + grid_w = pixel_values.shape[2] // patch + x = model.patch_embed(pixel_values).reshape(1, grid_h * grid_w, -1) + x = x + model._position_embedding(grid_w, grid_h)[None] + perm, segment = window_order(grid_w, grid_h, model.window) + perm, seg = mx.array(perm), mx.array(segment) + mask = (seg[:, None] == seg[None, :])[None, None] + x = mx.take(model.pre_layernorm(x), perm, axis=1) + half = (cfg.hidden_size // cfg.num_attention_heads) // 2 + tw = _rope_tables(perm % grid_w + 1, half, cfg.rope_theta) + th = _rope_tables(perm // grid_w + 1, half, cfg.rope_theta) + n_layer = len(model.layers) + for idx, layer in enumerate(model.layers): + is_global = (idx == n_layer - 1 + or (idx + 1) % cfg.sparse_factor == 0) + h = layer.layer_norm1(x) + a = layer.self_attn + B, L, _ = h.shape + shp = (B, L, a.n_heads, a.head_dim) + q = _rope_2d(a.q_proj(h).reshape(shp).transpose(0, 2, 1, 3), tw, th) + k = _rope_2d(a.k_proj(h).reshape(shp).transpose(0, 2, 1, 3), tw, th) + v = a.v_proj(h).reshape(shp).transpose(0, 2, 1, 3) + o = mx.fast.scaled_dot_product_attention( + q, k, v, scale=a.scale, mask=None if is_global else mask) + x = x + a.o_proj(o.transpose(0, 2, 1, 3).reshape(B, L, -1)) + x = x + layer.mlp(layer.layer_norm2(x)) + x = model.post_layernorm(x) + inverse = mx.zeros(perm.shape, dtype=mx.int32) + inverse[perm] = mx.arange(perm.size, dtype=mx.int32) + return mx.take(x, inverse, axis=1)[0] + + for grid_h, grid_w in [(4, 4), (7, 5), (2, 6), (1, 1)]: + img = mx.random.normal( + (1, grid_h * cfg.patch_size, grid_w * cfg.patch_size, 3)) + got, ref = model(img), masked_reference(img) + mx.eval(got, ref) + err = float(mx.abs(got - ref).max().item()) + assert err < 1e-4, (grid_h, grid_w, err) + + # --- pixel shuffle ------------------------------------------------------------ From b78b93e9384c3d18b617ec96d088e96d48ab37e3 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:15:43 -0700 Subject: [PATCH 16/17] perf(mtp): default muse-glimmer dflash to 2 drafts per round, the measured optimum --- docs/cli.md | 2 +- gmlx/mtp_load.py | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 5aafe67..c91a28c 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -542,7 +542,7 @@ gmlx serve Qwen3.6-27B-Q4_K_S.gguf --speculative | `--hf-source REPO` | Processor/config override for a single VLM model (rarely needed). | | `--speculative` | Serve a single positional model with MTP (native-head qwen3.5/3.6; gemma4 also needs `--draft-gguf`). | | `--draft-gguf PATH` | Companion drafter GGUF for assistant-shape MTP (gemma4); implies `--speculative`. | -| `--draft-block-size N` | MTP draft tokens per round (analogous to llama-server `--spec-draft-n-max`). Default: the drafter's own block size. Also via `GMLX_DRAFT_BLOCK_SIZE`. | +| `--draft-block-size N` | MTP draft tokens per round (analogous to llama-server `--spec-draft-n-max`). Default: the drafter's own block size; muse-glimmer defaults to 2 drafts and caps N at the loaded block (raise it with `GMLX_MUSE_DFLASH_BLOCK` at load). Also via `GMLX_DRAFT_BLOCK_SIZE`. | | `--adapter PATH` | GGUF LoRA adapter applied live over a single positional model at load (text only, not `--mmproj`/`--speculative`). In config mode set `adapter:` per model instead. | | `--stream-cpu` | Run a single positional model entirely on the CPU device: the over-RAM MoE path, same semantics as [`run --stream-cpu`](#loading). In config mode set `stream: cpu` per model instead; see [server-config.md](server-config.md#models). | | `--stream-experts` | Routed-expert stacks stream from disk while the every-token layers and KV cache stay on GPU; the decode feeder (default) serves decode from a wired expert arena and makes this the faster placement once warm. Config mode: `stream: experts`. Mutually exclusive with `--stream-cpu`. | diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index ef15bb1..f634063 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -83,9 +83,12 @@ # losing regime. Uncapped is earned by measurement, not inherited by default. _MTP_WIDTH_CAP_FALLBACK = 2 -# Drafted depth per DFlash round, matching llama.cpp's n_max=3. -# GMLX_MUSE_DFLASH_BLOCK overrides, up to the GGUF's dflash.block_size. -_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT = 4 +# Drafted depth per DFlash round. Verify cost on the 30B target rises ~36% +# from block 3 to 4 (the kquant small-M kernels hold near-flat only through +# M=3 at these projection shapes), which outweighs block 4's extra accepted +# tokens; llama.cpp defaults to 4 (n_max=3). GMLX_MUSE_DFLASH_BLOCK +# overrides, up to the GGUF's dflash.block_size. +_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT = 3 def _stamp_mtp_width_cap(drafter, model_type: str, *, target=None, From e35cac30ccb0234b7b708f0163ed9faf53905579 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:51:13 -0700 Subject: [PATCH 17/17] test(mtp): pin the muse-glimmer verify-walk draw and keep its draft ids in vocab --- tests/test_muse_glimmer_mtp.py | 40 ++++++++++++++++++++++++---------- 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/tests/test_muse_glimmer_mtp.py b/tests/test_muse_glimmer_mtp.py index 2e67f4d..6bf15a1 100644 --- a/tests/test_muse_glimmer_mtp.py +++ b/tests/test_muse_glimmer_mtp.py @@ -26,14 +26,19 @@ CAPTURE = (0, 2) N_GEN = 16 BLOCK = 4 -# A tiny random model's argmax is near-tied at many steps, and the verify path -# derives its tokens through a block SDPA (qL = drafts + 1) whose rounding -# differs from the 1-token decode path by ~1e-3 in logit space. Stop the -# identity claim at the first step whose top-2 margin sits under that floor. -GREEDY_TIE_TOL = 1e-2 +# Weight init is unseeded, and about a fifth of draws leave the tiny model's +# top-2 logits tied within the floor below at step 0, which makes the identity +# claim vacuous. Pin the draw: this one holds a 1.6e-2 minimum top-2 margin +# across all N_GEN reference steps. +SEED = 25 +# The verify path derives its tokens through a block SDPA (qL = drafts + 1) +# whose rounding differs from the 1-token decode path by at most 3e-7 in logit +# space. A gap above this floor is a real divergence, not rounding. +GREEDY_TIE_TOL = 1e-3 def _build(): + mx.random.seed(SEED) ensure_registered() cfg = synthesize_config(_muse_glimmer_meta(), tensor_shapes=_MUSE_GLIMMER_SHAPES) @@ -137,7 +142,8 @@ def test_verify_walk_is_token_identical_to_greedy(armed): Parametrized over the capture seam because the bug this guards only appeared with capture armed - the packed slice reached the head unevaluated. """ - lm, _ = _build() + lm, cfg = _build() + vocab = cfg["vocab_size"] prompt = mx.array([[1, 2, 3, 4, 5]]) cache = lm.make_cache() @@ -151,14 +157,24 @@ def test_verify_walk_is_token_identical_to_greedy(armed): ref.append(t) ref_logits = lm._spec_logits(lm.model(mx.array([[t]]), cache))[0, -1] + assert min(margins) > GREEDY_TIE_TOL, ( + "the pinned draw no longer has an unambiguous greedy chain; choose " + "another SEED rather than weakening the identity claim" + ) + if armed: lm.set_dflash_capture(CAPTURE) cache2 = lm.make_cache() hid, _ = lm.speculative_verify_hidden(prompt, cache2) tok = int(lm.speculative_argmax_from_hidden(hid)[0, -1].item()) - got = [tok] + got, accepts = [tok], [] while len(got) < N_GEN: - drafts = [900 + i for i in range(BLOCK - 1)] # never the target's pick + # One past the target's own pick, so position 0 always rejects. Ids stay + # in vocab: an out-of-range id is an out-of-bounds gather that reads + # uninitialized memory on the CPU backend, and a NaN landing in a + # rejected slot does not stay there - it propagates through the masked + # SDPA into the accepted row, whose argmax then collapses to 0. + drafts = [(ref[len(got)] + 1 + i) % vocab for i in range(BLOCK - 1)] hid, _ = lm.speculative_verify_hidden( mx.array([[tok] + drafts]), cache2) rows = lm.speculative_argmax_from_hidden(hid)[0].tolist() @@ -167,15 +183,15 @@ def test_verify_walk_is_token_identical_to_greedy(armed): if int(rows[i]) != d: break accepted += 1 + accepts.append(accepted) got.extend(drafts[:accepted]) tok = int(rows[accepted]) got.append(tok) lm.rollback_speculative_cache(cache2, None, accepted, BLOCK) - # compare over the prefix whose greedy pick is unambiguous - limit = next((i for i, m in enumerate(margins) if m < GREEDY_TIE_TOL), N_GEN) - assert limit > 0, "tiny model produced no unambiguous step" - assert got[:limit] == ref[:limit] + assert got == ref + assert accepts == [0] * len(accepts), ( + "every round was meant to reject at position 0 and roll back") # --- the drafter side of the same seam ---------------------------------------