From 4ac42717a793b661babcc203641133f7b11e3e11 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sat, 17 Oct 2026 23:25:54 -0700
Subject: [PATCH 01/16] feat(nemotron): nemotron 3.5 bolt text, with per-layer
experts and nvfp4 weight scales
---
gmlx/_vendor/manifest.json | 2 +-
gmlx/load/arch_table.py | 2 +-
gmlx/load/config_synth.py | 68 ++++-
gmlx/load/loader.py | 33 ++-
gmlx/load/modules.py | 27 +-
gmlx/load/mtp_target.py | 2 +-
gmlx/load/out_scale.py | 74 +++++
gmlx/load/remap.py | 29 +-
.../nemotron_h/model.py} | 142 +++++----
gmlx/models/nemotron_h/mtp.py | 11 +-
gmlx/models/owned.py | 3 +
gmlx/spec/mtp_load.py | 2 +-
gmlx/upstream/occupancy_fuse.py | 3 +-
tests/load/test_native_fp.py | 47 ++-
tests/load/test_out_scale.py | 98 +++++++
tests/load/test_remap.py | 17 +-
tests/models/test_nemotron_h.py | 269 ++++++++++++++++++
17 files changed, 715 insertions(+), 114 deletions(-)
create mode 100644 gmlx/load/out_scale.py
rename gmlx/{_vendor/mlx_lm/models/nemotron_h.py => models/nemotron_h/model.py} (80%)
create mode 100644 tests/load/test_out_scale.py
create mode 100644 tests/models/test_nemotron_h.py
diff --git a/gmlx/_vendor/manifest.json b/gmlx/_vendor/manifest.json
index 256ef5364..adfed61e2 100644
--- a/gmlx/_vendor/manifest.json
+++ b/gmlx/_vendor/manifest.json
@@ -99,7 +99,7 @@
"mlx_lm/models/nanochat.py": {"status": "verbatim", "upstream": "989d414c4c8c3f1ae0d2d9b06c06d45b7dc5fc0bd8b796f585e318586b27dfe4"},
"mlx_lm/models/nemotron-nas.py": {"status": "verbatim", "upstream": "05b40ddd35fd5b829172b2e40ce9674a787623e5bbb26c55d45ed31f5734a855"},
"mlx_lm/models/nemotron.py": {"status": "verbatim", "upstream": "1ca8e8bd88d450fb03ba1723b0376a3199f2a92db2e8a3ccdc512b6cea492bae"},
- "mlx_lm/models/nemotron_h.py": {"status": "verbatim", "upstream": "47143633f5ad663aa6834a18be69520ae4d588371e392bcdfdfad898b2571b1f"},
+ "mlx_lm/models/nemotron_h.py": {"status": "excluded", "upstream": "47143633f5ad663aa6834a18be69520ae4d588371e392bcdfdfad898b2571b1f"},
"mlx_lm/models/olmo.py": {"status": "verbatim", "upstream": "cc4cc1097d73449ae22ee2dcf637d6e68b5f11c0482e75261f563b46c41bd40e"},
"mlx_lm/models/olmo2.py": {"status": "verbatim", "upstream": "f14a7484ebf584fdd92faefcc394b1064ffe1828b3b7fd75267de38b1b50b4a9"},
"mlx_lm/models/olmo3.py": {"status": "verbatim", "upstream": "ede2b37d41cff6f73877e8ab4174e9eac10dc37b49962f03c297fe41d9a27393"},
diff --git a/gmlx/load/arch_table.py b/gmlx/load/arch_table.py
index 6b3a2d5ed..c6f04d944 100644
--- a/gmlx/load/arch_table.py
+++ b/gmlx/load/arch_table.py
@@ -76,7 +76,7 @@ class ArchEntry:
"qwen3vl": ("qwen3", "Qwen3-VL dense text tower (run on the 8B; Qwen3 dense layout, rope base from the file); the rotary position has three axes (time, height, width), and a text-only prompt puts the same position on all three, which reduces it to the one-axis rope of the qwen3 class; pairs with the Qwen3-VL mmproj (--mmproj) for image input, where the vision tower's deepstack features join the first decoder layers"),
"llama": ("llama", "Llama-2/3, Mistral-7B-as-llama, Vicuna; SPM merges reconstructed from scores. Sparse-MoE variants (Mixtral-8x7B/8x22B) ship under this arch with an expert count -> routed to model_type=mixtral (block_sparse_moe + SwitchGLU); legacy per-expert split weights are coalesced to the stacked form on load"),
"mistral3": ("llama", "llama.cpp 'mistral3' = Ministral-3 / Mistral-Small-3.1, Llama layout"),
- "nemotron_h_moe": ("nemotron", "NVIDIA Nemotron-H MoE hybrid (Nemotron-3-Super, Nemotron-3.5-Lightning): Mamba2 SSM layers + NoPE attention layers + MoE MLPs; layer_norm_epsilon (not rms_norm_eps). Lightning appends an in-file NextN/MTP block (blk.{block_count-1}: combined attention+MoE + nextn.* fusion, excluded from the trunk) wired for --speculative; llama.cpp also ships it as a same-arch mtp-*.gguf sidecar usable via --draft-gguf"),
+ "nemotron_h_moe": ("nemotron", "NVIDIA Nemotron-H MoE hybrid (Nemotron-3-Super, Nemotron-3.5-Lightning, Nemotron-3.5-Bolt): Mamba2 SSM layers + NoPE attention layers + MoE MLPs; layer_norm_epsilon (not rms_norm_eps). Bolt (Puzzle) varies the expert count, width and top-k per layer and ships dense NVFP4 Mamba projections with per-tensor .scale tensors. Lightning appends an in-file NextN/MTP block (blk.{block_count-1}: combined attention+MoE + nextn.* fusion, excluded from the trunk) wired for --speculative; llama.cpp also ships it as a same-arch mtp-*.gguf sidecar usable via --draft-gguf"),
"deepseek2": ("deepseek", "DeepSeek-V3/R1, Moonshot Kimi-K2.5/K2.7, + GLM-4.x MLA conversions (GLM-4.7-Flash); MLA attention (absorbed embed_q/unembed_out via KQuantMultiLinear, native Q8_0), fine-grained sigmoid-gated MoE + shared expert + group routing, leading dense block. The converter writes rope yarn_log_multiplier as 0.1 * mscale_all_dim, and the synth undoes the 0.1 factor (a raw read flattens attention ~1.85x). Kimi K2.x is ungrouped (384 experts, one group), and gmlx selects its own sampling family (t=1.0) from general.name. Its <|im_user|>/<|im_middle|> chat format leaves open at the generation prompt, and the stock kimi_k2 tool parser reads its tool calls. K2.x pairs with its MoonViT mmproj (--mmproj, projector 'kimik25') for vision. V2 (softmax gating) not yet supported"),
"deepseek4": ("deepseek", "DeepSeek V4 Flash (256x8.4B, dwarfstar 'deepseek4' arch, not a llama.cpp conversion; parity reference is the ds4 engine). MLA-lite attention (low-rank q, single shared 512-dim KV latent K=V, grouped low-rank output proj, per-head fp32 sinks, tail-64 NEOX rope with inverse-rope'd output) in three per-layer variants from compress_ratios: sliding-window(128), +compressed pool, +lightning-indexer top-512 sparse. Manifold-constrained hyper-connections (4-stream 4D hidden, Sinkhorn mixing). Every-layer 256-expert MoE, sqrt-softplus gating + selection-only correction bias, first 3 layers hash-routed (tid2eid), shared expert, clamped SwiGLU. QAT round-trips (fp8-E4M3 KV, Hadamard+fp4-E2M1 indexer) reproduced on-path for logit parity. Model class vendored from mlx-lm PR #1192; separate MTP drafter GGUF (--draft-gguf, or autodetected next to the target) for speculative decoding - the DSpark sidecar in either container (gmlx deepseek4-dspark or llama.cpp dflash, e.g. the unsloth dspark release), or the legacy deepseek4_mtp_support nextn file. The Vision-Exp checkpoint (per-layer exp_probs_b_vl image router bias, embedding_length_out) pairs with its mmproj (--mmproj, projector deepseek4v): a native-resolution DeepSeek ViT (2-D rope, bidirectional per image) + 3x3 unfold aligner; each image becomes a block of sentinel ids past the vocab (lead pads, START, grid rows, END) with the reference's in-block bidirectional sliding-window attention and the image router bias on the text side; image turns run single-row with unquantized KV and blocks are prefilled whole"),
"glm-dsa": ("deepseek", "GLM-5.2 (DeepSeek-V3.2): deepseek2 MLA + fine-grained sigmoid-gated MoE plus a per-layer DSA 'lightning indexer' (top-k sparse-attention key selection) + an MTP/nextn layer (dropped on load). Reuses the DEEPSEEK2 remap (indexer.* patterns appended) -> mlx-lm model_type glm_moe_dsa (subclasses deepseek_v32)"),
diff --git a/gmlx/load/config_synth.py b/gmlx/load/config_synth.py
index 992014071..a08544a16 100644
--- a/gmlx/load/config_synth.py
+++ b/gmlx/load/config_synth.py
@@ -2474,16 +2474,13 @@ def _synth_nemotron_h_moe(meta, shapes, config: dict) -> None:
if non_zero:
config["num_key_value_heads"] = non_zero[0]
+ ff_array = _read_int_array(meta, f"{arch}.feed_forward_length")
+
# MoE parameters
n_experts = _read_int(meta, f"{arch}.expert_count")
if n_experts is not None:
- config["n_routed_experts"] = n_experts
- config["num_experts_per_tok"] = _require(
- _read_int(meta, f"{arch}.expert_used_count"),
- arch=arch, gguf_field=f"{arch}.expert_used_count")
- config["moe_intermediate_size"] = _require(
- _read_int(meta, f"{arch}.expert_feed_forward_length"),
- arch=arch, gguf_field=f"{arch}.expert_feed_forward_length")
+ _synth_nemotron_h_experts(meta, shapes, config, arch, n_experts,
+ kv_array, ff_array)
shared_ffn = _read_int(meta, f"{arch}.expert_shared_feed_forward_length")
if shared_ffn is not None:
@@ -2500,8 +2497,9 @@ def _synth_nemotron_h_moe(meta, shapes, config: dict) -> None:
config["n_group"] = _read_int(meta, f"{arch}.expert_group_count") or 1
config["topk_group"] = _read_int(meta, f"{arch}.expert_group_used_count") or 1
- norm_topk = _read_int(meta, f"{arch}.expert_weights_norm")
- config["norm_topk_prob"] = bool(norm_topk) if norm_topk is not None else True
+ # llama.cpp leaves expert_weights_norm false when the key is absent.
+ config["norm_topk_prob"] = bool(
+ _read_int(meta, f"{arch}.expert_weights_norm") or 0)
scale = _read_float(meta, f"{arch}.expert_weights_scale")
if scale is None:
@@ -2511,13 +2509,14 @@ def _synth_nemotron_h_moe(meta, shapes, config: dict) -> None:
# intermediate_size for dense MLP layers (type "-"). Use
# feed_forward_length if uniform, else default to moe_intermediate_size.
- ff_array = _read_int_array(meta, f"{arch}.feed_forward_length")
if ff_array is not None:
non_zero_ff = [v for v in ff_array if v > 0]
if non_zero_ff:
config["intermediate_size"] = non_zero_ff[0]
if "intermediate_size" not in config:
- config["intermediate_size"] = config.get("moe_intermediate_size", 0)
+ moe_ff = config.get("moe_intermediate_size", 0)
+ config["intermediate_size"] = (
+ max(moe_ff) if isinstance(moe_ff, list) else moe_ff)
# Derive hybrid_override_pattern from per-layer metadata arrays.
# head_count_kv > 0 -> attention (*); feed_forward_length > 0 and kv==0
@@ -2540,6 +2539,53 @@ def _synth_nemotron_h_moe(meta, shapes, config: dict) -> None:
config["num_hidden_layers"] = n_layers
+def _synth_nemotron_h_experts(meta, shapes, config: dict, arch: str,
+ n_experts: int, kv_array, ff_array) -> None:
+ """The routed-expert sizes, per layer as llama.cpp's nemotron_h loader
+ reads them. Puzzle models (Nemotron-3.5-Bolt) vary all three by layer:
+ the expert count is the ``ffn_gate_inp`` row count (``expert_count`` is
+ only the largest), the width is ``expert_feed_forward_length`` (else
+ ``feed_forward_length / expert_used_count``), and the top-k is
+ ``expert_used_count``. Each is a scalar when every MoE layer agrees, else
+ a list over all blocks (MTP included) with 0 on the other layers."""
+ used = _require(
+ _read_int_array(meta, f"{arch}.expert_used_count"),
+ arch=arch, gguf_field=f"{arch}.expert_used_count")
+ ff_exp = _read_int_array(meta, f"{arch}.expert_feed_forward_length") or [0]
+ n_all = len(kv_array) if kv_array is not None else int(
+ config["num_hidden_layers"])
+
+ def at(arr, i):
+ if arr is None:
+ return 0
+ return arr[i] if len(arr) > 1 else arr[0]
+
+ counts, widths, topks = [0] * n_all, [0] * n_all, [0] * n_all
+ for i in range(n_all):
+ if at(kv_array, i) > 0 or at(ff_array, i) == 0:
+ continue
+ k = at(used, i)
+ router = shapes.get(f"blk.{i}.ffn_gate_inp.weight")
+ n = int(router[1]) if router is not None and len(router) == 2 \
+ else n_experts
+ if not k <= n <= n_experts:
+ raise ValueError(
+ f"{arch}: blk.{i}.ffn_gate_inp has {n} experts, outside "
+ f"[{k}, {n_experts}] (expert_used_count, expert_count)")
+ width = at(ff_exp, i) or (at(ff_array, i) // k if k else 0)
+ counts[i], widths[i], topks[i] = n, width, k
+
+ def collapse(vals, fallback):
+ moe = {v for v in vals if v}
+ if not moe:
+ return fallback
+ return moe.pop() if len(moe) == 1 else vals
+
+ config["n_routed_experts"] = collapse(counts, n_experts)
+ config["moe_intermediate_size"] = collapse(widths, at(ff_exp, 0))
+ config["num_experts_per_tok"] = collapse(topks, at(used, 0))
+
+
# DeepseekV2Model and its subclasses (deepseek2, glm-dsa) write the yarn log
# multiplier as `0.1 * mscale_all_dim`, not the value itself
# ([TAG_DEEPSEEK2_YARN_LOG_MUL_FIX] in upstream conversion/deepseek.py). Other
diff --git a/gmlx/load/loader.py b/gmlx/load/loader.py
index 822e20b03..9e953c571 100644
--- a/gmlx/load/loader.py
+++ b/gmlx/load/loader.py
@@ -51,6 +51,7 @@
dequantize_unattachable_leaves,
install_kquant_modules,
)
+from .out_scale import install_out_scales
from .populate import (
maybe_populate_for_load,
start_populate,
@@ -931,6 +932,11 @@ def _warm_touch_pass(
# class's cast_predicate.
"k2_horizon": (".mlp.gate.weight", ".e_score_correction_bias",
".selection_bias"),
+ # nemotron_h: llama.cpp routes in fp32 (F32 wire router and selection
+ # bias; Bolt picks top-9 of up to 496), and A_log feeds the fp32 scan.
+ # Kept in step with the model class's cast_predicate.
+ "nemotron_h": (".mixer.gate.weight", ".e_score_correction_bias",
+ ".A_log"),
}
# Params kept at their native f16 through the bf16 cast (no upcast). MLX
@@ -1059,22 +1065,23 @@ def _resolve_native_fp_wire(hf_weights, hf_kquant_meta, log) -> bool:
return False
-def _native_fp_multilinear_keys(model, hf_kquant_meta) -> set[str]:
- """Weight names of native-fp tensors destined for a MultiLinear leaf.
+def _native_fp_wire_only_keys(model, hf_kquant_meta) -> set[str]:
+ """Weight names of native-fp tensors destined for a dense Linear or a
+ MultiLinear leaf.
- These stay ggml wire bytes even in packed mode (kq.gather_qmm dispatch;
- no packed MultiLinear module exists) - the repack must skip them.
+ These stay ggml wire bytes even in packed mode (kq.quantized_matmul and
+ kq.gather_qmm dispatch; no packed module exists for either) - the repack
+ must skip them.
"""
from .modules import MultiLinear
from .native_fp import NATIVE_FP_CODECS
- if MultiLinear is None:
- return set()
+ wire_only = (nn.Linear,) if MultiLinear is None else (nn.Linear, MultiLinear)
keys = set()
for path, mod in tree_flatten(model.leaf_modules(),
is_leaf=nn.Module.is_module):
wk = f"{path}.weight"
- if (isinstance(mod, MultiLinear)
+ if (isinstance(mod, wire_only)
and hf_kquant_meta.get(wk) in NATIVE_FP_CODECS):
keys.add(wk)
return keys
@@ -1279,7 +1286,7 @@ def _install_and_load(
n_fp = repack_native_fp_weights(
hf_weights, hf_kquant_meta,
- skip=_native_fp_multilinear_keys(model, hf_kquant_meta))
+ skip=_native_fp_wire_only_keys(model, hf_kquant_meta))
if n_fp:
log(
f"[native-fp] de-interleaved {n_fp} mxfp4/nvfp4 tensors -> MLX packed layout"
@@ -1295,6 +1302,7 @@ def _install_and_load(
model, hf_kquant_meta, native_fp_wire=native_fp_wire)
log(f"[install] replaced {n_replaced} leaves with kquant modules")
_install_hadamard(model, hadamard, log)
+ install_out_scales(model, hf_weights, hf_kquant_meta)
from gmlx.stream.table_pread import install_deferred_tables
attached = install_deferred_tables(model, deferred_tables or {})
@@ -1336,7 +1344,8 @@ def _install_and_load(
n_fp32 += 1
continue
if v.dtype in (mx.float32, mx.float16) and k not in hf_kquant_meta:
- if fp32_keep and any(s in k for s in fp32_keep):
+ if k.endswith(".out_scale") or (
+ fp32_keep and any(s in k for s in fp32_keep)):
if v.dtype != mx.float32: # e.g. F16 ape tables
loadable[k] = v.astype(mx.float32)
continue
@@ -1737,7 +1746,7 @@ def load_model(
n_fp = repack_native_fp_weights(
hf_weights, hf_kquant_meta,
- skip=_native_fp_multilinear_keys(model, hf_kquant_meta))
+ skip=_native_fp_wire_only_keys(model, hf_kquant_meta))
if n_fp:
_log(
f"[native-fp] de-interleaved {n_fp} mxfp4/nvfp4 tensors -> MLX packed layout"
@@ -1757,6 +1766,7 @@ def load_model(
act_dtype=activation_dtype(config.get("model_type")))
_log(f"[install] replaced {n_replaced} leaves with kquant modules")
_install_hadamard(model, hadamard_targets, _log)
+ install_out_scales(model, hf_weights, hf_kquant_meta)
from gmlx.stream.table_pread import install_deferred_tables
attached = install_deferred_tables(model, deferred)
@@ -1839,7 +1849,8 @@ def load_model(
for k in list(loadable):
v = loadable[k]
if v.dtype in (mx.float32, mx.float16) and k not in hf_kquant_meta:
- if fp32_keep and any(s in k for s in fp32_keep):
+ if k.endswith(".out_scale") or (
+ fp32_keep and any(s in k for s in fp32_keep)):
if v.dtype != mx.float32: # e.g. F16 ape tables
loadable[k] = v.astype(mx.float32)
continue
diff --git a/gmlx/load/modules.py b/gmlx/load/modules.py
index 97d3fbb6d..ac262f7a5 100644
--- a/gmlx/load/modules.py
+++ b/gmlx/load/modules.py
@@ -375,7 +375,8 @@ def _kq_wire_k(w, codec):
def _kq_dense(proj, codecs):
return (isinstance(proj, KQuantLinear)
- and proj.kquant_type in codecs and "bias" not in proj)
+ and proj.kquant_type in codecs and "bias" not in proj
+ and "out_scale" not in proj)
def _kq_act_kind(m, caps):
@@ -2783,22 +2784,30 @@ def _replace(path: str, module):
bias, codec)
return NativeFPSwitchLinear(n_experts, out_dims, in_dims, bias,
codec)
- if MultiLinear is not None and isinstance(module, MultiLinear):
- # No packed MultiLinear module exists; both modes dispatch the
- # ggml wire bytes through kq.gather_qmm (the loader exempts
- # these tensors from the packed repack).
+ wire_only = isinstance(module, nn.Linear) or (
+ MultiLinear is not None and isinstance(module, MultiLinear))
+ if wire_only:
+ # No packed dense or MultiLinear module exists; both modes
+ # dispatch the ggml wire bytes through the kq kernels (the
+ # loader exempts these tensors from the packed repack).
import mlx_kquant as kq
if codec not in set(kq.codecs()):
raise NotImplementedError(
- f"native-fp codec {codec!r} on MultiLinear at "
- f"{path!r} needs an mlx-kquant build with that codec")
+ f"native-fp codec {codec!r} on "
+ f"{type(module).__name__} at {path!r} needs an "
+ f"mlx-kquant build with that codec")
+ if isinstance(module, nn.Linear):
+ out_dims, in_dims = module.weight.shape
+ n_replaced += 1
+ return KQuantLinear(in_dims, out_dims, "bias" in module, codec)
+ if MultiLinear is not None and isinstance(module, MultiLinear):
num_heads, out_dims, in_dims = module.weight.shape
n_replaced += 1
return KQuantMultiLinear(in_dims, out_dims, num_heads, codec)
raise NotImplementedError(
f"native-fp codec {codec!r} on {type(module).__name__} at "
- f"{path!r} - only SwitchLinear (MoE experts) and MultiLinear "
- f"(absorbed MLA) are wired so far")
+ f"{path!r} - only Linear, SwitchLinear (MoE experts) and "
+ f"MultiLinear (absorbed MLA) are wired so far")
if isinstance(module, nn.Linear):
out_dims, in_dims = module.weight.shape
bias = "bias" in module
diff --git a/gmlx/load/mtp_target.py b/gmlx/load/mtp_target.py
index 38cd32376..bf946aedc 100644
--- a/gmlx/load/mtp_target.py
+++ b/gmlx/load/mtp_target.py
@@ -230,7 +230,7 @@ def build(config: dict) -> nn.Module:
return glm5_next_mtp.Glm5NextSpecLM, build
if model_type == "nemotron_h":
import gmlx.models.nemotron_h.mtp as nemotron_h_mtp
- from gmlx._vendor.mlx_lm.models.nemotron_h import ModelArgs
+ from gmlx.models.nemotron_h.model import ModelArgs
def build(config: dict) -> nn.Module:
return nemotron_h_mtp.NemotronHSpecLM(ModelArgs.from_dict(config))
diff --git a/gmlx/load/out_scale.py b/gmlx/load/out_scale.py
new file mode 100644
index 000000000..fe1794378
--- /dev/null
+++ b/gmlx/load/out_scale.py
@@ -0,0 +1,74 @@
+"""Per-tensor weight scales (``.scale`` in a GGUF).
+
+llama.cpp multiplies a matmul's output by the weight's ``.scale`` tensor
+when the file has one (NVFP4 exports carry it as the global "scale2").
+Remap names it ``.out_scale``. ``install_out_scales`` swaps each
+such module's class for a subclass that applies the scale after its own
+``__call__``, so every caller that runs the module gets it. Fusions that
+read a projection's wire bytes directly skip a module with ``out_scale``,
+as they skip one with a bias.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+_SCALED: dict[type, type] = {}
+
+
+def _scaled_class(cls: type) -> type:
+ sub = _SCALED.get(cls)
+ if sub is None:
+ def __call__(self, x, *args, **kwargs):
+ if x.dtype == mx.float16:
+ # A small scale leaves an unscaled product past the float16
+ # range; bfloat16 holds it.
+ y = cls.__call__(self, x.astype(mx.bfloat16), *args, **kwargs)
+ y = y * self["out_scale"].astype(y.dtype)
+ return y.astype(mx.float16)
+ y = cls.__call__(self, x, *args, **kwargs)
+ return y * self["out_scale"].astype(y.dtype)
+
+ sub = type(f"OutScaled{cls.__name__}", (cls,), {
+ "__call__": __call__, "__module__": __name__})
+ _SCALED[cls] = sub
+ return sub
+
+
+def has_out_scale(module) -> bool:
+ return isinstance(module, nn.Module) and "out_scale" in module
+
+
+def install_out_scales(model: nn.Module, hf_weights: dict,
+ hf_kquant_meta: dict | None = None) -> int:
+ """Give every module with a ``.out_scale`` weight an ``out_scale``
+ parameter and the scaling subclass. Runs after ``install_kquant_modules``
+ and the Hadamard swap, before the weights load. An ``lm_head`` scale
+ applies only to an NVFP4 head, as in llama.cpp. Returns the count."""
+ suffix = ".out_scale"
+ paths = {k[:-len(suffix)] for k in hf_weights if k.endswith(suffix)}
+ if not paths:
+ return 0
+ meta = hf_kquant_meta or {}
+ n = 0
+ for path, m in model.named_modules():
+ if path not in paths:
+ continue
+ paths.discard(path)
+ key = path + suffix
+ scale = hf_weights[key]
+ if path.split(".")[-1] == "lm_head" and meta.get(
+ path + ".weight") != "nvfp4":
+ del hf_weights[key]
+ continue
+ if int(scale.size) != 1:
+ raise NotImplementedError(
+ f"{key}: per-expert weight scales ({int(scale.size)} values) "
+ f"are not supported yet")
+ if type(m) not in _SCALED.values():
+ m.__class__ = _scaled_class(type(m))
+ m.out_scale = mx.ones((1,), dtype=mx.float32)
+ m.freeze(recurse=False, keys=["out_scale"])
+ n += 1
+ return n
diff --git a/gmlx/load/remap.py b/gmlx/load/remap.py
index 66c3346d7..fe192c590 100644
--- a/gmlx/load/remap.py
+++ b/gmlx/load/remap.py
@@ -2144,11 +2144,22 @@ def _gemma_norm_transform(arch_string: str, hf_name: str | None,
})
-# Unsloth UD `.scale` tensors are metadata for Unsloth's MLX recipe (per-tensor
-# additional scaling on top of the base K-quant). They have no HF/MLX-stock
-# counterpart; skip with warning.
+# Weight scales llama.cpp loads for every arch (llama-model.cpp, the generic
+# ".scale" pass after load_arch_tensors): a scalar per matmul weight, or one
+# per expert for the stacked experts. The matmul output is multiplied by it
+# (NVFP4 "scale2"). `.scale` maps to `.out_scale`.
+_WEIGHT_SCALE_RE = re.compile(
+ r"^(blk\.\d+\.(?:attn_q|attn_k|attn_v|attn_output|attn_qkv|attn_gate"
+ r"|ffn_gate|ffn_down|ffn_up|ffn_gate_shexp|ffn_down_shexp|ffn_up_shexp"
+ r"|ffn_gate_exps|ffn_down_exps|ffn_up_exps"
+ r"|ssm_in|ssm_out|ssm_alpha|ssm_beta"
+ r"|nextn\.eh_proj|nextn\.shared_head_head)|output)\.scale$")
+
+
+# Other `.scale` tensors (Unsloth UD recipe metadata, `.input_scale`) have no
+# counterpart in llama.cpp's graph; skip with warning.
def _is_unsloth_ud_scale(name: str) -> bool:
- return name.endswith(".scale")
+ return name.endswith(".scale") or name.endswith(".input_scale")
# Module-level cache: {gguf_name_template: MODEL_TENSOR_enum}. The reverse of
@@ -2237,6 +2248,16 @@ def parse_gguf_name(arch_string: str, gguf_name: str) -> RemapDecision:
reason=f"no diffusion-gemma override for {gguf_name!r}; extend the "
f"DIFFUSION_GEMMA table")
+ # A weight scale follows its weight's mapping (after arch-priority claims).
+ m = _WEIGHT_SCALE_RE.match(gguf_name)
+ if m:
+ d = parse_gguf_name(arch_string, m.group(1) + ".weight")
+ if d.kind != RemapDecision.KIND_MAP or not d.hf_name.endswith(".weight"):
+ return d
+ return RemapDecision(RemapDecision.KIND_MAP,
+ hf_name=d.hf_name[:-len("weight")] + "out_scale",
+ transform="flatten", bid=d.bid)
+
# Universal Unsloth-UD .scale skip (after arch-priority claims).
if _is_unsloth_ud_scale(gguf_name):
return RemapDecision(RemapDecision.KIND_SKIP,
diff --git a/gmlx/_vendor/mlx_lm/models/nemotron_h.py b/gmlx/models/nemotron_h/model.py
similarity index 80%
rename from gmlx/_vendor/mlx_lm/models/nemotron_h.py
rename to gmlx/models/nemotron_h/model.py
index 353de36c9..3350cf3c5 100644
--- a/gmlx/_vendor/mlx_lm/models/nemotron_h.py
+++ b/gmlx/models/nemotron_h/model.py
@@ -1,22 +1,58 @@
+# SPDX-License-Identifier: MIT
# Copyright © 2025 Apple Inc.
+# Copyright (c) 2026 Asher Feldman
+"""Model for NVIDIA Nemotron-H (GGUF arch ``nemotron_h_moe``): Nano, Super,
+Lightning and Bolt.
+
+Absorbed from mlx-lm's ``nemotron_h``. The reference runtime is llama.cpp
+``src/models/nemotron-h.cpp`` (trunk graph) and ``src/llama-graph.cpp``
+``build_moe_ffn``. Changes from mlx-lm:
+
+- The MoE sizes may vary per layer: ``n_routed_experts``,
+ ``moe_intermediate_size`` and ``num_experts_per_tok`` each take a scalar
+ or a list indexed by layer (Bolt, Puzzle).
+- The router runs in fp32, as llama.cpp's F32 ``ffn_gate_inp`` matmul does.
+- The routing weights are renormalized over a sum clamped at
+ 6.103515625e-5, whatever the top-k, and scaled only when the scale is
+ neither 0 nor 1.
+
+Block kinds: ``M`` Mamba-2, ``*`` attention (no RoPE), ``E`` MoE, ``-``
+dense MLP. Every block is ``x + mixer(rms_norm(x))``.
+"""
+
+from __future__ import annotations
from dataclasses import dataclass
-from functools import partial
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
import mlx.core as mx
import mlx.nn as nn
-from .activations import swiglu
-from .base import (
- BaseModelArgs,
+from gmlx._vendor.mlx_lm.models.activations import swiglu
+from gmlx._vendor.mlx_lm.models.base import BaseModelArgs
+from gmlx._vendor.mlx_lm.models.ssm import ssm_update
+from gmlx._vendor.mlx_lm.models.switch_layers import SwitchMLP
+from gmlx.cache.attention import (
create_attention_mask,
create_ssm_mask,
scaled_dot_product_attention,
)
-from .cache import ArraysCache, KVCache
-from .ssm import ssm_update
-from .switch_layers import SwitchMLP
+from gmlx.cache.kv_cache import ArraysCache, KVCache
+
+# llama.cpp build_moe_ffn clamps the weight sum here before dividing.
+_WEIGHT_SUM_FLOOR = 6.103515625e-5
+
+PerLayer = Union[int, List[int], None]
+
+
+def per_layer(value: PerLayer, layer_idx: Optional[int]):
+ """``value`` at ``layer_idx`` when it is a list, else ``value``. An index
+ past the list (an appended MTP block) or None takes the largest entry."""
+ if not isinstance(value, (list, tuple)):
+ return value
+ if layer_idx is not None and 0 <= layer_idx < len(value):
+ return value[layer_idx]
+ return max(value) if value else None
@dataclass()
@@ -43,14 +79,14 @@ class ModelArgs(BaseModelArgs):
hybrid_override_pattern: Optional[List[str]] = None
layers_block_type: Optional[List[str]] = None
head_dim: Optional[int] = None
- moe_intermediate_size: Optional[int] = None
+ moe_intermediate_size: PerLayer = None
moe_shared_expert_intermediate_size: Optional[int] = None
moe_latent_size: Optional[int] = None
n_group: Optional[int] = None
- n_routed_experts: Optional[int] = None
+ n_routed_experts: PerLayer = None
n_shared_experts: Optional[int] = None
topk_group: Optional[int] = None
- num_experts_per_tok: Optional[int] = None
+ num_experts_per_tok: PerLayer = None
norm_topk_prob: Optional[bool] = None
routed_scaling_factor: Optional[float] = None
time_step_limit: Optional[Tuple[float, float]] = None
@@ -181,11 +217,7 @@ def _ssm(
)
B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
- if cache:
- state = cache[1]
- lengths = cache.lengths
- else:
- state, lengths = None, None
+ state = cache[1] if cache else None
y, state = ssm_update(
hidden_states,
@@ -336,30 +368,34 @@ def group_expert_select(
k = top_k
inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
scores = mx.take_along_axis(orig_scores, inds, axis=-1)
- if top_k > 1 and norm_topk_prob:
+ if norm_topk_prob:
denominator = scores.sum(axis=-1, keepdims=True)
- scores = scores / (denominator + 1e-20)
- scores = scores * routed_scaling_factor
+ scores = scores / mx.maximum(denominator, _WEIGHT_SUM_FLOOR)
+ if routed_scaling_factor not in (0.0, 1.0):
+ scores = scores * routed_scaling_factor
return inds, scores
class MoEGate(nn.Module):
- def __init__(self, config: ModelArgs):
+ def __init__(self, config: ModelArgs, layer_idx: Optional[int] = None):
super().__init__()
self.config = config
- self.top_k = config.num_experts_per_tok
- self.norm_topk_prob = config.norm_topk_prob
- self.n_routed_experts = config.n_routed_experts
- self.routed_scaling_factor = config.routed_scaling_factor
- self.n_group = config.n_group
- self.topk_group = config.topk_group
- self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
- self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+ self.top_k = int(per_layer(config.num_experts_per_tok, layer_idx))
+ self.norm_topk_prob = bool(config.norm_topk_prob)
+ self.n_routed_experts = int(per_layer(config.n_routed_experts, layer_idx))
+ self.routed_scaling_factor = float(config.routed_scaling_factor or 1.0)
+ self.n_group = config.n_group or 1
+ self.topk_group = config.topk_group or 1
+ # fp32 like the F32 wire router; llama.cpp routes in fp32.
+ self.weight = mx.zeros(
+ (self.n_routed_experts, config.hidden_size), dtype=mx.float32)
+ self.e_score_correction_bias = mx.zeros(
+ (self.n_routed_experts,), dtype=mx.float32)
def __call__(self, x):
return group_expert_select(
- x @ self.weight.T,
+ x.astype(mx.float32) @ self.weight.T.astype(mx.float32),
self.e_score_correction_bias,
self.top_k,
self.n_group,
@@ -370,10 +406,11 @@ def __call__(self, x):
class NemotronHMoE(nn.Module):
- def __init__(self, config: ModelArgs):
+ def __init__(self, config: ModelArgs, layer_idx: Optional[int] = None):
super().__init__()
self.config = config
- self.num_experts_per_tok = config.num_experts_per_tok
+ self.num_experts_per_tok = int(
+ per_layer(config.num_experts_per_tok, layer_idx))
self.moe_latent_size = config.moe_latent_size
# When latent projection is used, experts operate on the latent dim
@@ -384,12 +421,12 @@ def __init__(self, config: ModelArgs):
)
self.switch_mlp = SwitchMLP(
expert_input_dim,
- config.moe_intermediate_size,
- config.n_routed_experts,
+ int(per_layer(config.moe_intermediate_size, layer_idx)),
+ int(per_layer(config.n_routed_experts, layer_idx)),
activation=nn.ReLU2(),
)
- self.gate = MoEGate(config)
+ self.gate = MoEGate(config, layer_idx)
if config.n_shared_experts is not None:
intermediate_size = config.moe_shared_expert_intermediate_size
self.shared_experts = NemotronHMLP(
@@ -425,7 +462,8 @@ def __call__(self, x):
class NemotronHBlock(nn.Module):
- def __init__(self, args: ModelArgs, block_type: str):
+ def __init__(self, args: ModelArgs, block_type: str,
+ layer_idx: Optional[int] = None):
super().__init__()
self.norm = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
@@ -438,7 +476,7 @@ def __init__(self, args: ModelArgs, block_type: str):
elif self.block_type == "-":
self.mixer = NemotronHMLP(args)
elif self.block_type == "E":
- self.mixer = NemotronHMoE(args)
+ self.mixer = NemotronHMoE(args, layer_idx)
def __call__(
self,
@@ -460,8 +498,8 @@ def __init__(self, args: ModelArgs):
super().__init__()
self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
self.layers = [
- NemotronHBlock(args, block_type)
- for block_type in args.hybrid_override_pattern
+ NemotronHBlock(args, block_type, i)
+ for i, block_type in enumerate(args.hybrid_override_pattern)
]
self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
self.fa_idx = 0
@@ -528,35 +566,23 @@ def layers(self):
def make_cache(self):
caches = []
- for l in self.layers:
- if l.block_type == "M":
+ for layer in self.layers:
+ if layer.block_type == "M":
caches.append(ArraysCache(size=2))
- elif l.block_type == "*":
+ elif layer.block_type == "*":
caches.append(KVCache())
return caches
def sanitize(self, weights):
- weights = {k: v for (k, v) in weights.items() if not k.startswith("mtp.")}
- for k, v in weights.items():
- if "conv1d.weight" in k and v.shape[-1] != 1:
- weights[k] = v.moveaxis(2, 1)
-
- # Stack experts
- for l in range(self.args.num_hidden_layers):
- prefix = f"backbone.layers.{l}.mixer"
- for m, n in [("down_proj", "fc2"), ("up_proj", "fc1")]:
- if f"{prefix}.experts.0.{m}.weight" in weights:
- to_join = [
- weights.pop(f"{prefix}.experts.{e}.{m}.weight")
- for e in range(self.args.n_routed_experts)
- ]
- weights[f"{prefix}.switch_mlp.{n}.weight"] = mx.stack(to_join)
-
return weights
@property
def cast_predicate(self):
def predicate(k):
- return "e_score_correction_bias" not in k and "A_log" not in k
+ return not any(s in k for s in _FP32_PARAMS)
return predicate
+
+
+# Kept fp32 by the loader (loader._FP32_KEEP_BY_MODEL_TYPE["nemotron_h"]).
+_FP32_PARAMS = (".mixer.gate.weight", ".e_score_correction_bias", ".A_log")
diff --git a/gmlx/models/nemotron_h/mtp.py b/gmlx/models/nemotron_h/mtp.py
index 3d536f175..c28a79339 100644
--- a/gmlx/models/nemotron_h/mtp.py
+++ b/gmlx/models/nemotron_h/mtp.py
@@ -25,8 +25,8 @@
- ``shared_head_norm`` is a full LayerNorm (llama.cpp builds it LLM_NORM,
not LLM_NORM_RMS), weight-only.
-The trunk stays the stock mlx-lm ``nemotron_h`` classes; the vendored
-backbone subclass only adds the verify-time Mamba2 sink. Verify rollback is
+The trunk is the owned ``nemotron_h.model`` classes; the backbone
+subclass only adds the verify-time Mamba2 sink. Verify rollback is
two-part (the glm5_next recurrent precedent): the attention KV leaves trim,
and the Mamba2 layers replay the accepted prefix from the recorded
pre-verify (conv, ssm) state - ``ArraysCache`` state entries are replaced,
@@ -45,7 +45,7 @@
import mlx.core as mx
import mlx.nn as nn
-from gmlx._vendor.mlx_lm.models import nemotron_h as nh
+from gmlx.models.nemotron_h import model as nh
from gmlx.cache.attention import create_attention_mask, create_ssm_mask
from gmlx.envflags import env_bool
@@ -196,7 +196,7 @@ def rollback_verify_sink(sink: list, n: int) -> None:
class NemotronHSpecLM(nh.Model):
- """Stock nemotron_h ``Model`` + the ``speculative_*`` hooks the owned MTP
+ """The nemotron_h ``Model`` + the ``speculative_*`` hooks the owned MTP
engine probes on its target. The backbone is swapped for the sink-aware
subclass; numerics are byte-identical to stock."""
@@ -304,7 +304,8 @@ class NemotronHMTPLayer(nn.Module):
def __init__(self, args):
super().__init__()
self.self_attn = nh.NemotronHAttention(args)
- self.mlp = nh.NemotronHMoE(args)
+ # The MTP block follows the trunk in the per-layer arrays.
+ self.mlp = nh.NemotronHMoE(args, args.num_hidden_layers)
self.input_layernorm = nn.RMSNorm(
args.hidden_size, eps=args.layer_norm_epsilon)
self.post_attention_layernorm = nn.RMSNorm(
diff --git a/gmlx/models/owned.py b/gmlx/models/owned.py
index 29b0fb493..54adf2f51 100644
--- a/gmlx/models/owned.py
+++ b/gmlx/models/owned.py
@@ -45,6 +45,9 @@
# llama.cpp LLM_ARCH_K2_HORIZON; the mlx-vlm class is a plain Llama
# with none of the group norm, the output gate, MoE or MoVA.
"k2_horizon": "gmlx.models.k2_horizon.model",
+ # mlx-lm nemotron_h, with per-layer MoE sizes (Nemotron-3.5-Bolt) and
+ # llama.cpp's fp32 routing.
+ "nemotron_h": "gmlx.models.nemotron_h.model",
}
# mlx-vlm model_type -> the gmlx module that stands for the model package
diff --git a/gmlx/spec/mtp_load.py b/gmlx/spec/mtp_load.py
index 6314ed71f..355e89cea 100644
--- a/gmlx/spec/mtp_load.py
+++ b/gmlx/spec/mtp_load.py
@@ -284,7 +284,7 @@ def _load_mtp_drafter(
f"block_size={drafter.config.block_size}"
)
elif model_type == "nemotron_h":
- from gmlx._vendor.mlx_lm.models.nemotron_h import ModelArgs
+ from gmlx.models.nemotron_h.model import ModelArgs
from gmlx.models.nemotron_h.mtp import (
NemotronHMTPConfig,
NemotronHMTPDrafter,
diff --git a/gmlx/upstream/occupancy_fuse.py b/gmlx/upstream/occupancy_fuse.py
index 361b2d931..8191ed5c5 100644
--- a/gmlx/upstream/occupancy_fuse.py
+++ b/gmlx/upstream/occupancy_fuse.py
@@ -79,7 +79,7 @@ def _same_codec_kquant(projs) -> bool:
# fused wire would bypass that rotation.
if any(getattr(p, "_hadamard", None) is not None for p in projs):
return False
- if any("bias" in p for p in projs):
+ if any("bias" in p or "out_scale" in p for p in projs):
return False
return len({p.kquant_type for p in projs}) == 1
@@ -361,6 +361,7 @@ def _kq_build_fused(self):
if (
isinstance(d, KQuantLinear)
and "bias" not in d
+ and "out_scale" not in d
and wire_cols % 2 == 0
and k % 512 == 0
):
diff --git a/tests/load/test_native_fp.py b/tests/load/test_native_fp.py
index ebdf3621c..e98c7da25 100644
--- a/tests/load/test_native_fp.py
+++ b/tests/load/test_native_fp.py
@@ -275,7 +275,7 @@ def test_install_wire_dispatches_kquant_switchlinear(codec, bpr):
@pytest.mark.parametrize("native_fp_wire", [False, True])
-def test_install_native_fp_non_switch_fails_loud(native_fp_wire):
+def test_install_native_fp_embedding_fails_loud(native_fp_wire):
import mlx.nn as nn
from gmlx.load.modules import install_kquant_modules
@@ -283,13 +283,39 @@ def test_install_native_fp_non_switch_fails_loud(native_fp_wire):
class Tiny(nn.Module):
def __init__(self):
super().__init__()
- self.proj = nn.Linear(32, 32, bias=False)
+ self.embed = nn.Embedding(32, 32)
with pytest.raises(NotImplementedError, match="mxfp4"):
- install_kquant_modules(Tiny(), {"proj.weight": "mxfp4"},
+ install_kquant_modules(Tiny(), {"embed.weight": "mxfp4"},
native_fp_wire=native_fp_wire)
+@pytest.mark.parametrize("native_fp_wire", [False, True])
+@pytest.mark.parametrize("codec", ["mxfp4", "nvfp4"])
+def test_install_native_fp_linear_dispatches_wire(native_fp_wire, codec):
+ # Nemotron-3.5-Bolt ships dense NVFP4 Mamba projections. Both modes keep
+ # the wire bytes on a KQuantLinear; no packed dense module exists.
+ if not _kq_has_fp4():
+ pytest.skip("mlx-kquant build lacks the fp4 codecs")
+ import mlx.nn as nn
+ from mlx_kquant.codec_geometry import bytes_per_row
+ from mlx_kquant.nn import KQuantLinear
+
+ from gmlx.load.modules import install_kquant_modules
+
+ class Tiny(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.proj = nn.Linear(128, 64, bias=False)
+
+ model = Tiny()
+ assert install_kquant_modules(model, {"proj.weight": codec},
+ native_fp_wire=native_fp_wire) == 1
+ assert isinstance(model.proj, KQuantLinear)
+ assert model.proj.kquant_type == codec
+ assert model.proj.weight.shape == (64, bytes_per_row(codec, 128))
+
+
@pytest.mark.parametrize("native_fp_wire", [False, True])
def test_install_native_fp_multilinear_dispatches_wire(native_fp_wire):
# Kimi-K3 regression: llama-quantize's MXFP4_MOE ftype also assigns MXFP4
@@ -324,9 +350,10 @@ def __init__(self):
assert m.scales.dtype == mx.uint8 and m.scales.shape == (1,)
-def test_native_fp_multilinear_keys_exempt_from_repack():
- # The packed repack must leave MultiLinear-destined wire bytes untouched
- # (KQuantMultiLinear reads ggml wire, not MLX packed layout).
+def test_native_fp_wire_only_keys_exempt_from_repack():
+ # The packed repack must leave MultiLinear- and Linear-destined wire
+ # bytes untouched (KQuantMultiLinear and KQuantLinear read ggml wire,
+ # not MLX packed layout).
import mlx.nn as nn
import gmlx.load.loader as loader
@@ -340,11 +367,13 @@ class Tiny(nn.Module):
def __init__(self):
super().__init__()
self.embed_q = MultiLinear(64, 32, 2)
+ self.proj = nn.Linear(64, 32, bias=False)
model = Tiny()
- meta = {"embed_q.weight": "mxfp4", "experts.weight": "mxfp4"}
- keys = loader._native_fp_multilinear_keys(model, meta)
- assert keys == {"embed_q.weight"}
+ meta = {"embed_q.weight": "mxfp4", "proj.weight": "nvfp4",
+ "experts.weight": "mxfp4"}
+ keys = loader._native_fp_wire_only_keys(model, meta)
+ assert keys == {"embed_q.weight", "proj.weight"}
wire = np.arange(2 * 32 * 34, dtype=np.uint8).reshape(2, 32, 34)
hf = {"embed_q.weight": wire}
diff --git a/tests/load/test_out_scale.py b/tests/load/test_out_scale.py
new file mode 100644
index 000000000..c6eb55730
--- /dev/null
+++ b/tests/load/test_out_scale.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Asher Feldman
+"""Per-tensor weight scales: ``install_out_scales`` multiplies a module's
+output by its ``out_scale``, keeps float16 activations finite under a small
+scale, applies an lm_head scale only to an NVFP4 head, refuses per-expert
+scales, and the wire fusions skip a scaled projection."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from gmlx.load.out_scale import has_out_scale, install_out_scales
+
+
+class _Tiny(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.proj = nn.Linear(16, 8, bias=False)
+ self.other = nn.Linear(16, 8, bias=False)
+ self.lm_head = nn.Linear(16, 8, bias=False)
+
+
+def test_scales_the_output_of_named_modules_only():
+ model = _Tiny()
+ weights = {"proj.out_scale": mx.array([0.25])}
+ assert install_out_scales(model, weights) == 1
+ model.load_weights(list(weights.items()), strict=False)
+ x = mx.random.normal((2, 16))
+ np.testing.assert_allclose(
+ np.array(model.proj(x)), 0.25 * np.array(x @ model.proj.weight.T),
+ rtol=1e-6)
+ assert has_out_scale(model.proj) and not has_out_scale(model.other)
+ assert isinstance(model.proj, nn.Linear)
+ assert "out_scale" not in model.trainable_parameters()["proj"]
+
+
+def test_install_is_idempotent():
+ model = _Tiny()
+ weights = {"proj.out_scale": mx.array([0.5])}
+ install_out_scales(model, weights)
+ cls = type(model.proj)
+ install_out_scales(model, weights)
+ assert type(model.proj) is cls
+
+
+def test_float16_product_past_range_stays_finite():
+ # A 1e-4 scale with weights 1e4 larger: the unscaled product overflows
+ # float16, the scaled one does not.
+ model = _Tiny()
+ model.proj.weight = mx.full((8, 16), 1e4)
+ weights = {"proj.out_scale": mx.array([1e-4])}
+ install_out_scales(model, weights)
+ model.load_weights(list(weights.items()), strict=False)
+ y = model.proj(mx.full((1, 16), 1.0, dtype=mx.float16))
+ assert y.dtype == mx.float16
+ np.testing.assert_allclose(np.array(y.astype(mx.float32)), 16.0, rtol=1e-2)
+
+
+def test_lm_head_scale_needs_an_nvfp4_head():
+ model = _Tiny()
+ weights = {"lm_head.out_scale": mx.array([0.5])}
+ assert install_out_scales(model, weights, {"lm_head.weight": "q8_0"}) == 0
+ assert "lm_head.out_scale" not in weights
+ assert not has_out_scale(model.lm_head)
+ weights = {"lm_head.out_scale": mx.array([0.5])}
+ assert install_out_scales(model, weights, {"lm_head.weight": "nvfp4"}) == 1
+
+
+def test_per_expert_scale_fails_loud():
+ model = _Tiny()
+ with pytest.raises(NotImplementedError, match="per-expert"):
+ install_out_scales(model, {"proj.out_scale": mx.ones((4,))})
+
+
+def test_wire_fusions_skip_a_scaled_projection():
+ from mlx_kquant.nn import KQuantLinear
+
+ from gmlx.load.modules import _kq_dense
+ from gmlx.upstream.occupancy_fuse import _same_codec_kquant
+
+ class Mlp(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.gate_proj = KQuantLinear(64, 32, False, "q8_0")
+ self.up_proj = KQuantLinear(64, 32, False, "q8_0")
+
+ m = Mlp()
+ assert _kq_dense(m.gate_proj, {"q8_0"})
+ assert _same_codec_kquant([m.gate_proj, m.up_proj])
+ install_out_scales(m, {"gate_proj.out_scale": mx.array([0.5])})
+ assert isinstance(m.gate_proj, KQuantLinear)
+ assert not _kq_dense(m.gate_proj, {"q8_0"})
+ assert not _same_codec_kquant([m.gate_proj, m.up_proj])
diff --git a/tests/load/test_remap.py b/tests/load/test_remap.py
index 772975161..264c1c73b 100644
--- a/tests/load/test_remap.py
+++ b/tests/load/test_remap.py
@@ -1489,9 +1489,22 @@ def test_vlm_tower_tensors_skip():
assert d("qwen2", "mm.0.weight").kind == SKIP
+def test_weight_scale_follows_its_weight():
+ # llama.cpp's generic pass loads these for every arch and multiplies the
+ # matmul output; each lands next to its weight as out_scale.
+ r = d("qwen2", "blk.0.ffn_gate.scale")
+ assert r.kind == MAP
+ assert r.hf_name == "model.layers.0.mlp.gate_proj.out_scale"
+ r = d("nemotron_h_moe", "blk.3.ssm_in.scale")
+ assert r.hf_name == "backbone.layers.3.mixer.in_proj.out_scale"
+ assert d("llama", "output.scale").hf_name == "lm_head.out_scale"
+
+
def test_unsloth_ud_scale_skipped_when_not_claimed():
- # On qwen2 nothing claims a bare .scale -> universal Unsloth-UD skip.
- assert d("qwen2", "blk.0.ffn_gate.scale").kind == SKIP
+ # A .scale outside llama.cpp's weight-scale list, and every
+ # .input_scale, is skipped.
+ assert d("qwen2", "blk.0.attn_norm.scale").kind == SKIP
+ assert d("qwen2", "blk.0.ffn_gate.input_scale").kind == SKIP
def test_rope_freqs_skipped():
diff --git a/tests/models/test_nemotron_h.py b/tests/models/test_nemotron_h.py
new file mode 100644
index 000000000..464814178
--- /dev/null
+++ b/tests/models/test_nemotron_h.py
@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Asher Feldman
+"""Nemotron-H MoE block against a float64 reference of llama.cpp's
+``build_moe_ffn`` as the nemotron_h graph calls it, on tiny random weights.
+No GGUF and no download.
+
+The reference covers per-layer expert counts, widths and top-k (Bolt), the
+fp32 router, the selection-only bias, the weight-sum clamp at
+6.103515625e-5, renormalization at top-1, and the scale rule (0 and 1 leave
+the weights alone). The synth tests drive the per-layer derivation from
+router shapes.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import mlx.core as mx
+from mlx.utils import tree_flatten, tree_unflatten
+
+from gmlx.load.config_synth import synthesize_config
+from gmlx.models.nemotron_h.model import Model, ModelArgs, per_layer
+
+ARCH = "nemotron_h_moe"
+# M E * E E: experts 8/6/4 wide 24/16/32 top 3/2/1 on the MoE layers.
+_PATTERN = ["M", "E", "*", "E", "E"]
+_COUNTS = [0, 8, 0, 6, 4]
+_WIDTHS = [0, 24, 0, 16, 32]
+_TOPK = [0, 3, 0, 2, 1]
+
+
+def _tiny_args(**over) -> ModelArgs:
+ kw = dict(
+ model_type="nemotron_h",
+ vocab_size=40,
+ hidden_size=64,
+ intermediate_size=32,
+ num_hidden_layers=len(_PATTERN),
+ max_position_embeddings=1024,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ attention_bias=False,
+ mamba_num_heads=4,
+ mamba_head_dim=16,
+ mamba_proj_bias=False,
+ ssm_state_size=8,
+ conv_kernel=4,
+ n_groups=2,
+ mlp_bias=False,
+ layer_norm_epsilon=1e-5,
+ use_bias=False,
+ use_conv_bias=True,
+ hybrid_override_pattern=list(_PATTERN),
+ head_dim=16,
+ moe_intermediate_size=list(_WIDTHS),
+ moe_shared_expert_intermediate_size=48,
+ moe_latent_size=16,
+ n_routed_experts=list(_COUNTS),
+ n_shared_experts=1,
+ num_experts_per_tok=list(_TOPK),
+ norm_topk_prob=True,
+ routed_scaling_factor=5.0,
+ n_group=1,
+ topk_group=1,
+ )
+ kw.update(over)
+ return ModelArgs(**kw)
+
+
+# Random values here put the Mamba scan outside its working range.
+_KEEP_INIT = (".A_log", ".dt_bias", ".D", ".conv1d.")
+
+
+def _randomized(args: ModelArgs, seed: int = 0) -> Model:
+ model = Model(args)
+ rng = np.random.default_rng(seed)
+ params = [(k, mx.array(rng.normal(0, 0.3, v.shape).astype(np.float32)))
+ for k, v in tree_flatten(model.parameters())
+ if not any(s in k for s in _KEEP_INIT)]
+ model.update(tree_unflatten(params))
+ mx.eval(model.parameters())
+ return model
+
+
+def _np(a) -> np.ndarray:
+ return np.array(a.astype(mx.float32), dtype=np.float64)
+
+
+def _relu2(v):
+ return np.maximum(v, 0.0) ** 2
+
+
+def _ref_moe(moe, x: np.ndarray, top_k: int, norm: bool, scale: float):
+ """llama.cpp build_moe_ffn (sigmoid, selection bias, no gate, relu^2)
+ plus nemotron_h's latent projections and shared expert."""
+ w_r = _np(moe.gate.weight)
+ bias = _np(moe.gate.e_score_correction_bias)
+ probs = 1.0 / (1.0 + np.exp(-(x @ w_r.T)))
+ sel = probs + bias
+ idx = np.argsort(-sel, axis=-1, kind="stable")[..., :top_k]
+ w = np.take_along_axis(probs, idx, axis=-1)
+ if norm:
+ w = w / np.maximum(w.sum(-1, keepdims=True), 6.103515625e-5)
+ if scale not in (0.0, 1.0):
+ w = w * scale
+ xl = x @ _np(moe.fc1_latent_proj.weight).T
+ up = _np(moe.switch_mlp.fc1.weight)
+ down = _np(moe.switch_mlp.fc2.weight)
+ y = np.zeros_like(xl)
+ for t in range(x.shape[0]):
+ for j in range(top_k):
+ e = idx[t, j]
+ y[t] += w[t, j] * (down[e] @ _relu2(up[e] @ xl[t]))
+ out = y @ _np(moe.fc2_latent_proj.weight).T
+ sh = moe.shared_experts
+ return out + _relu2(x @ _np(sh.up_proj.weight).T) @ _np(sh.down_proj.weight).T
+
+
+@pytest.mark.parametrize("layer", [1, 3, 4])
+def test_per_layer_moe_sizes(layer):
+ model = _randomized(_tiny_args())
+ moe = model.layers[layer].mixer
+ assert moe.gate.weight.shape == (_COUNTS[layer], 64)
+ assert moe.gate.weight.dtype == mx.float32
+ assert moe.switch_mlp.fc1.weight.shape == (_COUNTS[layer], _WIDTHS[layer], 16)
+ assert moe.gate.top_k == _TOPK[layer]
+
+
+@pytest.mark.parametrize("layer", [1, 3, 4])
+@pytest.mark.parametrize("scale", [5.0, 1.0, 0.0])
+def test_moe_matches_reference(layer, scale):
+ model = _randomized(_tiny_args(routed_scaling_factor=scale), seed=layer)
+ moe = model.layers[layer].mixer
+ x = np.random.default_rng(7).normal(0, 1, (5, 64)).astype(np.float32)
+ got = _np(moe(mx.array(x)[None]))[0]
+ ref = _ref_moe(moe, x.astype(np.float64), _TOPK[layer], True, scale)
+ np.testing.assert_allclose(got, ref, rtol=2e-3, atol=2e-3)
+
+
+def test_weight_sum_clamp():
+ # Router logits far below zero: the top-k probabilities sum under the
+ # floor, so llama.cpp divides by the floor, not the sum.
+ model = _randomized(_tiny_args(), seed=3)
+ moe = model.layers[1].mixer
+ moe.gate.weight = mx.full(moe.gate.weight.shape, -0.5) # logits -32
+ moe.gate.e_score_correction_bias = mx.array(
+ np.linspace(0, 1, 8).astype(np.float32))
+ inds, scores = moe.gate(mx.ones((1, 2, 64)))
+ probs = 1.0 / (1.0 + np.exp(32.0))
+ np.testing.assert_allclose(
+ np.array(scores), probs / 6.103515625e-5 * 5.0, rtol=1e-3)
+ assert set(np.array(inds)[0, 0].tolist()) == {5, 6, 7}
+
+
+def test_top1_still_renormalizes():
+ model = _randomized(_tiny_args(), seed=4)
+ _, scores = model.layers[4].mixer.gate(mx.ones((1, 3, 64)))
+ np.testing.assert_allclose(np.array(scores), 5.0, rtol=1e-6)
+
+
+def test_forward_and_cached_decode_agree():
+ model = _randomized(_tiny_args(), seed=5)
+ tokens = mx.array([[1, 5, 9, 2, 7, 3]])
+ full = model(tokens)
+ cache = model.make_cache()
+ model(tokens[:, :-1], cache=cache)
+ step = model(tokens[:, -1:], cache=cache)
+ # The chunked scan and the one-step update round differently.
+ np.testing.assert_allclose(
+ _np(step[0, -1]), _np(full[0, -1]), rtol=5e-3, atol=5e-3)
+
+
+def test_per_layer_index_past_list_takes_largest():
+ assert per_layer([0, 8, 0, 6], 3) == 6
+ assert per_layer([0, 8, 0, 6], 4) == 8
+ assert per_layer([0, 8, 0, 6], None) == 8
+ assert per_layer(5, 2) == 5
+
+
+def test_cast_predicate_matches_loader_keep_row():
+ from gmlx.load.loader import _FP32_KEEP_BY_MODEL_TYPE
+
+ keep = _FP32_KEEP_BY_MODEL_TYPE["nemotron_h"]
+ model = Model(_tiny_args())
+ for name, _ in tree_flatten(model.parameters()):
+ assert model.cast_predicate(name) == (
+ not any(s in name for s in keep)), name
+ assert not model.cast_predicate("backbone.layers.1.mixer.gate.weight")
+ assert model.cast_predicate("backbone.layers.1.mixer.switch_mlp.fc1.weight")
+
+
+def _bolt_meta():
+ """Tiny Bolt-shaped KV, per-layer arrays over the 5 blocks above."""
+ kv = [0, 0, 2, 0, 0]
+ return {
+ "general.architecture": ARCH,
+ f"{ARCH}.embedding_length": 64,
+ f"{ARCH}.block_count": 5,
+ f"{ARCH}.attention.head_count": 4,
+ f"{ARCH}.attention.head_count_kv": kv,
+ f"{ARCH}.feed_forward_length": list(_WIDTHS),
+ f"{ARCH}.context_length": 1024,
+ f"{ARCH}.attention.layer_norm_rms_epsilon": 1e-5,
+ f"{ARCH}.attention.key_length": 16,
+ f"{ARCH}.ssm.inner_size": 64,
+ f"{ARCH}.ssm.time_step_rank": 4,
+ f"{ARCH}.ssm.state_size": 8,
+ f"{ARCH}.ssm.conv_kernel": 4,
+ f"{ARCH}.ssm.group_count": 2,
+ f"{ARCH}.expert_count": 8,
+ f"{ARCH}.expert_used_count": list(_TOPK),
+ f"{ARCH}.expert_feed_forward_length": list(_WIDTHS),
+ f"{ARCH}.expert_shared_count": 1,
+ f"{ARCH}.expert_shared_feed_forward_length": 48,
+ f"{ARCH}.expert_weights_norm": True,
+ f"{ARCH}.expert_weights_scale": 5.0,
+ f"{ARCH}.moe_latent_size": 16,
+ "tokenizer.ggml.tokens": ["t"] * 40,
+ }
+
+
+def _bolt_shapes():
+ return {f"blk.{i}.ffn_gate_inp.weight": [64, n]
+ for i, n in enumerate(_COUNTS) if n}
+
+
+def test_synth_reads_per_layer_sizes():
+ config = synthesize_config(_bolt_meta(), tensor_shapes=_bolt_shapes())
+ assert config["hybrid_override_pattern"] == _PATTERN
+ assert config["n_routed_experts"] == _COUNTS
+ assert config["moe_intermediate_size"] == _WIDTHS
+ assert config["num_experts_per_tok"] == _TOPK
+ model = Model(ModelArgs.from_dict(config))
+ assert model.layers[4].mixer.gate.weight.shape == (4, 64)
+
+
+def test_synth_uniform_sizes_stay_scalar():
+ meta = _bolt_meta()
+ meta[f"{ARCH}.expert_used_count"] = 2
+ meta[f"{ARCH}.expert_feed_forward_length"] = 24
+ shapes = {k: [64, 8] for k in _bolt_shapes()}
+ config = synthesize_config(meta, tensor_shapes=shapes)
+ assert config["n_routed_experts"] == 8
+ assert config["moe_intermediate_size"] == 24
+ assert config["num_experts_per_tok"] == 2
+
+
+def test_synth_width_falls_back_to_ff_over_topk():
+ # No expert_feed_forward_length: llama.cpp takes n_ff / n_expert_used.
+ meta = _bolt_meta()
+ del meta[f"{ARCH}.expert_feed_forward_length"]
+ config = synthesize_config(meta, tensor_shapes=_bolt_shapes())
+ assert config["moe_intermediate_size"] == [0, 8, 0, 8, 32]
+
+
+def test_synth_norm_defaults_off_like_llama_cpp():
+ meta = _bolt_meta()
+ del meta[f"{ARCH}.expert_weights_norm"]
+ assert synthesize_config(meta, tensor_shapes=_bolt_shapes())[
+ "norm_topk_prob"] is False
+
+
+def test_synth_rejects_router_past_expert_count():
+ shapes = _bolt_shapes()
+ shapes["blk.3.ffn_gate_inp.weight"] = [64, 9]
+ with pytest.raises(ValueError, match="blk.3.ffn_gate_inp"):
+ synthesize_config(_bolt_meta(), tensor_shapes=shapes)
From 2bdf26acd1965d61723aa07f427077a46ee2f4a2 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sat, 17 Oct 2026 23:41:26 -0700
Subject: [PATCH 02/16] feat(spec): lilicorr drafter for nemotron 3.5 bolt,
paired from its sibling repo
---
gmlx/load/arch_table.py | 3 +
gmlx/load/discovery.py | 34 ++-
gmlx/models/nemotron_h/mtp.py | 76 +++++-
gmlx/spec/dflash_drafter.py | 17 +-
gmlx/spec/lilicorr_drafter.py | 253 ++++++++++++++++++++
gmlx/spec/mtp_load.py | 199 ++++++++++++++--
tests/load/test_discovery.py | 51 ++++
tests/spec/test_lilicorr_drafter.py | 352 ++++++++++++++++++++++++++++
8 files changed, 949 insertions(+), 36 deletions(-)
create mode 100644 gmlx/spec/lilicorr_drafter.py
create mode 100644 tests/spec/test_lilicorr_drafter.py
diff --git a/gmlx/load/arch_table.py b/gmlx/load/arch_table.py
index c6f04d944..1af96d65a 100644
--- a/gmlx/load/arch_table.py
+++ b/gmlx/load/arch_table.py
@@ -180,6 +180,9 @@ class UnsupportedArchError(Exception):
# the HF mtp.* tensors, or the llama.cpp mtp-*.gguf NextN sidecar, which
# keeps the target's arch.
"qwen4_exp": ("qwen4exp-mtp", "qwen4exp"),
+ # LiLiCorr DFlash drafters (Nemotron-3.5-Bolt). The in-file NextN head
+ # and the same-arch mtp sidecar need no row.
+ "nemotron_h": ("dflash",),
}
diff --git a/gmlx/load/discovery.py b/gmlx/load/discovery.py
index 0de5315ed..358ef4c87 100644
--- a/gmlx/load/discovery.py
+++ b/gmlx/load/discovery.py
@@ -1165,11 +1165,13 @@ def _emit_dir(classified, spec, used_ids, out, *, configured=None,
f"{dr.path}", file=sys.stderr)
else:
unpaired.append(dr)
- if dr.declared and (not targets or dr.drafter_kind == "dflash2"):
+ if (dr.declared or _sibling_drafter_repo(dr.path)) and (
+ not targets or dr.drafter_kind == "dflash2"):
cross.append(dr)
# Cross-directory pass: a drafter that declares its base model (a DFlash 2
- # header does) pairs with that model wherever the scan found it, and a
+ # header does), or sits in a marked sibling repo of its target's, pairs
+ # with that model wherever the scan found it, and a
# DFlash 2 drafter takes over a target its same-directory DFlash v1
# sibling claimed above. Declared hits only; the filename rule stays
# same-directory.
@@ -1289,6 +1291,29 @@ def _configured_in_dir(dirpath: str, configured) -> list:
return out
+# A drafter repo named for its target's repo plus one of these markers
+# (NVIDIA's ``-LiLiCorr-GGUF`` beside ``-GGUF``) pairs with the
+# models of that repo wherever the scan found them.
+_SIBLING_DRAFTER_MARKER_RE = re.compile(r"-lilicorr(?=-gguf$|$)", re.IGNORECASE)
+
+
+def _repo_name(path: str) -> str:
+ """The repo a GGUF sits in: ``--`` of an HF cache entry, else
+ its directory's name (the ``__`` local layout)."""
+ parts = os.path.normpath(os.path.abspath(path)).split(os.sep)
+ for part in reversed(parts[:-1]):
+ if part.startswith("models--"):
+ return part[len("models--"):].replace("--", "__", 1)
+ return parts[-2] if len(parts) > 1 else ""
+
+
+def _sibling_drafter_repo(drafter_path: str) -> str | None:
+ """The target repo a marked drafter repo names, or None."""
+ repo = _repo_name(drafter_path)
+ target = _SIBLING_DRAFTER_MARKER_RE.sub("", repo)
+ return target.lower() if target != repo else None
+
+
def _target_ids(c: ClassifiedGguf) -> set:
"""Ids a target answers to: its declared base ids plus ``general.name``."""
ids = set(c.declared)
@@ -1331,6 +1356,11 @@ def _open(mc) -> bool:
and not (dr.n_embd and c.n_embd and dr.n_embd != c.n_embd)]
if not cands:
return []
+ sibling = _sibling_drafter_repo(dr.path)
+ if sibling is not None:
+ hits = [mc for mc, c in cands if _repo_name(c.path).lower() == sibling]
+ if hits:
+ return hits
if dr.declared:
hits = [mc for mc, c in cands if set(dr.declared) & _target_ids(c)]
if hits:
diff --git a/gmlx/models/nemotron_h/mtp.py b/gmlx/models/nemotron_h/mtp.py
index c28a79339..cc24e5858 100644
--- a/gmlx/models/nemotron_h/mtp.py
+++ b/gmlx/models/nemotron_h/mtp.py
@@ -92,8 +92,16 @@ def __call__(
inputs,
cache: Optional[Any] = None,
mamba_sink: Optional[list] = None,
+ capture: Optional[tuple] = None,
+ exact: bool = True,
):
+ """``capture`` is ``(layer_ids, out)``: the raw output of each listed
+ layer is appended to ``out`` in layer order (the DFlash seam).
+ ``exact`` off runs a sinked verify batched, as llama.cpp does: the
+ block drafters verify 8 positions, where stepping every layer costs
+ 8 decode steps."""
hidden_states = self.embeddings(inputs)
+ capture_ids, captured = capture if capture is not None else ((), None)
if cache is None:
cache = [None] * len(self.layers)
@@ -101,7 +109,7 @@ def __call__(
ssm_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
cache_counter = 0
- for layer in self.layers:
+ for i, layer in enumerate(self.layers):
if layer.block_type == "M" or layer.block_type == "*":
c = cache[cache_counter]
cache_counter += 1
@@ -120,15 +128,22 @@ def __call__(
mamba_sink.append({
"layer": layer, "cache": c,
"pre": (c[0], c[1]),
- "inputs": hidden_states, "mask": mask,
+ "inputs": hidden_states, "mask": mask, "exact": exact,
})
- hidden_states = _mamba_stepwise(layer, hidden_states, mask, c)
+ if exact:
+ hidden_states = _mamba_stepwise(layer, hidden_states, mask, c)
+ else:
+ hidden_states = layer(hidden_states, mask=mask, cache=c)
+ elif mamba_sink is not None and not exact:
+ hidden_states = layer(hidden_states, mask=mask, cache=c)
elif mamba_sink is not None and layer.block_type == "*" and c is not None:
hidden_states = _attn_stepwise(layer, hidden_states, c)
elif mamba_sink is not None and layer.block_type == "E":
hidden_states = _moe_rowgate(layer, hidden_states)
else:
hidden_states = layer(hidden_states, mask=mask, cache=c)
+ if i in capture_ids:
+ captured.append(hidden_states)
return self.norm_f(hidden_states)
@@ -185,14 +200,20 @@ def rollback_verify_sink(sink: list, n: int) -> None:
to the state after its first ``n`` (the accepted prefix): restore the
recorded pre-verify (conv, ssm) state, then replay the block over the
accepted prefix. O(n <= block) per layer; the attention KV leaves are
- trimmed by the caller. The replay steps position-wise like the verify
- (see :func:`_mamba_stepwise`), so the rewound state is bit-identical to
- a plain decode over the accepted tokens."""
+ trimmed by the caller. The replay runs the way the verify ran: stepped
+ position-wise for an exact verify (see :func:`_mamba_stepwise`), so the
+ rewound state is bit-identical to a plain decode over the accepted
+ tokens, else batched."""
for e in sink:
cache = e["cache"]
cache[0], cache[1] = e["pre"]
mask = e["mask"]
- _mamba_stepwise(e["layer"], e["inputs"][:, :n], mask, cache)
+ if isinstance(mask, mx.array):
+ mask = mask[:, :n]
+ if e.get("exact", True):
+ _mamba_stepwise(e["layer"], e["inputs"][:, :n], mask, cache)
+ else:
+ e["layer"](e["inputs"][:, :n], mask=mask, cache=cache)
class NemotronHSpecLM(nh.Model):
@@ -200,10 +221,40 @@ class NemotronHSpecLM(nh.Model):
engine probes on its target. The backbone is swapped for the sink-aware
subclass; numerics are byte-identical to stock."""
+ _dflash_capture: Optional[tuple] = None
+
def __init__(self, args):
super().__init__(args)
self.backbone = NemotronHSpecBackbone(args)
+ def set_dflash_capture(self, layer_ids) -> None:
+ """Arm the DFlash capture: every engine-facing hidden becomes
+ ``[trunk | cap ...]``, the post-norm trunk followed by the raw
+ outputs of ``layer_ids`` in layer order."""
+ # Around nn.Module.__setattr__, which would route a tuple into the
+ # parameter tree.
+ object.__setattr__(self, "_dflash_capture",
+ tuple(int(i) for i in layer_ids) or None)
+
+ def _dflash_trunk(self, hidden: mx.array) -> mx.array:
+ if self._dflash_capture is None:
+ return hidden
+ # The head is a quantized kernel that reads the buffer directly; a
+ # lazy slice would hand it the packed strides.
+ return mx.contiguous(hidden[..., : int(self.args.hidden_size)])
+
+ def _hidden(self, inputs, cache, mamba_sink=None) -> mx.array:
+ ids = self._dflash_capture
+ if ids is None:
+ return self.backbone(inputs, cache=cache, mamba_sink=mamba_sink)
+ caps: list = []
+ hidden = self.backbone(inputs, cache=cache, mamba_sink=mamba_sink,
+ capture=(ids, caps), exact=False)
+ if len(caps) != len(ids):
+ raise RuntimeError(
+ f"dflash capture collected {len(caps)} of {len(ids)} layers")
+ return mx.concatenate([hidden, *caps], axis=-1)
+
@property
def model(self):
# mlx-vlm-facing alias: MTPTextTarget and the drafter bind walk read
@@ -234,12 +285,13 @@ def __call__(
# keyword; the token ids are authoritative (a GGUF text target has
# no vision features) and shared_kv is never used.
del inputs_embeds, n_to_process, kwargs
- hidden = self.backbone(inputs, cache=cache)
- logits = self.lm_head(hidden)
if not (return_hidden or return_shared_kv):
from gmlx._vendor.mlx_vlm.models.base import LanguageModelOutput
- return LanguageModelOutput(logits=logits)
+ return LanguageModelOutput(
+ logits=self.lm_head(self.backbone(inputs, cache=cache)))
+ hidden = self._hidden(inputs, cache)
+ logits = self.lm_head(self._dflash_trunk(hidden))
return _SpecOutput(logits=logits, hidden_states=[hidden])
# --- speculative hooks (owned engine contract) -------------------------
@@ -248,7 +300,7 @@ def __call__(
# from-hidden hooks apply no further norm.
def speculative_logits_from_hidden(self, hidden: mx.array) -> mx.array:
- return self.lm_head(hidden)
+ return self.lm_head(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)
@@ -257,7 +309,7 @@ def speculative_verify_hidden(self, verify_input: mx.array, prompt_cache):
"""The single verify forward (qL = drafts + 1): backbone only, with
the Mamba2 rollback sink recorded."""
sink: list = []
- hidden = self.backbone(verify_input, cache=prompt_cache, mamba_sink=sink)
+ hidden = self._hidden(verify_input, prompt_cache, mamba_sink=sink)
return hidden, {}, sink
def rollback_speculative_cache(
diff --git a/gmlx/spec/dflash_drafter.py b/gmlx/spec/dflash_drafter.py
index 88fe4f51b..36507d8b3 100644
--- a/gmlx/spec/dflash_drafter.py
+++ b/gmlx/spec/dflash_drafter.py
@@ -117,6 +117,17 @@ class DFlashConfig:
log_snr_range: Optional[tuple] = None
# The anchor row proposes a token as well (llama.cpp's default).
sample_from_anchor: bool = True
+ # Per-head attention sinks in every draft layer (Nemotron drafters).
+ attention_sinks: bool = False
+ # The LiLiCorr lattice head; lilicorr_layers == 0 means none.
+ lilicorr_layers: int = 0
+ lilicorr_hidden_size: int = 0
+ lilicorr_top_k: int = 0
+ lilicorr_heads: int = 0
+ lilicorr_factor_dim: int = 0
+ lilicorr_mlp_size: int = 0
+ lilicorr_vector_eps: float = 1e-4
+ lilicorr_logit_scale: float = 1.0
def __post_init__(self):
n = int(self.num_hidden_layers)
@@ -415,6 +426,8 @@ def __init__(self, config: DFlashConfig, layer_idx: int):
self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ if config.attention_sinks:
+ self.sinks = mx.zeros((self.n_heads,))
def inject(self, h_ctx: mx.array, rope, cache) -> None:
"""Write the context rows' K/V into the ring at the ring's positions."""
@@ -440,7 +453,9 @@ def draft(self, x: mx.array, rope, cache) -> mx.array:
k = mx.concatenate([ck, k], axis=2)
v = mx.concatenate([cv, v], axis=2)
mask = block_attention_mask(ctx_len, L, self.sliding_window, self.is_causal)
- o = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
+ sinks = self["sinks"].astype(q.dtype) if "sinks" in self else None
+ o = mx.fast.scaled_dot_product_attention(
+ q, k, v, scale=self.scale, mask=mask, sinks=sinks)
return self.o_proj(o.transpose(0, 2, 1, 3).reshape(B, L, -1))
diff --git a/gmlx/spec/lilicorr_drafter.py b/gmlx/spec/lilicorr_drafter.py
new file mode 100644
index 000000000..a99e5840d
--- /dev/null
+++ b/gmlx/spec/lilicorr_drafter.py
@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Asher Feldman
+"""LiLiCorr drafter: a DFlash backbone with a lattice head that rescores
+the top-k candidates of every draft slot (NVIDIA Nemotron-3.5 drafters).
+
+The backbone is DFlash 2's conv-wrapped layers with per-head attention
+sinks and no candidate selector. The head follows llama.cpp
+``src/models/dflash.cpp`` ``build_lilicorr_scores`` and the walk in
+``common/speculative.cpp``:
+
+- Each of the ``block - 1`` draft slots keeps its top-k candidates in
+ descending logit order. A node per (slot, candidate) sums the projected
+ token embedding, the projected draft hidden of its slot, an MLP over five
+ features (log-prob, prob, log-prob minus the slot's best, rank / (k - 1),
+ first-rank flag) and slot and rank embeddings.
+- A small pre-norm transformer runs over all nodes with a learned per-head
+ attention bias, then factor heads give unit ``out`` and ``in`` vectors per
+ node and an ``anchor_out`` vector from the encoder row of the last
+ committed position.
+- Slot 0 scores ``scale * ``; slot ``s`` scores candidate
+ ``n`` after predecessor ``p`` as ``scale * ``.
+ Drafts are the greedy walk through that lattice, at any temperature.
+
+The head only runs on a full block, so a round always denoises the trained
+block and walks as many slots as the engine asked for.
+"""
+
+from __future__ import annotations
+
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .dflash_drafter import (
+ DFlash2DecoderLayer,
+ DFlashConfig,
+ DFlashDrafter,
+ greedy_path,
+ path_estimate,
+)
+from .drafter_protocol import DraftStash
+
+_FEATURES = 5
+_FEATURE_NORM_EPS = 1e-5
+
+
+class _Ring(list):
+ """A drafter's per-layer KV ring plus the encoder row of its last
+ committed position, which anchors the lattice."""
+
+ anchor: Optional[mx.array] = None
+
+
+def _l2_normalize(x: mx.array, eps: float) -> mx.array:
+ # ggml_l2_norm: x / max(||x||, eps).
+ norm = mx.sqrt((x * x).sum(axis=-1, keepdims=True))
+ return x / mx.maximum(norm, eps)
+
+
+class LiLiCorrLayer(nn.Module):
+ def __init__(self, hidden: int, heads: int, mlp: int, eps: float):
+ super().__init__()
+ self.heads = int(heads)
+ self.eps = float(eps)
+ self.attn_norm = mx.ones((hidden,))
+ self.in_proj = nn.Linear(hidden, 3 * hidden, bias=True)
+ self.out_proj = nn.Linear(hidden, hidden, bias=True)
+ self.mlp_norm = mx.ones((hidden,))
+ self.mlp_up = nn.Linear(hidden, mlp, bias=True)
+ self.mlp_down = nn.Linear(mlp, hidden, bias=True)
+
+ def __call__(self, x: mx.array, bias: mx.array) -> mx.array:
+ """x [N, H] nodes; bias [heads, N, N] added to the scores."""
+ n, hidden = x.shape
+ hd = hidden // self.heads
+ a = mx.fast.rms_norm(x, self.attn_norm, self.eps)
+ q, k, v = mx.split(self.in_proj(a).astype(mx.float32), 3, axis=-1)
+ q, k, v = (t.reshape(n, self.heads, hd).transpose(1, 0, 2)
+ for t in (q, k, v))
+ scores = (q @ k.transpose(0, 2, 1)) * (hd ** -0.5) + bias
+ o = mx.softmax(scores, axis=-1) @ v
+ x = x + self.out_proj(o.transpose(1, 0, 2).reshape(n, hidden))
+ m = mx.fast.rms_norm(x, self.mlp_norm, self.eps)
+ return x + self.mlp_down(nn.silu(self.mlp_up(m)))
+
+
+class LiLiCorrHead(nn.Module):
+ """The lattice scorer over ``slots`` x ``top_k`` candidate nodes."""
+
+ def __init__(self, config: DFlashConfig, embed_dim: int):
+ super().__init__()
+ hidden = int(config.lilicorr_hidden_size)
+ k = int(config.lilicorr_top_k)
+ slots = int(config.native_block_size or config.block_size) - 1
+ factor = int(config.lilicorr_factor_dim)
+ self.top_k = k
+ self.eps = float(config.rms_norm_eps)
+ self.vector_eps = float(config.lilicorr_vector_eps)
+ self.logit_scale = float(config.lilicorr_logit_scale)
+ if hidden != embed_dim:
+ self.token_proj = nn.Linear(embed_dim, hidden, bias=True)
+ self.pass_hidden_proj = nn.Linear(embed_dim, hidden, bias=True)
+ self.context_proj = nn.Linear(embed_dim, hidden, bias=True)
+ self.feature_norm = nn.LayerNorm(_FEATURES, eps=_FEATURE_NORM_EPS)
+ self.feature_up = nn.Linear(_FEATURES, hidden, bias=True)
+ self.feature_down = nn.Linear(hidden, hidden, bias=True)
+ self.slot_embedding = mx.zeros((slots, hidden))
+ self.rank_embedding = mx.zeros((k, hidden))
+ self.attention_bias = mx.zeros(
+ (config.lilicorr_heads, slots * k, slots * k))
+ self.layers = [
+ LiLiCorrLayer(hidden, config.lilicorr_heads,
+ config.lilicorr_mlp_size, config.rms_norm_eps)
+ for _ in range(int(config.lilicorr_layers))]
+ self.output_norm = mx.ones((hidden,))
+ self.anchor_norm = mx.ones((hidden,))
+ self.factor_input_proj = nn.Linear(3 * hidden, hidden, bias=True)
+ self.out_head = nn.Linear(hidden, factor, bias=True)
+ self.in_head = nn.Linear(hidden, factor, bias=True)
+ self.anchor_out_head = nn.Linear(hidden, factor, bias=True)
+
+ def candidates(self, logits: mx.array):
+ """Top-k ids per slot in descending logit order, with their
+ log-probabilities under the full-vocab softmax. ``logits`` [S, V]."""
+ lg = logits.astype(mx.float32)
+ k = self.top_k
+ cands = mx.argpartition(-lg, kth=k - 1, axis=-1)[:, :k]
+ vals = mx.take_along_axis(lg, cands, axis=-1)
+ order = mx.argsort(-vals, axis=-1)
+ cands = mx.take_along_axis(cands, order, axis=-1)
+ lp = mx.take_along_axis(lg, cands, axis=-1) - mx.logsumexp(
+ lg, axis=-1, keepdims=True)
+ return cands, lp
+
+ def lattice(self, passes: mx.array, logits: mx.array, embed,
+ anchor: mx.array):
+ """Lattice of one block.
+
+ ``passes`` [S, D] are the normed draft hiddens of the slots,
+ ``logits`` [S, V] their head outputs, ``embed`` the target's token
+ embedding and ``anchor`` [D] the encoder row of the last committed
+ position. Returns ``(cands [S, k], first [k], edges [S-1, k, k])``:
+ ``edges[s-1][p][n]`` scores candidate ``n`` of slot ``s`` after
+ candidate ``p`` of slot ``s-1``. All lazy.
+ """
+ f32 = mx.float32
+ slots, k = int(logits.shape[0]), self.top_k
+ cands, lp = self.candidates(logits)
+ rank = mx.arange(k, dtype=f32)
+ rank_feat = rank / (k - 1) if k > 1 else mx.zeros_like(rank)
+ first_feat = (rank == 0).astype(f32)
+ feat = mx.stack([
+ lp, mx.exp(lp), lp - lp[:, :1],
+ mx.broadcast_to(rank_feat, lp.shape),
+ mx.broadcast_to(first_feat, lp.shape)], axis=-1) # [S, k, 5]
+ feat = self.feature_down(nn.silu(self.feature_up(
+ self.feature_norm(feat)))).astype(f32)
+
+ tokens = embed(cands) # [S, k, D]
+ if "token_proj" in self:
+ tokens = self.token_proj(tokens)
+ passes = self.pass_hidden_proj(passes).astype(f32)[:, None, :]
+ pos = self.slot_embedding[:, None, :] + self.rank_embedding[None]
+ x = (tokens.astype(f32) + passes + feat + pos.astype(f32))
+ x = x.reshape(slots * k, -1)
+ bias = self.attention_bias.astype(f32)
+ for layer in self.layers:
+ x = layer(x, bias).astype(f32)
+ x = mx.fast.rms_norm(x, self.output_norm, self.eps)
+
+ a = self.context_proj(anchor.reshape(1, -1)).astype(f32)
+ a = mx.fast.rms_norm(a, self.anchor_norm, self.eps)
+ ab = mx.broadcast_to(a, x.shape)
+ factors = nn.silu(self.factor_input_proj(
+ mx.concatenate([x, ab, x * ab], axis=-1)).astype(f32))
+ eps = self.vector_eps
+ out = _l2_normalize(self.out_head(factors).astype(f32), eps)
+ inn = _l2_normalize(self.in_head(factors).astype(f32), eps)
+ a_out = _l2_normalize(self.anchor_out_head(a).astype(f32), eps)[0]
+ out = out.reshape(slots, k, -1)
+ inn = inn.reshape(slots, k, -1)
+ first = self.logit_scale * (inn[0] @ a_out)
+ if slots > 1:
+ edges = self.logit_scale * (out[:-1] @ inn[1:].transpose(0, 2, 1))
+ else:
+ edges = mx.zeros((0, k, k), dtype=f32)
+ return cands, first, edges
+
+
+class LiLiCorrDrafter(DFlashDrafter):
+ """DFlash backbone (conv layers, sinks) + the LiLiCorr lattice head."""
+
+ layer_class = DFlash2DecoderLayer
+ kind_label = "lilicorr"
+ # The walk is greedy at every temperature (llama.cpp's reading), so the
+ # verify stays exact-match.
+ stochastic_draft = False
+ supports_q_stash = False
+
+ def __init__(self, config: DFlashConfig):
+ if config.conv_kernel_size <= 0 or config.conv_group_size <= 0:
+ raise ValueError("LiLiCorrDrafter needs conv_kernel_size and conv_group_size")
+ if not config.lilicorr_layers:
+ raise ValueError("LiLiCorrDrafter needs the lilicorr head metadata")
+ super().__init__(config)
+ self.lilicorr = LiLiCorrHead(config, int(config.hidden_size))
+ object.__setattr__(self, "_round_path", None)
+
+ @property
+ def round_estimate(self) -> Optional[mx.array]:
+ path = self._round_path
+ return None if path is None else path_estimate(*path)
+
+ def make_cache(self, left_padding: Optional[List[int]] = None) -> list:
+ return _Ring(super().make_cache(left_padding))
+
+ def append_context(self, captures: mx.array, cache=None) -> None:
+ ring = self._cache if cache is None else cache
+ h_ctx = self.hidden_norm(self.fc(captures))
+ for layer, c in zip(self.layers, ring):
+ layer.inject(h_ctx, self.rope, c)
+ ring.anchor = h_ctx[0, -1]
+
+ def draft_block(
+ self,
+ last_bonus,
+ hidden: mx.array,
+ cache,
+ block_size: int,
+ sampler,
+ token_dtype: mx.Dtype = mx.int32,
+ greedy: bool = False,
+ stash: Optional[DraftStash] = None,
+ ) -> mx.array:
+ """One round: the full trained block through the backbone, the
+ lattice over its slots, and the greedy walk over the first
+ ``block_size - 1`` of them."""
+ del hidden, cache, sampler, greedy, stash
+ native = self._native_block_size
+ n = min(int(block_size), native) - 1
+ block = self._block_tokens(last_bonus, native, token_dtype)
+ h = self._draft_hidden(block)[0, 1:]
+ logits = self._logits(h)
+ anchor: Any = getattr(self._cache, "anchor", None)
+ if anchor is None:
+ anchor = mx.zeros((int(self.config.hidden_size),), dtype=h.dtype)
+ cands, first, edges = self.lilicorr.lattice(
+ h, logits, self.embed_tokens, anchor)
+ edges = edges[: n - 1]
+ sel = greedy_path(first, edges)
+ object.__setattr__(self, "_round_path", (first, edges, sel))
+ return mx.take_along_axis(cands[:n], sel[:, None], axis=-1).reshape(1, -1)
diff --git a/gmlx/spec/mtp_load.py b/gmlx/spec/mtp_load.py
index 355e89cea..d7191f545 100644
--- a/gmlx/spec/mtp_load.py
+++ b/gmlx/spec/mtp_load.py
@@ -111,6 +111,7 @@
"muse_glimmer": _MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT,
"dflash2": None,
"dflash_dspark": None,
+ "lilicorr": None,
}
@@ -1046,12 +1047,15 @@ def dflash_container(arrays: dict) -> str:
Muse Glimmer DFlash one and the DFlash 2 drafters under ``dflash``, and
picks its graph on header keys. Tensor presence is the equivalent split
here - the DeepSeek DSpark carries MLA's ``attn_q_a`` or the
- hyper-connection head, the DFlash-backbone DSpark the markov head on
- plain ``attn_q`` layers, DFlash 2 the candidate selector, Muse Glimmer
- plain ``attn_q`` with per-head QK-norms and no hyper-connections.
+ hyper-connection head, LiLiCorr its ``lilicorr.*`` lattice head, the
+ DFlash-backbone DSpark the markov head on plain ``attn_q`` layers,
+ DFlash 2 the candidate selector, Muse Glimmer plain ``attn_q`` with
+ per-head QK-norms and no hyper-connections.
"""
if any(n.startswith("output_hc_") or ".attn_q_a" in n for n in arrays):
return "dspark"
+ if any(n.startswith("lilicorr.") for n in arrays):
+ return "lilicorr"
if any(n.startswith("markov_w1") for n in arrays):
return "dflash_dspark"
if any(n.startswith("selector_hidden") for n in arrays):
@@ -1111,10 +1115,41 @@ def dflash_container(arrays: dict) -> str:
"log_snr_fc1": "log_snr_fc1.weight",
"log_snr_fc2": "log_snr_fc2.weight",
}
+# LiLiCorr (llama.cpp dflash.cpp): DFlash 2 conv layers with attention sinks
+# and no selector, plus the lattice head. The head's own layers sit at
+# blk.{n_layers + i} and its feature MLP at blk.{n_layers + head_layers};
+# the rest of the head is under lilicorr.* (the drafter's ``lilicorr``).
+_LILICORR_BLK = {**_DFLASH_BLK, **_DFLASH2_BLK,
+ "attn_sinks": "self_attn.sinks"}
+_LILICORR_HEAD_BLK = {
+ "attn_norm": "attn_norm",
+ "attn_output": "out_proj.weight",
+ "ffn_norm": "mlp_norm",
+ "ffn_up": "mlp_up.weight",
+ "ffn_down": "mlp_down.weight",
+}
+_LILICORR_FEATURE_BLK = {
+ "ffn_norm": "feature_norm.weight",
+ "ffn_up": "feature_up.weight",
+ "ffn_down": "feature_down.weight",
+}
+_LILICORR_ROOT = {
+ **_DFLASH_ROOT,
+ **{f"lilicorr.{n}": f"lilicorr.{n}.weight" for n in (
+ "token_proj", "pass_hidden_proj", "context_proj", "factor_input_proj",
+ "out_head", "in_head", "anchor_out_head")},
+ "lilicorr.slot_embedding": "lilicorr.slot_embedding",
+ "lilicorr.rank_embedding": "lilicorr.rank_embedding",
+ "lilicorr.attention_bias": "lilicorr.attention_bias",
+ "lilicorr.output_norm": "lilicorr.output_norm",
+ "lilicorr.anchor_norm": "lilicorr.anchor_norm",
+}
+_LILICORR_HEAD_LAYER_RE = re.compile(r"^lilicorr\.layers\.(\d+)\.attn\.in_proj$")
_DFLASH_CONTAINER_MAPS = {
"muse_glimmer": (_DFLASH_BLK, _DFLASH_ROOT),
"dflash2": ({**_DFLASH_BLK, **_DFLASH2_BLK}, {**_DFLASH_ROOT, **_DFLASH2_ROOT}),
"dflash_dspark": (_DFLASH_BLK, {**_DFLASH_ROOT, **_DFLASH_DSPARK_ROOT}),
+ "lilicorr": (_LILICORR_BLK, _LILICORR_ROOT),
}
# Which target families a container can drive. A load-time check; the
# arch-tag filter (arch_table.drafter_serves) is discovery's pairing-time one.
@@ -1122,14 +1157,48 @@ def dflash_container(arrays: dict) -> str:
"dflash2": ("qwen3_5", "qwen3_5_text", "muse_glimmer"),
"dflash_dspark": ("qwen3_5", "qwen3_5_text"),
"muse_glimmer": ("muse_glimmer",),
+ "lilicorr": ("nemotron_h",),
}
-def remap_dflash_arrays(arrays: dict, kquant_meta: dict, container: str):
+def _lilicorr_blk_target(idx: int, leaf: str, n_layers: int,
+ head_layers: int) -> str | None:
+ """The head-side target of LiLiCorr's ``blk.{idx}.{leaf}``, past the
+ backbone's ``n_layers`` blocks."""
+ i = idx - n_layers
+ if i == head_layers:
+ target = _LILICORR_FEATURE_BLK.get(leaf)
+ return None if target is None else f"lilicorr.{target}"
+ target = _LILICORR_HEAD_BLK.get(leaf)
+ if target is None or i > head_layers:
+ return None
+ return f"lilicorr.layers.{i}.{target}"
+
+
+def _dflash_target(container: str, base: str, n_layers: int | None,
+ head_layers: int) -> str | None:
+ blk_map, root_map = _DFLASH_CONTAINER_MAPS[container]
+ if base.startswith("blk."):
+ _, idx, leaf = base.split(".", 2)
+ if container == "lilicorr" and n_layers is not None \
+ and int(idx) >= n_layers:
+ return _lilicorr_blk_target(int(idx), leaf, n_layers, head_layers)
+ target = blk_map.get(leaf)
+ return None if target is None else f"layers.{idx}.{target}"
+ if container == "lilicorr":
+ m = _LILICORR_HEAD_LAYER_RE.match(base)
+ if m:
+ return f"lilicorr.layers.{m.group(1)}.in_proj.weight"
+ return root_map.get(base)
+
+
+def remap_dflash_arrays(arrays: dict, kquant_meta: dict, container: str,
+ *, n_layers: int | None = None, head_layers: int = 0):
"""Remap a ``dflash`` GGUF of ``container`` onto the drafter param tree.
Closed tensor set: unknown names are hard errors (converter drift must
- surface at load, not as an unfilled param)."""
- blk_map, root_map = _DFLASH_CONTAINER_MAPS[container]
+ surface at load, not as an unfilled param). A weight's ``.scale`` (NVFP4
+ exports) lands on its module's ``out_scale``. ``n_layers`` and
+ ``head_layers`` place LiLiCorr's head blocks after the backbone's."""
hf_weights: dict[str, mx.array] = {}
hf_kquant_meta: dict[str, str] = {}
stats = {"mapped": 0}
@@ -1137,24 +1206,23 @@ def remap_dflash_arrays(arrays: dict, kquant_meta: dict, container: str):
if name.endswith((".scales", ".biases")):
continue
base, suffix = name, ""
- for s in (".weight", ".bias"):
+ for s in (".weight", ".bias", ".scale"):
if name.endswith(s):
base, suffix = name[: -len(s)], s
break
- if base.startswith("blk."):
- _, idx, leaf = base.split(".", 2)
- target = blk_map.get(leaf)
- if target is not None:
- target = f"layers.{idx}.{target}"
- else:
- target = root_map.get(base)
- if target is None:
+ target = _dflash_target(container, base, n_layers, head_layers)
+ if target is None or (suffix == ".scale"
+ and not target.endswith(".weight")):
raise RuntimeError(
f"{container} dflash remap: unknown tensor {name!r} "
f"(the drafter tensor set is closed)"
)
if suffix == ".bias":
- target = _strip_weight(target) + ".bias"
+ target = (_strip_weight(target) + ".bias"
+ if target.endswith(".weight") else target)
+ elif suffix == ".scale":
+ target = _strip_weight(target) + ".out_scale"
+ arr = arr.reshape(-1)
if base == "conf_proj" and suffix == ".weight" \
and getattr(arr, "ndim", 2) == 1:
arr = arr.reshape(1, -1)
@@ -1229,7 +1297,8 @@ def _dflash_config_from_meta(
f"{draft_gguf_path}: {name} has {ne[-1]} rows, the "
f"target vocab is {vocab}"
)
- n_layers = 1 + max(
+ # block_count is the backbone; LiLiCorr's head blocks follow it.
+ n_layers = int(meta.get("dflash.block_count") or 0) or 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 = [
@@ -1301,10 +1370,40 @@ def _dflash_config_from_meta(
confidence_head=confidence,
log_snr_range=snr_range,
sample_from_anchor=anchor,
+ attention_sinks=any(n.endswith(".attn_sinks.weight") for n in arrays),
+ **(_lilicorr_head_keys(draft_gguf_path, meta)
+ if container == "lilicorr" else {}),
)
return config, layer_ids
+def _lilicorr_head_keys(draft_gguf_path: str, meta: dict) -> dict:
+ """The LiLiCorr head sizes, every key required as llama.cpp's dflash
+ loader requires them once ``lilicorr.hidden_size`` is present."""
+ keys = {
+ "lilicorr_hidden_size": "lilicorr.hidden_size",
+ "lilicorr_top_k": "lilicorr.candidate_topk",
+ "lilicorr_heads": "lilicorr.num_heads",
+ "lilicorr_layers": "lilicorr.num_layers",
+ "lilicorr_factor_dim": "lilicorr.factor_dim",
+ "lilicorr_mlp_size": "lilicorr.mlp_hidden_size",
+ "lilicorr_vector_eps": "lilicorr.vector_eps",
+ "lilicorr_logit_scale": "lilicorr.logit_scale",
+ }
+ missing = [k for k in keys.values() if meta.get(k) is None]
+ if missing:
+ raise ValueError(f"{draft_gguf_path}: LiLiCorr keys missing: {missing}")
+ out = {f: meta[k] for f, k in keys.items()}
+ for f in ("lilicorr_vector_eps", "lilicorr_logit_scale"):
+ out[f] = float(out[f])
+ for f in set(out) - {"lilicorr_vector_eps", "lilicorr_logit_scale"}:
+ out[f] = int(out[f])
+ hidden, heads = out["lilicorr_hidden_size"], out["lilicorr_heads"]
+ if not hidden or not heads or hidden % heads or not out["lilicorr_top_k"]:
+ raise ValueError(f"{draft_gguf_path}: invalid LiLiCorr head sizes {out}")
+ return out
+
+
def _dspark_head_keys(draft_gguf_path: str, meta: dict, arrays: dict):
"""The DSpark head facts of a DFlash-backbone drafter, read the way
llama.cpp's dflash loader reads them: the rank from the markov table,
@@ -1510,6 +1609,61 @@ def _load_dflash_dspark_drafter(
return drafter
+def _load_lilicorr_drafter(
+ draft_gguf_path: str,
+ target,
+ target_config_dict: dict,
+ *,
+ arrays: dict,
+ kquant_meta: dict,
+ meta: dict,
+ shapes: dict | None = None,
+ active_before: float | None = None,
+ log=loadlog.verbose_print,
+):
+ """Build + load + bind a LiLiCorr drafter and wire the target's
+ ``_dflash_capture``."""
+ from .lilicorr_drafter import LiLiCorrDrafter
+
+ config, layer_ids = _dflash_config_from_meta(
+ draft_gguf_path, meta, target_config_dict, "lilicorr",
+ arrays=arrays, shapes=shapes)
+ drafter = LiLiCorrDrafter(config)
+ log(
+ f"[mtp] drafter: lilicorr layers={config.num_hidden_layers} "
+ f"targets={layer_ids} block_total={config.block_size} "
+ f"(native {config.native_block_size}) window={config.sliding_window} "
+ f"sinks={config.attention_sinks} head top_k={config.lilicorr_top_k} "
+ f"layers={config.lilicorr_layers}"
+ )
+ d_weights, d_meta, d_stats = remap_dflash_arrays(
+ arrays, kquant_meta, "lilicorr", n_layers=config.num_hidden_layers,
+ head_layers=config.lilicorr_layers)
+ 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,
+ # The head runs in fp32, as llama.cpp's F32 head tensors do.
+ fp32_keep=("lilicorr.",),
+ )
+ drafter.bind(target)
+ _wire_dflash_capture(target, layer_ids)
+
+ from .drafter_protocol import validate_drafter
+
+ validate_drafter(drafter)
+ log("[mtp] lilicorr drafter bound; target capture layers wired")
+ _stamp_mtp_width_cap(
+ drafter, str(target_config_dict.get("model_type") or "lilicorr"),
+ target=target, hard_limit=1, log=log)
+ return drafter
+
+
def _load_dflash_drafter(
draft_gguf_path: str,
target,
@@ -1537,6 +1691,7 @@ def _load_dflash_drafter(
"dflash2": _load_dflash2_drafter,
"dflash_dspark": _load_dflash_dspark_drafter,
"muse_glimmer": _load_muse_glimmer_dflash_drafter,
+ "lilicorr": _load_lilicorr_drafter,
}
loader = loaders.get(container)
if loader is None:
@@ -1568,15 +1723,17 @@ def _drafter_header_arch(draft_gguf_path: str) -> str | None:
def _assistant_kind(model_type: str | None, draft_gguf_path: str) -> str:
"""Which companion loader a ``--draft-gguf`` takes: ``deepseek4`` for a
deepseek_v4 target (its drafters share the ``dflash`` tag with DSpark),
- ``dflash`` for a muse_glimmer target or any ``dflash`` header, else
- ``gemma4``. Each loader validates the pairing it is handed."""
+ ``nemotron`` for a nemotron_h MTP sidecar, ``dflash`` for a muse_glimmer
+ target or any ``dflash`` header, else ``gemma4``. Each loader validates
+ the pairing it is handed."""
if model_type == "deepseek_v4":
return "deepseek4"
if model_type == "qwen4_exp":
return "qwen4exp"
- if model_type == "nemotron_h":
+ header_arch = _drafter_header_arch(draft_gguf_path)
+ if model_type == "nemotron_h" and header_arch != "dflash":
return "nemotron"
- if model_type == "muse_glimmer" or _drafter_header_arch(draft_gguf_path) == "dflash":
+ if model_type == "muse_glimmer" or header_arch == "dflash":
return "dflash"
return "gemma4"
diff --git a/tests/load/test_discovery.py b/tests/load/test_discovery.py
index 764341be0..5caa79ddb 100644
--- a/tests/load/test_discovery.py
+++ b/tests/load/test_discovery.py
@@ -1172,3 +1172,54 @@ def test_supported_main_arch_with_backbone_field_is_model():
"deepseek4.embedding_length_out": 16384},
"DeepSeek-V4-Flash-Vision-Exp-UD-IQ3_XXS-00001-of-00004.gguf")
assert c.kind == "model"
+
+
+def _nemotron_classify(monkeypatch):
+ def _f(path):
+ name = os.path.basename(path)
+ if "LiLiCorr" in name:
+ meta = {"general.architecture": "dflash",
+ "dflash.embedding_length": 4096}
+ else:
+ meta = {"general.architecture": "nemotron_h_moe",
+ "nemotron_h_moe.embedding_length": 4096}
+ return disc._classify_meta(meta, basename=name, path=path)
+
+ monkeypatch.setattr(disc, "classify_gguf", _f)
+
+
+def test_lilicorr_sibling_repo_pairs_its_target(tmp_path, monkeypatch):
+ # NVIDIA ships the drafter as -LiLiCorr-GGUF beside -GGUF,
+ # with no base-model id in its header.
+ _nemotron_classify(monkeypatch)
+ repo = "nvidia__NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB"
+ (tmp_path / f"{repo}-GGUF").mkdir()
+ (tmp_path / f"{repo}-LiLiCorr-GGUF").mkdir()
+ _write(tmp_path / f"{repo}-GGUF", "NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB.gguf")
+ _write(tmp_path / f"{repo}-LiLiCorr-GGUF",
+ "NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB-LiLiCorr.gguf")
+ models = _scan(tmp_path, recursive=True)
+ assert len(models) == 1
+ assert models[0].draft_gguf.endswith("-LiLiCorr.gguf")
+ assert models[0].speculative is True
+
+
+def test_lilicorr_sibling_rule_needs_the_matching_repo(tmp_path, monkeypatch):
+ _nemotron_classify(monkeypatch)
+ (tmp_path / "nvidia__Other-GGUF").mkdir()
+ (tmp_path / "nvidia__Bolt-LiLiCorr-GGUF").mkdir()
+ _write(tmp_path / "nvidia__Other-GGUF", "Other-Q4_K_M.gguf")
+ _write(tmp_path / "nvidia__Bolt-LiLiCorr-GGUF", "Bolt-LiLiCorr.gguf")
+ models = _scan(tmp_path, recursive=True)
+ assert [m.draft_gguf for m in models] == [None]
+
+
+def test_repo_name_reads_both_layouts():
+ assert disc._repo_name("/m/nvidia__X-GGUF/X.gguf") == "nvidia__X-GGUF"
+ assert disc._repo_name(
+ "/hf/models--nvidia--X-LiLiCorr-GGUF/snapshots/abc/X.gguf"
+ ) == "nvidia__X-LiLiCorr-GGUF"
+ assert disc._sibling_drafter_repo(
+ "/hf/models--nvidia--X-LiLiCorr-GGUF/snapshots/abc/X.gguf"
+ ) == "nvidia__x-gguf"
+ assert disc._sibling_drafter_repo("/m/nvidia__X-GGUF/X.gguf") is None
diff --git a/tests/spec/test_lilicorr_drafter.py b/tests/spec/test_lilicorr_drafter.py
new file mode 100644
index 000000000..8734c5ba7
--- /dev/null
+++ b/tests/spec/test_lilicorr_drafter.py
@@ -0,0 +1,352 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Asher Feldman
+"""LiLiCorr drafter units: the lattice head against a float64 reference of
+llama.cpp's ``build_lilicorr_scores`` and its greedy walk, the closed
+remap of the GGUF tensor set, container and loader routing, and the
+nemotron_h target's packed-hidden capture with its batched verify
+rollback. Tiny random weights; no GGUF and no download."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from gmlx.spec.dflash_drafter import DFlashConfig, greedy_path
+from gmlx.spec.lilicorr_drafter import LiLiCorrDrafter
+
+H, V, BLOCK, K = 64, 48, 5, 4 # 4 slots x 4 candidates = 16 nodes
+
+
+def _config(**over) -> DFlashConfig:
+ kw = dict(
+ hidden_size=H, intermediate_size=96, num_hidden_layers=2,
+ num_attention_heads=4, num_key_value_heads=2, head_dim=16,
+ rms_norm_eps=1e-6, vocab_size=V, max_position_embeddings=4096,
+ rope_theta=10000.0, block_size=BLOCK, native_block_size=BLOCK,
+ mask_token_id=V - 1, target_layer_ids=[0, 2], num_target_layers=4,
+ layer_types=["sliding_attention"] * 2, sliding_window=16,
+ conv_kernel_size=2, conv_group_size=16, attention_sinks=True,
+ lilicorr_layers=2, lilicorr_hidden_size=32, lilicorr_top_k=K,
+ lilicorr_heads=4, lilicorr_factor_dim=24, lilicorr_mlp_size=40,
+ lilicorr_vector_eps=1e-4, lilicorr_logit_scale=8.0)
+ kw.update(over)
+ return DFlashConfig(**kw)
+
+
+def _randomize(module: nn.Module, seed: int, scale: float = 0.3):
+ rng = np.random.default_rng(seed)
+ module.update(tree_unflatten([
+ (k, mx.array(rng.normal(0, scale, v.shape).astype(np.float32)))
+ for k, v in tree_flatten(module.parameters())]))
+ mx.eval(module.parameters())
+
+
+def _np(a) -> np.ndarray:
+ return np.array(a.astype(mx.float32), dtype=np.float64)
+
+
+def _lin(m, x):
+ y = x @ _np(m.weight).T
+ return y + _np(m.bias) if "bias" in m else y
+
+
+def _rms(x, w, eps):
+ return x / np.sqrt((x * x).mean(-1, keepdims=True) + eps) * w
+
+
+def _silu(x):
+ return x / (1.0 + np.exp(-x))
+
+
+def _l2n(x, eps):
+ return x / np.maximum(np.sqrt((x * x).sum(-1, keepdims=True)), eps)
+
+
+def _ref_lattice(head, passes, logits, embed_w, anchor):
+ """build_lilicorr_scores in float64, node order slot-major."""
+ slots, k = logits.shape[0], head.top_k
+ probs = np.exp(logits - logits.max(-1, keepdims=True))
+ probs /= probs.sum(-1, keepdims=True)
+ cands = np.argsort(-logits, axis=-1, kind="stable")[:, :k]
+ p = np.take_along_axis(probs, cands, -1)
+ lp = np.log(p)
+ rank = np.arange(k, dtype=np.float64)
+ feat = np.stack([lp, p, lp - lp[:, :1],
+ np.broadcast_to(rank / (k - 1), lp.shape),
+ np.broadcast_to((rank == 0) * 1.0, lp.shape)], -1)
+ mu = feat.mean(-1, keepdims=True)
+ var = ((feat - mu) ** 2).mean(-1, keepdims=True)
+ fn = (feat - mu) / np.sqrt(var + 1e-5) * _np(head.feature_norm.weight) \
+ + _np(head.feature_norm.bias)
+ fs = _lin(head.feature_down, _silu(_lin(head.feature_up, fn)))
+ tok = _lin(head.token_proj, embed_w[cands])
+ ps = _lin(head.pass_hidden_proj, passes)[:, None, :]
+ pos = _np(head.slot_embedding)[:, None, :] + _np(head.rank_embedding)[None]
+ x = (tok + ps + fs + pos).reshape(slots * k, -1)
+ n, hidden = x.shape
+ bias = _np(head.attention_bias)
+ for layer in head.layers:
+ heads = layer.heads
+ hd = hidden // heads
+ a = _rms(x, _np(layer.attn_norm), head.eps)
+ q, kk, v = np.split(_lin(layer.in_proj, a), 3, -1)
+ q, kk, v = (t.reshape(n, heads, hd).transpose(1, 0, 2) for t in (q, kk, v))
+ s = q @ kk.transpose(0, 2, 1) / np.sqrt(hd) + bias
+ s = np.exp(s - s.max(-1, keepdims=True))
+ s /= s.sum(-1, keepdims=True)
+ o = (s @ v).transpose(1, 0, 2).reshape(n, hidden)
+ x = x + _lin(layer.out_proj, o)
+ m = _rms(x, _np(layer.mlp_norm), head.eps)
+ x = x + _lin(layer.mlp_down, _silu(_lin(layer.mlp_up, m)))
+ x = _rms(x, _np(head.output_norm), head.eps)
+ a = _rms(_lin(head.context_proj, anchor[None]), _np(head.anchor_norm), head.eps)
+ ab = np.broadcast_to(a, x.shape)
+ f = _silu(_lin(head.factor_input_proj, np.concatenate([x, ab, x * ab], -1)))
+ eps = head.vector_eps
+ out = _l2n(_lin(head.out_head, f), eps).reshape(slots, k, -1)
+ inn = _l2n(_lin(head.in_head, f), eps).reshape(slots, k, -1)
+ a_out = _l2n(_lin(head.anchor_out_head, a), eps)[0]
+ first = head.logit_scale * inn[0] @ a_out
+ edges = head.logit_scale * np.einsum("spf,snf->spn", out[:-1], inn[1:])
+ return cands, first, edges
+
+
+def _ref_walk(cands, first, edges, n_max):
+ """common/speculative.cpp's lattice loop over the packed rows."""
+ rows = [first] + [e.reshape(-1) for e in edges]
+ pred, out = 0, []
+ for i in range(n_max):
+ if i == 0:
+ scores = rows[0]
+ else:
+ scores = rows[i][pred * len(first):(pred + 1) * len(first)]
+ pred = int(np.argmax(scores))
+ out.append(int(cands[i][pred]))
+ return out
+
+
+def _head_inputs(seed=1):
+ rng = np.random.default_rng(seed)
+ slots = BLOCK - 1
+ passes = rng.normal(0, 1, (slots, H)).astype(np.float32)
+ logits = rng.normal(0, 3, (slots, V)).astype(np.float32)
+ anchor = rng.normal(0, 1, (H,)).astype(np.float32)
+ embed = nn.Embedding(V, H)
+ _randomize(embed, seed + 7, 1.0)
+ return passes, logits, anchor, embed
+
+
+def test_lattice_matches_reference():
+ drafter = LiLiCorrDrafter(_config())
+ head = drafter.lilicorr
+ _randomize(head, 3)
+ passes, logits, anchor, embed = _head_inputs()
+ cands, first, edges = head.lattice(
+ mx.array(passes), mx.array(logits), embed, mx.array(anchor))
+ rc, rf, re_ = _ref_lattice(head, passes.astype(np.float64),
+ logits.astype(np.float64),
+ _np(embed.weight), anchor.astype(np.float64))
+ np.testing.assert_array_equal(np.array(cands), rc)
+ np.testing.assert_allclose(_np(first), rf, rtol=1e-3, atol=1e-3)
+ np.testing.assert_allclose(_np(edges), re_, rtol=1e-3, atol=1e-3)
+ assert edges.shape == (BLOCK - 2, K, K)
+
+
+@pytest.mark.parametrize("n_max", [1, 2, BLOCK - 1])
+def test_greedy_walk_matches_llama_cpp_loop(n_max):
+ drafter = LiLiCorrDrafter(_config())
+ _randomize(drafter.lilicorr, 5)
+ passes, logits, anchor, embed = _head_inputs(2)
+ cands, first, edges = drafter.lilicorr.lattice(
+ mx.array(passes), mx.array(logits), embed, mx.array(anchor))
+ sel = greedy_path(first, edges[: n_max - 1])
+ got = np.take_along_axis(np.array(cands)[:n_max], np.array(sel)[:, None], -1)
+ want = _ref_walk(np.array(cands), _np(first), _np(edges), n_max)
+ assert got[:, 0].tolist() == want
+
+
+# Random values here put the Mamba scan outside its working range.
+_KEEP_INIT = (".A_log", ".dt_bias", ".D", ".conv1d.")
+
+
+def _tiny_target(seed=0):
+ from gmlx.models.nemotron_h.model import ModelArgs
+ from gmlx.models.nemotron_h.mtp import NemotronHSpecLM
+
+ args = ModelArgs(
+ model_type="nemotron_h", vocab_size=40, hidden_size=64,
+ intermediate_size=32, num_hidden_layers=4,
+ max_position_embeddings=1024, num_attention_heads=4,
+ num_key_value_heads=2, attention_bias=False, mamba_num_heads=4,
+ mamba_head_dim=16, mamba_proj_bias=False, ssm_state_size=8,
+ conv_kernel=4, n_groups=2, mlp_bias=False, layer_norm_epsilon=1e-5,
+ use_bias=False, use_conv_bias=True,
+ hybrid_override_pattern=["M", "*", "E", "M"], head_dim=16,
+ moe_intermediate_size=24, moe_shared_expert_intermediate_size=48,
+ moe_latent_size=16, n_routed_experts=8, n_shared_experts=1,
+ num_experts_per_tok=2, norm_topk_prob=True, routed_scaling_factor=5.0,
+ n_group=1, topk_group=1)
+ model = NemotronHSpecLM(args)
+ rng = np.random.default_rng(seed)
+ model.update(tree_unflatten([
+ (k, mx.array(rng.normal(0, 0.3, v.shape).astype(np.float32)))
+ for k, v in tree_flatten(model.parameters())
+ if not any(s in k for s in _KEEP_INIT)]))
+ mx.eval(model.parameters())
+ return model
+
+
+def test_draft_block_runs_full_block_and_walks_requested_slots():
+ model = _tiny_target()
+ cfg = _config(hidden_size=64, vocab_size=40, mask_token_id=39,
+ target_layer_ids=[0, 2], num_target_layers=4)
+ drafter = LiLiCorrDrafter(cfg)
+ _randomize(drafter, 9, 0.1)
+ drafter.reset(model)
+ model.set_dflash_capture([0, 2])
+ out = model(mx.array([[1, 5, 9, 2, 7]]), cache=model.make_cache(),
+ return_hidden=True)
+ drafter.prefill_from_target_hidden(None, out.hidden_states[-1], 3, None)
+ assert drafter._cache.anchor is not None
+ for bs in (2, 3, BLOCK):
+ toks = drafter.draft_block(4, None, None, bs, None, greedy=True)
+ assert toks.shape == (1, bs - 1)
+
+
+def test_capture_packs_trunk_and_layers():
+ model = _tiny_target(1)
+ ids = mx.array([[1, 5, 9, 2]])
+ plain = model(ids, cache=model.make_cache(), return_hidden=True)
+ model.set_dflash_capture([0, 2])
+ packed = model(ids, cache=model.make_cache(), return_hidden=True)
+ h = packed.hidden_states[-1]
+ assert h.shape[-1] == 64 * 3
+ np.testing.assert_allclose(_np(h[..., :64]), _np(plain.hidden_states[-1]),
+ rtol=1e-5, atol=1e-5)
+ np.testing.assert_allclose(_np(packed.logits), _np(plain.logits),
+ rtol=1e-5, atol=1e-5)
+ np.testing.assert_allclose(
+ _np(model.speculative_logits_from_hidden(h)), _np(plain.logits),
+ rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize("accepted", [0, 2, 4])
+def test_batched_verify_rollback_matches_a_fresh_prefix(accepted):
+ # The DFlash verify runs batched; rolling back to the accepted prefix
+ # must leave the caches as running that prefix would.
+ model = _tiny_target(2)
+ model.set_dflash_capture([0, 2])
+ prompt = mx.array([[1, 5, 9, 2]])
+ block = mx.array([[7, 3, 8, 6, 4]])
+ cache = model.make_cache()
+ model(prompt, cache=cache)
+ _, _, sink = model.speculative_verify_hidden(block, cache)
+ model.rollback_speculative_cache(cache, sink, accepted, block.shape[1])
+ ref = model.make_cache()
+ model(prompt, cache=ref)
+ model(block[:, : accepted + 1], cache=ref)
+ nxt = mx.array([[2]])
+ np.testing.assert_allclose(
+ _np(model(nxt, cache=cache).logits), _np(model(nxt, cache=ref).logits),
+ rtol=5e-3, atol=5e-3)
+
+
+# The real header's names and shapes (GGUF order), at tiny width.
+def _lilicorr_arrays(cfg: DFlashConfig) -> dict:
+ n, hd, hl = cfg.num_hidden_layers, cfg.lilicorr_hidden_size, cfg.lilicorr_layers
+ ff, k, slots = cfg.intermediate_size, cfg.lilicorr_top_k, BLOCK - 1
+ kv = cfg.num_key_value_heads * cfg.head_dim
+ groups = H // cfg.conv_group_size
+ shapes = {"fc.weight": [2 * H, H], "fc.scale": [1],
+ "enc.output_norm.weight": [H], "output_norm.weight": [H]}
+ for i in range(n):
+ b = f"blk.{i}."
+ shapes.update({
+ b + "attn_norm.weight": [H], b + "ffn_norm.weight": [H],
+ b + "attn_q.weight": [H, H], b + "attn_k.weight": [H, kv],
+ b + "attn_v.weight": [H, kv], b + "attn_output.weight": [H, H],
+ b + "attn_q_norm.weight": [cfg.head_dim],
+ b + "attn_k_norm.weight": [cfg.head_dim],
+ b + "attn_sinks.weight": [cfg.num_attention_heads],
+ b + "ffn_gate.weight": [H, ff], b + "ffn_up.weight": [H, ff],
+ b + "ffn_down.weight": [ff, H], b + "ffn_down.scale": [1],
+ b + "attn_conv_base": [H, 2, 2], b + "ffn_conv_base": [H, 2, 2],
+ b + "attn_conv_proj.weight": [H, 2 * 2 * groups],
+ b + "ffn_conv_proj.weight": [H, 2 * 2 * groups],
+ })
+ for i in range(hl):
+ b = f"blk.{n + i}."
+ shapes.update({
+ b + "attn_norm.weight": [hd], b + "ffn_norm.weight": [hd],
+ b + "attn_output.weight": [hd, hd], b + "attn_output.bias": [hd],
+ b + "attn_output.scale": [1],
+ b + "ffn_up.weight": [hd, cfg.lilicorr_mlp_size],
+ b + "ffn_up.bias": [cfg.lilicorr_mlp_size],
+ b + "ffn_down.weight": [cfg.lilicorr_mlp_size, hd],
+ b + "ffn_down.bias": [hd],
+ f"lilicorr.layers.{i}.attn.in_proj.weight": [hd, 3 * hd],
+ f"lilicorr.layers.{i}.attn.in_proj.bias": [3 * hd],
+ })
+ f = f"blk.{n + hl}."
+ shapes.update({f + "ffn_norm.weight": [5], f + "ffn_norm.bias": [5],
+ f + "ffn_up.weight": [5, hd], f + "ffn_up.bias": [hd],
+ f + "ffn_down.weight": [hd, hd], f + "ffn_down.bias": [hd]})
+ for name in ("token_proj", "pass_hidden_proj", "context_proj"):
+ shapes.update({f"lilicorr.{name}.weight": [H, hd],
+ f"lilicorr.{name}.bias": [hd],
+ f"lilicorr.{name}.scale": [1]})
+ for name in ("out_head", "in_head", "anchor_out_head"):
+ shapes.update({f"lilicorr.{name}.weight": [hd, cfg.lilicorr_factor_dim],
+ f"lilicorr.{name}.bias": [cfg.lilicorr_factor_dim]})
+ shapes.update({
+ "lilicorr.factor_input_proj.weight": [3 * hd, hd],
+ "lilicorr.factor_input_proj.bias": [hd],
+ "lilicorr.output_norm.weight": [hd], "lilicorr.anchor_norm.weight": [hd],
+ "lilicorr.slot_embedding.weight": [hd, slots],
+ "lilicorr.rank_embedding.weight": [hd, k],
+ "lilicorr.attention_bias.bias": [slots * k, slots * k, cfg.lilicorr_heads],
+ })
+ return {name: mx.zeros(tuple(reversed(ne))) for name, ne in shapes.items()}
+
+
+def test_remap_fills_every_param_from_the_closed_set():
+ from gmlx.load.out_scale import install_out_scales
+ from gmlx.spec.mtp_load import dflash_container, remap_dflash_arrays
+
+ cfg = _config(target_layer_ids=[0, 1])
+ arrays = _lilicorr_arrays(cfg)
+ assert dflash_container(arrays) == "lilicorr"
+ weights, _, stats = remap_dflash_arrays(
+ arrays, {}, "lilicorr", n_layers=cfg.num_hidden_layers,
+ head_layers=cfg.lilicorr_layers)
+ assert stats["mapped"] == len(arrays)
+ drafter = LiLiCorrDrafter(cfg)
+ install_out_scales(drafter, weights)
+ params = dict(tree_flatten(drafter.parameters()))
+ assert set(weights) == set(params)
+ for name, arr in weights.items():
+ assert tuple(arr.shape) == tuple(params[name].shape), name
+
+
+def test_remap_refuses_an_unknown_tensor():
+ from gmlx.spec.mtp_load import remap_dflash_arrays
+
+ cfg = _config()
+ arrays = _lilicorr_arrays(cfg)
+ arrays["lilicorr.mystery.weight"] = mx.zeros((4,))
+ with pytest.raises(RuntimeError, match="lilicorr.mystery"):
+ remap_dflash_arrays(arrays, {}, "lilicorr", n_layers=2, head_layers=2)
+
+
+def test_nemotron_dflash_header_routes_to_the_dflash_loader(monkeypatch):
+ import gmlx.spec.mtp_load as ml
+
+ monkeypatch.setattr(ml, "_drafter_header_arch", lambda p: "dflash")
+ assert ml._assistant_kind("nemotron_h", "x.gguf") == "dflash"
+ monkeypatch.setattr(ml, "_drafter_header_arch", lambda p: "nemotron_h_moe")
+ assert ml._assistant_kind("nemotron_h", "x.gguf") == "nemotron"
From 04696d6982f90795415d71c70cfe17a8d3cac61f Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sat, 17 Oct 2026 23:43:49 -0700
Subject: [PATCH 03/16] feat(profiles): nemotron instruct and reasoning-medium
intents
---
docs/family-defaults.md | 2 +-
gmlx/gen/profiles.py | 8 +++++++-
tests/gen/test_profiles.py | 10 ++++++++++
3 files changed, 18 insertions(+), 2 deletions(-)
diff --git a/docs/family-defaults.md b/docs/family-defaults.md
index 28ea6e5ae..14ba56d72 100644
--- a/docs/family-defaults.md
+++ b/docs/family-defaults.md
@@ -31,7 +31,7 @@ a model ends up with.
| `deepseek` | `deepseek2`, `deepseek4` | temperature=0.6 top_p=0.95 | - |
| `deepseek41` | `deepseek41` | temperature=1.0 top_p=0.95 | `@reasoning-high`: reasoning_effort=high; `@reasoning-low`: reasoning_effort=low; `@reasoning-max`: reasoning_effort=max |
| `minimax` | `minimax-m2`, `minimax-m3` | temperature=1.0 top_p=0.95 top_k=40 | - |
-| `nemotron` | `nemotron_h_moe` | temperature=1.0 top_p=0.95 | - |
+| `nemotron` | `nemotron_h_moe` | temperature=1.0 top_p=0.95 | `@instruct`: enable_thinking=False; `@reasoning-medium`: medium_effort=True |
| `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=<think:opensource> thinking_end_token=</think:opensource> | `@reasoning-high`: reasoning_effort=high; `@reasoning-low`: reasoning_effort=low |
| `hy4` | `hyv4` | temperature=0.9 top_p=1.0 thinking_start_token=<think:6124c78e> thinking_end_token=</think:6124c78e> | `@reasoning-high`: reasoning_effort=high; `@reasoning-low`: reasoning_effort=low |
diff --git a/gmlx/gen/profiles.py b/gmlx/gen/profiles.py
index 99ff11a82..bf34a5655 100644
--- a/gmlx/gen/profiles.py
+++ b/gmlx/gen/profiles.py
@@ -142,11 +142,17 @@
"intents": {},
},
# https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16, 2026-07.
+ # Thinking is on by default. Nemotron 3.5 Bolt's template also takes
+ # medium_effort, which asks for a shorter reasoning pass; older templates
+ # ignore it.
"nemotron": {
"label": "Nemotron",
"arches": ("nemotron_h_moe",),
"base": {"sampling": {"temperature": 1.0, "top_p": 0.95}},
- "intents": {},
+ "intents": {
+ "instruct": {"chat_template_kwargs": {"enable_thinking": False}},
+ "reasoning-medium": {"chat_template_kwargs": {"medium_effort": True}},
+ },
},
# https://huggingface.co/tencent/Hunyuan-A13B-Instruct, 2026-07.
"hunyuan": {
diff --git a/tests/gen/test_profiles.py b/tests/gen/test_profiles.py
index bc620599d..6f90e2dd0 100644
--- a/tests/gen/test_profiles.py
+++ b/tests/gen/test_profiles.py
@@ -122,6 +122,16 @@ def test_gpt_oss_reasoning_intents():
assert hi["sampling"] == base["sampling"]
+def test_nemotron_intents():
+ base = profiles.groups_for("nemotron")
+ off = profiles.groups_for("nemotron", "instruct")
+ assert off["chat_template_kwargs"] == {"enable_thinking": False}
+ assert off["sampling"] == base["sampling"]
+ med = profiles.groups_for("nemotron", "reasoning-medium")
+ assert med["chat_template_kwargs"] == {"medium_effort": True}
+ assert profiles.groups_for("nemotron", "reasoning-high") == base
+
+
def test_kimi_reasoning_intents_and_thinking_tokens():
base = profiles.groups_for("kimi")
assert base["sampling"]["temperature"] == 1.0
From e698cb6d0c6ae0f5f9e93d62d06e4958115740ce Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sat, 17 Oct 2026 23:45:02 -0700
Subject: [PATCH 04/16] docs(load): native-fp dispatch comment covers the dense
linear
---
gmlx/load/modules.py | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/gmlx/load/modules.py b/gmlx/load/modules.py
index ac262f7a5..c81ad6f3c 100644
--- a/gmlx/load/modules.py
+++ b/gmlx/load/modules.py
@@ -2771,10 +2771,9 @@ def _replace(path: str, module):
if codec in NATIVE_FP_CODECS:
# Micro-scaling float codec. The SwitchGLU expert path (gpt-oss,
# DS-V4) and the absorbed-MLA MultiLinear (Kimi-K3, whose
- # MXFP4_MOE ftype also quantizes attn_k_b/attn_v_b) are wired.
- # A native-fp Linear / Embedding would need its own module
- # (packed) or its own wire Metal leaves; fail loud rather than
- # mis-dispatch onto the kquant path.
+ # MXFP4_MOE ftype also quantizes attn_k_b/attn_v_b) are wired,
+ # and so is a dense Linear. An Embedding has no native-fp module;
+ # fail loud rather than mis-dispatch onto the kquant path.
if isinstance(module, _switch_linear_types):
n_experts, out_dims, in_dims = module.weight.shape
bias = "bias" in module
From 7a4ddc4dba81f9b60f6a14cb7369efdee6c2ceb2 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sat, 17 Oct 2026 23:45:03 -0700
Subject: [PATCH 05/16] docs(changelog): nemotron 3.5 bolt, weight scales and
dense nvfp4
---
CHANGELOG.md | 9 +++++++++
1 file changed, 9 insertions(+)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9dab8da62..ce53e89cd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
scales it, where it was refused.
- K2 Horizon 7B and K2 Horizon MoVA 36B-A4B (GGUF arch `k2-horizon`), with
`reasoning-low`, `-medium` and `-high` intents and tool calls.
+- NVIDIA Nemotron 3.5 Bolt 74B-A9B, with image input through its `--mmproj`
+ companion and its LiLiCorr drafter, which gmlx pairs from the sibling
+ `-LiLiCorr-GGUF` repo.
+- Nemotron models take `@instruct` to turn thinking off, and
+ `@reasoning-medium` for a shorter reasoning pass where the template has one.
### Changed
@@ -61,6 +66,10 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
when its template has the word thinking in its text. It now warns.
- A Krea 2 GGUF whose stored projector shape is not a list of integers
crashed the load. It is refused with the reason.
+- A GGUF whose NVFP4 weights carry a per-tensor `.scale` gave wrong output,
+ as gmlx dropped the scale that llama.cpp applies.
+- A dense NVFP4 or MXFP4 projection outside an expert block loads, where it
+ was refused.
- A streamed image request refused while it ran, such as one too large for
memory, got a server error frame where the same request unstreamed got a 400.
- The `gmlx pull` plan of an image model could say its parts fit in stages
From cebe890ab78df4f1b890961ea4738d2ff911705e Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 00:22:10 -0700
Subject: [PATCH 06/16] fix(nemotron): ascii copyright line in the owned model
---
gmlx/models/nemotron_h/model.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/gmlx/models/nemotron_h/model.py b/gmlx/models/nemotron_h/model.py
index 3350cf3c5..805fd7af8 100644
--- a/gmlx/models/nemotron_h/model.py
+++ b/gmlx/models/nemotron_h/model.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: MIT
-# Copyright © 2025 Apple Inc.
+# Portions copyright (c) 2025 Apple Inc. (mlx-lm nemotron_h, MIT)
# Copyright (c) 2026 Asher Feldman
"""Model for NVIDIA Nemotron-H (GGUF arch ``nemotron_h_moe``): Nano, Super,
Lightning and Bolt.
From 7612e47296fa89cc4710e3f83a919c5c0fda68cf Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 00:02:07 -0700
Subject: [PATCH 07/16] feat(nemotron): bolt image input through the c-radio
tower
---
gmlx/_vendor/manifest.json | 2 +-
gmlx/_vendor/mlx_vlm/prompt_utils.py | 1 +
gmlx/load/loader.py | 4 +
gmlx/load/vlm.py | 339 +++++++++++++++++
gmlx/models/nemotron_h/model.py | 4 +-
gmlx/models/nemotron_h/vision.py | 196 ++++++++++
gmlx/models/nemotron_h/vlm_model.py | 162 ++++++++
gmlx/models/owned.py | 1 +
gmlx/serve/patches/render.py | 4 +-
tests/fixtures/nemotron_bolt_template.jinja | 240 ++++++++++++
tests/models/test_nemotron_h_vision.py | 385 ++++++++++++++++++++
tests/serve/test_chat_image_placement.py | 4 +-
tests/serve/test_messages_route.py | 4 +-
13 files changed, 1339 insertions(+), 7 deletions(-)
create mode 100644 gmlx/models/nemotron_h/vision.py
create mode 100644 gmlx/models/nemotron_h/vlm_model.py
create mode 100644 tests/fixtures/nemotron_bolt_template.jinja
create mode 100644 tests/models/test_nemotron_h_vision.py
diff --git a/gmlx/_vendor/manifest.json b/gmlx/_vendor/manifest.json
index adfed61e2..ae94ca4aa 100644
--- a/gmlx/_vendor/manifest.json
+++ b/gmlx/_vendor/manifest.json
@@ -1501,7 +1501,7 @@
"mlx_vlm/models/zaya1_vl/zaya1_vl.py": {"status": "verbatim", "upstream": "255a5c86005d38c160b11bb228a2822ec784465d001e46da95ae0f67498989fd"},
"mlx_vlm/moe_offload.py": {"sha256": "2918037c2f0cb7e34f8a92a9e497fbf5c433f38f3846246d4d73f205b36a6eea", "status": "edited", "upstream": "c5bc9fc7e76af1d4514076f1e4d9e259099decf452b6f3d50fff0c8ffa31ec5d"},
"mlx_vlm/privacy_filter.py": {"status": "verbatim", "upstream": "5f45a1b183f6ee0c1c7d33ed1fcc9fa2d85f837a964baa4f4f29bc9fec1411ad"},
- "mlx_vlm/prompt_utils.py": {"sha256": "b9fb62ef0701743dc8d5ebffa0aaa3576d309c1b0a84d26899c2149fb21e0ff7", "status": "edited", "upstream": "5e4e1889ca5a8b8b8ec1efd54a41b282a671b7cb9fef927777e1427be5735134"},
+ "mlx_vlm/prompt_utils.py": {"sha256": "da312cf0c1daae4a5b0690cd79246b15f51b75ca65a8c1a1cae066122814d197", "status": "edited", "upstream": "5e4e1889ca5a8b8b8ec1efd54a41b282a671b7cb9fef927777e1427be5735134"},
"mlx_vlm/py.typed": {"status": "verbatim", "upstream": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},
"mlx_vlm/quant/__init__.py": {"status": "verbatim", "upstream": "fc84bc2cce35428cc1a30c3ce5eedac3f227f328b117e4bb5936f1211edcfaa4"},
"mlx_vlm/quant/awq.py": {"status": "verbatim", "upstream": "75da9639b48141ffb67cd14b6cb6be6c56ec9f011fbc10325f604eaba1d3ac21"},
diff --git a/gmlx/_vendor/mlx_vlm/prompt_utils.py b/gmlx/_vendor/mlx_vlm/prompt_utils.py
index 5cecb77a8..f6f58771a 100644
--- a/gmlx/_vendor/mlx_vlm/prompt_utils.py
+++ b/gmlx/_vendor/mlx_vlm/prompt_utils.py
@@ -125,6 +125,7 @@ class MessageFormat(Enum):
"deepseek_v4_vl": MessageFormat.LIST_WITH_IMAGE_FIRST,
"deepseek_v41_vl": MessageFormat.LIST_WITH_IMAGE_FIRST,
"muse_spark": MessageFormat.LIST_WITH_IMAGE_FIRST,
+ "nemotron_h_omni": MessageFormat.LIST_WITH_IMAGE_FIRST,
}
# Models that don't support multi-image
diff --git a/gmlx/load/loader.py b/gmlx/load/loader.py
index 9e953c571..06a431055 100644
--- a/gmlx/load/loader.py
+++ b/gmlx/load/loader.py
@@ -937,6 +937,10 @@ def _warm_touch_pass(
# Kept in step with the model class's cast_predicate.
"nemotron_h": (".mixer.gate.weight", ".e_score_correction_bias",
".A_log"),
+ # nemotron_h_omni: the same set under language_model.*, and the vision
+ # tower's F32 tensors, since llama.cpp runs the tower in fp32.
+ "nemotron_h_omni": (".mixer.gate.weight", ".e_score_correction_bias",
+ ".A_log", "vision_tower."),
}
# Params kept at their native f16 through the bf16 cast (no upcast). MLX
diff --git a/gmlx/load/vlm.py b/gmlx/load/vlm.py
index 712bbf883..4259cb957 100644
--- a/gmlx/load/vlm.py
+++ b/gmlx/load/vlm.py
@@ -73,6 +73,15 @@ def resolve_vlm_model_type(llm_arch: str, mm_meta: dict) -> str:
return "deepseek_v41_vl"
proj = (mm_meta.get("clip.projector_type")
or _mm(mm_meta, "clip.vision.projector_type"))
+ if proj == "nemotron_omni":
+ # NVIDIA Nemotron-3.5-Bolt: a C-RADIO ViT and a relu^2 projector onto
+ # the nemotron_h text tower (gmlx.models.nemotron_h.vlm_model).
+ if llm_arch != "nemotron_h_moe":
+ raise UnsupportedVLMError(
+ f"mmproj projector 'nemotron_omni' on LLM arch {llm_arch!r} "
+ "is not supported (only Nemotron-H text GGUFs, arch "
+ "'nemotron_h_moe')")
+ return "nemotron_h_omni"
if proj == "pixtral":
# 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.
@@ -808,6 +817,31 @@ def _qwen3vl_deepstack_blocks(mm_meta: dict) -> list[int]:
}
+# Nemotron-3.5-Bolt (nemotron_omni): the C-RADIO ViT and projector, mapped
+# onto gmlx.models.nemotron_h.vision. The patch embed is a 16x16 conv stored
+# [out, C, ph, pw] that the tower applies as a Linear over the
+# (C, ph, pw)-flattened patch. The 6-channel patch embed (frame pairs) serves
+# video only, which this tower does not take.
+_NEMOTRON_OMNI_BLK_SUBMAP = {
+ "attn_q": "attn.q_proj",
+ "attn_k": "attn.k_proj",
+ "attn_v": "attn.v_proj",
+ "attn_out": "attn.out_proj",
+ "ln1": "norm1",
+ "ln2": "norm2",
+ "ffn_up": "mlp.fc1",
+ "ffn_down": "mlp.fc2",
+}
+_NEMOTRON_OMNI_TOP_MAP = {
+ "v.position_embd.weight": "vision_tower.position_embedding",
+ "v.class_embd": "vision_tower.registers",
+ "mm.model.mlp.0.weight": "vision_tower.projector.norm.weight",
+ "mm.model.mlp.1.weight": "vision_tower.projector.fc1.weight",
+ "mm.model.mlp.3.weight": "vision_tower.projector.fc2.weight",
+}
+_NEMOTRON_OMNI_VIDEO_ONLY = frozenset({"v.patch_embd.weight.1"})
+
+
# DeepSeek-V4-Flash-Vision-Exp (deepseek4v): the DeepSeek ViT + 3x3 unfold
# aligner, mapped onto the vendored gmlx.models.deepseek_v4.vision tower. The
# converter split the fused wqkv and w1 (llama.cpp PR 28133), so q/k/v and
@@ -1070,6 +1104,33 @@ def remap_vision_arrays(
out[f"vision_tower.blocks.{bid}.{tgt}.{leaf}"] = arr
return out, skipped, vis_kqmeta
+ if model_type == "nemotron_h_omni":
+ for name, arr in arrays.items():
+ if name.endswith(".scales") or name.endswith(".biases"):
+ continue
+ if name in _NEMOTRON_OMNI_VIDEO_ONLY:
+ skipped.append(name)
+ continue
+ if name == "v.patch_embd.weight":
+ out["vision_tower.patch_embed.weight"] = arr.reshape(
+ arr.shape[0], -1)
+ continue
+ hit = _NEMOTRON_OMNI_TOP_MAP.get(name)
+ if hit is not None:
+ out[hit] = arr
+ continue
+ m = _VISION_BLK_RE.match(name)
+ if m is None:
+ skipped.append(name)
+ continue
+ sub, _, leaf = m.group(2).rpartition(".")
+ tgt = _NEMOTRON_OMNI_BLK_SUBMAP.get(sub)
+ if tgt is None:
+ skipped.append(name)
+ continue
+ out[f"vision_tower.blocks.{m.group(1)}.{tgt}.{leaf}"] = arr
+ return out, skipped, vis_kqmeta
+
if model_type == "deepseek_v4_vl":
for name, arr in arrays.items():
if name.endswith(".scales") or name.endswith(".biases"):
@@ -1816,6 +1877,75 @@ def _synthesize_glm5next_vlm_config(
return config
+def _synthesize_nemotron_omni_vlm_config(
+ text_config: dict, mm_meta: dict, llm_meta: dict,
+ mm_shapes: dict | None = None,
+) -> dict:
+ """Nemotron-3.5-Bolt VLM config: the nemotron_h text synth plus the
+ C-RADIO tower read from ``clip.vision.*``. The register count, the
+ position grid and the projector width have no metadata key, so they are
+ read off ``v.class_embd``, ``v.position_embd.weight`` and
+ ``mm.model.mlp.1.weight``."""
+ import math
+
+ shapes = mm_shapes or {}
+ hidden = _mm_int(mm_meta, "clip.vision.embedding_length")
+ vision_config: dict = {
+ "model_type": "nemotron_h_omni",
+ "hidden_size": hidden,
+ "intermediate_size": _mm_int(mm_meta, "clip.vision.feed_forward_length"),
+ "num_hidden_layers": _mm_int(mm_meta, "clip.vision.block_count"),
+ "num_attention_heads": _mm_int(
+ mm_meta, "clip.vision.attention.head_count"),
+ "patch_size": _mm_int(mm_meta, "clip.vision.patch_size"),
+ "merge": _mm_int(mm_meta, "clip.vision.projector.scale_factor"),
+ "projection_dim": _mm_int(mm_meta, "clip.vision.projection_dim"),
+ }
+ eps = _mm(mm_meta, "clip.vision.attention.layer_norm_epsilon")
+ if eps is not None:
+ vision_config["layer_norm_eps"] = float(eps)
+ regs = shapes.get("v.class_embd")
+ if regs:
+ vision_config["num_registers"] = _other_dim(regs, hidden)
+ pos = shapes.get("v.position_embd.weight")
+ if pos:
+ n = _other_dim(pos, hidden)
+ side = math.isqrt(n)
+ if side * side != n:
+ raise ValueError(
+ f"nemotron_omni position table has {n} rows, not a square grid")
+ vision_config["position_grid"] = side
+ up = shapes.get("mm.model.mlp.1.weight")
+ if up:
+ vision_config["projector_hidden_size"] = _other_dim(
+ up, hidden * vision_config["merge"] ** 2)
+ if vision_config["projection_dim"] != int(text_config["hidden_size"]):
+ raise ValueError(
+ f"mmproj projection width {vision_config['projection_dim']} does "
+ f"not match the nemotron_h hidden size {text_config['hidden_size']}")
+
+ config: dict = {
+ "model_type": "nemotron_h_omni",
+ "text_config": text_config,
+ "vision_config": vision_config,
+ "vocab_size": int(text_config.get("vocab_size", 131072)),
+ }
+ img_id = _gguf_token_id(llm_meta, NEMOTRON_OMNI_IMAGE_TOKEN)
+ if img_id is None:
+ raise ValueError(
+ f"the text GGUF has no {NEMOTRON_OMNI_IMAGE_TOKEN} token")
+ config["image_token_id"] = img_id
+ config["image_token_index"] = img_id
+ return config
+
+
+# Nemotron-3.5-Bolt: the chat template writes per image; mtmd
+# brackets the embeddings with ![]()
and .
+NEMOTRON_OMNI_IMAGE_TOKEN = ""
+NEMOTRON_OMNI_IMAGE_START = "![]()
"
+NEMOTRON_OMNI_IMAGE_END = ""
+
+
# DeepSeek-V4-Flash-Vision-Exp image placeholder (fullwidth bars, as the
# chat template emits it once per image).
DEEPSEEK4V_IMAGE_TOKEN = "<\uff5cdeepseek_image\uff5c>"
@@ -1975,6 +2105,10 @@ def synthesize_vlm_config(
return _synthesize_glm5next_vlm_config(
text_config, mm_meta, llm_meta, mm_shapes)
+ if model_type == "nemotron_h_omni":
+ return _synthesize_nemotron_omni_vlm_config(
+ text_config, mm_meta, llm_meta, mm_shapes)
+
if model_type == "deepseek_v4_vl":
return _synthesize_deepseek4v_vlm_config(text_config, mm_meta, llm_meta)
@@ -2285,6 +2419,8 @@ def _synthesize_vlm_processor(model_type: str, tokenizer, mm_meta: dict,
tokenizer, mm_meta, audio=bool(mm_meta.get("clip.has_audio_encoder")))
if model_type == "glm5_next":
return _synthesize_glm5next_processor(tokenizer, mm_meta)
+ if model_type == "nemotron_h_omni":
+ return _synthesize_nemotron_omni_processor(tokenizer, mm_meta)
if model_type == "deepseek_v4_vl":
return _synthesize_deepseek4v_processor(tokenizer, mm_meta, config or {})
if model_type == "deepseek_v41_vl":
@@ -3305,6 +3441,209 @@ def model_input_names(self):
return _attach_streaming_helpers(processor, tokenizer)
+class _NemotronOmniGgufImageProcessor(ImageProcessingMixin):
+ """Torch-free Nemotron-3.5-Bolt image preprocessing (numpy + PIL only).
+
+ Ports llama.cpp's ``mtmd_image_preprocessor_dyn_size`` with
+ ``img_tool::calc_size_preserved_ratio`` (``tools/mtmd/mtmd-image.cpp``):
+ round each edge to a multiple of ``patch_size * merge`` (32), then, past
+ the pixel budget, scale both edges by the sqrt area ratio and floor them
+ to the multiple, or below it, scale up and ceil them. The image is
+ stretched to that size BICUBIC with no padding, scaled to [0, 1] and
+ normalized with the mmproj's mean and std. Patches flatten channel-major
+ in row-major patch order, as the tower's patch embed reads them.
+ """
+
+ model_input_names = ["pixel_values", "image_grid_thw"]
+
+ def __init__(self, image_mean, image_std, patch_size=16, merge=2,
+ min_pixels=262144, max_pixels=3407872):
+ super().__init__()
+ self.image_mean = list(image_mean)
+ self.image_std = list(image_std)
+ self.patch_size = int(patch_size)
+ self.merge = int(merge)
+ self.factor = self.patch_size * self.merge
+ self.min_pixels = int(min_pixels)
+ self.max_pixels = int(max_pixels)
+
+ def soft_tokens(self, grid) -> int:
+ _t, gh, gw = (int(v) for v in grid)
+ return (gh // self.merge) * (gw // self.merge)
+
+ def target_size(self, height: int, width: int) -> tuple[int, int]:
+ """(height, width) after the resize. float32 like the C++."""
+ import math
+
+ import numpy as np
+ f = self.factor
+ f32 = np.float32
+
+ def round_by(x):
+ return int(math.floor(float(f32(x) / f32(f)) + 0.5)) * f
+
+ def ceil_by(x):
+ return int(math.ceil(float(f32(x) / f32(f)))) * f
+
+ def floor_by(x):
+ return int(math.floor(float(f32(x) / f32(f)))) * f
+
+ w_bar = max(f, round_by(width))
+ h_bar = max(f, round_by(height))
+ if self.max_pixels > 0 and h_bar * w_bar > self.max_pixels:
+ beta = np.sqrt(f32(height) * f32(width) / f32(self.max_pixels))
+ h_bar = max(f, floor_by(f32(height) / beta))
+ w_bar = max(f, floor_by(f32(width) / beta))
+ elif self.min_pixels > 0 and h_bar * w_bar < self.min_pixels:
+ beta = np.sqrt(f32(self.min_pixels) / (f32(height) * f32(width)))
+ h_bar = ceil_by(f32(height) * beta)
+ w_bar = ceil_by(f32(width) * beta)
+ return h_bar, w_bar
+
+ 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, w = self.target_size(img.height, img.width)
+ if (h, w) != (img.height, img.width):
+ img = img.resize((w, h), Image.Resampling.BICUBIC)
+ arr = np.asarray(img, dtype=np.float32) / np.float32(255.0)
+ mean = np.array(self.image_mean, dtype=np.float32)
+ std = np.array(self.image_std, dtype=np.float32)
+ arr = np.transpose((arr - mean) / std, (2, 0, 1)) # [C, H, W]
+ ps = self.patch_size
+ c = arr.shape[0]
+ gh, gw = h // ps, w // ps
+ patches = arr.reshape(c, gh, ps, gw, ps).transpose(1, 3, 0, 2, 4)
+ return patches.reshape(gh * gw, c * ps * ps), (1, gh, gw)
+
+ def __call__(self, images, **kwargs):
+ import numpy as np
+ flat = _PixtralGgufImageProcessor._flatten(images)
+ vecs, grids = [], []
+ for img in flat:
+ v, g = self._one(img)
+ vecs.append(v)
+ grids.append(g)
+ return {
+ "pixel_values": np.concatenate(vecs, axis=0),
+ "image_grid_thw": np.array(grids, dtype=np.int64),
+ }
+
+
+def _nemotron_omni_image_processor(
+ mm_meta: dict) -> _NemotronOmniGgufImageProcessor:
+ """The Nemotron-3.5-Bolt image processor from the mmproj's
+ ``clip.vision.*`` keys, which llama.cpp requires for this projector."""
+ return _NemotronOmniGgufImageProcessor(
+ image_mean=_mm_floats(mm_meta, "clip.vision.image_mean"),
+ image_std=_mm_floats(mm_meta, "clip.vision.image_std"),
+ patch_size=_mm_int(mm_meta, "clip.vision.patch_size"),
+ merge=_mm_int(mm_meta, "clip.vision.projector.scale_factor"),
+ min_pixels=_mm_int(mm_meta, "clip.vision.image_min_pixels"),
+ max_pixels=_mm_int(mm_meta, "clip.vision.image_max_pixels"))
+
+
+def _synthesize_nemotron_omni_processor(tokenizer, mm_meta: dict):
+ """Build the Nemotron-3.5-Bolt processor from the GGUFs alone.
+
+ The chat template writes one ```` per image item. The processor
+ turns each into ``![]()
``, one ```` per soft token and
+ ````, the stream llama.cpp's mtmd builds around the embeddings.
+ Video input is refused: the 6-channel frame-pair path is not wired."""
+ from transformers.feature_extraction_utils import BatchFeature
+ from transformers.processing_utils import ProcessorMixin
+
+ from gmlx._vendor.mlx_vlm.models.base import to_mlx
+
+ from gmlx.upstream.single_bos_fix import KWARGS_REACH_TOKENIZER
+
+ image_processor = _nemotron_omni_image_processor(mm_meta)
+
+ class NemotronOmniProcessor(ProcessorMixin):
+ attributes = ["image_processor", "tokenizer"]
+ image_processor_class = "AutoImageProcessor"
+ tokenizer_class = "AutoTokenizer"
+
+ image_token = NEMOTRON_OMNI_IMAGE_TOKEN
+
+ # Images and text only; the base signature also names videos and
+ # audio, which reach the tokenizer through kwargs here.
+ def __call__( # pyright: ignore[reportIncompatibleMethodOverride]
+ self, images=None, text=None, **kwargs):
+ if kwargs.get("videos") is not None:
+ raise ValueError(
+ "Nemotron-3.5-Bolt video input is not supported yet; "
+ "pass still images")
+ kwargs.pop("videos", None)
+ if kwargs.get("audio") is not None:
+ raise ValueError(
+ "Nemotron-3.5-Bolt takes no audio input")
+ kwargs.pop("audio", None)
+ 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_grid_thw"])
+
+ 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, grids):
+ """Wrap each ```` in ``![]()
`` ... ```` and repeat
+ it once per soft token of its image. Grids 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(grids):
+ raise ValueError(
+ f"the prompt has more {self.image_token} "
+ f"placeholders than the {len(grids)} images given")
+ n = self.image_processor.soft_tokens(grids[index])
+ rebuilt += (NEMOTRON_OMNI_IMAGE_START
+ + self.image_token * n
+ + NEMOTRON_OMNI_IMAGE_END + part)
+ index += 1
+ 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)))
+
+ # __call__ hands **kwargs to self.tokenizer (single_bos_fix).
+ setattr(NemotronOmniProcessor, KWARGS_REACH_TOKENIZER, True)
+ processor = NemotronOmniProcessor(
+ image_processor=image_processor, tokenizer=tokenizer,
+ chat_template=getattr(tokenizer, "chat_template", None))
+ return _attach_streaming_helpers(processor, tokenizer)
+
+
class _DeepseekV4VGgufImageProcessor(ImageProcessingMixin):
"""Torch-free DeepSeek-V4-Flash-Vision-Exp image preprocessing (numpy +
PIL only), a port of the checkpoint's ``inference/image_processor.py``
diff --git a/gmlx/models/nemotron_h/model.py b/gmlx/models/nemotron_h/model.py
index 805fd7af8..33ffe4f1a 100644
--- a/gmlx/models/nemotron_h/model.py
+++ b/gmlx/models/nemotron_h/model.py
@@ -519,8 +519,10 @@ def __call__(
self,
inputs,
cache: Optional[Any] = None,
+ input_embeddings: Optional[mx.array] = None,
):
- hidden_states = self.embeddings(inputs)
+ hidden_states = (self.embeddings(inputs) if input_embeddings is None
+ else input_embeddings)
if cache is None:
cache = [None] * len(self.layers)
diff --git a/gmlx/models/nemotron_h/vision.py b/gmlx/models/nemotron_h/vision.py
new file mode 100644
index 000000000..726a0d7c9
--- /dev/null
+++ b/gmlx/models/nemotron_h/vision.py
@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Asher Feldman
+"""Vision tower for Nemotron-3.5-Bolt (mmproj projector ``nemotron_omni``):
+a C-RADIO ViT and the RMSNorm / relu^2 projector.
+
+The reference is llama.cpp ``tools/mtmd/models/nemotron-omni.cpp`` with
+``clip_graph::build_vit``, ``resize_position_embeddings`` and
+``build_patch_merge_permute`` from ``clip.cpp``:
+
+- a 16x16 patch embed with no bias, plus the 128x128 position table resized
+ to the patch grid (bilinear, half-pixel centers, edge clamp, no
+ antialias, no crop);
+- the register tokens prepended after the position add, and dropped after
+ the last block;
+- pre-norm blocks with biased LayerNorms, biased q/k/v/out projections,
+ full attention over the image, and an exact (erf) GELU MLP;
+- a 2x2 pixel shuffle (channel ``c + C*dx + 2*C*dy``), then RMSNorm and a
+ bias-free relu^2 MLP into the text width.
+
+The tower runs in float32 against the mmproj's BF16 weights, as llama.cpp
+does. Each image runs on its own.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from gmlx._vendor.mlx_vlm.models.base import BaseModelConfig
+
+
+@dataclass
+class VisionConfig(BaseModelConfig):
+ model_type: str = "nemotron_h_omni"
+ hidden_size: int = 1280
+ intermediate_size: int = 5120
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 16
+ patch_size: int = 16
+ num_channels: int = 3
+ layer_norm_eps: float = 1e-6
+ num_registers: int = 10
+ position_grid: int = 128
+ merge: int = 2
+ projector_hidden_size: int = 20480
+ projection_dim: int = 4096
+ projector_norm_eps: float = 1e-5
+
+
+def _axis_weights(n_out: int, n_in: int) -> np.ndarray:
+ """ggml's bilinear upscale weights along one axis, ``[n_out, n_in]``:
+ half-pixel sample points, edge indices clamped, the fraction clamped to
+ [0, 1], in float32 as ggml computes them."""
+ w = np.zeros((n_out, n_in), dtype=np.float32)
+ sf = np.float32(n_out) / np.float32(n_in)
+ for i in range(n_out):
+ x = (np.float32(i) + np.float32(0.5)) / sf - np.float32(0.5)
+ x0 = int(np.floor(x))
+ x1 = x0 + 1
+ x0 = min(max(x0, 0), n_in - 1)
+ x1 = min(max(x1, 0), n_in - 1)
+ d = float(np.clip(x - np.float32(x0), 0.0, 1.0))
+ w[i, x0] += 1.0 - d
+ w[i, x1] += d
+ return w
+
+
+@lru_cache(maxsize=32)
+def _resize_weights(grid_h: int, grid_w: int, side: int):
+ return (mx.array(_axis_weights(grid_h, side)),
+ mx.array(_axis_weights(grid_w, side)))
+
+
+def resize_positions(table: mx.array, grid_h: int, grid_w: int,
+ side: int) -> mx.array:
+ """The ``[side*side, C]`` table (row ``y*side + x``) resized to
+ ``[grid_h*grid_w, C]`` in patch order."""
+ c = table.shape[-1]
+ t = table.astype(mx.float32)
+ if grid_h == side and grid_w == side:
+ return t
+ ry, rx = _resize_weights(grid_h, grid_w, side)
+ t = (ry @ t.reshape(side, side * c)).reshape(grid_h, side, c)
+ t = mx.swapaxes(t, 1, 2) @ rx.T # [gh, C, gw]
+ return mx.swapaxes(t, 1, 2).reshape(grid_h * grid_w, c)
+
+
+class Attention(nn.Module):
+ def __init__(self, dims: int, num_heads: int):
+ super().__init__()
+ self.num_heads = num_heads
+ self.scale = (dims // num_heads) ** -0.5
+ self.q_proj = nn.Linear(dims, dims, bias=True)
+ self.k_proj = nn.Linear(dims, dims, bias=True)
+ self.v_proj = nn.Linear(dims, dims, bias=True)
+ self.out_proj = nn.Linear(dims, dims, bias=True)
+
+ def __call__(self, x: mx.array) -> mx.array:
+ n, d = x.shape
+ h = self.num_heads
+
+ def heads(t):
+ return t.reshape(1, n, h, d // h).transpose(0, 2, 1, 3)
+
+ o = mx.fast.scaled_dot_product_attention(
+ heads(self.q_proj(x)), heads(self.k_proj(x)),
+ heads(self.v_proj(x)), scale=self.scale)
+ return self.out_proj(o.transpose(0, 2, 1, 3).reshape(n, d))
+
+
+class MLP(nn.Module):
+ def __init__(self, dims: int, hidden: int):
+ super().__init__()
+ self.fc1 = nn.Linear(dims, hidden, bias=True)
+ self.fc2 = nn.Linear(hidden, dims, bias=True)
+
+ def __call__(self, x: mx.array) -> mx.array:
+ return self.fc2(nn.gelu(self.fc1(x)))
+
+
+class Block(nn.Module):
+ def __init__(self, config: VisionConfig):
+ super().__init__()
+ d = config.hidden_size
+ self.norm1 = nn.LayerNorm(d, eps=config.layer_norm_eps)
+ self.attn = Attention(d, config.num_attention_heads)
+ self.norm2 = nn.LayerNorm(d, eps=config.layer_norm_eps)
+ self.mlp = MLP(d, config.intermediate_size)
+
+ def __call__(self, x: mx.array) -> mx.array:
+ x = x + self.attn(self.norm1(x))
+ return x + self.mlp(self.norm2(x))
+
+
+class Projector(nn.Module):
+ def __init__(self, config: VisionConfig):
+ super().__init__()
+ merged = config.hidden_size * config.merge * config.merge
+ self.norm = nn.RMSNorm(merged, eps=config.projector_norm_eps)
+ self.fc1 = nn.Linear(merged, config.projector_hidden_size, bias=False)
+ self.fc2 = nn.Linear(config.projector_hidden_size,
+ config.projection_dim, bias=False)
+
+ def __call__(self, x: mx.array) -> mx.array:
+ return self.fc2(nn.relu2(self.fc1(self.norm(x))))
+
+
+def pixel_shuffle(x: mx.array, grid_h: int, grid_w: int,
+ merge: int) -> mx.array:
+ """``[gh*gw, C]`` patch rows to ``[(gh/m)*(gw/m), m*m*C]``, each row
+ holding a m x m block with channel ``c + C*dx + m*C*dy``. An odd edge
+ pads with zeros, as ``build_patch_merge_permute`` does."""
+ c = x.shape[-1]
+ x = x.reshape(grid_h, grid_w, c)
+ ph, pw = (-grid_h) % merge, (-grid_w) % merge
+ if ph or pw:
+ x = mx.pad(x, [(0, ph), (0, pw), (0, 0)])
+ h, w = (grid_h + ph) // merge, (grid_w + pw) // merge
+ x = x.reshape(h, merge, w, merge, c).transpose(0, 2, 1, 3, 4)
+ return x.reshape(h * w, merge * merge * c)
+
+
+class VisionModel(nn.Module):
+ def __init__(self, config: VisionConfig):
+ super().__init__()
+ self.config = config
+ d = config.hidden_size
+ p = config.patch_size
+ self.patch_embed = nn.Linear(config.num_channels * p * p, d, bias=False)
+ self.position_embedding = mx.zeros(
+ (config.position_grid * config.position_grid, d))
+ self.registers = mx.zeros((config.num_registers, d))
+ self.blocks = [Block(config) for _ in range(config.num_hidden_layers)]
+ self.projector = Projector(config)
+
+ def __call__(self, patches: mx.array,
+ grid: Tuple[int, int]) -> mx.array:
+ """``patches``: ``[gh*gw, C*p*p]`` in row-major patch order, each
+ flattened channel-major (``c, y, x``). Returns the projected
+ ``[(gh/2)*(gw/2), projection_dim]`` features."""
+ gh, gw = (int(v) for v in grid)
+ cfg = self.config
+ x = self.patch_embed(patches.astype(mx.float32))
+ x = x + resize_positions(self.position_embedding, gh, gw,
+ cfg.position_grid)
+ x = mx.concatenate(
+ [self.registers.astype(x.dtype), x], axis=0)
+ for block in self.blocks:
+ x = block(x)
+ x = x[cfg.num_registers:]
+ return self.projector(pixel_shuffle(x, gh, gw, cfg.merge))
diff --git a/gmlx/models/nemotron_h/vlm_model.py b/gmlx/models/nemotron_h/vlm_model.py
new file mode 100644
index 000000000..b31c1057a
--- /dev/null
+++ b/gmlx/models/nemotron_h/vlm_model.py
@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Asher Feldman
+"""Nemotron-3.5-Bolt with image input (mmproj projector ``nemotron_omni``).
+
+The text half is the class the text-only path uses
+(:mod:`gmlx.models.nemotron_h.model`), in the ``language_model`` shape
+mlx-vlm's generate stack expects. The vision half is the C-RADIO tower and
+projector (:mod:`gmlx.models.nemotron_h.vision`).
+
+The processor expands each ```` the chat template writes into
+``![]()
``, one ```` per soft token and ````, the stream
+llama.cpp's mtmd builds around the embeddings. The features replace the
+```` rows of the token embeddings, and the text model runs
+unchanged: it has no position encoding to adjust.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from gmlx._vendor.mlx_vlm.models.base import (
+ BaseModelConfig,
+ InputEmbeddingsFeatures,
+ LanguageModelOutput,
+)
+
+from .model import Model as _TextModel
+from .model import ModelArgs
+from .vision import VisionConfig, VisionModel
+
+TextConfig = ModelArgs
+
+__all__ = ["LanguageModel", "Model", "ModelConfig", "TextConfig",
+ "VisionConfig"]
+
+
+@dataclass
+class ModelConfig(BaseModelConfig):
+ text_config: TextConfig = None # pyright: ignore[reportAssignmentType]
+ vision_config: VisionConfig = field(default_factory=VisionConfig)
+ model_type: str = "nemotron_h_omni"
+ image_token_id: int = 18
+ image_token_index: Optional[int] = None
+ vocab_size: int = 131072
+ 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
+
+
+class LanguageModel(_TextModel):
+ def __init__(self, config: TextConfig):
+ super().__init__(config)
+ # The serve engine and the APC guard read the text config here.
+ self.config = config
+
+ @property
+ def model(self):
+ # mlx-vlm and the drafter binders reach the trunk as ``.model``.
+ return self.backbone
+
+ def __call__( # pyright: ignore[reportIncompatibleMethodOverride]
+ self,
+ inputs: mx.array,
+ inputs_embeds: Optional[mx.array] = None,
+ mask: Optional[mx.array] = None,
+ cache=None,
+ **kwargs,
+ ):
+ # The backbone builds its own attention and SSM masks from the caches.
+ skip_logits = kwargs.pop("skip_logits", False)
+ del mask, kwargs
+ h = self.backbone(inputs, cache=cache, input_embeddings=inputs_embeds)
+ return LanguageModelOutput(
+ logits=None if skip_logits else self.lm_head(h))
+
+
+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.language_model = LanguageModel(config.text_config)
+
+ def _image_features(self, pixel_values: mx.array, grid_thw) -> mx.array:
+ """Flattened patches of every image, in image order, to their
+ projected features, concatenated in the same order."""
+ feats = []
+ start = 0
+ for _t, gh, gw in (tuple(int(v) for v in row) for row in grid_thw):
+ n = gh * gw
+ feats.append(self.vision_tower(
+ pixel_values[start:start + n], (gh, gw)))
+ start += n
+ return feats[0] if len(feats) == 1 else 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.backbone.embeddings(input_ids)
+ if pixel_values is None:
+ return InputEmbeddingsFeatures(inputs_embeds=embeds)
+ features = kwargs.get("cached_image_features")
+ if features is None:
+ grid_thw = kwargs.get("image_grid_thw")
+ if grid_thw is None:
+ raise ValueError(
+ "nemotron_h_omni vision needs image_grid_thw alongside "
+ "pixel_values")
+ features = self._image_features(pixel_values, grid_thw)
+ 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 rows, 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.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/models/owned.py b/gmlx/models/owned.py
index 54adf2f51..f17eaf469 100644
--- a/gmlx/models/owned.py
+++ b/gmlx/models/owned.py
@@ -59,4 +59,5 @@
"muse_glimmer": "gmlx.models.muse_glimmer.vlm_model",
"muse_spark": "gmlx.models.muse_spark.vlm_model",
"qwen4_exp": "gmlx.models.qwen4_exp.vlm_model",
+ "nemotron_h_omni": "gmlx.models.nemotron_h.vlm_model",
}
diff --git a/gmlx/serve/patches/render.py b/gmlx/serve/patches/render.py
index 0217d237c..e0c72f976 100644
--- a/gmlx/serve/patches/render.py
+++ b/gmlx/serve/patches/render.py
@@ -76,7 +76,7 @@
"qwen3_5", "qwen3_5_moe", "qwen4_exp", "qwen3_vl", "qwen3_vl_moe",
"qwen3_omni_moe",
"kimi_k25", "deepseek_v4_vl", "deepseek_v41_vl", "muse_glimmer",
- "muse_spark", "glm5_next", "mistral3", "pixtral",
+ "muse_spark", "glm5_next", "mistral3", "pixtral", "nemotron_h_omni",
))
# Image items as the client and the stock rebuild write them.
_IMAGE_ITEMS = ("image", "image_url", "input_image")
@@ -84,7 +84,7 @@
# message in that tool turn.
_TOOL_IMAGES = frozenset((
"qwen3_5", "qwen3_5_moe", "kimi_k25", "muse_glimmer", "muse_spark",
- "glm5_next", "gemma4_unified",
+ "glm5_next", "gemma4_unified", "nemotron_h_omni",
))
# The model types whose chat template raises on a JSON-string
# ``tool_calls[*].function.arguments``.
diff --git a/tests/fixtures/nemotron_bolt_template.jinja b/tests/fixtures/nemotron_bolt_template.jinja
new file mode 100644
index 000000000..3c7d68e2f
--- /dev/null
+++ b/tests/fixtures/nemotron_bolt_template.jinja
@@ -0,0 +1,240 @@
+{% macro render_extra_keys(json_dict, handled_keys) %}
+ {%- if json_dict is mapping %}
+ {%- for json_key in json_dict if json_key not in handled_keys %}
+ {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+ {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '' }}
+ {%- else %}
+ {{-'\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '' }}
+ {%- endif %}
+ {%- endfor %}
+ {%- endif %}
+{% endmacro %}
+{%- macro render_content(content) -%}
+ {%- if content is none -%}
+ {{- '' -}}
+ {%- elif content is string -%}
+ {{- content | string -}}
+ {%- else -%}
+ {%- set text_ns = namespace(val='') -%}
+ {%- set counters = namespace(images=0, videos=0) -%}
+ {%- for part in content -%}
+ {%- if part['type'] == 'image' or part['type'] == 'image_url' or part['type'] == 'input_image' -%}
+ {%- set counters.images = counters.images + 1 -%}
+ {%- elif part['type'] == 'video' or part['type'] == 'video_url' or part['type'] == 'input_video' -%}
+ {%- set counters.videos = counters.videos + 1 -%}
+ {%- elif part['type'] == 'text' or part['type'] == 'input_text' -%}
+ {%- set text_ns.val = text_ns.val + (part['text'] | default('', true)) -%}
+ {%- endif -%}
+ {%- endfor -%}
+ {%- if '' in text_ns.val -%}
+ {%- set counters.images = 0 -%}
+ {%- endif -%}
+ {%- if '' in c %}
+ {%- set c = c.split('')[-1] %}
+ {%- elif '' in c %}
+ {%- set c = c.split('')[0] %}
+ {%- endif %}
+ {%- set c = "" ~ c %}
+ {%- if c | length > 0 %}
+ {{- c ~ '\n' -}}
+ {%- endif %}
+ {%- endif %}
+ {%- else %}
+ {{- "" -}}
+ {%- endif %}
+ {%- for tool_call in message.tool_calls %}
+ {%- if tool_call.function is defined %}
+ {%- set tool_call = tool_call.function %}
+ {%- endif %}
+ {{- '\n\n' -}}
+ {%- if tool_call.arguments is defined %}
+ {%- for args_name, args_value in tool_call.arguments|items %}
+ {{- '\n' -}}
+ {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+ {{- args_value ~ '\n\n' -}}
+ {%- endfor %}
+ {%- endif %}
+ {{- '\n\n' -}}
+ {%- endfor %}
+ {{- '<|im_end|>\n' }}
+ {%- else %}
+ {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+ {{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
+ {%- else %}
+ {%- set c = (content | default('', true) | string) %}
+ {%- if '' in c and '' in c %}
+ {%- set c = "" ~ c.split('')[-1] %}
+ {%- endif %}
+ {%- set c = c | trim %}
+ {%- if c | length > 0 %}
+ {{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
+ {%- else %}
+ {{- '<|im_start|>assistant\n<|im_end|>\n' }}
+ {%- endif %}
+ {%- endif %}
+ {%- endif %}
+ {%- elif message.role == "user" or message.role == "system" %}
+ {{- '<|im_start|>' + message.role + '\n' }}
+ {#- Render structured image/video content in its original order. -#}
+ {%- if message.role == "user" and loop.index0 == ns.last_user_idx and medium_effort %}
+ {{- content + '\n\n{reasoning effort: efficient}' }}
+ {%- else %}
+ {{- content }}
+ {%- endif %}
+ {{- '<|im_end|>\n' }}
+ {%- elif message.role == "tool" %}
+ {%- if loop.previtem and loop.previtem.role != "tool" %}
+ {{- '<|im_start|>user\n' }}
+ {%- endif %}
+ {{- '\n' }}
+ {{- content }}
+ {{- '\n\n' }}
+ {%- if not loop.last and loop.nextitem.role != "tool" %}
+ {{- '<|im_end|>\n' }}
+ {%- elif loop.last %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+ {%- else %}
+ {{- '<|im_start|>' ~ message.role ~ '\n' ~ content ~ '<|im_end|>\n' }}
+ {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+ {%- if enable_thinking %}
+ {{- '<|im_start|>assistant\n\n' }}
+ {%- else %}
+ {{- '<|im_start|>assistant\n' }}
+ {%- endif %}
+{%- endif %}
diff --git a/tests/models/test_nemotron_h_vision.py b/tests/models/test_nemotron_h_vision.py
new file mode 100644
index 000000000..225ba3619
--- /dev/null
+++ b/tests/models/test_nemotron_h_vision.py
@@ -0,0 +1,385 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Asher Feldman
+"""Nemotron-3.5-Bolt vision (mmproj ``nemotron_omni``) against a float64
+reference of llama.cpp's nemotron-omni graph, on tiny random weights. No
+GGUF and no download.
+
+The reference follows the C++ step by step: the 16x16 conv patch embed on
+the image, ggml's bilinear upscale of the position table, the registers
+prepended after the position add, ``build_vit`` with exact GELU, the
+ggml reshape/permute sequence of ``build_patch_merge_permute``, and the
+RMSNorm / relu^2 projector. The other tests cover the image processor's
+resize and token counts, the ``![]()
`` wrapping, the resolver, the remap
+and the config synth.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pytest
+
+import mlx.core as mx
+from mlx.utils import tree_flatten, tree_unflatten
+
+from gmlx.load import vlm
+from gmlx.models.nemotron_h.vision import VisionConfig, VisionModel
+
+_CFG = dict(hidden_size=32, intermediate_size=64, num_hidden_layers=2,
+ num_attention_heads=4, patch_size=4, num_registers=3,
+ position_grid=6, merge=2, projector_hidden_size=48,
+ projection_dim=24, layer_norm_eps=1e-6, projector_norm_eps=1e-5)
+
+
+def _tower(seed=0) -> VisionModel:
+ tower = VisionModel(VisionConfig(**_CFG))
+ rng = np.random.default_rng(seed)
+ params = [(k, mx.array(rng.normal(0, 0.3, v.shape).astype(np.float32)))
+ for k, v in tree_flatten(tower.parameters())]
+ tower.update(tree_unflatten(params))
+ mx.eval(tower.parameters())
+ return tower
+
+
+def _w(tower, path) -> np.ndarray:
+ return np.array(dict(tree_flatten(tower.parameters()))[path],
+ dtype=np.float64)
+
+
+def _ggml_bilinear(table: np.ndarray, side: int, gh: int, gw: int):
+ """ggml_compute_forward_upscale_f32, GGML_SCALE_MODE_BILINEAR, no
+ flags, on the ``[side*side, C]`` table (row y*side + x)."""
+ src = table.reshape(side, side, -1)
+ out = np.zeros((gh, gw, src.shape[-1]))
+ sf0 = np.float32(gw) / np.float32(side)
+ sf1 = np.float32(gh) / np.float32(side)
+ for i1 in range(gh):
+ y = (np.float32(i1) + np.float32(0.5)) / sf1 - np.float32(0.5)
+ y0 = int(math.floor(y))
+ y1 = min(max(y0 + 1, 0), side - 1)
+ y0 = min(max(y0, 0), side - 1)
+ dy = min(max(float(y) - y0, 0.0), 1.0)
+ for i0 in range(gw):
+ x = (np.float32(i0) + np.float32(0.5)) / sf0 - np.float32(0.5)
+ x0 = int(math.floor(x))
+ x1 = min(max(x0 + 1, 0), side - 1)
+ x0 = min(max(x0, 0), side - 1)
+ dx = min(max(float(x) - x0, 0.0), 1.0)
+ out[i1, i0] = (src[y0, x0] * (1 - dx) * (1 - dy)
+ + src[y0, x1] * dx * (1 - dy)
+ + src[y1, x0] * (1 - dx) * dy
+ + src[y1, x1] * dx * dy)
+ return out.reshape(gh * gw, -1)
+
+
+def _ggml_patch_merge_permute(x: np.ndarray, gh: int, gw: int, m: int):
+ """build_patch_merge_permute, step by step. numpy axes are ggml's ne
+ reversed, so ggml permute(0, 2, 1, 3) is a swap of numpy axes 0, 1."""
+ c = x.shape[-1]
+ a = x.reshape(gh, gw, c) # ne (C, W, H)
+ a = a.reshape(gh, gw // m, c * m) # ne (C*m, W/m, H)
+ a = np.transpose(a, (1, 0, 2)) # ne (C*m, H, W/m)
+ a = a.reshape(gw // m, gh // m, c * m * m) # ne (C*m*m, H/m, W/m)
+ a = np.transpose(a, (1, 0, 2)) # ne (C*m*m, W/m, H/m)
+ return a.reshape(-1, c * m * m)
+
+
+def _layer_norm(x, w, b, eps):
+ mu = x.mean(-1, keepdims=True)
+ var = ((x - mu) ** 2).mean(-1, keepdims=True)
+ return (x - mu) / np.sqrt(var + eps) * w + b
+
+
+def _gelu_erf(x):
+ return 0.5 * x * (1.0 + np.vectorize(math.erf)(x / math.sqrt(2.0)))
+
+
+def _ref_tower(tower, image: np.ndarray) -> np.ndarray:
+ """image: [C, H, W] normalized pixels."""
+ cfg = tower.config
+ p, d = cfg.patch_size, cfg.hidden_size
+ c, h, w = image.shape
+ gh, gw = h // p, w // p
+ kernel = _w(tower, "patch_embed.weight").reshape(d, c, p, p)
+ x = np.zeros((gh, gw, d))
+ for y in range(gh):
+ for xx in range(gw):
+ patch = image[:, y * p:(y + 1) * p, xx * p:(xx + 1) * p]
+ x[y, xx] = np.tensordot(kernel, patch, axes=([1, 2, 3], [0, 1, 2]))
+ x = x.reshape(gh * gw, d)
+ x = x + _ggml_bilinear(_w(tower, "position_embedding"),
+ cfg.position_grid, gh, gw)
+ x = np.concatenate([_w(tower, "registers"), x], axis=0)
+ nh = cfg.num_attention_heads
+ hd = d // nh
+ for i in range(cfg.num_hidden_layers):
+ b = f"blocks.{i}."
+
+ def lin(name, v, b=b):
+ return v @ _w(tower, b + name + ".weight").T + _w(
+ tower, b + name + ".bias")
+
+ hn = _layer_norm(x, _w(tower, b + "norm1.weight"),
+ _w(tower, b + "norm1.bias"), cfg.layer_norm_eps)
+ q, k, v = (lin(f"attn.{n}_proj", hn).reshape(-1, nh, hd)
+ for n in "qkv")
+ s = np.einsum("qhd,khd->hqk", q, k) / math.sqrt(hd)
+ s = np.exp(s - s.max(-1, keepdims=True))
+ s = s / s.sum(-1, keepdims=True)
+ o = np.einsum("hqk,khd->qhd", s, v).reshape(-1, d)
+ x = x + lin("attn.out_proj", o)
+ hn = _layer_norm(x, _w(tower, b + "norm2.weight"),
+ _w(tower, b + "norm2.bias"), cfg.layer_norm_eps)
+ x = x + lin("mlp.fc2", _gelu_erf(lin("mlp.fc1", hn)))
+ x = x[cfg.num_registers:]
+ x = _ggml_patch_merge_permute(x, gh, gw, cfg.merge)
+ x = x / np.sqrt((x * x).mean(-1, keepdims=True)
+ + cfg.projector_norm_eps) * _w(tower, "projector.norm.weight")
+ x = np.maximum(x @ _w(tower, "projector.fc1.weight").T, 0.0) ** 2
+ return x @ _w(tower, "projector.fc2.weight").T
+
+
+def _patchify(image: np.ndarray, p: int) -> np.ndarray:
+ """The processor's patch layout: row-major patches, channel-major."""
+ c, h, w = image.shape
+ a = image.reshape(c, h // p, p, w // p, p).transpose(1, 3, 0, 2, 4)
+ return a.reshape((h // p) * (w // p), c * p * p)
+
+
+@pytest.mark.parametrize("gh,gw", [(4, 6), (6, 6), (8, 2)])
+def test_tower_matches_llama_cpp_graph(gh, gw):
+ tower = _tower(seed=gh * 10 + gw)
+ p = _CFG["patch_size"]
+ image = np.random.default_rng(1).normal(0, 1, (3, gh * p, gw * p))
+ got = tower(mx.array(_patchify(image, p).astype(np.float32)), (gh, gw))
+ ref = _ref_tower(tower, image)
+ assert got.shape == ((gh // 2) * (gw // 2), _CFG["projection_dim"])
+ np.testing.assert_allclose(np.array(got), ref, rtol=2e-3, atol=2e-3)
+
+
+def test_pixel_shuffle_channel_order():
+ # Channel c + C*dx + 2*C*dy of a merged row reads patch (2y+dy, 2x+dx).
+ from gmlx.models.nemotron_h.vision import pixel_shuffle
+ gh, gw, c = 4, 6, 3
+ x = np.arange(gh * gw * c, dtype=np.float32).reshape(gh * gw, c)
+ got = np.array(pixel_shuffle(mx.array(x), gh, gw, 2))
+ np.testing.assert_array_equal(got, _ggml_patch_merge_permute(x, gh, gw, 2))
+ for y in range(gh // 2):
+ for xx in range(gw // 2):
+ for dy in range(2):
+ for dx in range(2):
+ patch = (2 * y + dy) * gw + 2 * xx + dx
+ np.testing.assert_array_equal(
+ got[y * (gw // 2) + xx, c * dx + 2 * c * dy:][:c],
+ x[patch])
+
+
+def test_position_table_identity_at_native_grid():
+ tower = _tower()
+ from gmlx.models.nemotron_h.vision import resize_positions
+ side = _CFG["position_grid"]
+ got = resize_positions(tower.position_embedding, side, side, side)
+ np.testing.assert_array_equal(np.array(got),
+ np.array(tower.position_embedding))
+
+
+def _processor():
+ return vlm._NemotronOmniGgufImageProcessor(
+ image_mean=[0.48145467, 0.4578275, 0.40821072],
+ image_std=[0.26862955, 0.26130259, 0.27577711],
+ patch_size=16, merge=2, min_pixels=262144, max_pixels=3407872)
+
+
+@pytest.mark.parametrize("hw,target", [
+ ((480, 640), (480, 640)), # already aligned, inside the budget
+ ((500, 700), (512, 704)), # rounded to the nearest 32
+ ((100, 100), (512, 512)), # below the minimum: scaled up, ceiled
+ ((1800, 2600), (1536, 2208)), # above the maximum: scaled down, floored
+ ((3000, 40), (4448, 64)), # thin: rounded under the minimum, ceiled
+])
+def test_resize_follows_calc_size_preserved_ratio(hw, target):
+ assert _processor().target_size(*hw) == target
+
+
+def test_tokens_and_patches_of_one_image():
+ from PIL import Image
+ proc = _processor()
+ out = proc([Image.new("RGB", (640, 480), (255, 0, 0))])
+ assert out["image_grid_thw"].tolist() == [[1, 30, 40]]
+ assert out["pixel_values"].shape == (30 * 40, 3 * 16 * 16)
+ assert proc.soft_tokens(out["image_grid_thw"][0]) == 300
+ # Pure red, CLIP-normalized, channel-major within each patch.
+ red = (1.0 - 0.48145467) / 0.26862955
+ np.testing.assert_allclose(out["pixel_values"][0, :256], red, rtol=1e-5)
+ np.testing.assert_allclose(
+ out["pixel_values"][0, 256:512], -0.4578275 / 0.26130259, rtol=1e-5)
+
+
+_SPECIALS = ["", "![]()
", ""]
+
+
+def _Tok():
+ """A word-level tokenizer with Bolt's image ids: 18, ![]()
19,
+ 20."""
+ from tokenizers import Tokenizer, models, pre_tokenizers
+ from transformers import PreTrainedTokenizerFast
+
+ vocab = {"[UNK]": 0, "hi": 1, **{t: 18 + i for i, t in enumerate(_SPECIALS)}}
+ tok = Tokenizer(models.WordLevel(vocab, unk_token="[UNK]"))
+ tok.pre_tokenizer = pre_tokenizers.Whitespace()
+ tok.add_special_tokens(_SPECIALS)
+ return PreTrainedTokenizerFast(tokenizer_object=tok, unk_token="[UNK]")
+
+
+def _omni_meta():
+ return {
+ "clip.vision.image_mean": [0.48145467, 0.4578275, 0.40821072],
+ "clip.vision.image_std": [0.26862955, 0.26130259, 0.27577711],
+ "clip.vision.patch_size": 16,
+ "clip.vision.projector.scale_factor": 2,
+ "clip.vision.image_min_pixels": 262144,
+ "clip.vision.image_max_pixels": 3407872,
+ "clip.vision.embedding_length": 1280,
+ "clip.vision.feed_forward_length": 5120,
+ "clip.vision.block_count": 32,
+ "clip.vision.attention.head_count": 16,
+ "clip.vision.projection_dim": 4096,
+ "clip.vision.attention.layer_norm_epsilon": 1e-6,
+ "clip.projector_type": "nemotron_omni",
+ }
+
+
+def test_processor_wraps_each_image():
+ proc = vlm._synthesize_nemotron_omni_processor(_Tok(), _omni_meta())
+ grids = np.array([[1, 30, 40], [1, 16, 16]])
+ text = proc._expand(["AB"], grids)[0]
+ assert text == ("![]()
" + "" * 300 + "A"
+ "![]()
" + "" * 64 + "B")
+ with pytest.raises(ValueError, match="more placeholders"):
+ proc._expand([""], grids[:1])
+
+
+def test_processor_stream_brackets_the_image_rows():
+ from PIL import Image
+ proc = vlm._synthesize_nemotron_omni_processor(_Tok(), _omni_meta())
+ out = proc(images=[Image.new("RGB", (512, 512))], text=[" hi"])
+ ids = np.array(out["input_ids"])[0].tolist()
+ assert ids == [19] + [18] * 256 + [20, 1]
+ assert np.array(out["pixel_values"]).shape == (32 * 32, 768)
+
+
+def test_processor_refuses_video():
+ proc = vlm._synthesize_nemotron_omni_processor(_Tok(), _omni_meta())
+ with pytest.raises(ValueError, match="video input is not supported"):
+ proc(text="hi", videos=["clip.mp4"])
+
+
+def test_resolver():
+ meta = {"clip.projector_type": "nemotron_omni"}
+ assert vlm.resolve_vlm_model_type("nemotron_h_moe", meta) == \
+ "nemotron_h_omni"
+ with pytest.raises(vlm.UnsupportedVLMError, match="nemotron_omni"):
+ vlm.resolve_vlm_model_type("llama", meta)
+
+
+def _bolt_mmproj_shapes(n_layers=2):
+ """The Bolt mmproj's tensors, in GGUF ne order, with fewer blocks."""
+ shapes = {
+ "mm.model.mlp.0.weight": [5120],
+ "mm.model.mlp.1.weight": [5120, 20480],
+ "mm.model.mlp.3.weight": [20480, 4096],
+ "v.class_embd": [1280, 10],
+ "v.patch_embd.weight": [16, 16, 3, 1280],
+ "v.patch_embd.weight.1": [16, 16, 6, 1280],
+ "v.position_embd.weight": [1280, 16384],
+ }
+ for i in range(n_layers):
+ for t in ("attn_q", "attn_k", "attn_v", "attn_out"):
+ shapes[f"v.blk.{i}.{t}.weight"] = [1280, 1280]
+ shapes[f"v.blk.{i}.{t}.bias"] = [1280]
+ shapes[f"v.blk.{i}.ffn_up.weight"] = [1280, 5120]
+ shapes[f"v.blk.{i}.ffn_up.bias"] = [5120]
+ shapes[f"v.blk.{i}.ffn_down.weight"] = [5120, 1280]
+ shapes[f"v.blk.{i}.ffn_down.bias"] = [1280]
+ for t in ("ln1", "ln2"):
+ shapes[f"v.blk.{i}.{t}.weight"] = [1280]
+ shapes[f"v.blk.{i}.{t}.bias"] = [1280]
+ return shapes
+
+
+def test_config_synth_reads_tensor_dims():
+ meta = {**_omni_meta(), "clip.vision.block_count": 2}
+ llm_meta = {"tokenizer.ggml.tokens": [""] * 18 + [""]}
+ cfg = vlm._synthesize_nemotron_omni_vlm_config(
+ {"hidden_size": 4096, "vocab_size": 131072}, meta, llm_meta,
+ _bolt_mmproj_shapes())
+ v = cfg["vision_config"]
+ assert (v["num_registers"], v["position_grid"],
+ v["projector_hidden_size"]) == (10, 128, 20480)
+ assert cfg["image_token_id"] == 18
+ with pytest.raises(ValueError, match="hidden size"):
+ vlm._synthesize_nemotron_omni_vlm_config(
+ {"hidden_size": 2688}, meta, llm_meta, _bolt_mmproj_shapes())
+
+
+def test_remap_fills_every_tower_param():
+ shapes = _bolt_mmproj_shapes()
+ arrays = {k: mx.zeros(tuple(reversed(s)), dtype=mx.float32)
+ for k, s in shapes.items()}
+ out, skipped, kq = vlm.remap_vision_arrays(arrays, "nemotron_h_omni")
+ assert skipped == ["v.patch_embd.weight.1"] and not kq
+ meta = {**_omni_meta(), "clip.vision.block_count": 2}
+ cfg = vlm._synthesize_nemotron_omni_vlm_config(
+ {"hidden_size": 4096}, meta,
+ {"tokenizer.ggml.tokens": [""] * 18 + [""]}, shapes)
+ tower = VisionModel(VisionConfig.from_dict(cfg["vision_config"]))
+ params = dict(tree_flatten(tower.parameters()))
+ assert {k[len("vision_tower."):] for k in out} == set(params)
+ for k, v in out.items():
+ assert v.shape == params[k[len("vision_tower."):]].shape, k
+
+
+def test_container_scatters_features_on_image_rows():
+ from gmlx.models.nemotron_h.vlm_model import Model, ModelConfig
+ from gmlx.models.nemotron_h.model import ModelArgs
+
+ text = ModelArgs(
+ model_type="nemotron_h", vocab_size=40, hidden_size=24,
+ intermediate_size=32, num_hidden_layers=2, max_position_embeddings=64,
+ num_attention_heads=4, num_key_value_heads=2, attention_bias=False,
+ mamba_num_heads=4, mamba_head_dim=6, mamba_proj_bias=False,
+ ssm_state_size=4, conv_kernel=4, n_groups=2, mlp_bias=False,
+ layer_norm_epsilon=1e-5, use_bias=False, use_conv_bias=True,
+ hybrid_override_pattern=["M", "*"], head_dim=6)
+ model = Model(ModelConfig(text_config=text,
+ vision_config=VisionConfig(**_CFG),
+ image_token_id=18))
+ p = _CFG["patch_size"]
+ pixels = mx.zeros((4 * 6, 3 * p * p))
+ feats = model.vision_tower(pixels, (4, 6))
+ ids = mx.array([[19] + [18] * 6 + [20, 5]])
+ emb = model.get_input_embeddings(
+ ids, pixels, image_grid_thw=np.array([[1, 4, 6]])).inputs_embeds
+ np.testing.assert_allclose(np.array(emb[0, 1:7]), np.array(feats),
+ rtol=1e-6)
+ np.testing.assert_array_equal(
+ np.array(emb[0, 7]),
+ np.array(model.language_model.backbone.embeddings(mx.array([20]))[0]))
+ logits = model(ids, pixels, image_grid_thw=np.array([[1, 4, 6]])).logits
+ assert logits.shape == (1, 9, 40)
+ with pytest.raises(ValueError, match="6 image placeholder"):
+ model.get_input_embeddings(
+ ids, mx.zeros((8 * 6, 3 * p * p)),
+ image_grid_thw=np.array([[1, 8, 6]]))
+
+
+def test_fp32_keep_row_covers_tower_and_router():
+ from gmlx.load.loader import _FP32_KEEP_BY_MODEL_TYPE
+ keep = _FP32_KEEP_BY_MODEL_TYPE["nemotron_h_omni"]
+ for name in ("vision_tower.blocks.0.norm1.weight",
+ "vision_tower.position_embedding",
+ "language_model.backbone.layers.1.mixer.gate.weight"):
+ assert any(s in name for s in keep), name
+ assert set(_FP32_KEEP_BY_MODEL_TYPE["nemotron_h"]) <= set(keep)
diff --git a/tests/serve/test_chat_image_placement.py b/tests/serve/test_chat_image_placement.py
index 08b6492ca..58681c721 100644
--- a/tests/serve/test_chat_image_placement.py
+++ b/tests/serve/test_chat_image_placement.py
@@ -298,8 +298,8 @@ def test_the_gemma4_processor_keeps_the_client_order(route, model_type):
# order.
_IN_PLACE = ("deepseek_v41_vl", "deepseek_v4_vl", "gemma4", "gemma4_unified",
"glm5_next", "kimi_k25", "mistral3", "muse_glimmer", "muse_spark",
- "pixtral", "qwen3_5", "qwen3_5_moe", "qwen3_omni_moe", "qwen3_vl",
- "qwen3_vl_moe", "qwen4_exp")
+ "nemotron_h_omni", "pixtral", "qwen3_5", "qwen3_5_moe",
+ "qwen3_omni_moe", "qwen3_vl", "qwen3_vl_moe", "qwen4_exp")
# The rows mlx-vlm's table lacks, which the loader adds on load.
_LOADER_ROWS = ("deepseek_v41_vl", "deepseek_v4_vl")
diff --git a/tests/serve/test_messages_route.py b/tests/serve/test_messages_route.py
index 3bee51ada..1c5e31cc3 100644
--- a/tests/serve/test_messages_route.py
+++ b/tests/serve/test_messages_route.py
@@ -270,10 +270,12 @@ def test_a_message_of_tool_results_renders_as_before(route, name, case):
"glm-5.3": ("glm5_next", "glm5_template.jinja", "<|begin_of_image|>"),
"muse": ("muse_glimmer", "muse_template.jinja", "<|patch|>"),
"muse-spark": ("muse_spark", "muse_spark_template.jinja", "<|patch|>"),
+ "nemotron-bolt": ("nemotron_h_omni", "nemotron_bolt_template.jinja",
+ ""),
}
# The templates that show an image in a tool turn.
_IN_TOOL_TURN = ("gemma4-12b", "glm-5.3", "kimi-k2.5", "muse", "muse-spark",
- "qwen3.5", "qwen3.6", "qwen3.8")
+ "nemotron-bolt", "qwen3.5", "qwen3.6", "qwen3.8")
def test_the_tool_image_types_are_the_checked_ones():
From 7e89d8504fc6f9053976ee1b4cb8a97170d96799 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 00:30:06 -0700
Subject: [PATCH 08/16] feat(nemotron): lilicorr drafts text turns on the bolt
image container
---
gmlx/load/mtp_target.py | 5 ++-
gmlx/models/nemotron_h/mtp.py | 4 ++-
gmlx/models/nemotron_h/vlm_model.py | 18 +++++------
tests/models/test_nemotron_h_vision.py | 42 +++++++++++++++++++++++---
4 files changed, 54 insertions(+), 15 deletions(-)
diff --git a/gmlx/load/mtp_target.py b/gmlx/load/mtp_target.py
index bf946aedc..6322ee780 100644
--- a/gmlx/load/mtp_target.py
+++ b/gmlx/load/mtp_target.py
@@ -252,6 +252,8 @@ def build(config: dict) -> nn.Module:
# The Vision-Exp container's language model carries the deepseek_v4
# text hooks (DeepseekV4SpecHooks).
"deepseek_v4_vl": "deepseek_v4",
+ # The Bolt image container's language model is NemotronHSpecLM.
+ "nemotron_h_omni": "nemotron_h",
}
@@ -269,7 +271,8 @@ def _vlm_spec_language_model(model_type: str):
image turns stays correct - unlike ``_mtp_target_classes``'s text build,
which passes an empty one). None means the built ``.language_model``
already carries the hooks (muse_glimmer / glm5_next / qwen4_exp wire
- their mixins in their own vlm_model, and the owned gemma4 containers
+ their mixins in their own vlm_model, nemotron_h_omni holds
+ NemotronHSpecLM, and the owned gemma4 containers
hold the owned language model), or the arch has no spec support
(the per-arch hook check downstream fails loud)."""
if model_type in ("qwen3_5", "qwen3_5_moe"):
diff --git a/gmlx/models/nemotron_h/mtp.py b/gmlx/models/nemotron_h/mtp.py
index cc24e5858..6f239b40d 100644
--- a/gmlx/models/nemotron_h/mtp.py
+++ b/gmlx/models/nemotron_h/mtp.py
@@ -94,13 +94,15 @@ def __call__(
mamba_sink: Optional[list] = None,
capture: Optional[tuple] = None,
exact: bool = True,
+ input_embeddings: Optional[mx.array] = None,
):
"""``capture`` is ``(layer_ids, out)``: the raw output of each listed
layer is appended to ``out`` in layer order (the DFlash seam).
``exact`` off runs a sinked verify batched, as llama.cpp does: the
block drafters verify 8 positions, where stepping every layer costs
8 decode steps."""
- hidden_states = self.embeddings(inputs)
+ hidden_states = (self.embeddings(inputs) if input_embeddings is None
+ else input_embeddings)
capture_ids, captured = capture if capture is not None else ((), None)
if cache is None:
diff --git a/gmlx/models/nemotron_h/vlm_model.py b/gmlx/models/nemotron_h/vlm_model.py
index b31c1057a..4e27a62ce 100644
--- a/gmlx/models/nemotron_h/vlm_model.py
+++ b/gmlx/models/nemotron_h/vlm_model.py
@@ -28,8 +28,8 @@
LanguageModelOutput,
)
-from .model import Model as _TextModel
from .model import ModelArgs
+from .mtp import NemotronHSpecLM
from .vision import VisionConfig, VisionModel
TextConfig = ModelArgs
@@ -53,17 +53,15 @@ def __post_init__(self):
self.image_token_index = self.image_token_id
-class LanguageModel(_TextModel):
+class LanguageModel(NemotronHSpecLM):
+ """The text model with the speculative hooks, so a drafter runs on text
+ turns. Image turns take the plain path with the merged embeddings."""
+
def __init__(self, config: TextConfig):
super().__init__(config)
# The serve engine and the APC guard read the text config here.
self.config = config
- @property
- def model(self):
- # mlx-vlm and the drafter binders reach the trunk as ``.model``.
- return self.backbone
-
def __call__( # pyright: ignore[reportIncompatibleMethodOverride]
self,
inputs: mx.array,
@@ -73,8 +71,10 @@ def __call__( # pyright: ignore[reportIncompatibleMethodOverride]
**kwargs,
):
# The backbone builds its own attention and SSM masks from the caches.
- skip_logits = kwargs.pop("skip_logits", False)
- del mask, kwargs
+ del mask
+ if kwargs.get("return_hidden") or kwargs.get("return_shared_kv"):
+ return super().__call__(inputs, cache=cache, **kwargs)
+ skip_logits = kwargs.get("skip_logits", False)
h = self.backbone(inputs, cache=cache, input_embeddings=inputs_embeds)
return LanguageModelOutput(
logits=None if skip_logits else self.lm_head(h))
diff --git a/tests/models/test_nemotron_h_vision.py b/tests/models/test_nemotron_h_vision.py
index 225ba3619..57123aa60 100644
--- a/tests/models/test_nemotron_h_vision.py
+++ b/tests/models/test_nemotron_h_vision.py
@@ -341,7 +341,7 @@ def test_remap_fills_every_tower_param():
assert v.shape == params[k[len("vision_tower."):]].shape, k
-def test_container_scatters_features_on_image_rows():
+def _container():
from gmlx.models.nemotron_h.vlm_model import Model, ModelConfig
from gmlx.models.nemotron_h.model import ModelArgs
@@ -353,9 +353,13 @@ def test_container_scatters_features_on_image_rows():
ssm_state_size=4, conv_kernel=4, n_groups=2, mlp_bias=False,
layer_norm_epsilon=1e-5, use_bias=False, use_conv_bias=True,
hybrid_override_pattern=["M", "*"], head_dim=6)
- model = Model(ModelConfig(text_config=text,
- vision_config=VisionConfig(**_CFG),
- image_token_id=18))
+ return Model(ModelConfig(text_config=text,
+ vision_config=VisionConfig(**_CFG),
+ image_token_id=18))
+
+
+def test_container_scatters_features_on_image_rows():
+ model = _container()
p = _CFG["patch_size"]
pixels = mx.zeros((4 * 6, 3 * p * p))
feats = model.vision_tower(pixels, (4, 6))
@@ -375,6 +379,36 @@ def test_container_scatters_features_on_image_rows():
image_grid_thw=np.array([[1, 8, 6]]))
+def test_container_language_model_carries_the_spec_hooks():
+ # A drafter runs on text turns through the container's language model;
+ # image turns keep their features on the same class.
+ from gmlx.load.mtp_target import (
+ _MTP_TARGET_HOOKS_BY_TYPE,
+ _spec_hook_key,
+ _vlm_spec_language_model,
+ )
+
+ model = _container()
+ lm = model.language_model
+ for hook in _MTP_TARGET_HOOKS_BY_TYPE[_spec_hook_key("nemotron_h_omni")]:
+ assert callable(getattr(lm, hook)), hook
+ assert callable(lm.set_dflash_capture)
+ assert lm.model.embed_tokens is lm.backbone.embeddings
+ assert _vlm_spec_language_model("nemotron_h_omni") is None
+
+ p = _CFG["patch_size"]
+ pixels = mx.random.normal((4 * 6, 3 * p * p), key=mx.random.key(1))
+ ids = mx.array([[19] + [18] * 6 + [20, 5]])
+ grid = np.array([[1, 4, 6]])
+ emb = model.get_input_embeddings(ids, pixels, image_grid_thw=grid)
+ with_image = model(ids, pixels, image_grid_thw=grid).logits
+ np.testing.assert_allclose(
+ np.array(with_image),
+ np.array(lm(ids, inputs_embeds=emb.inputs_embeds).logits), rtol=1e-6)
+ assert not np.allclose(np.array(with_image[0, -1]),
+ np.array(lm(ids).logits[0, -1]))
+
+
def test_fp32_keep_row_covers_tower_and_router():
from gmlx.load.loader import _FP32_KEEP_BY_MODEL_TYPE
keep = _FP32_KEEP_BY_MODEL_TYPE["nemotron_h_omni"]
From da2e3bf805baa7678d4aba0f8669f221c737d22b Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 00:30:07 -0700
Subject: [PATCH 09/16] feat(spec): run and chat pair the drafter of a sibling
lilicorr repo
---
gmlx/commands/cli.py | 21 +++++++++-----
gmlx/load/discovery.py | 53 +++++++++++++++++++++++++++++++++++-
gmlx/spec/mtp_load.py | 16 +++++++----
tests/load/test_discovery.py | 46 +++++++++++++++++++++++++++++++
4 files changed, 123 insertions(+), 13 deletions(-)
diff --git a/gmlx/commands/cli.py b/gmlx/commands/cli.py
index 6f2e51fa9..adc27755b 100644
--- a/gmlx/commands/cli.py
+++ b/gmlx/commands/cli.py
@@ -839,20 +839,27 @@ def _sibling_drafter(gguf_path: str) -> str | None:
def _companion_drafter(gguf_path: str) -> str | None:
"""The companion drafter of a companion-only family
- (``MTP_COMPANION_AUTO_MODEL_TYPES``), which auto-enables MTP on the text
- path and VLM text-only MTP alike. The loader re-resolves the same path
- when draft_gguf_path is not given. Narrower than :func:`_sibling_drafter`
- on purpose: qwen3_5's native head wins over a sidecar, so a DFlash2
- companion stays explicit ``--draft-gguf``. Header-cache peeks only."""
+ (``MTP_COMPANION_AUTO_MODEL_TYPES``), or the drafter in a marked sibling
+ repo of a target with no native head (``-LiLiCorr-GGUF``). Either
+ auto-enables MTP on the text path and VLM text-only MTP alike. The loader
+ re-resolves the same path when draft_gguf_path is not given. Narrower
+ than :func:`_sibling_drafter` on purpose: qwen3_5's native head wins over
+ a sidecar, so a DFlash2 companion stays explicit ``--draft-gguf``.
+ Header-cache peeks only."""
try:
import gmlx.load.arch_table as arch_table
import gmlx.load.config_synth as config_synth
- from gmlx.load.discovery import find_mtp_companion, header_meta
+ from gmlx.load.discovery import (
+ find_mtp_companion,
+ find_sibling_drafter,
+ header_meta,
+ )
meta = header_meta(gguf_path)
mt = config_synth.GGUF_ARCH_TO_MODEL_TYPE.get((meta or {}).get("arch") or "")
if mt not in arch_table.MTP_COMPANION_AUTO_MODEL_TYPES:
- return None
+ arches = arch_table.drafter_arches(mt) if mt else ()
+ return find_sibling_drafter(gguf_path, arches) if arches else None
return find_mtp_companion(gguf_path, arch_table.drafter_arches(mt))
except Exception:
return None
diff --git a/gmlx/load/discovery.py b/gmlx/load/discovery.py
index 358ef4c87..583f134e9 100644
--- a/gmlx/load/discovery.py
+++ b/gmlx/load/discovery.py
@@ -465,6 +465,12 @@ def find_mtp_companion(
parent = os.path.dirname(ap)
folders = [parent, os.path.join(parent, "MTP"),
os.path.join(os.path.dirname(parent), "MTP")]
+ return _pick_drafter(folders, arches, ap)
+
+
+def _pick_drafter(folders, arches: tuple, ap: str) -> str | None:
+ """The best drafter GGUF of ``arches`` in ``folders``: earlier arches
+ win, then earlier folders, then the lexically first name."""
best: tuple[int, str] | None = None
for folder in dict.fromkeys(folders):
try:
@@ -491,6 +497,49 @@ def find_mtp_companion(
return best[1] if best else None
+def find_sibling_drafter(
+ path: str,
+ drafter_arch: str | tuple,
+) -> str | None:
+ """Path of a drafter GGUF in the sibling repo of ``path``'s repo that
+ carries a drafter marker (``-LiLiCorr-GGUF`` beside
+ ``-GGUF``), or ``None``. Covers the ``__`` local layout
+ and the Hugging Face cache. A target with a native MTP head keeps it, so
+ this returns ``None`` for one."""
+ meta = header_meta(path)
+ if meta and meta.get("mtp"):
+ return None
+ arches = (drafter_arch,) if isinstance(drafter_arch, str) else tuple(drafter_arch)
+ ap = os.path.abspath(os.path.expanduser(path))
+ return _pick_drafter(_sibling_drafter_folders(ap), arches, ap)
+
+
+def _sibling_drafter_folders(ap: str) -> list:
+ parts = os.path.normpath(ap).split(os.sep)
+ for i in range(len(parts) - 2, -1, -1):
+ if parts[i].startswith("models--"):
+ root = os.sep.join(parts[:i])
+ out = []
+ for repo in _marked_repo_names(parts[i]):
+ snaps = os.path.join(root, repo, "snapshots")
+ try:
+ revs = sorted(os.listdir(snaps))
+ except OSError:
+ continue
+ out += [os.path.join(snaps, r) for r in revs]
+ return out
+ parent = os.path.dirname(ap)
+ return [os.path.join(os.path.dirname(parent), repo)
+ for repo in _marked_repo_names(os.path.basename(parent))]
+
+
+def _marked_repo_names(repo: str) -> list:
+ """``repo`` with each drafter marker before its ``-GGUF`` suffix."""
+ m = re.search(r"-gguf$", repo, re.IGNORECASE)
+ stem, suffix = (repo[:m.start()], repo[m.start():]) if m else (repo, "")
+ return [f"{stem}-{marker}{suffix}" for marker in _SIBLING_DRAFTER_MARKERS]
+
+
def fill_families(cfg) -> None:
"""Fill each config model's sampling ``family`` (profiles.py key) from its
GGUF header, in place. Shared by serve registration and the run/chat config
@@ -1294,7 +1343,9 @@ def _configured_in_dir(dirpath: str, configured) -> list:
# A drafter repo named for its target's repo plus one of these markers
# (NVIDIA's ``-LiLiCorr-GGUF`` beside ``-GGUF``) pairs with the
# models of that repo wherever the scan found them.
-_SIBLING_DRAFTER_MARKER_RE = re.compile(r"-lilicorr(?=-gguf$|$)", re.IGNORECASE)
+_SIBLING_DRAFTER_MARKERS = ("LiLiCorr",)
+_SIBLING_DRAFTER_MARKER_RE = re.compile(
+ "-(?:" + "|".join(_SIBLING_DRAFTER_MARKERS) + ")(?=-gguf$|$)", re.IGNORECASE)
def _repo_name(path: str) -> str:
diff --git a/gmlx/spec/mtp_load.py b/gmlx/spec/mtp_load.py
index d7191f545..0196c5e3b 100644
--- a/gmlx/spec/mtp_load.py
+++ b/gmlx/spec/mtp_load.py
@@ -1995,14 +1995,20 @@ def _resolve_companion_drafter(model_type: str | None, gguf_path: str, *, log):
For model_types in ``arch_table.MTP_COMPANION_AUTO_MODEL_TYPES`` the
native-head extraction cannot serve the target, so a missing companion
raises the family error rather than falling into a misleading
- no-native-head message. Returns None for every other family (their
- drafts come from the in-GGUF head, or an explicit ``--draft-gguf``).
- Shared by the text and VLM loaders."""
+ no-native-head message. Every other family takes the drafter of a marked
+ sibling repo when the target has no native head, else None (drafts come
+ from the in-GGUF head, or an explicit ``--draft-gguf``). Shared by the
+ text and VLM loaders."""
import gmlx.load.arch_table as arch_table
- from gmlx.load.discovery import find_mtp_companion
+ from gmlx.load.discovery import find_mtp_companion, find_sibling_drafter
if model_type not in arch_table.MTP_COMPANION_AUTO_MODEL_TYPES:
- return None
+ arches = arch_table.drafter_arches(model_type) if model_type else ()
+ path = find_sibling_drafter(gguf_path, arches) if arches else None
+ if path is not None:
+ loadlog.fact("mtp_companion", os.path.basename(path))
+ log(f"[mtp] sibling drafter autodetected: {path}")
+ return path
path = find_mtp_companion(gguf_path, arch_table.drafter_arches(model_type))
if path is None:
raise ValueError(
diff --git a/tests/load/test_discovery.py b/tests/load/test_discovery.py
index 5caa79ddb..411bd9b09 100644
--- a/tests/load/test_discovery.py
+++ b/tests/load/test_discovery.py
@@ -1223,3 +1223,49 @@ def test_repo_name_reads_both_layouts():
"/hf/models--nvidia--X-LiLiCorr-GGUF/snapshots/abc/X.gguf"
) == "nvidia__x-gguf"
assert disc._sibling_drafter_repo("/m/nvidia__X-GGUF/X.gguf") is None
+
+
+def _bolt_pair(target_dir, drafter_dir, monkeypatch, *, mtp=False):
+ target_dir.mkdir(parents=True)
+ drafter_dir.mkdir(parents=True)
+ target = target_dir / "Bolt.gguf"
+ drafter = drafter_dir / "Bolt-LiLiCorr.gguf"
+ for p in (target, drafter):
+ p.write_bytes(b"GGUF")
+ metas = {str(target): {"arch": "nemotron_h_moe", "kind": "model", "mtp": mtp},
+ str(drafter): {"arch": "dflash", "kind": "drafter"}}
+ monkeypatch.setattr(disc, "header_meta", lambda p: metas.get(str(p)))
+ return str(target), str(drafter)
+
+
+def test_sibling_drafter_in_the_local_layout(tmp_path, monkeypatch):
+ target, drafter = _bolt_pair(tmp_path / "nvidia__Bolt-GGUF",
+ tmp_path / "nvidia__Bolt-LiLiCorr-GGUF",
+ monkeypatch)
+ assert disc.find_sibling_drafter(target, ("dflash",)) == drafter
+ assert disc.find_mtp_companion(target, ("dflash",)) is None
+
+
+def test_sibling_drafter_in_the_hf_cache(tmp_path, monkeypatch):
+ hub = tmp_path / "hub"
+ target, drafter = _bolt_pair(
+ hub / "models--nvidia--Bolt-GGUF" / "snapshots" / "r1",
+ hub / "models--nvidia--Bolt-LiLiCorr-GGUF" / "snapshots" / "r2",
+ monkeypatch)
+ assert disc.find_sibling_drafter(target, ("dflash",)) == drafter
+
+
+def test_sibling_drafter_leaves_a_native_head_alone(tmp_path, monkeypatch):
+ target, _ = _bolt_pair(tmp_path / "nvidia__Lightning-GGUF",
+ tmp_path / "nvidia__Lightning-LiLiCorr-GGUF",
+ monkeypatch, mtp=True)
+ assert disc.find_sibling_drafter(target, ("dflash",)) is None
+
+
+def test_run_auto_enables_the_sibling_drafter(tmp_path, monkeypatch):
+ from gmlx.commands import cli
+
+ target, drafter = _bolt_pair(tmp_path / "nvidia__Bolt-GGUF",
+ tmp_path / "nvidia__Bolt-LiLiCorr-GGUF",
+ monkeypatch)
+ assert cli._companion_drafter(target) == drafter
From 8de441de33a73aa9c3f66d7ccd146801d1fb74d6 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 00:31:21 -0700
Subject: [PATCH 10/16] docs: nemotron 3.5 bolt image input and its sibling
drafter repo
---
docs/speculative-decoding.md | 5 +++--
docs/vlm.md | 1 +
2 files changed, 4 insertions(+), 2 deletions(-)
diff --git a/docs/speculative-decoding.md b/docs/speculative-decoding.md
index bcf729813..4b0622579 100644
--- a/docs/speculative-decoding.md
+++ b/docs/speculative-decoding.md
@@ -19,8 +19,9 @@ extra layer that drafts tokens, as the Qwen3.5, Qwen3.6 and Qwen3.8 models
do. `run` and `chat` turn speculation on for these without a flag. DeepSeek-V4,
Qwen3.8-Flash-Next and Muse Glimmer turn it on the same way when their
companion drafter sits in the model's folder, or in an `MTP` folder inside
-or beside it. Other families, such as gemma-4 with its assistant drafter,
-need `--draft-gguf`.
+or beside it. Nemotron 3.5 Bolt turns it on when its LiLiCorr drafter repo,
+`-LiLiCorr-GGUF`, sits beside the model's repo. Other families, such
+as gemma-4 with its assistant drafter, need `--draft-gguf`.
| Flag | Config key | Effect |
|------|------------|--------|
diff --git a/docs/vlm.md b/docs/vlm.md
index 6a2f853a4..da9850f4e 100644
--- a/docs/vlm.md
+++ b/docs/vlm.md
@@ -84,6 +84,7 @@ projector and architecture that it found. gmlx supports these families:
| gemma-4 unified | `gemma4uv`, `gemma4ua` | gemma-4-12B | It reads images, and audio when the companion has an audio encoder. |
| GLM-5.3-Flash | `glm5v` with `glm5-next` | GLM-5.3-Flash | The first uploads name both files `glm5next`, and those load too. |
| Muse Glimmer | `muse-glimmer` | Muse-Glimmer-30B | |
+| Nemotron 3.5 Bolt | `nemotron_omni` with `nemotron_h_moe` | NVIDIA-Nemotron-3.5-Bolt-74B-A9B | Each image becomes 256 to 3328 tokens. Video input is refused. |
| Muse Spark | `muse-glimmer` and `muse-spark-audio` with `muse-spark` | Muse-Spark-Open-816B-A42B | It reads images and audio. It is larger than RAM on most Macs, so it needs `--stream-experts`. |
| Kimi K2.5 and K2.7 | `kimik25` with `deepseek2` | Kimi-K2.5, Kimi-K2.7-Code | It is larger than RAM on most Macs, so it needs `--stream-experts`. |
| DeepSeek-V4-Flash-Vision-Exp | `deepseek4v` with `deepseek4` | The unsloth UD builds | |
From da00af6a6610f458afb69fabea2b314a7fbf3daf Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 01:52:20 -0700
Subject: [PATCH 11/16] fix(nemotron): bolt image turns keep their features
under server speculation
---
gmlx/models/nemotron_h/mtp.py | 9 ++++++---
gmlx/models/nemotron_h/vlm_model.py | 10 ++++++----
tests/models/test_nemotron_h_vision.py | 8 ++++++++
3 files changed, 20 insertions(+), 7 deletions(-)
diff --git a/gmlx/models/nemotron_h/mtp.py b/gmlx/models/nemotron_h/mtp.py
index 6f239b40d..c9526c2f5 100644
--- a/gmlx/models/nemotron_h/mtp.py
+++ b/gmlx/models/nemotron_h/mtp.py
@@ -245,13 +245,16 @@ def _dflash_trunk(self, hidden: mx.array) -> mx.array:
# lazy slice would hand it the packed strides.
return mx.contiguous(hidden[..., : int(self.args.hidden_size)])
- def _hidden(self, inputs, cache, mamba_sink=None) -> mx.array:
+ def _hidden(self, inputs, cache, mamba_sink=None,
+ input_embeddings=None) -> mx.array:
ids = self._dflash_capture
if ids is None:
- return self.backbone(inputs, cache=cache, mamba_sink=mamba_sink)
+ return self.backbone(inputs, cache=cache, mamba_sink=mamba_sink,
+ input_embeddings=input_embeddings)
caps: list = []
hidden = self.backbone(inputs, cache=cache, mamba_sink=mamba_sink,
- capture=(ids, caps), exact=False)
+ capture=(ids, caps), exact=False,
+ input_embeddings=input_embeddings)
if len(caps) != len(ids):
raise RuntimeError(
f"dflash capture collected {len(caps)} of {len(ids)} layers")
diff --git a/gmlx/models/nemotron_h/vlm_model.py b/gmlx/models/nemotron_h/vlm_model.py
index 4e27a62ce..92de95e07 100644
--- a/gmlx/models/nemotron_h/vlm_model.py
+++ b/gmlx/models/nemotron_h/vlm_model.py
@@ -29,7 +29,7 @@
)
from .model import ModelArgs
-from .mtp import NemotronHSpecLM
+from .mtp import NemotronHSpecLM, _SpecOutput
from .vision import VisionConfig, VisionModel
TextConfig = ModelArgs
@@ -54,8 +54,8 @@ def __post_init__(self):
class LanguageModel(NemotronHSpecLM):
- """The text model with the speculative hooks, so a drafter runs on text
- turns. Image turns take the plain path with the merged embeddings."""
+ """The text model with the speculative hooks. Both paths take the
+ merged image embeddings, so a drafter runs on image turns too."""
def __init__(self, config: TextConfig):
super().__init__(config)
@@ -73,7 +73,9 @@ def __call__( # pyright: ignore[reportIncompatibleMethodOverride]
# The backbone builds its own attention and SSM masks from the caches.
del mask
if kwargs.get("return_hidden") or kwargs.get("return_shared_kv"):
- return super().__call__(inputs, cache=cache, **kwargs)
+ hidden = self._hidden(inputs, cache, input_embeddings=inputs_embeds)
+ logits = self.lm_head(self._dflash_trunk(hidden))
+ return _SpecOutput(logits=logits, hidden_states=[hidden])
skip_logits = kwargs.get("skip_logits", False)
h = self.backbone(inputs, cache=cache, input_embeddings=inputs_embeds)
return LanguageModelOutput(
diff --git a/tests/models/test_nemotron_h_vision.py b/tests/models/test_nemotron_h_vision.py
index 57123aa60..5db5e6615 100644
--- a/tests/models/test_nemotron_h_vision.py
+++ b/tests/models/test_nemotron_h_vision.py
@@ -407,6 +407,14 @@ def test_container_language_model_carries_the_spec_hooks():
np.array(lm(ids, inputs_embeds=emb.inputs_embeds).logits), rtol=1e-6)
assert not np.allclose(np.array(with_image[0, -1]),
np.array(lm(ids).logits[0, -1]))
+ # The speculative prefill asks for hiddens and passes the embeddings.
+ spec = lm(ids, inputs_embeds=emb.inputs_embeds, return_hidden=True)
+ np.testing.assert_allclose(np.array(spec.logits), np.array(with_image),
+ rtol=1e-6)
+ lm.set_dflash_capture([0])
+ packed = lm(ids, inputs_embeds=emb.inputs_embeds, return_hidden=True)
+ np.testing.assert_allclose(np.array(packed.logits), np.array(with_image),
+ rtol=1e-5, atol=1e-5)
def test_fp32_keep_row_covers_tower_and_router():
From ae817ef0827086ec9cb34fe228d24904e9fa9d74 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 01:52:55 -0700
Subject: [PATCH 12/16] docs: the server speculates on media turns, run and
chat do not
---
docs/speculative-decoding.md | 3 ++-
docs/vlm.md | 2 +-
2 files changed, 3 insertions(+), 2 deletions(-)
diff --git a/docs/speculative-decoding.md b/docs/speculative-decoding.md
index 4b0622579..80ec6adb3 100644
--- a/docs/speculative-decoding.md
+++ b/docs/speculative-decoding.md
@@ -41,7 +41,8 @@ Speculation works with some other features and stays off with others:
- `--stream-cpu`: `run` and `chat` refuse `--speculative`.
- The [lossy MoE settings](streaming.md#the-lossy-settings): off, and
`chat` refuses `--speculative`.
-- `--mmproj`: text turns speculate, and turns with images or audio do not.
+- `--mmproj`: text turns speculate. Turns with images or audio speculate on
+ the server, and decode without it in `run` and `chat`.
- `--adapter`: the adapted model checks each draft, so the output matches
plain decoding with the adapter.
diff --git a/docs/vlm.md b/docs/vlm.md
index da9850f4e..330e9f74b 100644
--- a/docs/vlm.md
+++ b/docs/vlm.md
@@ -109,7 +109,7 @@ have limits:
| Combination | Result |
|-------------|--------|
-| [Speculative decoding](speculative-decoding.md) | Text turns use speculation, and media turns decode without it. The drafter is a native head, a `--draft-gguf` companion, or a companion found beside the model. |
+| [Speculative decoding](speculative-decoding.md) | Text turns use speculation. Media turns use it on the server, and decode without it in `run` and `chat`. |
| [`--stream-experts`](streaming.md) | `run` and `serve` stream the language model and keep the vision encoder on the GPU. Chat refuses it, and a served model cannot combine it with speculation. |
| `--stream-cpu` | The command refuses it, because it would move the vision encoder to the CPU too. |
| The lossy MoE settings, such as `--moe-experts` | Chat refuses them with `--mmproj`. |
From e8e26348aecbf4b37d2cd409cdb2264201384e23 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 09:39:39 -0700
Subject: [PATCH 13/16] fix(nemotron): norm eps falls back to
layer_norm_epsilon, as llama.cpp reads it
---
gmlx/load/config_synth.py | 6 +++++-
tests/models/test_nemotron_h.py | 9 +++++++++
2 files changed, 14 insertions(+), 1 deletion(-)
diff --git a/gmlx/load/config_synth.py b/gmlx/load/config_synth.py
index a08544a16..24a46a8a4 100644
--- a/gmlx/load/config_synth.py
+++ b/gmlx/load/config_synth.py
@@ -2426,8 +2426,12 @@ def _synth_mistral3(meta, shapes, config: dict) -> None:
def _synth_nemotron_h_moe(meta, shapes, config: dict) -> None:
arch = "nemotron_h_moe"
- # nemotron_h model uses layer_norm_epsilon (not rms_norm_eps)
+ # nemotron_h model uses layer_norm_epsilon (not rms_norm_eps). As in
+ # llama.cpp, a file with no rms eps (Super VL) takes attention's
+ # layer_norm_epsilon.
eps = _read_float(meta, f"{arch}.attention.layer_norm_rms_epsilon")
+ if eps is None:
+ eps = _read_float(meta, f"{arch}.attention.layer_norm_epsilon")
if eps is not None:
config["layer_norm_epsilon"] = eps
config.pop("rms_norm_eps", None)
diff --git a/tests/models/test_nemotron_h.py b/tests/models/test_nemotron_h.py
index 464814178..67d26a1af 100644
--- a/tests/models/test_nemotron_h.py
+++ b/tests/models/test_nemotron_h.py
@@ -267,3 +267,12 @@ def test_synth_rejects_router_past_expert_count():
shapes["blk.3.ffn_gate_inp.weight"] = [64, 9]
with pytest.raises(ValueError, match="blk.3.ffn_gate_inp"):
synthesize_config(_bolt_meta(), tensor_shapes=shapes)
+
+
+def test_synth_eps_falls_back_to_layer_norm_epsilon():
+ # Super VL carries attention.layer_norm_epsilon and no rms eps.
+ meta = _bolt_meta()
+ del meta[f"{ARCH}.attention.layer_norm_rms_epsilon"]
+ meta[f"{ARCH}.attention.layer_norm_epsilon"] = 2e-5
+ config = synthesize_config(meta, tensor_shapes=_bolt_shapes())
+ assert config["layer_norm_epsilon"] == pytest.approx(2e-5)
From 0a14db43ccda842f91075f02b5ea7f597ef63f87 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 09:39:39 -0700
Subject: [PATCH 14/16] feat(nemotron): super vl image input, with its fused
qkv, post norm and q8_0 tower
---
docs/vlm.md | 2 +-
gmlx/load/loader.py | 10 ++-
gmlx/load/vlm.py | 48 ++++++++++---
gmlx/models/nemotron_h/vision.py | 38 ++++++----
tests/models/test_nemotron_h_vision.py | 96 ++++++++++++++++++++++----
5 files changed, 155 insertions(+), 39 deletions(-)
diff --git a/docs/vlm.md b/docs/vlm.md
index 330e9f74b..9f001d689 100644
--- a/docs/vlm.md
+++ b/docs/vlm.md
@@ -84,7 +84,7 @@ projector and architecture that it found. gmlx supports these families:
| gemma-4 unified | `gemma4uv`, `gemma4ua` | gemma-4-12B | It reads images, and audio when the companion has an audio encoder. |
| GLM-5.3-Flash | `glm5v` with `glm5-next` | GLM-5.3-Flash | The first uploads name both files `glm5next`, and those load too. |
| Muse Glimmer | `muse-glimmer` | Muse-Glimmer-30B | |
-| Nemotron 3.5 Bolt | `nemotron_omni` with `nemotron_h_moe` | NVIDIA-Nemotron-3.5-Bolt-74B-A9B | Each image becomes 256 to 3328 tokens. Video input is refused. |
+| Nemotron 3.5 Bolt and Super VL | `nemotron_omni` with `nemotron_h_moe` | NVIDIA-Nemotron-3.5-Bolt-74B-A9B, NVIDIA-Nemotron-3.5-Super-VL-120B-A12B | Each image becomes 256 to 3328 tokens. Video input is refused. |
| Muse Spark | `muse-glimmer` and `muse-spark-audio` with `muse-spark` | Muse-Spark-Open-816B-A42B | It reads images and audio. It is larger than RAM on most Macs, so it needs `--stream-experts`. |
| Kimi K2.5 and K2.7 | `kimik25` with `deepseek2` | Kimi-K2.5, Kimi-K2.7-Code | It is larger than RAM on most Macs, so it needs `--stream-experts`. |
| DeepSeek-V4-Flash-Vision-Exp | `deepseek4v` with `deepseek4` | The unsloth UD builds | |
diff --git a/gmlx/load/loader.py b/gmlx/load/loader.py
index 06a431055..7a0e7ee41 100644
--- a/gmlx/load/loader.py
+++ b/gmlx/load/loader.py
@@ -937,10 +937,10 @@ def _warm_touch_pass(
# Kept in step with the model class's cast_predicate.
"nemotron_h": (".mixer.gate.weight", ".e_score_correction_bias",
".A_log"),
- # nemotron_h_omni: the same set under language_model.*, and the vision
- # tower's F32 tensors, since llama.cpp runs the tower in fp32.
+ # nemotron_h_omni: the same set under language_model.*. The vision tower
+ # keeps its file dtypes (_F16_KEEP_BY_MODEL_TYPE).
"nemotron_h_omni": (".mixer.gate.weight", ".e_score_correction_bias",
- ".A_log", "vision_tower."),
+ ".A_log"),
}
# Params kept at their native f16 through the bf16 cast (no upcast). MLX
@@ -958,6 +958,10 @@ def _warm_touch_pass(
# so the audio encoder takes the plain cast and gives the same output.
"muse_spark": ("vision_tower.", "vision_adapter.", "vision_projection."),
"kimi_k25": ("vision_tower.", "mm_projector."),
+ # nemotron_h_omni's tower runs in float32 as llama.cpp's does: its F32
+ # norms and tables stay F32, Bolt's BF16 weights stay BF16, and Super
+ # VL's Q8_0 weights, dequantized to F16 at remap, stay F16.
+ "nemotron_h_omni": ("vision_tower.",),
}
# Params cast to float32 at load, whatever float dtype the file stores them
diff --git a/gmlx/load/vlm.py b/gmlx/load/vlm.py
index 4259cb957..78c8d6ff3 100644
--- a/gmlx/load/vlm.py
+++ b/gmlx/load/vlm.py
@@ -817,12 +817,17 @@ def _qwen3vl_deepstack_blocks(mm_meta: dict) -> list[int]:
}
-# Nemotron-3.5-Bolt (nemotron_omni): the C-RADIO ViT and projector, mapped
-# onto gmlx.models.nemotron_h.vision. The patch embed is a 16x16 conv stored
-# [out, C, ph, pw] that the tower applies as a Linear over the
-# (C, ph, pw)-flattened patch. The 6-channel patch embed (frame pairs) serves
-# video only, which this tower does not take.
+# Nemotron-3.5 Bolt and Super VL (nemotron_omni): the C-RADIO ViT and
+# projector, mapped onto gmlx.models.nemotron_h.vision. The patch embed is a
+# 16x16 conv stored [out, C, ph, pw] that the tower applies as a Linear over
+# the (C, ph, pw)-flattened patch. The 6-channel patch embed (frame pairs)
+# serves video only, which this tower does not take. Super VL fuses q/k/v,
+# adds post_ln and ships Q8_0 matmul weights. Those dequantize to float16 at
+# load: the kq kernels round the tower's float32 activations to bfloat16,
+# which puts the features 8e-3 off llama.cpp's, where float16 weights under
+# float32 activations land 6e-4 off.
_NEMOTRON_OMNI_BLK_SUBMAP = {
+ "attn_qkv": "attn.qkv_proj",
"attn_q": "attn.q_proj",
"attn_k": "attn.k_proj",
"attn_v": "attn.v_proj",
@@ -838,6 +843,8 @@ def _qwen3vl_deepstack_blocks(mm_meta: dict) -> list[int]:
"mm.model.mlp.0.weight": "vision_tower.projector.norm.weight",
"mm.model.mlp.1.weight": "vision_tower.projector.fc1.weight",
"mm.model.mlp.3.weight": "vision_tower.projector.fc2.weight",
+ "v.post_ln.weight": "vision_tower.post_norm.weight",
+ "v.post_ln.bias": "vision_tower.post_norm.bias",
}
_NEMOTRON_OMNI_VIDEO_ONLY = frozenset({"v.patch_embd.weight.1"})
@@ -1105,6 +1112,18 @@ def remap_vision_arrays(
return out, skipped, vis_kqmeta
if model_type == "nemotron_h_omni":
+ codecs = mm_codecs or {}
+
+ def _emit_omni(tgt: str, arr: mx.array, src: str) -> None:
+ codec = codecs.get(src)
+ if codec is not None:
+ import mlx_kquant as kq
+
+ scales = arrays[src[: -len(".weight")] + ".scales"]
+ arr = kq.dequantize(arr, scales, codec, mx.float16).reshape(
+ arr.shape[0], -1)
+ out[tgt] = arr
+
for name, arr in arrays.items():
if name.endswith(".scales") or name.endswith(".biases"):
continue
@@ -1115,9 +1134,14 @@ def remap_vision_arrays(
out["vision_tower.patch_embed.weight"] = arr.reshape(
arr.shape[0], -1)
continue
+ if name == "v.position_embd.weight":
+ # Super VL stores it [1, rows, C].
+ out["vision_tower.position_embedding"] = arr.reshape(
+ -1, arr.shape[-1])
+ continue
hit = _NEMOTRON_OMNI_TOP_MAP.get(name)
if hit is not None:
- out[hit] = arr
+ _emit_omni(hit, arr, name)
continue
m = _VISION_BLK_RE.match(name)
if m is None:
@@ -1128,7 +1152,8 @@ def remap_vision_arrays(
if tgt is None:
skipped.append(name)
continue
- out[f"vision_tower.blocks.{m.group(1)}.{tgt}.{leaf}"] = arr
+ _emit_omni(f"vision_tower.blocks.{m.group(1)}.{tgt}.{leaf}",
+ arr, name)
return out, skipped, vis_kqmeta
if model_type == "deepseek_v4_vl":
@@ -1881,11 +1906,12 @@ def _synthesize_nemotron_omni_vlm_config(
text_config: dict, mm_meta: dict, llm_meta: dict,
mm_shapes: dict | None = None,
) -> dict:
- """Nemotron-3.5-Bolt VLM config: the nemotron_h text synth plus the
- C-RADIO tower read from ``clip.vision.*``. The register count, the
+ """Nemotron-3.5 Bolt and Super VL config: the nemotron_h text synth plus
+ the C-RADIO tower read from ``clip.vision.*``. The register count, the
position grid and the projector width have no metadata key, so they are
read off ``v.class_embd``, ``v.position_embd.weight`` and
- ``mm.model.mlp.1.weight``."""
+ ``mm.model.mlp.1.weight``; the fused q/k/v and the post-norm follow
+ their tensors."""
import math
shapes = mm_shapes or {}
@@ -1900,6 +1926,8 @@ def _synthesize_nemotron_omni_vlm_config(
"patch_size": _mm_int(mm_meta, "clip.vision.patch_size"),
"merge": _mm_int(mm_meta, "clip.vision.projector.scale_factor"),
"projection_dim": _mm_int(mm_meta, "clip.vision.projection_dim"),
+ "fused_qkv": "v.blk.0.attn_qkv.weight" in shapes,
+ "post_norm": "v.post_ln.weight" in shapes,
}
eps = _mm(mm_meta, "clip.vision.attention.layer_norm_epsilon")
if eps is not None:
diff --git a/gmlx/models/nemotron_h/vision.py b/gmlx/models/nemotron_h/vision.py
index 726a0d7c9..c4f534afc 100644
--- a/gmlx/models/nemotron_h/vision.py
+++ b/gmlx/models/nemotron_h/vision.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Asher Feldman
-"""Vision tower for Nemotron-3.5-Bolt (mmproj projector ``nemotron_omni``):
-a C-RADIO ViT and the RMSNorm / relu^2 projector.
+"""Vision tower for Nemotron-3.5 Bolt and Super VL (mmproj projector
+``nemotron_omni``): a C-RADIO ViT and the RMSNorm / relu^2 projector.
The reference is llama.cpp ``tools/mtmd/models/nemotron-omni.cpp`` with
``clip_graph::build_vit``, ``resize_position_embeddings`` and
@@ -12,13 +12,15 @@
antialias, no crop);
- the register tokens prepended after the position add, and dropped after
the last block;
-- pre-norm blocks with biased LayerNorms, biased q/k/v/out projections,
- full attention over the image, and an exact (erf) GELU MLP;
+- pre-norm blocks with biased LayerNorms, biased q/k/v/out projections
+ (one fused ``[q | k | v]`` projection in the Super VL file), full
+ attention over the image, and an exact (erf) GELU MLP;
+- a final biased LayerNorm over every token when the file has ``post_ln``;
- a 2x2 pixel shuffle (channel ``c + C*dx + 2*C*dy``), then RMSNorm and a
bias-free relu^2 MLP into the text width.
The tower runs in float32 against the mmproj's BF16 weights, as llama.cpp
-does. Each image runs on its own.
+does. Q8_0 weights run on the kquant modules. Each image runs on its own.
"""
from __future__ import annotations
@@ -50,6 +52,8 @@ class VisionConfig(BaseModelConfig):
projector_hidden_size: int = 20480
projection_dim: int = 4096
projector_norm_eps: float = 1e-5
+ fused_qkv: bool = False
+ post_norm: bool = False
def _axis_weights(n_out: int, n_in: int) -> np.ndarray:
@@ -91,13 +95,16 @@ def resize_positions(table: mx.array, grid_h: int, grid_w: int,
class Attention(nn.Module):
- def __init__(self, dims: int, num_heads: int):
+ def __init__(self, dims: int, num_heads: int, fused_qkv: bool = False):
super().__init__()
self.num_heads = num_heads
self.scale = (dims // num_heads) ** -0.5
- self.q_proj = nn.Linear(dims, dims, bias=True)
- self.k_proj = nn.Linear(dims, dims, bias=True)
- self.v_proj = nn.Linear(dims, dims, bias=True)
+ if fused_qkv:
+ self.qkv_proj = nn.Linear(dims, 3 * dims, bias=True)
+ else:
+ self.q_proj = nn.Linear(dims, dims, bias=True)
+ self.k_proj = nn.Linear(dims, dims, bias=True)
+ self.v_proj = nn.Linear(dims, dims, bias=True)
self.out_proj = nn.Linear(dims, dims, bias=True)
def __call__(self, x: mx.array) -> mx.array:
@@ -107,9 +114,12 @@ def __call__(self, x: mx.array) -> mx.array:
def heads(t):
return t.reshape(1, n, h, d // h).transpose(0, 2, 1, 3)
+ if "qkv_proj" in self:
+ q, k, v = mx.split(self.qkv_proj(x), 3, axis=-1)
+ else:
+ q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
o = mx.fast.scaled_dot_product_attention(
- heads(self.q_proj(x)), heads(self.k_proj(x)),
- heads(self.v_proj(x)), scale=self.scale)
+ heads(q), heads(k), heads(v), scale=self.scale)
return self.out_proj(o.transpose(0, 2, 1, 3).reshape(n, d))
@@ -128,7 +138,7 @@ def __init__(self, config: VisionConfig):
super().__init__()
d = config.hidden_size
self.norm1 = nn.LayerNorm(d, eps=config.layer_norm_eps)
- self.attn = Attention(d, config.num_attention_heads)
+ self.attn = Attention(d, config.num_attention_heads, config.fused_qkv)
self.norm2 = nn.LayerNorm(d, eps=config.layer_norm_eps)
self.mlp = MLP(d, config.intermediate_size)
@@ -176,6 +186,8 @@ def __init__(self, config: VisionConfig):
(config.position_grid * config.position_grid, d))
self.registers = mx.zeros((config.num_registers, d))
self.blocks = [Block(config) for _ in range(config.num_hidden_layers)]
+ if config.post_norm:
+ self.post_norm = nn.LayerNorm(d, eps=config.layer_norm_eps)
self.projector = Projector(config)
def __call__(self, patches: mx.array,
@@ -192,5 +204,7 @@ def __call__(self, patches: mx.array,
[self.registers.astype(x.dtype), x], axis=0)
for block in self.blocks:
x = block(x)
+ if "post_norm" in self:
+ x = self.post_norm(x)
x = x[cfg.num_registers:]
return self.projector(pixel_shuffle(x, gh, gw, cfg.merge))
diff --git a/tests/models/test_nemotron_h_vision.py b/tests/models/test_nemotron_h_vision.py
index 5db5e6615..2dfbf29ad 100644
--- a/tests/models/test_nemotron_h_vision.py
+++ b/tests/models/test_nemotron_h_vision.py
@@ -33,8 +33,12 @@
projection_dim=24, layer_norm_eps=1e-6, projector_norm_eps=1e-5)
-def _tower(seed=0) -> VisionModel:
- tower = VisionModel(VisionConfig(**_CFG))
+# The Super VL file fuses q/k/v and adds post_ln; Bolt's has neither.
+_SUPER = dict(fused_qkv=True, post_norm=True)
+
+
+def _tower(seed=0, **over) -> VisionModel:
+ tower = VisionModel(VisionConfig(**{**_CFG, **over}))
rng = np.random.default_rng(seed)
params = [(k, mx.array(rng.normal(0, 0.3, v.shape).astype(np.float32)))
for k, v in tree_flatten(tower.parameters())]
@@ -123,8 +127,13 @@ def lin(name, v, b=b):
hn = _layer_norm(x, _w(tower, b + "norm1.weight"),
_w(tower, b + "norm1.bias"), cfg.layer_norm_eps)
- q, k, v = (lin(f"attn.{n}_proj", hn).reshape(-1, nh, hd)
- for n in "qkv")
+ if cfg.fused_qkv:
+ # build_vit views [q | k | v] off one projection.
+ q, k, v = (t.reshape(-1, nh, hd) for t in np.split(
+ lin("attn.qkv_proj", hn), 3, axis=-1))
+ else:
+ q, k, v = (lin(f"attn.{n}_proj", hn).reshape(-1, nh, hd)
+ for n in "qkv")
s = np.einsum("qhd,khd->hqk", q, k) / math.sqrt(hd)
s = np.exp(s - s.max(-1, keepdims=True))
s = s / s.sum(-1, keepdims=True)
@@ -133,6 +142,9 @@ def lin(name, v, b=b):
hn = _layer_norm(x, _w(tower, b + "norm2.weight"),
_w(tower, b + "norm2.bias"), cfg.layer_norm_eps)
x = x + lin("mlp.fc2", _gelu_erf(lin("mlp.fc1", hn)))
+ if cfg.post_norm:
+ x = _layer_norm(x, _w(tower, "post_norm.weight"),
+ _w(tower, "post_norm.bias"), cfg.layer_norm_eps)
x = x[cfg.num_registers:]
x = _ggml_patch_merge_permute(x, gh, gw, cfg.merge)
x = x / np.sqrt((x * x).mean(-1, keepdims=True)
@@ -148,9 +160,10 @@ def _patchify(image: np.ndarray, p: int) -> np.ndarray:
return a.reshape((h // p) * (w // p), c * p * p)
+@pytest.mark.parametrize("variant", [{}, _SUPER], ids=["bolt", "super-vl"])
@pytest.mark.parametrize("gh,gw", [(4, 6), (6, 6), (8, 2)])
-def test_tower_matches_llama_cpp_graph(gh, gw):
- tower = _tower(seed=gh * 10 + gw)
+def test_tower_matches_llama_cpp_graph(gh, gw, variant):
+ tower = _tower(seed=gh * 10 + gw, **variant)
p = _CFG["patch_size"]
image = np.random.default_rng(1).normal(0, 1, (3, gh * p, gw * p))
got = tower(mx.array(_patchify(image, p).astype(np.float32)), (gh, gw))
@@ -309,6 +322,21 @@ def _bolt_mmproj_shapes(n_layers=2):
return shapes
+def _super_mmproj_shapes(n_layers=2):
+ """The Super VL mmproj: one fused attn_qkv, post_ln, and a 3-D
+ position table."""
+ shapes = _bolt_mmproj_shapes(n_layers)
+ for i in range(n_layers):
+ for t in ("attn_q", "attn_k", "attn_v"):
+ del shapes[f"v.blk.{i}.{t}.weight"], shapes[f"v.blk.{i}.{t}.bias"]
+ shapes[f"v.blk.{i}.attn_qkv.weight"] = [1280, 3840]
+ shapes[f"v.blk.{i}.attn_qkv.bias"] = [3840]
+ shapes["v.position_embd.weight"] = [1280, 16384, 1]
+ shapes["v.post_ln.weight"] = [1280]
+ shapes["v.post_ln.bias"] = [1280]
+ return shapes
+
+
def test_config_synth_reads_tensor_dims():
meta = {**_omni_meta(), "clip.vision.block_count": 2}
llm_meta = {"tokenizer.ggml.tokens": [""] * 18 + [""]}
@@ -319,13 +347,20 @@ def test_config_synth_reads_tensor_dims():
assert (v["num_registers"], v["position_grid"],
v["projector_hidden_size"]) == (10, 128, 20480)
assert cfg["image_token_id"] == 18
+ assert not v["fused_qkv"] and not v["post_norm"]
+ sv = vlm._synthesize_nemotron_omni_vlm_config(
+ {"hidden_size": 4096}, meta, llm_meta,
+ _super_mmproj_shapes())["vision_config"]
+ assert sv["fused_qkv"] and sv["post_norm"] and sv["position_grid"] == 128
with pytest.raises(ValueError, match="hidden size"):
vlm._synthesize_nemotron_omni_vlm_config(
{"hidden_size": 2688}, meta, llm_meta, _bolt_mmproj_shapes())
-def test_remap_fills_every_tower_param():
- shapes = _bolt_mmproj_shapes()
+@pytest.mark.parametrize("shapes_of", [_bolt_mmproj_shapes,
+ _super_mmproj_shapes])
+def test_remap_fills_every_tower_param(shapes_of):
+ shapes = shapes_of()
arrays = {k: mx.zeros(tuple(reversed(s)), dtype=mx.float32)
for k, s in shapes.items()}
out, skipped, kq = vlm.remap_vision_arrays(arrays, "nemotron_h_omni")
@@ -341,6 +376,33 @@ def test_remap_fills_every_tower_param():
assert v.shape == params[k[len("vision_tower."):]].shape, k
+def test_remap_dequantizes_q8_0_weights_to_float16():
+ # Super VL ships Q8_0 matmul weights. They load as float16, not through
+ # the kq kernels, whose bfloat16 activations cost the tower accuracy.
+ import mlx_kquant as kq
+
+ shapes = _super_mmproj_shapes(1)
+ arrays = {k: mx.zeros(tuple(reversed(s)), dtype=mx.float32)
+ for k, s in shapes.items()}
+ quant = {"v.blk.0.attn_qkv.weight": "vision_tower.blocks.0.attn.qkv_proj.weight",
+ "mm.model.mlp.1.weight": "vision_tower.projector.fc1.weight"}
+ want = {}
+ for src, tgt in quant.items():
+ w = mx.random.normal(arrays[src].shape, key=mx.random.key(len(src)))
+ wire, scales = kq.quantize(w, "q8_0")
+ arrays[src] = wire
+ arrays[src[: -len(".weight")] + ".scales"] = scales
+ want[tgt] = np.array(kq.dequantize(wire, scales, "q8_0", mx.float32))
+ out, _, kq_meta = vlm.remap_vision_arrays(
+ arrays, "nemotron_h_omni", mm_codecs={k: "q8_0" for k in quant})
+ assert not kq_meta
+ assert not any(k.endswith(".scales") for k in out)
+ for tgt, ref in want.items():
+ assert out[tgt].dtype == mx.float16 and out[tgt].shape == ref.shape
+ np.testing.assert_allclose(np.array(out[tgt].astype(mx.float32)), ref,
+ rtol=1e-3, atol=1e-3)
+
+
def _container():
from gmlx.models.nemotron_h.vlm_model import Model, ModelConfig
from gmlx.models.nemotron_h.model import ModelArgs
@@ -417,11 +479,19 @@ def test_container_language_model_carries_the_spec_hooks():
rtol=1e-5, atol=1e-5)
-def test_fp32_keep_row_covers_tower_and_router():
- from gmlx.load.loader import _FP32_KEEP_BY_MODEL_TYPE
+def test_keep_rows_cover_tower_and_router():
+ from gmlx.load.loader import (
+ _F16_KEEP_BY_MODEL_TYPE,
+ _FP32_KEEP_BY_MODEL_TYPE,
+ )
keep = _FP32_KEEP_BY_MODEL_TYPE["nemotron_h_omni"]
+ assert any(s in "language_model.backbone.layers.1.mixer.gate.weight"
+ for s in keep)
+ assert set(_FP32_KEEP_BY_MODEL_TYPE["nemotron_h"]) <= set(keep)
+ # The tower keeps its file dtypes, F16 included.
+ tower = _F16_KEEP_BY_MODEL_TYPE["nemotron_h_omni"]
for name in ("vision_tower.blocks.0.norm1.weight",
"vision_tower.position_embedding",
- "language_model.backbone.layers.1.mixer.gate.weight"):
- assert any(s in name for s in keep), name
- assert set(_FP32_KEEP_BY_MODEL_TYPE["nemotron_h"]) <= set(keep)
+ "vision_tower.blocks.0.attn.qkv_proj.weight"):
+ assert any(s in name for s in tower), name
+ assert not any(s in name for s in keep), name
From 6490c9b8663646d78cd3e72e29381f363e9ddd32 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 09:39:40 -0700
Subject: [PATCH 15/16] feat(spec): a headless nemotron target pairs its mtp
sidecar in run and chat
---
CHANGELOG.md | 2 ++
docs/speculative-decoding.md | 6 +++---
gmlx/commands/cli.py | 9 +++++----
gmlx/load/arch_table.py | 6 +++---
gmlx/load/discovery.py | 20 ++++++++++++++++++++
gmlx/spec/mtp_load.py | 14 +++++++-------
tests/load/test_discovery.py | 28 ++++++++++++++++++++++++++++
7 files changed, 68 insertions(+), 17 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ce53e89cd..fd83f2869 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -38,6 +38,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
- NVIDIA Nemotron 3.5 Bolt 74B-A9B, with image input through its `--mmproj`
companion and its LiLiCorr drafter, which gmlx pairs from the sibling
`-LiLiCorr-GGUF` repo.
+- NVIDIA Nemotron 3.5 Super VL 120B-A12B, with image input through its
+ `--mmproj` companion and an MTP drafter from the `mtp-*.gguf` file beside it.
- Nemotron models take `@instruct` to turn thinking off, and
`@reasoning-medium` for a shorter reasoning pass where the template has one.
diff --git a/docs/speculative-decoding.md b/docs/speculative-decoding.md
index 80ec6adb3..0b00ebb3f 100644
--- a/docs/speculative-decoding.md
+++ b/docs/speculative-decoding.md
@@ -17,9 +17,9 @@ gmlx run model.gguf --bench-depths "0,4096" --speculative # measure the gain
Some GGUF files carry a [native head](glossary.md#native-head), a small
extra layer that drafts tokens, as the Qwen3.5, Qwen3.6 and Qwen3.8 models
do. `run` and `chat` turn speculation on for these without a flag. DeepSeek-V4,
-Qwen3.8-Flash-Next and Muse Glimmer turn it on the same way when their
-companion drafter sits in the model's folder, or in an `MTP` folder inside
-or beside it. Nemotron 3.5 Bolt turns it on when its LiLiCorr drafter repo,
+Qwen3.8-Flash-Next, Muse Glimmer and Nemotron 3.5 Super VL turn it on the
+same way when their companion drafter sits in the model's folder, or in an
+`MTP` folder inside or beside it. Nemotron 3.5 Bolt turns it on when its LiLiCorr drafter repo,
`-LiLiCorr-GGUF`, sits beside the model's repo. Other families, such
as gemma-4 with its assistant drafter, need `--draft-gguf`.
diff --git a/gmlx/commands/cli.py b/gmlx/commands/cli.py
index adc27755b..2d32ee398 100644
--- a/gmlx/commands/cli.py
+++ b/gmlx/commands/cli.py
@@ -839,8 +839,9 @@ def _sibling_drafter(gguf_path: str) -> str | None:
def _companion_drafter(gguf_path: str) -> str | None:
"""The companion drafter of a companion-only family
- (``MTP_COMPANION_AUTO_MODEL_TYPES``), or the drafter in a marked sibling
- repo of a target with no native head (``-LiLiCorr-GGUF``). Either
+ (``MTP_COMPANION_AUTO_MODEL_TYPES``), or for a target with no native
+ head its NextN sidecar or the drafter in a marked sibling repo
+ (``-LiLiCorr-GGUF``). Each
auto-enables MTP on the text path and VLM text-only MTP alike. The loader
re-resolves the same path when draft_gguf_path is not given. Narrower
than :func:`_sibling_drafter` on purpose: qwen3_5's native head wins over
@@ -850,8 +851,8 @@ def _companion_drafter(gguf_path: str) -> str | None:
import gmlx.load.arch_table as arch_table
import gmlx.load.config_synth as config_synth
from gmlx.load.discovery import (
+ find_headless_drafter,
find_mtp_companion,
- find_sibling_drafter,
header_meta,
)
@@ -859,7 +860,7 @@ def _companion_drafter(gguf_path: str) -> str | None:
mt = config_synth.GGUF_ARCH_TO_MODEL_TYPE.get((meta or {}).get("arch") or "")
if mt not in arch_table.MTP_COMPANION_AUTO_MODEL_TYPES:
arches = arch_table.drafter_arches(mt) if mt else ()
- return find_sibling_drafter(gguf_path, arches) if arches else None
+ return find_headless_drafter(gguf_path, arches) if arches else None
return find_mtp_companion(gguf_path, arch_table.drafter_arches(mt))
except Exception:
return None
diff --git a/gmlx/load/arch_table.py b/gmlx/load/arch_table.py
index 1af96d65a..cca9c348c 100644
--- a/gmlx/load/arch_table.py
+++ b/gmlx/load/arch_table.py
@@ -180,9 +180,9 @@ class UnsupportedArchError(Exception):
# the HF mtp.* tensors, or the llama.cpp mtp-*.gguf NextN sidecar, which
# keeps the target's arch.
"qwen4_exp": ("qwen4exp-mtp", "qwen4exp"),
- # LiLiCorr DFlash drafters (Nemotron-3.5-Bolt). The in-file NextN head
- # and the same-arch mtp sidecar need no row.
- "nemotron_h": ("dflash",),
+ # LiLiCorr DFlash drafters (Nemotron-3.5-Bolt) and the llama.cpp mtp-*.gguf
+ # NextN sidecar (Nemotron-3.5-Super-VL), which keeps the target's arch.
+ "nemotron_h": ("dflash", "nemotron_h_moe"),
}
diff --git a/gmlx/load/discovery.py b/gmlx/load/discovery.py
index 583f134e9..16ef2f6b1 100644
--- a/gmlx/load/discovery.py
+++ b/gmlx/load/discovery.py
@@ -514,6 +514,26 @@ def find_sibling_drafter(
return _pick_drafter(_sibling_drafter_folders(ap), arches, ap)
+def find_headless_drafter(
+ path: str,
+ drafter_arch: str | tuple,
+) -> str | None:
+ """The drafter of a target with no native MTP head: a NextN sidecar of
+ the target's own arch beside it (the head in its own file, as
+ llama.cpp's ``mtp-*.gguf``) when ``drafter_arch`` lists that arch, else
+ :func:`find_sibling_drafter`. ``None`` for a target with a head."""
+ meta = header_meta(path)
+ if not meta or meta.get("mtp"):
+ return None
+ arches = (drafter_arch,) if isinstance(drafter_arch, str) else tuple(drafter_arch)
+ arch = meta.get("arch")
+ if arch in arches:
+ hit = find_mtp_companion(path, (arch,))
+ if hit is not None:
+ return hit
+ return find_sibling_drafter(path, arches)
+
+
def _sibling_drafter_folders(ap: str) -> list:
parts = os.path.normpath(ap).split(os.sep)
for i in range(len(parts) - 2, -1, -1):
diff --git a/gmlx/spec/mtp_load.py b/gmlx/spec/mtp_load.py
index 0196c5e3b..549c74f38 100644
--- a/gmlx/spec/mtp_load.py
+++ b/gmlx/spec/mtp_load.py
@@ -1995,19 +1995,19 @@ def _resolve_companion_drafter(model_type: str | None, gguf_path: str, *, log):
For model_types in ``arch_table.MTP_COMPANION_AUTO_MODEL_TYPES`` the
native-head extraction cannot serve the target, so a missing companion
raises the family error rather than falling into a misleading
- no-native-head message. Every other family takes the drafter of a marked
- sibling repo when the target has no native head, else None (drafts come
- from the in-GGUF head, or an explicit ``--draft-gguf``). Shared by the
- text and VLM loaders."""
+ no-native-head message. In every other family, a target with no native
+ head takes its NextN sidecar or the drafter of a marked sibling repo;
+ else None (drafts come from the in-GGUF head, or ``--draft-gguf``).
+ Shared by the text and VLM loaders."""
import gmlx.load.arch_table as arch_table
- from gmlx.load.discovery import find_mtp_companion, find_sibling_drafter
+ from gmlx.load.discovery import find_headless_drafter, find_mtp_companion
if model_type not in arch_table.MTP_COMPANION_AUTO_MODEL_TYPES:
arches = arch_table.drafter_arches(model_type) if model_type else ()
- path = find_sibling_drafter(gguf_path, arches) if arches else None
+ path = find_headless_drafter(gguf_path, arches) if arches else None
if path is not None:
loadlog.fact("mtp_companion", os.path.basename(path))
- log(f"[mtp] sibling drafter autodetected: {path}")
+ log(f"[mtp] companion drafter autodetected: {path}")
return path
path = find_mtp_companion(gguf_path, arch_table.drafter_arches(model_type))
if path is None:
diff --git a/tests/load/test_discovery.py b/tests/load/test_discovery.py
index 411bd9b09..004e186fd 100644
--- a/tests/load/test_discovery.py
+++ b/tests/load/test_discovery.py
@@ -1269,3 +1269,31 @@ def test_run_auto_enables_the_sibling_drafter(tmp_path, monkeypatch):
tmp_path / "nvidia__Bolt-LiLiCorr-GGUF",
monkeypatch)
assert cli._companion_drafter(target) == drafter
+
+
+def _super_vl(tmp_path, monkeypatch, *, mtp=False):
+ target = tmp_path / "Super-VL-Q4_0.gguf"
+ sidecar = tmp_path / "mtp-Super-VL-Q4_0.gguf"
+ for p in (target, sidecar):
+ p.write_bytes(b"GGUF")
+ metas = {str(target): {"arch": "nemotron_h_moe", "kind": "model", "mtp": mtp},
+ str(sidecar): {"arch": "nemotron_h_moe", "kind": "drafter"}}
+ monkeypatch.setattr(disc, "header_meta", lambda p: metas.get(str(p)))
+ return str(target), str(sidecar)
+
+
+def test_headless_target_takes_its_nextn_sidecar(tmp_path, monkeypatch):
+ # llama.cpp ships Super VL's MTP head as mtp-*.gguf beside the model.
+ from gmlx.commands import cli
+
+ target, sidecar = _super_vl(tmp_path, monkeypatch)
+ arches = ("dflash", "nemotron_h_moe")
+ assert disc.find_headless_drafter(target, arches) == sidecar
+ assert cli._companion_drafter(target) == sidecar
+ # Only a family whose row lists its own arch loads such a sidecar.
+ assert disc.find_headless_drafter(target, ("dflash",)) is None
+
+
+def test_headless_lookup_leaves_a_native_head_alone(tmp_path, monkeypatch):
+ target, _ = _super_vl(tmp_path, monkeypatch, mtp=True)
+ assert disc.find_headless_drafter(target, ("nemotron_h_moe",)) is None
From e82b59d18e05ca22d998d2f54d354f46370a6e9f Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Sun, 18 Oct 2026 10:01:54 -0700
Subject: [PATCH 16/16] perf(nemotron): mtp verify runs batched and rollback
reruns only the mamba scan
---
CHANGELOG.md | 3 +
gmlx/models/nemotron_h/model.py | 28 ++++++-
gmlx/models/nemotron_h/mtp.py | 117 +++++++---------------------
tests/spec/test_lilicorr_drafter.py | 34 +++++++-
4 files changed, 86 insertions(+), 96 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd83f2869..ce510072d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -57,6 +57,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
family from another source, and names it, where it downloaded the VAE again.
- `presence_penalty` and `frequency_penalty` count every generated token, as
vLLM does, where they counted the last 20 tokens, the prompt's included.
+- MTP speculative decoding on Nemotron 3.5 checks each draft in one batched
+ forward, as llama.cpp does, and runs faster. Greedy output can now differ
+ from plain decoding where two tokens tie.
### Fixed
diff --git a/gmlx/models/nemotron_h/model.py b/gmlx/models/nemotron_h/model.py
index 33ffe4f1a..09acbdd0b 100644
--- a/gmlx/models/nemotron_h/model.py
+++ b/gmlx/models/nemotron_h/model.py
@@ -172,6 +172,7 @@ def _conv(
conv_input: mx.array,
cache: Optional[ArraysCache],
mask: Optional[mx.array],
+ record: Optional[dict] = None,
) -> mx.array:
if mask is not None:
conv_input = mx.where(mask[..., None], conv_input, 0)
@@ -193,6 +194,8 @@ def _conv(
cache[0] = mx.take_along_axis(padded_input, positions, axis=1)
else:
cache[0] = padded_input[:, -n_keep:, :]
+ if record is not None:
+ record["conv_window"] = padded_input
else:
padded_input = mx.pad(
conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)]
@@ -241,8 +244,10 @@ def __call__(
hidden_states: mx.array,
mask: Optional[mx.array],
cache: Optional[ArraysCache] = None,
+ record: Optional[dict] = None,
) -> mx.array:
-
+ """``record`` (a dict) receives the conv window and the SSM inputs
+ :meth:`replay_prefix` needs to rewind ``cache``."""
projected = self.in_proj(hidden_states)
gate, conv_input, dt = mx.split(
@@ -250,7 +255,7 @@ def __call__(
[self.intermediate_size, self.intermediate_size + self.conv_dim],
axis=-1,
)
- conv_output = self._conv(conv_input, cache, mask)
+ conv_output = self._conv(conv_input, cache, mask, record)
hidden_states_ssm, B, C = mx.split(
conv_output,
[
@@ -259,12 +264,25 @@ def __call__(
],
axis=-1,
)
+ if record is not None and "conv_window" in record:
+ record["ssm_inputs"] = (hidden_states_ssm, B, C, dt)
y = self._ssm(hidden_states_ssm, B, C, dt, cache, mask)
if cache:
cache.advance(y.shape[1])
y = self.norm(y, gate)
return self.out_proj(y)
+ def replay_prefix(self, record: dict, n: int, cache: ArraysCache,
+ mask: Optional[mx.array]) -> None:
+ """Set ``cache``, holding the state before a recorded call, to the
+ state after the call's first ``n`` positions. Only the scan runs."""
+ n_keep = self.conv_kernel_size - 1
+ cache[0] = mx.contiguous(record["conv_window"][:, n:n + n_keep])
+ x, B, C, dt = (a[:, :n] for a in record["ssm_inputs"])
+ if isinstance(mask, mx.array):
+ mask = mask[:, :n]
+ self._ssm(x, B, C, dt, cache, mask)
+
class NemotronHAttention(nn.Module):
def __init__(self, args: ModelArgs):
@@ -483,9 +501,13 @@ def __call__(
x,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
+ record: Optional[dict] = None,
):
hidden_states = self.norm(x)
- if self.block_type == "M" or self.block_type == "*":
+ if self.block_type == "M":
+ hidden_states = self.mixer(hidden_states, mask=mask, cache=cache,
+ record=record)
+ elif self.block_type == "*":
hidden_states = self.mixer(hidden_states, mask=mask, cache=cache)
else:
hidden_states = self.mixer(hidden_states)
diff --git a/gmlx/models/nemotron_h/mtp.py b/gmlx/models/nemotron_h/mtp.py
index c9526c2f5..04b47e8d1 100644
--- a/gmlx/models/nemotron_h/mtp.py
+++ b/gmlx/models/nemotron_h/mtp.py
@@ -28,13 +28,14 @@
The trunk is the owned ``nemotron_h.model`` classes; the backbone
subclass only adds the verify-time Mamba2 sink. Verify rollback is
two-part (the glm5_next recurrent precedent): the attention KV leaves trim,
-and the Mamba2 layers replay the accepted prefix from the recorded
-pre-verify (conv, ssm) state - ``ArraysCache`` state entries are replaced,
-not mutated, so the recorded pre-arrays stay valid.
-
-Correctness is drafter-independent: the verify walk emits the target's own
-greedy/sampled tokens, so the drafter affects speed (acceptance), never
-output. The losslessness gate is the greedy A/B vs plain decode.
+and the Mamba2 layers rerun their scan over the accepted prefix from the
+recorded pre-verify (conv, ssm) state and SSM inputs - ``ArraysCache``
+state entries are replaced, not mutated, so the recorded pre-arrays stay
+valid.
+
+The verify forward runs batched, as llama.cpp does. Its M>1 kernels round
+differently from decode, so greedy output can split from plain decode at a
+near-tie logit; otherwise the drafter affects speed, never output.
"""
from __future__ import annotations
@@ -77,8 +78,9 @@ class NemotronHSpecBackbone(nh.NemotronHModel):
The loop is the stock ``NemotronHModel.__call__`` with one addition:
when ``mamba_sink`` is given, each Mamba block records its cache, the
- pre-update (conv, ssm) state arrays, and its own input, so
- :func:`rollback_verify_sink` can rewind to the accepted prefix.
+ pre-update (conv, ssm) state arrays, its input, conv window and SSM
+ inputs, so :func:`rollback_verify_sink` can rewind to the accepted
+ prefix.
"""
@property
@@ -93,14 +95,10 @@ def __call__(
cache: Optional[Any] = None,
mamba_sink: Optional[list] = None,
capture: Optional[tuple] = None,
- exact: bool = True,
input_embeddings: Optional[mx.array] = None,
):
"""``capture`` is ``(layer_ids, out)``: the raw output of each listed
- layer is appended to ``out`` in layer order (the DFlash seam).
- ``exact`` off runs a sinked verify batched, as llama.cpp does: the
- block drafters verify 8 positions, where stepping every layer costs
- 8 decode steps."""
+ layer is appended to ``out`` in layer order (the DFlash seam)."""
hidden_states = (self.embeddings(inputs) if input_embeddings is None
else input_embeddings)
capture_ids, captured = capture if capture is not None else ((), None)
@@ -127,21 +125,11 @@ def __call__(
# The mixer replaces c[0]/c[1] (conv, ssm) rather than
# mutating them, so holding the pre-call arrays is a
# zero-copy snapshot.
- mamba_sink.append({
- "layer": layer, "cache": c,
- "pre": (c[0], c[1]),
- "inputs": hidden_states, "mask": mask, "exact": exact,
- })
- if exact:
- hidden_states = _mamba_stepwise(layer, hidden_states, mask, c)
- else:
- hidden_states = layer(hidden_states, mask=mask, cache=c)
- elif mamba_sink is not None and not exact:
- hidden_states = layer(hidden_states, mask=mask, cache=c)
- elif mamba_sink is not None and layer.block_type == "*" and c is not None:
- hidden_states = _attn_stepwise(layer, hidden_states, c)
- elif mamba_sink is not None and layer.block_type == "E":
- hidden_states = _moe_rowgate(layer, hidden_states)
+ entry = {"layer": layer, "cache": c, "pre": (c[0], c[1]),
+ "inputs": hidden_states, "mask": mask}
+ mamba_sink.append(entry)
+ hidden_states = layer(hidden_states, mask=mask, cache=c,
+ record=entry)
else:
hidden_states = layer(hidden_states, mask=mask, cache=c)
if i in capture_ids:
@@ -150,72 +138,23 @@ def __call__(
return self.norm_f(hidden_states)
-def _mamba_stepwise(layer, x: mx.array, mask, cache) -> mx.array:
- """Mamba block one position at a time: ``ssm_update`` runs T=1 on the
- sequential kernel and T>1 through the segsum scan, whose rounding
- differs. Stepping keeps verify and replay bit-identical to decode."""
- T = int(x.shape[1])
- if T <= 1:
- return layer(x, mask=mask, cache=cache)
- outs = []
- for t in range(T):
- m = mask[:, t:t + 1] if isinstance(mask, mx.array) else mask
- outs.append(layer(x[:, t:t + 1], mask=m, cache=cache))
- return mx.concatenate(outs, axis=1)
-
-
-def _attn_stepwise(layer, x: mx.array, cache) -> mx.array:
- """Attention block one position at a time: SDPA's qL=1 and qL>1 kernels
- round differently. A qL=1 causal step over full history needs no mask;
- rollback stays trim-based."""
- T = int(x.shape[1])
- if T <= 1:
- return layer(x, mask=None, cache=cache)
- outs = []
- for t in range(T):
- outs.append(layer(x[:, t:t + 1], mask=None, cache=cache))
- return mx.concatenate(outs, axis=1)
-
-
-def _moe_rowgate(layer, x: mx.array) -> mx.array:
- """MoE block with the router gate run per row: the gate matmul's
- M=1-vs-M>1 rounding shifts scores enough to flip near-tie argmaxes
- downstream. switch_mlp and shared_experts are M-invariant and stay
- batched (stock ``NemotronHMoE.__call__`` body otherwise)."""
- moe = layer.mixer
- if getattr(moe, "moe_latent_size", None) is not None:
- outs = [layer(x[:, t:t + 1]) for t in range(int(x.shape[1]))]
- return mx.concatenate(outs, axis=1)
- xn = layer.norm(x)
- pairs = [moe.gate(xn[:, t:t + 1]) for t in range(int(xn.shape[1]))]
- inds = mx.concatenate([p[0] for p in pairs], axis=1)
- scores = mx.concatenate([p[1] for p in pairs], axis=1)
- y = moe.switch_mlp(xn, inds)
- y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
- if moe.config.n_shared_experts is not None:
- y = y + moe.shared_experts(xn)
- return x + y
-
-
def rollback_verify_sink(sink: list, n: int) -> None:
- """Rewind the Mamba2 caches after an MTP verify forward over S positions
- to the state after its first ``n`` (the accepted prefix): restore the
- recorded pre-verify (conv, ssm) state, then replay the block over the
- accepted prefix. O(n <= block) per layer; the attention KV leaves are
- trimmed by the caller. The replay runs the way the verify ran: stepped
- position-wise for an exact verify (see :func:`_mamba_stepwise`), so the
- rewound state is bit-identical to a plain decode over the accepted
- tokens, else batched."""
+ """Rewind the Mamba2 caches after a verify forward over S positions to
+ the state after its first ``n`` (the accepted prefix): restore the
+ recorded pre-verify (conv, ssm) state, then rerun the scan over the
+ prefix's recorded SSM inputs. A layer whose cache tracks per-row lengths
+ records none and replays the whole block. The attention KV leaves are
+ trimmed by the caller."""
for e in sink:
cache = e["cache"]
cache[0], cache[1] = e["pre"]
+ if "ssm_inputs" in e:
+ e["layer"].mixer.replay_prefix(e, n, cache, e["mask"])
+ continue
mask = e["mask"]
if isinstance(mask, mx.array):
mask = mask[:, :n]
- if e.get("exact", True):
- _mamba_stepwise(e["layer"], e["inputs"][:, :n], mask, cache)
- else:
- e["layer"](e["inputs"][:, :n], mask=mask, cache=cache)
+ e["layer"](e["inputs"][:, :n], mask=mask, cache=cache)
class NemotronHSpecLM(nh.Model):
@@ -253,7 +192,7 @@ def _hidden(self, inputs, cache, mamba_sink=None,
input_embeddings=input_embeddings)
caps: list = []
hidden = self.backbone(inputs, cache=cache, mamba_sink=mamba_sink,
- capture=(ids, caps), exact=False,
+ capture=(ids, caps),
input_embeddings=input_embeddings)
if len(caps) != len(ids):
raise RuntimeError(
diff --git a/tests/spec/test_lilicorr_drafter.py b/tests/spec/test_lilicorr_drafter.py
index 8734c5ba7..499abecec 100644
--- a/tests/spec/test_lilicorr_drafter.py
+++ b/tests/spec/test_lilicorr_drafter.py
@@ -235,12 +235,14 @@ def test_capture_packs_trunk_and_layers():
rtol=1e-5, atol=1e-5)
+@pytest.mark.parametrize("capture", [None, [0, 2]], ids=["mtp", "dflash"])
@pytest.mark.parametrize("accepted", [0, 2, 4])
-def test_batched_verify_rollback_matches_a_fresh_prefix(accepted):
- # The DFlash verify runs batched; rolling back to the accepted prefix
- # must leave the caches as running that prefix would.
+def test_batched_verify_rollback_matches_a_fresh_prefix(accepted, capture):
+ # The verify runs batched; rolling back to the accepted prefix must
+ # leave the caches as running that prefix would.
model = _tiny_target(2)
- model.set_dflash_capture([0, 2])
+ if capture:
+ model.set_dflash_capture(capture)
prompt = mx.array([[1, 5, 9, 2]])
block = mx.array([[7, 3, 8, 6, 4]])
cache = model.make_cache()
@@ -256,6 +258,30 @@ def test_batched_verify_rollback_matches_a_fresh_prefix(accepted):
rtol=5e-3, atol=5e-3)
+def test_rollback_reruns_only_the_scan():
+ model = _tiny_target(3)
+ cache = model.make_cache()
+ model(mx.array([[1, 5, 9, 2]]), cache=cache)
+ _, _, sink = model.speculative_verify_hidden(mx.array([[7, 3, 8, 6]]), cache)
+ mx.eval(sink[0]["ssm_inputs"])
+ want = model.make_cache()
+ model(mx.array([[1, 5, 9, 2]]), cache=want)
+ model(mx.array([[7, 3]]), cache=want)
+
+ def refuse(*_):
+ raise AssertionError("rollback read a projection")
+
+ for e in sink:
+ e["layer"].mixer.in_proj = refuse
+ e["layer"].mixer.out_proj = refuse
+ model.rollback_speculative_cache(cache, sink, 1, 4)
+ for got, ref in zip(cache, want):
+ if got is None or not hasattr(got, "cache"):
+ continue
+ for g, r in zip(got.cache, ref.cache):
+ np.testing.assert_allclose(_np(g), _np(r), rtol=2e-3, atol=2e-3)
+
+
# The real header's names and shapes (GGUF order), at tiny width.
def _lilicorr_arrays(cfg: DFlashConfig) -> dict:
n, hd, hl = cfg.num_hidden_layers, cfg.lilicorr_hidden_size, cfg.lilicorr_layers