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/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/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..c91a28c 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
@@ -541,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/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 5f692ac..ff75aa8 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 |
@@ -406,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
diff --git a/gmlx/arch_table.py b/gmlx/arch_table.py
index adbe582..981e38d 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)"),
@@ -127,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",
})
@@ -157,6 +159,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/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/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/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/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/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 bec74d9..3e08447 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(
@@ -907,6 +925,15 @@ 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_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)
@@ -2690,6 +2717,19 @@ def _warm_touch_pass(
".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. 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."),
+}
+
def preset_native_fp_wire_env(args) -> None:
"""Pre-set wire mode when a streaming placement is coming.
@@ -2789,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:
@@ -2809,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
@@ -2916,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/mtp_load.py b/gmlx/mtp_load.py
index 4e4ece9..f634063 100644
--- a/gmlx/mtp_load.py
+++ b/gmlx/mtp_load.py
@@ -77,11 +77,19 @@
"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.
_MTP_WIDTH_CAP_FALLBACK = 2
+# 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,
log=loadlog.verbose_print):
@@ -486,6 +494,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 +656,237 @@ 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)
+ default_total = min(_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT, native_total)
+ block_total = max(
+ 2, min(env_int("GMLX_MUSE_DFLASH_BLOCK", default_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 +1201,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 +1320,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
@@ -1192,14 +1456,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_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_model.py b/gmlx/muse_glimmer_model.py
new file mode 100644
index 0000000..f6fb3c2
--- /dev/null
+++ b/gmlx/muse_glimmer_model.py
@@ -0,0 +1,285 @@
+# 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,
+ inputs_embeds: Optional[mx.array] = None,
+ ):
+ # 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:
+ 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)
+
+
+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__()
+ 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."""
+ if self.args.tie_word_embeddings:
+ out = self.model.embed_tokens.as_linear(h)
+ else:
+ out = self.lm_head(h)
+ 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))
+
+ @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/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])
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/muse_glimmer_vlm_model.py b/gmlx/muse_glimmer_vlm_model.py
new file mode 100644
index 0000000..92f59da
--- /dev/null
+++ b/gmlx/muse_glimmer_vlm_model.py
@@ -0,0 +1,532 @@
+# 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
+from .muse_glimmer_mtp import SpecHooks, _SpecOutput
+
+
+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 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``)."""
+ 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]`` (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]
+ 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)
+
+
+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, 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
+ 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):
+ 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, 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))
+
+
+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)
+ # 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)
+
+ 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, groups = window_partition(grid_w, grid_h, self.window)
+ perm = mx.array(perm)
+
+ 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 groups)
+
+ 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(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
+ 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,
+ 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)
+ logits = self._spec_logits(h)
+ if not want_hidden:
+ return LanguageModelOutput(logits=logits)
+ return _SpecOutput(logits=logits, hidden_states=[h])
+
+ @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]
+ # 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)
+ 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/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/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/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
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..1d55f31 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),
@@ -493,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..2a7ae0e 100644
--- a/gmlx/vlm.py
+++ b/gmlx/vlm.py
@@ -34,6 +34,8 @@
)
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,
load_gguf_wire_bytes,
@@ -66,6 +68,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 +224,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 +665,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 +1060,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 +1139,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 +1354,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 +1764,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 +2013,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 +2052,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 +2086,8 @@ 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, ()),
+ 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/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 = [
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):
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_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_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"))
diff --git a/tests/test_muse_glimmer_mtp.py b/tests/test_muse_glimmer_mtp.py
new file mode 100644
index 0000000..6bf15a1
--- /dev/null
+++ b/tests/test_muse_glimmer_mtp.py
@@ -0,0 +1,274 @@
+"""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
+# 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)
+ 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, cfg = _build()
+ vocab = cfg["vocab_size"]
+ 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]
+
+ 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, accepts = [tok], []
+ while len(got) < N_GEN:
+ # 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()
+ accepted = 0
+ for i, d in enumerate(drafts):
+ 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)
+
+ 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 ---------------------------------------
+
+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
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_muse_glimmer_vision.py b/tests/test_muse_glimmer_vision.py
new file mode 100644
index 0000000..6e437f1
--- /dev/null
+++ b/tests/test_muse_glimmer_vision.py
@@ -0,0 +1,343 @@
+"""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,
+ window_partition,
+)
+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))
+
+
+# --- 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 ------------------------------------------------------------
+
+
+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_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_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"
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"
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_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.
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