diff --git a/CHANGELOG.md b/CHANGELOG.md index 9dab8da62..ce510072d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,13 @@ 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. +- 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. ### Changed @@ -50,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 @@ -61,6 +71,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 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/docs/speculative-decoding.md b/docs/speculative-decoding.md index bcf729813..0b00ebb3f 100644 --- a/docs/speculative-decoding.md +++ b/docs/speculative-decoding.md @@ -17,10 +17,11 @@ 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. Other families, such as gemma-4 with its assistant drafter, -need `--draft-gguf`. +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`. | Flag | Config key | Effect | |------|------------|--------| @@ -40,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 6a2f853a4..9f001d689 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 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 | | @@ -108,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`. | diff --git a/gmlx/_vendor/manifest.json b/gmlx/_vendor/manifest.json index 256ef5364..ae94ca4aa 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"}, @@ -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/commands/cli.py b/gmlx/commands/cli.py index 6f2e51fa9..2d32ee398 100644 --- a/gmlx/commands/cli.py +++ b/gmlx/commands/cli.py @@ -839,20 +839,28 @@ 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 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 + 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_headless_drafter, + find_mtp_companion, + 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_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/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/gmlx/load/arch_table.py b/gmlx/load/arch_table.py index 6b3a2d5ed..cca9c348c 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)"), @@ -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) 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/config_synth.py b/gmlx/load/config_synth.py index 992014071..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) @@ -2474,16 +2478,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 +2501,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 +2513,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 +2543,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/discovery.py b/gmlx/load/discovery.py index 0de5315ed..16ef2f6b1 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,69 @@ 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 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): + 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 @@ -1165,11 +1234,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 +1360,31 @@ 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_MARKERS = ("LiLiCorr",) +_SIBLING_DRAFTER_MARKER_RE = re.compile( + "-(?:" + "|".join(_SIBLING_DRAFTER_MARKERS) + ")(?=-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 +1427,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/load/loader.py b/gmlx/load/loader.py index 822e20b03..7a0e7ee41 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,15 @@ 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"), + # 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"), } # Params kept at their native f16 through the bf16 cast (no upcast). MLX @@ -948,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 @@ -1059,22 +1073,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 +1294,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 +1310,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 +1352,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 +1754,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 +1774,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 +1857,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..c81ad6f3c 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): @@ -2770,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 @@ -2783,22 +2783,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..6322ee780 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)) @@ -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/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/load/vlm.py b/gmlx/load/vlm.py index 712bbf883..78c8d6ff3 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,38 @@ def _qwen3vl_deepstack_blocks(mm_meta: dict) -> list[int]: } +# 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", + "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", + "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"}) + + # 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 +1111,51 @@ def remap_vision_arrays( out[f"vision_tower.blocks.{bid}.{tgt}.{leaf}"] = arr 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 + 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 + 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: + _emit_omni(hit, arr, name) + 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 + _emit_omni(f"vision_tower.blocks.{m.group(1)}.{tgt}.{leaf}", + arr, name) + 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 +1902,78 @@ 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 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``; the fused q/k/v and the post-norm follow + their tensors.""" + 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"), + "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: + 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 +2133,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 +2447,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 +3469,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/_vendor/mlx_lm/models/nemotron_h.py b/gmlx/models/nemotron_h/model.py similarity index 74% rename from gmlx/_vendor/mlx_lm/models/nemotron_h.py rename to gmlx/models/nemotron_h/model.py index 353de36c9..09acbdd0b 100644 --- a/gmlx/_vendor/mlx_lm/models/nemotron_h.py +++ b/gmlx/models/nemotron_h/model.py @@ -1,22 +1,58 @@ -# Copyright © 2025 Apple Inc. +# SPDX-License-Identifier: MIT +# 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. + +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 @@ -136,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) @@ -157,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)] @@ -181,11 +220,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, @@ -209,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( @@ -218,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, [ @@ -227,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): @@ -336,30 +386,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 +424,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 +439,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 +480,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,16 +494,20 @@ 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, 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) @@ -460,8 +520,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 @@ -481,8 +541,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) @@ -528,35 +590,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..04b47e8d1 100644 --- a/gmlx/models/nemotron_h/mtp.py +++ b/gmlx/models/nemotron_h/mtp.py @@ -25,16 +25,17 @@ - ``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, -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 @@ -45,7 +46,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 @@ -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 @@ -92,8 +94,14 @@ def __call__( inputs, cache: Optional[Any] = None, mamba_sink: Optional[list] = None, + capture: Optional[tuple] = None, + input_embeddings: Optional[mx.array] = None, ): - hidden_states = self.embeddings(inputs) + """``capture`` is ``(layer_ids, out)``: the raw output of each listed + 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) 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 @@ -117,93 +125,80 @@ 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, - }) - hidden_states = _mamba_stepwise(layer, hidden_states, mask, 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: + captured.append(hidden_states) 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 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.""" + """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"] - _mamba_stepwise(e["layer"], e["inputs"][:, :n], mask, cache) + if isinstance(mask, mx.array): + mask = mask[:, :n] + e["layer"](e["inputs"][:, :n], mask=mask, cache=cache) 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.""" + _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, + input_embeddings=None) -> mx.array: + ids = self._dflash_capture + if ids is None: + 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), + input_embeddings=input_embeddings) + 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 +229,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 +244,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 +253,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( @@ -304,7 +300,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/nemotron_h/vision.py b/gmlx/models/nemotron_h/vision.py new file mode 100644 index 000000000..c4f534afc --- /dev/null +++ b/gmlx/models/nemotron_h/vision.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Asher Feldman +"""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 +``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 + (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. Q8_0 weights run on the kquant modules. 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 + fused_qkv: bool = False + post_norm: bool = False + + +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, fused_qkv: bool = False): + super().__init__() + self.num_heads = num_heads + self.scale = (dims // num_heads) ** -0.5 + 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: + n, d = x.shape + h = self.num_heads + + 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(q), heads(k), heads(v), 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, config.fused_qkv) + 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)] + 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, + 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) + 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/gmlx/models/nemotron_h/vlm_model.py b/gmlx/models/nemotron_h/vlm_model.py new file mode 100644 index 000000000..92de95e07 --- /dev/null +++ b/gmlx/models/nemotron_h/vlm_model.py @@ -0,0 +1,164 @@ +# 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 ModelArgs +from .mtp import NemotronHSpecLM, _SpecOutput +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(NemotronHSpecLM): + """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) + # The serve engine and the APC guard read the text config here. + self.config = config + + 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. + del mask + if kwargs.get("return_hidden") or kwargs.get("return_shared_kv"): + 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( + 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 29b0fb493..f17eaf469 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 @@ -56,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/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 6314ed71f..549c74f38 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, } @@ -284,7 +285,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, @@ -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" @@ -1838,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``). + 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 + from gmlx.load.discovery import find_headless_drafter, find_mtp_companion 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_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] companion 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/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/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/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 diff --git a/tests/load/test_discovery.py b/tests/load/test_discovery.py index 764341be0..004e186fd 100644 --- a/tests/load/test_discovery.py +++ b/tests/load/test_discovery.py @@ -1172,3 +1172,128 @@ 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 + + +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 + + +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 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..67d26a1af --- /dev/null +++ b/tests/models/test_nemotron_h.py @@ -0,0 +1,278 @@ +#!/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) + + +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) diff --git a/tests/models/test_nemotron_h_vision.py b/tests/models/test_nemotron_h_vision.py new file mode 100644 index 000000000..2dfbf29ad --- /dev/null +++ b/tests/models/test_nemotron_h_vision.py @@ -0,0 +1,497 @@ +#!/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) + + +# 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())] + 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) + 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) + 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))) + 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) + + 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("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, 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)) + 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 _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 + [""]} + 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 + 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()) + + +@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") + 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_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 + + 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) + 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)) + 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_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])) + # 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_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", + "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 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(): diff --git a/tests/spec/test_lilicorr_drafter.py b/tests/spec/test_lilicorr_drafter.py new file mode 100644 index 000000000..499abecec --- /dev/null +++ b/tests/spec/test_lilicorr_drafter.py @@ -0,0 +1,378 @@ +#!/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("capture", [None, [0, 2]], ids=["mtp", "dflash"]) +@pytest.mark.parametrize("accepted", [0, 2, 4]) +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) + 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() + 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) + + +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 + 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"