From bbe16fa56b5dc354da20caabd0fd03d8e20e5380 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 27 Aug 2026 22:43:13 -0700 Subject: [PATCH 1/4] Add remaining dense GGUF graph routes Implement exact explicit-float MiniMax-M2, Mistral4 latent-cache, and GLM-DSA import paths. Correct the pinned Mistral4 DeepSeek-V2 inventory, preserve packed/runtime fail-closed boundaries, and record immutable oversized artifact evidence. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0ee8718a-33b9-4015-9e20-3745bad54686 Signed-off-by: Justin Chu --- docs/api/build_from_gguf.md | 16 +- src/mobius/__main__.py | 8 + src/mobius/_registry.py | 13 + src/mobius/components/_deepseek_mla.py | 60 +- .../integrations/gguf/_arch_registry.py | 129 ++- .../integrations/gguf/_arch_registry_test.py | 5 + .../gguf/_artifact_blocker_evidence.py | 169 ++++ src/mobius/integrations/gguf/_builder.py | 56 +- .../integrations/gguf/_config_mapping.py | 345 +++++++- .../gguf/_conventional_decoder_test.py | 1 - src/mobius/integrations/gguf/_docs.py | 14 +- .../integrations/gguf/_glm_moe_dsa_test.py | 165 +++- .../integrations/gguf/_remaining_dense.py | 492 +++++++++++ .../gguf/_remaining_dense_test.py | 807 ++++++++++++++++++ .../integrations/gguf/_route_census_test.py | 3 + .../integrations/gguf/_tensor_mapping.py | 81 ++ .../integrations/gguf/_tensor_processors.py | 80 ++ .../gguf/_tensor_processors_test.py | 27 + .../gguf/_upstream_data/llamacpp_pin.json | 2 +- src/mobius/models/__init__.py | 4 + src/mobius/models/gguf_minimax_m2.py | 104 +++ src/mobius/models/gguf_mistral4.py | 327 +++++++ src/mobius/models/glm_moe_dsa.py | 50 +- src/mobius/models/glm_moe_dsa_test.py | 73 +- src/mobius/tasks/__init__.py | 3 + src/mobius/tasks/_glm_moe_dsa.py | 17 +- src/mobius/tasks/_mistral4_gguf.py | 86 ++ tests/_test_configs.py | 19 + tests/build_graph_test.py | 2 + tests/gguf_remaining_dense_artifact_test.py | 50 ++ tests/model_coverage_test.py | 6 + tests/synthetic_parity_test.py | 2 + 32 files changed, 3106 insertions(+), 110 deletions(-) create mode 100644 src/mobius/integrations/gguf/_artifact_blocker_evidence.py create mode 100644 src/mobius/integrations/gguf/_remaining_dense.py create mode 100644 src/mobius/integrations/gguf/_remaining_dense_test.py create mode 100644 src/mobius/models/gguf_minimax_m2.py create mode 100644 src/mobius/models/gguf_mistral4.py create mode 100644 src/mobius/tasks/_mistral4_gguf.py create mode 100644 tests/gguf_remaining_dense_artifact_test.py diff --git a/docs/api/build_from_gguf.md b/docs/api/build_from_gguf.md index 4d039bcf3..5812664db 100644 --- a/docs/api/build_from_gguf.md +++ b/docs/api/build_from_gguf.md @@ -9,7 +9,7 @@ Support is capability-specific: graph import does not imply runtime packaging. | Census | Total | Closure | |---|---:|---| -| Architectures | 148 | graph verdicts: {'deferred': 43, 'rejected': 2, 'supported': 103}; importable: 102; quantized import: {'rejected': 44, 'supported': 104}; runtime: {'deferred': 136, 'rejected': 2, 'supported': 10} | +| Architectures | 148 | graph verdicts: {'deferred': 41, 'rejected': 2, 'supported': 105}; importable: 105; quantized import: {'rejected': 47, 'supported': 101}; runtime: {'deferred': 136, 'rejected': 2, 'supported': 10} | | Active stored qtypes | 25 | 24 have an import route; 1 are explicitly deferred with no route | | Serialized projector strings | 60 | {'graph-importable': 9, 'runtime-supported': 0} | | Tokenizer pre identifiers | 87 | 56 semantic groups; route dispositions: {'deferred-compiled-semantics': 45, 'deferred-pinned-artifact-evidence': 4, 'deferred-pinned-artifact-mismatch': 7, 'validated-pinned-source': 31} | @@ -88,12 +88,14 @@ network-free selection, budget, exclusions, and fail-closed candidate reasons ar | `tiny-starcoder2-q2-k-ort-genai-0.15.2` | `tensorblock/tiny-random-starcoder2-GGUF@82c9eb61d1af6ea00dff834f1ff0620144b333e8`
`tiny-random-starcoder2-Q2_K.gguf`
68,039,904 B
`ab0a4b4e79c906520808db065a00a317c8b097e2176638b34088975eada6e0ed` | `hyper-accel/tiny-random-starcoder2@193576733055d2108dc9906d0da6e0806ad9be57` | `hyper-accel/tiny-random-starcoder2@193576733055d2108dc9906d0da6e0806ad9be57`
`special_tokens_map.json` 1,300 B `0fc9ac706a35e6337b19d484bc8f866b6a0ee7ad2509b03f7305dc0838c82d2c`, `tokenizer.json` 2,060,947 B `17fa145258b20c18287f1e3bd804e074cc13333f11984a2f5a2f11c5110437aa`, `tokenizer_config.json` 7,877 B `8149d8e6b21275ad2cc346885ad92c0e9b5aa3e28a78bb5d39b7febc7e52545d`
metadata `38608a4848dc8535113cff0312858853a0c8e29a5cdc145f258e4a9d84b35113` | ONNX Runtime 1.29.0 `CPUExecutionProvider`; ort-genai 0.15.2; result=passed; full-logit; dynamic KV cache prefill, full-sequence replay, rollback, reorder, and 20 decode steps; Explicit-float correctness route; source GGUF blocks are dequantized. The tiny source config declares BOS/EOS ID 50256 outside its 49152-token model vocabulary. | Runtime support above is independent from tokenizer materialization support below. - ### Fail-closed runtime evidence | Evidence ID | Pinned candidate | Bounded result | Withheld runtime claims | |---|---|---|---| | `nemotron-h-moe-30b-iq2-xxs-runtime-blocker` | `bartowski/nvidia_Nemotron-3-Nano-30B-A3B-GGUF@1fc64d5b160654ec892df2708aa893b0e96e6491`
`nvidia_Nemotron-3-Nano-30B-A3B-IQ2_XXS.gguf`
18,010,755,296 B
`f3da710c046ce7cc6ff28a9b5f1a9153ac72e3f60603e51c7bb679d80716b58a` | config/tokenizer `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16@bf77c3174f68ad409e1c2aa60daeb46e32d1c606`; GGUF tokenizer metadata `6089bcaf08b3fe0d49379ca7e85bd3c93e8705bac6130636425c159212971225`; result=blocked; 401 tensors / 31,577,940,288 parameters; ORT 1.29.0 / CPUExecutionProvider; ort-genai 0.15.2; The smallest model GGUF in the pinned 29-file repository revision is 18,010,755,296 bytes, above the 16 GiB bounded-artifact policy; explicit float16/float32 weights require 63,155,880,576/126,311,761,152 bytes.
The normal optimized CPU export has 37,142 nodes (40,167 before Mobius optimization) and 6,028 MatMul nodes because its truthful ReLU2 routed experts remain an ONNX loop. ORT 1.29 MoE/QMoE exposes ReLU but not ReLU2. QMoE has separate router_probs/router_weights for correction-biased selection with unbiased sigmoid mixing, while shared experts and optional latent projections can surround it; those are not fused-op blockers. The pinned 30B has no latent projection. QMoE does not support GGUF IQ2_XXS storage, and Mobius has no proven IQ2_XXS packer/kernel path.
ORT GenAI 0.15.2 discovers sparse/nonconsecutive KV and conv/recurrent slots, but derives recurrent_state names while this export uses ssm_state, does not beam-reorder recurrent state, and rejects nonzero recurrent-state rewind.
The GGUF tokenizer declares pre=pixtral, whose compiled llama.cpp behavior is not serialized; exact ORT tokenizer materialization is unavailable. | cached decode and deterministic generation, full-logit parity, package and report roundtrip, state replay, rollback, and reorder, tensor value closure | +| `glm-dsa-iq1-s-artifact-budget-blocker` | `unsloth/GLM-5.2-GGUF@abc55e72527792c6e77069c99b4cb7de16fa9f23`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00001-of-00006.gguf` 9,423,744 B `46b6148389219ae45167cb8124fbb18ef7d432daf619b4faf9e06ea80d3f4777`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00002-of-00006.gguf` 49,208,128,256 B `f2180207285e04fcaa5b8c53ba6e77ad5cc58666b6e7c6b04a5eded3fe8bef09`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00003-of-00006.gguf` 49,684,417,024 B `b1c0c5a302cc8d5d9ea0bcd4467c01db72c26839f820f7e882079582ea0a8d2b`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00004-of-00006.gguf` 49,396,052,864 B `a6a42da6975e29f89866dcde2956e9e50e6ea26635fb5063b74f3973f4f863b6`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00005-of-00006.gguf` 49,246,275,936 B `a4a9851a50db533f21ef824e5d8038f04e6782e7d602d18e5fdd6643f68ccccb`
`UD-IQ1_S/GLM-5.2-UD-IQ1_S-00006-of-00006.gguf` 19,171,063,136 B `3b767f55df64e0432d52fcf1a14eb47a1ef3bbc91339e2ae220f38602237d7d7`
total 216,715,360,960 B | blocked; The smallest mainline GLM-5.2 GGUF is a 216,715,360,960-byte split set whose trunk also declares an unsupported routed DSA/MLA MTP block. | real-weight full-logit parity, runtime packaging, deterministic generation | +| `minimax-m2-iq1-s-artifact-budget-blocker` | `mradermacher/MiniMax-M2-i1-GGUF@2d4f9b1a86d32ce4dfc47db312c8d6fcae8d7b37`
`MiniMax-M2.i1-IQ1_S.gguf` 46,514,882,176 B `7bae986e3cd380c28c6177d612fce1d52373241a0dfa13a6fca25de79abf15fb`
total 46,514,882,176 B | blocked; The smallest immutable public MiniMax-M2 GGUF is 46,514,882,176 bytes; real-weight parity cannot enter the 16 GiB bounded evidence set. | real-weight full-logit parity, runtime packaging, deterministic generation | +| `mistral4-iq1-m-artifact-budget-blocker` | `unsloth/Mistral-Small-4-119B-2603-GGUF@bd93c721735aa32c035c0f19e738cb3371fd56ff`
`Mistral-Small-4-119B-2603-UD-IQ1_M.gguf` 32,306,941,632 B `40fcdee4869110938638c6b8bac253f442b196518d4623f4afdf4b885cd961c7`
total 32,306,941,632 B | blocked; The smallest immutable public Mistral4 GGUF is 32,306,941,632 bytes; real-weight parity cannot enter the 16 GiB bounded evidence set. | real-weight full-logit parity, runtime packaging, deterministic generation | ## Remaining route work @@ -111,12 +113,10 @@ remain machine-readable in `_route_census.py`; this table groups only shared nex | `dependency-or-runtime-abi-blocked` | `projector-runtime-abi` | `projector:resampler` | dynamic processor-to-graph media shape ABI | | `dependency-or-runtime-abi-blocked` | `tokenizer-compiled-semantics` | `tokenizer:afmoe`, `tokenizer:bloom`, `tokenizer:chameleon`, `tokenizer:codeshell`, `tokenizer:command-r`, `tokenizer:dbrx`, `tokenizer:deepseek-coder`, `tokenizer:deepseek-llm`, `tokenizer:deepseek-v3`, `tokenizer:default`, `tokenizer:exaone`, `tokenizer:exaone-moe`, `tokenizer:falcon`, `tokenizer:gpt3-finnish`, `tokenizer:granite-docling`, `tokenizer:granite-embed-multi-97m`, `tokenizer:grok-2`, `tokenizer:hunyuan`, `tokenizer:hunyuan-dense`, `tokenizer:jais`, `tokenizer:jais-2`, `tokenizer:joyai-llm`, `tokenizer:kimi-k2`, `tokenizer:laguna`, `tokenizer:megrez`, `tokenizer:mellum2`, `tokenizer:minerva-7b`, `tokenizer:minicpm5`, `tokenizer:minimax-m2`, `tokenizer:mpt`, `tokenizer:olmo`, `tokenizer:poro-chat`, `tokenizer:refact`, `tokenizer:sarvam-moe`, `tokenizer:seed-coder`, `tokenizer:smaug-bpe`, `tokenizer:solar-open`, `tokenizer:stablelm2`, `tokenizer:starcoder`, `tokenizer:superbpe`, `tokenizer:tekken`, `tokenizer:trillion`, `tokenizer:viking`, `tokenizer:whitespace`, `tokenizer:youtu` | compiled pinned llama.cpp oracle; dispatch-equivalence fixture | | `dependency-or-runtime-abi-blocked` | `tokenizer-compiled-semantics` | `tokenizer:bailingmoe`, `tokenizer:bailingmoe2`, `tokenizer:chatglm-bpe`, `tokenizer:cohere2moe`, `tokenizer:glm4`, `tokenizer:llada-moe`, `tokenizer:tiny_aya` | upstream tokenizer semantic parity; independently proven replacement reconstruction | -| `evidence-only` | `architecture-runtime-evidence` | `architecture:apertus`, `architecture:arcee`, `architecture:arctic`, `architecture:baichuan`, `architecture:bailingmoe`, `architecture:bert`, `architecture:bitnet`, `architecture:bloom`, `architecture:chatglm`, `architecture:codeshell`, `architecture:cohere2`, `architecture:command-r`, `architecture:dbrx`, `architecture:deci`, `architecture:deepseek`, `architecture:dflash`, `architecture:dots1`, `architecture:dream`, `architecture:eagle3`, `architecture:ernie4_5`, `architecture:ernie4_5-moe`, `architecture:eurobert`, `architecture:exaone`, `architecture:falcon`, `architecture:gemma`, `architecture:gemma-embedding`, `architecture:gemma2`, `architecture:gemma3`, `architecture:gemma4`, `architecture:granite`, `architecture:granitemoe`, `architecture:grok`, `architecture:grovemoe`, `architecture:hunyuan-dense`, `architecture:hunyuan-moe`, `architecture:hy_v3`, `architecture:internlm2`, `architecture:jais`, `architecture:jais2`, `architecture:jina-bert-v2`, `architecture:jina-bert-v3`, `architecture:lfm2moe`, `architecture:llada`, `architecture:llada-moe`, `architecture:llama-embed`, `architecture:maincoder`, `architecture:mamba`, `architecture:mamba2`, `architecture:minicpm`, `architecture:minicpm3`, `architecture:modern-bert`, `architecture:muse-glimmer`, `architecture:nemotron`, `architecture:nemotron_h`, `architecture:neo-bert`, `architecture:nomic-bert`, `architecture:nomic-bert-moe`, `architecture:olmo2`, `architecture:olmoe`, `architecture:openelm`, `architecture:orion`, `architecture:pangu-embedded`, `architecture:phi2`, `architecture:phi3`, `architecture:phimoe`, `architecture:plamo`, `architecture:plm`, `architecture:qwen`, `architecture:qwen2moe`, `architecture:qwen2vl`, `architecture:qwen3`, `architecture:qwen35`, `architecture:qwen3moe`, `architecture:qwen3next`, `architecture:refact`, `architecture:rnd1`, `architecture:seed_oss`, `architecture:smallthinker`, `architecture:smollm3`, `architecture:stablelm`, `architecture:t5`, `architecture:t5encoder`, `architecture:talkie`, `architecture:xverse` | immutable representative GGUF; full-logit prefill and cached-decode parity; deterministic generation/state evidence | +| `evidence-only` | `architecture-runtime-evidence` | `architecture:apertus`, `architecture:arcee`, `architecture:arctic`, `architecture:baichuan`, `architecture:bailingmoe`, `architecture:bert`, `architecture:bitnet`, `architecture:bloom`, `architecture:chatglm`, `architecture:codeshell`, `architecture:cohere2`, `architecture:command-r`, `architecture:dbrx`, `architecture:deci`, `architecture:deepseek`, `architecture:dflash`, `architecture:dots1`, `architecture:dream`, `architecture:eagle3`, `architecture:ernie4_5`, `architecture:ernie4_5-moe`, `architecture:eurobert`, `architecture:exaone`, `architecture:falcon`, `architecture:gemma`, `architecture:gemma-embedding`, `architecture:gemma2`, `architecture:gemma3`, `architecture:gemma4`, `architecture:glm-dsa`, `architecture:granite`, `architecture:granitemoe`, `architecture:grok`, `architecture:grovemoe`, `architecture:hunyuan-dense`, `architecture:hunyuan-moe`, `architecture:hy_v3`, `architecture:internlm2`, `architecture:jais`, `architecture:jais2`, `architecture:jina-bert-v2`, `architecture:jina-bert-v3`, `architecture:lfm2moe`, `architecture:llada`, `architecture:llada-moe`, `architecture:llama-embed`, `architecture:maincoder`, `architecture:mamba`, `architecture:mamba2`, `architecture:minicpm`, `architecture:minicpm3`, `architecture:minimax-m2`, `architecture:mistral4`, `architecture:modern-bert`, `architecture:muse-glimmer`, `architecture:nemotron`, `architecture:nemotron_h`, `architecture:neo-bert`, `architecture:nomic-bert`, `architecture:nomic-bert-moe`, `architecture:olmo2`, `architecture:olmoe`, `architecture:openelm`, `architecture:orion`, `architecture:pangu-embedded`, `architecture:phi2`, `architecture:phi3`, `architecture:phimoe`, `architecture:plamo`, `architecture:plm`, `architecture:qwen`, `architecture:qwen2moe`, `architecture:qwen2vl`, `architecture:qwen3`, `architecture:qwen35`, `architecture:qwen3moe`, `architecture:qwen3next`, `architecture:refact`, `architecture:rnd1`, `architecture:seed_oss`, `architecture:smallthinker`, `architecture:smollm3`, `architecture:stablelm`, `architecture:t5`, `architecture:t5encoder`, `architecture:talkie`, `architecture:xverse` | immutable representative GGUF; full-logit prefill and cached-decode parity; deterministic generation/state evidence | | `evidence-only` | `draft-runtime-evidence` | `draft:dflash`, `draft:eagle3` | target acceptance loop; draft cache orchestration; deterministic speedup parity | | `evidence-only` | `mtp-runtime-evidence` | `mtp:hy_v3`, `mtp:qwen35` | target acceptance loop; cache-threaded draft/target parity | | `evidence-only` | `projector-runtime-evidence` | `projector:adapter`, `projector:gemma3`, `projector:gemma4v`, `projector:ldp`, `projector:ldpv2`, `projector:mlp`, `projector:muse-glimmer`, `projector:qwen2.5vl_merger`, `projector:qwen2vl_merger` | paired text target; processor boundary; deterministic multimodal package execution | -| `immediately-implementable` | `architecture-implementation` | `architecture:minimax-m2`, `architecture:mistral4` | exact metadata extraction; tensor closure; dedicated graph and parity | -| `immediately-implementable` | `architecture-implementation` | `architecture:glm-dsa` | suffix-exact tensor mapping; packed-value transform proof; graph closure | | `immediately-implementable` | `projector-implementation` | `projector:cogvlm`, `projector:deepseekocr`, `projector:deepseekocr2`, `projector:dots3note_a`, `projector:dots3note_v`, `projector:dots_ocr`, `projector:exaone4_5`, `projector:gemma3na`, `projector:gemma3nv`, `projector:gemma4a`, `projector:gemma4ua`, `projector:gemma4uv`, `projector:glm4v`, `projector:glma`, `projector:granite4_vision`, `projector:granite_speech`, `projector:hunyuanvl`, `projector:idefics3`, `projector:internvl`, `projector:janus_pro`, `projector:kimik25`, `projector:kimivl`, `projector:lfm2`, `projector:lfm2a`, `projector:lightonocr`, `projector:llama4`, `projector:meralion`, `projector:mimo_audio`, `projector:mimovl`, `projector:minicpmv4_6`, `projector:minimax_m3`, `projector:musicflamingo`, `projector:nemotron_v2_vl`, `projector:paddleocr`, `projector:parakeet`, `projector:pixtral`, `projector:pockettts_spkenc`, `projector:qwen2.5o`, `projector:qwen2a`, `projector:qwen3a`, `projector:qwen3tts_spkenc`, `projector:qwen3vl_merger`, `projector:step3vl`, `projector:ultravox`, `projector:voxtral`, `projector:yasa2`, `projector:youtuvl` | metadata schema; tensor closure; component graph parity | | `intentionally-rejected` | `policy-rejections` | `architecture:bailingmoe2`, `architecture:clip`, `architecture:dots3note`, `architecture:exaone-moe`, `architecture:exaone4`, `architecture:glm4`, `architecture:glm4moe`, `architecture:gptj` | policy change plus independent correctness proof | | `intentionally-rejected` | `policy-rejections` | `projector:pockettts_gen`, `projector:qwen3tts_gen` | sidecar role must become a valid projector contract | @@ -214,7 +214,7 @@ Reason codes are concise user-facing categories; detailed architecture audits re | `gemma3n` | — | none (fails before config extraction) | exact-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — Gemma3n GGUF is the text member of a vision-and-audio package whose gemma3nv and gemma3na clip companions carry distinct encoders and projectors. | | `gemma4` | — | model=`gemma4_text`; tensor=`llama`+`gemma4_extras`; mmproj=`gemma4` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | RUNTIME_EVIDENCE_PENDING — Config extraction, exact tensor-name closure, and a full synthetic GGUF graph build are covered, but no representative real-weight GGUF has yet passed ORT parity or generation validation. | | `gemma4-assistant` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — Gemma4 Assistant is a standalone target-coupled model with pre/post projections, masked embeddings, scalar layer scales, its own KV cache, and a live target-model context. | -| `glm-dsa` | `glm_dsa` | none (no tensor mapping route) | audited-direct-loader-conditional-union | config=supported; tensor_map=deferred; graph=supported; runtime=deferred; quantized_import=supported | TENSOR_MAP_DEFERRED — Config extraction and the glm_moe_dsa graph are both available, but no GGUF→HuggingFace tensor-name mapping has been written for GLM-5.2's MLA + DSA-indexer tensor families yet, so weights cannot be routed into the graph. | +| `glm-dsa` | `glm_dsa` | model=`glm_moe_dsa`; tensor=`glm_dsa` | audited-direct-loader-conditional-union | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | RUNTIME_EVIDENCE_PENDING / FLOAT_IMPORT_ONLY — Exact explicit-float MLA, DSA indexer, dense-prefix/routed-MoE tensor routing, and packed dynamic-cache graph construction are covered. | | `glm4` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — GLM4 serializes complete fused-FFN trailing blocks and NextN tensors, but the pinned loader skips appended blocks; GLM-OCR converter transforms also permute Q/K for M-RoPE. | | `glm4moe` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — GLM4-MoE serializes biased attention and periodic dense/routed expert trailing blocks with mandatory router bias, but the pinned loader skips them. | | `gpt-oss` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — The pinned GPT-OSS converter splits interleaved gate/up expert rows and repacks checkpoint block+scale tensors into expert-major MXFP4 values. | @@ -256,10 +256,10 @@ Reason codes are concise user-facing categories; detailed architecture audits re | `minicpm` | — | model=`minicpm`; module=`minicpm_gguf`; tensor=`llama`; mmproj=`generic_projector` | audited-direct-loader-conditional-union | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | RUNTIME_EVIDENCE_PENDING — Exact dense graph import, tensor closure, scaling, LongRoPE factors, Q/K transforms, and expanded dynamic K/V cache are covered. | | `minicpm3` | — | model=`minicpm3`; module=`minicpm3_gguf`; tensor=`minicpm3` | audited-direct-loader-conditional-union | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | RUNTIME_EVIDENCE_PENDING — Exact Q/KV-LoRA MLA graph import with expanded K/V cache is covered. | | `minimax-01` | — | model=`minimax`; tensor=`minimax` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | RUNTIME_EVIDENCE_PENDING — Graph import is exact, but released ORT GenAI packaging cannot represent the heterogeneous KV/recurrent state slots or bounded rollback snapshots; runtime packaging remains tracked by #605. | -| `minimax-m2` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — MiniMax-M2 uses full-vector Q/K norms, partial RoPE, and all-layer correction-biased routed experts under metadata-selected gating. | +| `minimax-m2` | — | model=`minimax_m2_gguf`; module=`minimax_m2_gguf`; tensor=`minimax_m2` | audited-direct-loader-conditional-union | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | RUNTIME_EVIDENCE_PENDING / FLOAT_IMPORT_ONLY — Exact explicit-float import owns MiniMax-M2's non-square Q/O geometry, full-vector Q/K RMSNorm, partial NeoX RoPE, standard dynamic/static KV cache, and F32 selection-biased sigmoid MoE routing. | | `minimax-m3` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — MiniMax-M3 adds F32 sparse-indexer tensors and a second index-key cache with position/cell maps, block masks, rollback, and reorder semantics alongside main K/V state. | | `mistral3` | — | none (fails before config extraction) | exact-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — The pinned Mistral3 loader selects dense or routed-expert text blocks from metadata and applies architecture-specific output temperature scaling. | -| `mistral4` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | CONFIG_DEFERRED — Mistral4 has no NextN metadata or MTP graph; it inherits Mistral3's conditional dense/MoE tensor loader and overrides graph construction. | +| `mistral4` | — | model=`mistral4_gguf`; module=`mistral4_gguf`; tensor=`mistral4` | audited-direct-loader-conditional-union | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | RUNTIME_EVIDENCE_PENDING / FLOAT_IMPORT_ONLY — Exact explicit-float import owns Mistral4's DeepSeek-V2 MLA projections, dense-prefix plus mandatory shared/routed MoE blocks, YaRN scaling, and one graph-visible latent-plus-RoPE K cache per layer. | | `modern-bert` | — | model=`modernbert`; tensor=`modern_bert` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | RUNTIME_EVIDENCE_PENDING — Config extraction, exact pinned tensor closure, encoder-only task dispatch, and synthetic ORT execution are covered, but no pinned real GGUF artifact has passed independent embedding parity. | | `mpt` | — | model=`mpt`; module=`gguf_legacy`; tensor=`legacy_layernorm` | audited-direct-loader-conditional-union | config=supported; tensor_map=supported; graph=supported; runtime=supported; quantized_import=rejected | FLOAT_IMPORT_ONLY — Runtime support is restricted to the exact evidenced tiny MPT Q2_K artifact, explicit-float portable graph, pinned tokenizer, and ORT GenAI 0.15.2. | | `muse-glimmer` | `muse_glimmer` | model=`muse_glimmer_text`; tensor=`llama`+`muse_glimmer_extras`; mmproj=`muse_glimmer` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | RUNTIME_EVIDENCE_PENDING — Config extraction, exact tensor-name closure, and a full synthetic GGUF graph build are covered, but no representative real-weight GGUF has yet passed ORT parity or generation validation. | diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 05c477064..16da5afa6 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -180,6 +180,14 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: "--static-cache cannot represent Falcon-H1's per-layer K, V, " "convolution, and SSM states" ) + if model_type == "glm_moe_dsa": + raise ValueError( + "--static-cache cannot represent GLM-DSA's per-layer-varying packed DSA cache" + ) + if model_type == "mistral4_gguf": + raise ValueError( + "--static-cache is not implemented for Mistral4's latent K-only cache" + ) if model_type == "gemma4": from mobius.tasks._gemma4 import Gemma4Task diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 14d846b15..b20f2a66a 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -111,6 +111,8 @@ Llama4CausalLMModel, MageVLForConditionalGeneration, MaincoderCausalLMModel, + MiniMaxM2GGUFCausalLMModel, + Mistral4GGUFCausalLMModel, MoECausalLMModel, MoonshineForConditionalGeneration, NanoChatCausalLMModel, @@ -492,6 +494,17 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: test_model_id="Maincode/Maincoder-1B", test_revision="088ec98640bdeb105f46a9ef6a1370ed5d0d2ea5", ), + "minimax_m2_gguf": ModelRegistration( + MiniMaxM2GGUFCausalLMModel, + family="minimax", + variant="gguf", + ), + "mistral4_gguf": ModelRegistration( + Mistral4GGUFCausalLMModel, + task="mistral4-gguf-text-generation", + family="mistral", + variant="gguf", + ), "bloom": ModelRegistration(BloomCausalLMModel), "orion": ModelRegistration(LayerNormCausalLMModel), "chatglm": ModelRegistration(ChatGLMCausalLMModel), diff --git a/src/mobius/components/_deepseek_mla.py b/src/mobius/components/_deepseek_mla.py index d5d62329a..0271478f4 100644 --- a/src/mobius/components/_deepseek_mla.py +++ b/src/mobius/components/_deepseek_mla.py @@ -45,6 +45,7 @@ def __init__( config: ArchitectureConfig, scale: float | None = None, linear_class: type | None = None, + split_kv_b: bool = False, ): super().__init__() if linear_class is None: @@ -82,12 +83,25 @@ def __init__( bias=False, ) self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) - # Decompresses latent KV into per-head k_nope + v - self.kv_b_proj = linear_class( - self.kv_lora_rank, - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), - bias=False, - ) + self._split_kv_b = split_kv_b + if split_kv_b: + self.k_b_proj = linear_class( + self.kv_lora_rank, + self.num_heads * self.qk_nope_head_dim, + bias=False, + ) + self.v_b_proj = linear_class( + self.kv_lora_rank, + self.num_heads * self.v_head_dim, + bias=False, + ) + else: + # Decompresses latent KV into per-head k_nope + v + self.kv_b_proj = linear_class( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) self.o_proj = linear_class( self.num_heads * self.v_head_dim, @@ -154,18 +168,28 @@ def forward( # Decompress latent KV → per-head k_nope + v k_pass = self.kv_a_layernorm(op, k_pass) - kv_decompressed = self.kv_b_proj(op, k_pass) - # (B, S, num_heads * (nope + v_dim)) → (B, S, num_heads, nope + v_dim) - kv_decompressed = op.Reshape( - kv_decompressed, - [0, 0, self.num_heads, self.qk_nope_head_dim + self.v_head_dim], - ) - k_nope, value_states = op.Split( - kv_decompressed, - [self.qk_nope_head_dim, self.v_head_dim], - axis=-1, - _outputs=2, - ) + if self._split_kv_b: + k_nope = op.Reshape( + self.k_b_proj(op, k_pass), + [0, 0, self.num_heads, self.qk_nope_head_dim], + ) + value_states = op.Reshape( + self.v_b_proj(op, k_pass), + [0, 0, self.num_heads, self.v_head_dim], + ) + else: + kv_decompressed = self.kv_b_proj(op, k_pass) + # (B, S, num_heads * (nope + v_dim)) → (B, S, num_heads, nope + v_dim) + kv_decompressed = op.Reshape( + kv_decompressed, + [0, 0, self.num_heads, self.qk_nope_head_dim + self.v_head_dim], + ) + k_nope, value_states = op.Split( + kv_decompressed, + [self.qk_nope_head_dim, self.v_head_dim], + axis=-1, + _outputs=2, + ) # k_nope: (B, S, H, nope_dim) → (B, S, H*nope_dim)... not needed yet # value_states: (B, S, H, v_dim) → (B, S, H*v_dim) for Attention op value_states = op.Reshape(value_states, [0, 0, -1]) diff --git a/src/mobius/integrations/gguf/_arch_registry.py b/src/mobius/integrations/gguf/_arch_registry.py index 15c6a9707..5082eef9d 100644 --- a/src/mobius/integrations/gguf/_arch_registry.py +++ b/src/mobius/integrations/gguf/_arch_registry.py @@ -505,11 +505,6 @@ "Q/K permutation, optional long/short RoPE tensors, and a conditional dense-or-MoE " "loader. The existing MiniCPM graph does not prove this complete GGUF contract." ), - "minimax-m2": ( - "MiniMax-M2 uses full-vector Q/K norms, partial RoPE, and all-layer " - "correction-biased routed experts under metadata-selected gating. Mobius has no " - "exact graph or suffix-safe expert import for that topology." - ), "minimax-m3": ( "MiniMax-M3 adds F32 sparse-indexer tensors and a second index-key cache with " "position/cell maps, block masks, rollback, and reorder semantics alongside main " @@ -584,11 +579,6 @@ "KV cache, and three chained heads selected by offsets. Mobius permits one head " "and cannot preserve that state or FP8 converter transform." ), - "mistral4": ( - "Mistral4 has no NextN metadata or MTP graph; it inherits Mistral3's conditional " - "dense/MoE tensor loader and overrides graph construction. It is not llama, " - "mistral alias text, or any DeepSeek/Qwen MTP family." - ), "step35": ( "Step3.5 executes one or more interleaved-SWA NextN heads with optional gates, " "routed/shared experts, centered-norm transforms, per-layer head geometry, and " @@ -1334,6 +1324,42 @@ "runtime packaging remains tracked by #605." ), ), + GGUFArchitectureSpec( + gguf_arch="minimax-m2", + model_type="minimax_m2_gguf", + module_type="minimax_m2_gguf", + config_key_map="minimax_m2", + config_postprocessor="minimax_m2", + tensor_map_recipe=("minimax_m2",), + required_metadata=( + "context_length", + "embedding_length", + "feed_forward_length", + "block_count", + "attention.head_count", + "attention.head_count_kv", + "attention.key_length", + "attention.value_length", + "attention.layer_norm_rms_epsilon", + "rope.freq_base", + "rope.dimension_count", + "expert_count", + "expert_used_count", + "expert_feed_forward_length", + "expert_gating_func", + ), + runtime=Support.DEFERRED, + quantized_import=Support.REJECTED, + reason=( + "Exact explicit-float import owns MiniMax-M2's non-square Q/O geometry, " + "full-vector Q/K RMSNorm, partial NeoX RoPE, standard dynamic/static KV cache, " + "and F32 selection-biased sigmoid MoE routing. Loader-accepted Q/K/V biases " + "and fused QKV are rejected because the pinned graph ignores or cannot execute " + "them. Packed projection/expert import and runtime packaging remain deferred; " + "the smallest immutable public GGUF is 46,514,882,176 bytes, above the " + "16 GiB evidence budget. Use keep_quantized=False." + ), + ), GGUFArchitectureSpec( gguf_arch="plamo", model_type="plamo", @@ -2131,6 +2157,48 @@ runtime=Support.DEFERRED, reason=_DEEPSEEK4_GGUF_GRAPH_REASON, ), + GGUFArchitectureSpec( + gguf_arch="mistral4", + model_type="mistral4_gguf", + module_type="mistral4_gguf", + config_key_map="mistral4", + config_postprocessor="mistral4", + tensor_map_recipe=("mistral4",), + tensor_processor="mistral4", + required_metadata=( + "context_length", + "embedding_length", + "feed_forward_length", + "block_count", + "attention.head_count", + "attention.head_count_kv", + "attention.key_length", + "attention.value_length", + "attention.key_length_mla", + "attention.value_length_mla", + "attention.q_lora_rank", + "attention.kv_lora_rank", + "attention.layer_norm_rms_epsilon", + "rope.freq_base", + "rope.dimension_count", + "expert_count", + "expert_used_count", + "expert_feed_forward_length", + "expert_shared_count", + ), + rope_interleave=True, + runtime=Support.DEFERRED, + quantized_import=Support.REJECTED, + reason=( + "Exact explicit-float import owns Mistral4's DeepSeek-V2 MLA projections, " + "dense-prefix plus mandatory shared/routed MoE blocks, YaRN scaling, and one " + "graph-visible latent-plus-RoPE K cache per layer. Legacy non-MLA, Q-LoRA-free, " + "temperature-scaled, and NextN layouts fail closed. Packed import and runtime " + "packaging remain deferred; the smallest immutable public GGUF is " + "32,306,941,632 bytes, above the 16 GiB evidence budget. " + "Use keep_quantized=False." + ), + ), # GLM-5.2 GGUFs (e.g. unsloth/GLM-5.2-GGUF) tag the architecture 'glm-dsa' # (MLA + DeepSeek Sparse Attention + MoE) and mobius's registry key is # 'glm_moe_dsa'. The format bridge is keyed on the authoritative @@ -2142,12 +2210,43 @@ model_type="glm_moe_dsa", aliases=frozenset({"glm_dsa"}), config_key_map="glm_dsa", - tensor_map=Support.DEFERRED, + config_postprocessor="glm_dsa", + tensor_map_recipe=("glm_dsa",), + tensor_processor="glm_dsa", + required_metadata=( + "context_length", + "embedding_length", + "feed_forward_length", + "block_count", + "attention.head_count", + "attention.head_count_kv", + "attention.key_length", + "attention.value_length", + "attention.key_length_mla", + "attention.value_length_mla", + "attention.q_lora_rank", + "attention.kv_lora_rank", + "attention.layer_norm_rms_epsilon", + "rope.freq_base", + "rope.dimension_count", + "expert_count", + "expert_used_count", + "expert_feed_forward_length", + "expert_shared_count", + "attention.indexer.head_count", + "attention.indexer.key_length", + "attention.indexer.top_k", + ), + rope_interleave=True, + quantized_import=Support.REJECTED, reason=( - "Config extraction and the glm_moe_dsa graph are both available, but " - "no GGUF→HuggingFace tensor-name mapping has been written for GLM-5.2's " - "MLA + DSA-indexer tensor families yet, so weights cannot be routed " - "into the graph. " + _NO_TENSOR_MAP + "Exact explicit-float MLA, DSA indexer, dense-prefix/routed-MoE tensor routing, " + "and packed dynamic-cache graph construction are covered. Packed import remains " + "rejected because GLM-5.2's rank-3 K/V-B transforms and correction-biased routed " + "experts do not have an independently proven storage-preserving runtime path. " + "The released GGUF also carries a routed DSA/MLA MTP block outside the current " + "sidecar ABI, and representative real-weight runtime evidence exceeds the 16 GiB " + "artifact budget. Use keep_quantized=False." ), ), GGUFArchitectureSpec( diff --git a/src/mobius/integrations/gguf/_arch_registry_test.py b/src/mobius/integrations/gguf/_arch_registry_test.py index 766e90a9e..6edb68ce5 100644 --- a/src/mobius/integrations/gguf/_arch_registry_test.py +++ b/src/mobius/integrations/gguf/_arch_registry_test.py @@ -85,6 +85,8 @@ "gemma-embedding", "llama-embed", "maincoder", + "minimax-m2", + "mistral4", "starcoder", "xverse", } @@ -328,6 +330,8 @@ def test_every_float_importable_architecture_has_a_quantized_verdict(self) -> No "mamba", "mamba2", "maincoder", + "minimax-m2", + "mistral4", "nemotron_h", "nemotron_h_moe", "neo-bert", @@ -343,6 +347,7 @@ def test_every_float_importable_architecture_has_a_quantized_verdict(self) -> No "grok", "grovemoe", "hunyuan-moe", + "glm-dsa", "hy_v3", "jais", "mpt", diff --git a/src/mobius/integrations/gguf/_artifact_blocker_evidence.py b/src/mobius/integrations/gguf/_artifact_blocker_evidence.py new file mode 100644 index 000000000..d5e1563a2 --- /dev/null +++ b/src/mobius/integrations/gguf/_artifact_blocker_evidence.py @@ -0,0 +1,169 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Immutable oversized-artifact evidence for graph routes using synthetic parity.""" + +from __future__ import annotations + +import dataclasses +import re +from pathlib import PurePosixPath +from types import MappingProxyType + +__all__ = [ + "GGUFArtifactBlockerEvidence", + "GGUFArtifactFile", + "MAX_BOUNDED_ARTIFACT_BYTES", + "artifact_blocker_evidence", + "iter_artifact_blocker_evidence", +] + +MAX_BOUNDED_ARTIFACT_BYTES = 16 * 1024**3 + + +@dataclasses.dataclass(frozen=True, slots=True) +class GGUFArtifactFile: + """One immutable Hub path in a complete logical GGUF candidate.""" + + path: str + size: int + lfs_sha256: str + + def __post_init__(self) -> None: + if ( + not self.path + or PurePosixPath(self.path).is_absolute() + or ".." in PurePosixPath(self.path).parts + ): + raise ValueError("GGUF artifact evidence paths must be safe Hub-relative paths") + if self.size <= 0 or re.fullmatch(r"[0-9a-f]{64}", self.lfs_sha256) is None: + raise ValueError("GGUF artifact evidence requires a positive size and LFS SHA-256") + + +@dataclasses.dataclass(frozen=True, slots=True) +class GGUFArtifactBlockerEvidence: + """The smallest pinned real candidate for one synthetic-only graph route.""" + + evidence_id: str + architecture: str + repository: str + revision: str + files: tuple[GGUFArtifactFile, ...] + blocker: str + + def __post_init__(self) -> None: + if ( + not self.evidence_id + or not self.architecture + or "/" not in self.repository + or not self.blocker.strip() + ): + raise ValueError("GGUF artifact blocker evidence fields must be non-empty") + if re.fullmatch(r"[0-9a-f]{40}", self.revision) is None: + raise ValueError("GGUF artifact blocker evidence requires an immutable revision") + paths = tuple(file.path for file in self.files) + if not paths or paths != tuple(sorted(paths)) or len(paths) != len(set(paths)): + raise ValueError("GGUF artifact blocker files must be sorted and unique") + if self.total_size <= MAX_BOUNDED_ARTIFACT_BYTES: + raise ValueError("Artifact blocker candidates must exceed the 16 GiB policy") + + @property + def total_size(self) -> int: + """Complete logical GGUF size.""" + return sum(file.size for file in self.files) + + +_MINIMAX_M2 = GGUFArtifactBlockerEvidence( + evidence_id="minimax-m2-iq1-s-artifact-budget-blocker", + architecture="minimax-m2", + repository="mradermacher/MiniMax-M2-i1-GGUF", + revision="2d4f9b1a86d32ce4dfc47db312c8d6fcae8d7b37", + files=( + GGUFArtifactFile( + "MiniMax-M2.i1-IQ1_S.gguf", + 46_514_882_176, + "7bae986e3cd380c28c6177d612fce1d52373241a0dfa13a6fca25de79abf15fb", + ), + ), + blocker=( + "The smallest immutable public MiniMax-M2 GGUF is 46,514,882,176 bytes; " + "real-weight parity cannot enter the 16 GiB bounded evidence set." + ), +) + +_MISTRAL4 = GGUFArtifactBlockerEvidence( + evidence_id="mistral4-iq1-m-artifact-budget-blocker", + architecture="mistral4", + repository="unsloth/Mistral-Small-4-119B-2603-GGUF", + revision="bd93c721735aa32c035c0f19e738cb3371fd56ff", + files=( + GGUFArtifactFile( + "Mistral-Small-4-119B-2603-UD-IQ1_M.gguf", + 32_306_941_632, + "40fcdee4869110938638c6b8bac253f442b196518d4623f4afdf4b885cd961c7", + ), + ), + blocker=( + "The smallest immutable public Mistral4 GGUF is 32,306,941,632 bytes; " + "real-weight parity cannot enter the 16 GiB bounded evidence set." + ), +) + +_GLM_DSA = GGUFArtifactBlockerEvidence( + evidence_id="glm-dsa-iq1-s-artifact-budget-blocker", + architecture="glm-dsa", + repository="unsloth/GLM-5.2-GGUF", + revision="abc55e72527792c6e77069c99b4cb7de16fa9f23", + files=( + GGUFArtifactFile( + "UD-IQ1_S/GLM-5.2-UD-IQ1_S-00001-of-00006.gguf", + 9_423_744, + "46b6148389219ae45167cb8124fbb18ef7d432daf619b4faf9e06ea80d3f4777", + ), + GGUFArtifactFile( + "UD-IQ1_S/GLM-5.2-UD-IQ1_S-00002-of-00006.gguf", + 49_208_128_256, + "f2180207285e04fcaa5b8c53ba6e77ad5cc58666b6e7c6b04a5eded3fe8bef09", + ), + GGUFArtifactFile( + "UD-IQ1_S/GLM-5.2-UD-IQ1_S-00003-of-00006.gguf", + 49_684_417_024, + "b1c0c5a302cc8d5d9ea0bcd4467c01db72c26839f820f7e882079582ea0a8d2b", + ), + GGUFArtifactFile( + "UD-IQ1_S/GLM-5.2-UD-IQ1_S-00004-of-00006.gguf", + 49_396_052_864, + "a6a42da6975e29f89866dcde2956e9e50e6ea26635fb5063b74f3973f4f863b6", + ), + GGUFArtifactFile( + "UD-IQ1_S/GLM-5.2-UD-IQ1_S-00005-of-00006.gguf", + 49_246_275_936, + "a4a9851a50db533f21ef824e5d8038f04e6782e7d602d18e5fdd6643f68ccccb", + ), + GGUFArtifactFile( + "UD-IQ1_S/GLM-5.2-UD-IQ1_S-00006-of-00006.gguf", + 19_171_063_136, + "3b767f55df64e0432d52fcf1a14eb47a1ef3bbc91339e2ae220f38602237d7d7", + ), + ), + blocker=( + "The smallest mainline GLM-5.2 GGUF is a 216,715,360,960-byte split set " + "whose trunk also declares an unsupported routed DSA/MLA MTP block." + ), +) + +_ARTIFACT_BLOCKERS = MappingProxyType( + {evidence.evidence_id: evidence for evidence in (_GLM_DSA, _MINIMAX_M2, _MISTRAL4)} +) + + +def artifact_blocker_evidence( + evidence_id: str, +) -> GGUFArtifactBlockerEvidence | None: + """Return one immutable oversized-artifact evidence record.""" + return _ARTIFACT_BLOCKERS.get(evidence_id) + + +def iter_artifact_blocker_evidence() -> tuple[GGUFArtifactBlockerEvidence, ...]: + """Return all oversized graph-route candidates in stable order.""" + return tuple(_ARTIFACT_BLOCKERS[key] for key in sorted(_ARTIFACT_BLOCKERS)) diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 6fcc50e47..114c193c4 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -713,6 +713,11 @@ def _validate_gguf_model( ) validate_remaining_moe_tensor_contract(gguf_model) + from mobius.integrations.gguf._remaining_dense import ( + validate_remaining_dense_tensor_contract, + ) + + validate_remaining_dense_tensor_contract(gguf_model) from mobius.integrations.gguf._hy_v3 import validate_hy_v3_tensor_contract validate_hy_v3_tensor_contract(gguf_model) @@ -6679,7 +6684,15 @@ def _logical_source_filename(reference: str | Path, resolved_path: str | Path) - "encoder_fused_qkv": frozenset({"jina-bert-v3"}), "moe_layer_frequency": frozenset({"ernie4_5-moe", "nomic-bert-moe"}), "routing_weight_normalization_floor": frozenset( - {"dots1", "ernie4_5-moe", "hy_v3", "smallthinker"} + { + "dots1", + "ernie4_5-moe", + "glm-dsa", + "hy_v3", + "minimax-m2", + "mistral4", + "smallthinker", + } ), "router_logit_softcapping": frozenset(), } @@ -7105,6 +7118,33 @@ def build_from_gguf( "smallthinker GGUF only supports the dedicated " "'smallthinker-gguf-text-generation' task" ) + if gguf_arch == "glm-dsa": + from mobius.tasks import GlmMoeDsaTask + + if static_cache: + raise ValueError( + "static_cache=True is not supported for GLM-DSA GGUF; " + "dsa_kv_cache_specs() describes a per-layer-varying packed dynamic cache" + ) + if task is not None and task != "glm-moe-dsa" and not isinstance(task, GlmMoeDsaTask): + raise ValueError("glm-dsa GGUF only supports the dedicated 'glm-moe-dsa' task") + if gguf_arch == "mistral4": + from mobius.tasks import Mistral4GGUFCausalLMTask + + if static_cache: + raise ValueError( + "static_cache=True is not supported for Mistral4 GGUF; " + "the dedicated graph owns a dynamic latent K-only cache" + ) + if ( + task is not None + and task != "mistral4-gguf-text-generation" + and not isinstance(task, Mistral4GGUFCausalLMTask) + ): + raise ValueError( + "mistral4 GGUF only supports the dedicated " + "'mistral4-gguf-text-generation' task" + ) if gguf_arch == "qwen4exp": from mobius.tasks import Qwen4ExpCausalLMTask @@ -8417,11 +8457,11 @@ def _preflight_quantization_report( target_bits=target_bits, target_block_size=target_block_size, ) - is_kimi_reshaped_projection = gguf_arch in { + is_reshaped_mla_projection = gguf_arch in { "kimi-linear", "kimi-k3", } and module_hf_name.endswith((".k_b_proj.weight", ".v_b_proj.weight")) - if is_kimi_reshaped_projection and route is not QuantImportRoute.REJECTED: + if is_reshaped_mla_projection and route is not QuantImportRoute.REJECTED: if quant_spec.dequantize is not Support.SUPPORTED: route = QuantImportRoute.REJECTED exactness = None @@ -8430,7 +8470,7 @@ def _preflight_quantization_report( route = QuantImportRoute.DEQUANTIZE_REQUANTIZE exactness = RepackExactness.LOSSY reason = ( - "The Kimi MLA layout transform changes affine block groups and " + "The MLA layout transform changes affine block groups and " "requires lossy dequantization/requantization." ) disposition = disposition_for_import_route(route, exactness) @@ -9309,11 +9349,11 @@ def _load_quantized_state_dict( # through the custom 130-byte-block parser below rather than the # generic target-splitting path, which assumes mainline Q1_0 bytes. affine_targets = [] - is_kimi_reshaped_projection = gguf_arch in { + is_reshaped_mla_projection = gguf_arch in { "kimi-linear", "kimi-k3", } and module_hf_name.endswith((".k_b_proj.weight", ".v_b_proj.weight")) - if is_kimi_reshaped_projection: + if is_reshaped_mla_projection: # These tensors are rank-3 in GGUF. They target one flattened # projection rather than an expert-major collection. affine_targets = [] @@ -9368,7 +9408,7 @@ def _load_quantized_state_dict( ) if explicitly_dequantized and quant_spec.dequantize is Support.SUPPORTED: route = QuantImportRoute.DEQUANTIZE_FLOAT - if is_kimi_reshaped_projection: + if is_reshaped_mla_projection: if quant_spec.dequantize is not Support.SUPPORTED: raise ValueError( f"Cannot reshape quantized {quant_spec.name} tensor {hf_name}: " @@ -9635,7 +9675,7 @@ def _load_quantized_state_dict( f"block-{repacked.block_size} for {hf_name}, but the graph " f"expects INT{target_bits} block-{target_block_size}." ) - elif is_kimi_reshaped_projection: + elif is_reshaped_mla_projection: values = gguf_model.dequantize_raw_tensor(raw, qtype, np_shape) if hf_name.endswith(".k_b_proj.weight"): values = values.transpose(0, 2, 1).reshape( diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index a4f84c272..0a58f0d06 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -138,10 +138,11 @@ # extracted config matches what GlmMoeDsaCausalLMModel (a DeepSeek-V3 subclass) # consumes. Both spellings of the architecture string are accepted. _GLM_DSA_KEY_MAP = { - "attention.key_length": "head_dim", + "attention.key_length_mla": "head_dim", "rope.dimension_count": "qk_rope_head_dim", "attention.q_lora_rank": "q_lora_rank", "attention.kv_lora_rank": "kv_lora_rank", + "attention.value_length_mla": "v_head_dim", "attention.sliding_window": "sliding_window", "expert_count": "num_local_experts", "expert_used_count": "num_experts_per_tok", @@ -149,11 +150,27 @@ "expert_shared_count": "n_shared_experts", "expert_weights_scale": "routed_scaling_factor", "expert_weights_norm": "norm_topk_prob", + "expert_group_count": "n_group", + "expert_group_used_count": "topk_group", + "leading_dense_block_count": "first_k_dense_replace", "attention.indexer.head_count": "index_n_heads", "attention.indexer.key_length": "index_head_dim", "attention.indexer.top_k": "index_topk", } +_MINIMAX_M2_KEY_MAP = { + "attention.key_length": "head_dim", + "expert_feed_forward_length": "moe_intermediate_size", +} + +_MISTRAL4_KEY_MAP = { + "attention.key_length_mla": "head_dim", + "attention.q_lora_rank": "q_lora_rank", + "attention.kv_lora_rank": "kv_lora_rank", + "expert_feed_forward_length": "moe_intermediate_size", + "leading_dense_block_count": "first_k_dense_replace", +} + _MAMBA_KEY_MAP = { "attention.layer_norm_rms_epsilon": "layer_norm_epsilon", "ssm.conv_kernel": "conv_kernel", @@ -307,6 +324,8 @@ "draft": _DRAFT_KEY_MAP, "muse_glimmer": _MUSE_GLIMMER_KEY_MAP, "glm_dsa": _GLM_DSA_KEY_MAP, + "minimax_m2": _MINIMAX_M2_KEY_MAP, + "mistral4": _MISTRAL4_KEY_MAP, "mamba": _MAMBA_KEY_MAP, "falcon_h1": _FALCON_H1_KEY_MAP, "plamo2": _PLAMO2_KEY_MAP, @@ -4134,6 +4153,322 @@ def _kimi_linear_postprocess( return KimiLinearConfig(**fields) +def _glm_dsa_postprocess( + config: ArchitectureConfig, + metadata: dict[str, Any], + model: Any, +) -> ArchitectureConfig: + """Restore the exact pinned GLM-5.2 MLA, DSA, and routed-MoE config.""" + arch = model.architecture + _validate_closed_rope_scaling_metadata(metadata, arch) + raw_gating = metadata.get(f"{arch}.expert_gating_func") + gating = int(raw_gating) if raw_gating is not None else 2 + if gating != 2: + raise ValueError(f"{arch}.expert_gating_func must be SIGMOID (2), got {gating}") + + qk_dim = int(metadata[f"{arch}.attention.key_length_mla"]) + rope_dim = int(metadata[f"{arch}.rope.dimension_count"]) + nope_dim = qk_dim - rope_dim + value_dim = int(metadata[f"{arch}.attention.value_length_mla"]) + raw_kv_rank = metadata.get(f"{arch}.attention.kv_lora_rank") + kv_rank = int(raw_kv_rank) if raw_kv_rank is not None else None + compressed_key_dim = int(metadata[f"{arch}.attention.key_length"]) + if min(nope_dim, rope_dim, value_dim) <= 0 or (kv_rank is not None and kv_rank <= 0): + raise ValueError("GLM-5.2 requires positive NoPE, RoPE, value, and KV-LoRA dimensions") + if kv_rank is not None and compressed_key_dim != kv_rank + rope_dim: + raise ValueError( + f"{arch}.attention.key_length must equal kv_lora_rank + rope.dimension_count " + f"({kv_rank + rope_dim}), got {compressed_key_dim}" + ) + + dense_prefix = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) + if not 0 <= dense_prefix <= config.num_hidden_layers: + raise ValueError( + f"{arch}.leading_dense_block_count must be in [0, " + f"{config.num_hidden_layers}], got {dense_prefix}" + ) + + indexer_types = _glm_dsa_indexer_types(config, metadata) + names = set(model.tensor_names) + routed_layers = set(range(dense_prefix, config.num_hidden_layers)) + bias_layers = { + layer for layer in routed_layers if f"blk.{layer}.exp_probs_b.bias" in names + } + if bias_layers and bias_layers != routed_layers: + raise ValueError( + f"{arch} correction bias must be present for every routed layer or none; " + f"found {sorted(bias_layers)}, expected {sorted(routed_layers)}" + ) + route_scale = float(metadata.get(f"{arch}.expert_weights_scale", 0.0)) + if math.isclose(route_scale, 0.0): + route_scale = 1.0 + if not math.isfinite(route_scale) or route_scale <= 0: + raise ValueError(f"{arch}.expert_weights_scale must resolve to a positive value") + + fields = _shallow_fields(config) + fields.update( + model_type="glm_moe_dsa", + num_key_value_heads=config.num_attention_heads, + head_dim=nope_dim, + q_lora_rank=( + int(metadata[f"{arch}.attention.q_lora_rank"]) + if f"{arch}.attention.q_lora_rank" in metadata + else None + ), + kv_lora_rank=kv_rank, + qk_nope_head_dim=nope_dim, + qk_rope_head_dim=rope_dim, + v_head_dim=value_dim, + intermediate_size=int(metadata[f"{arch}.feed_forward_length"]), + moe_intermediate_size=config.moe_intermediate_size, + n_shared_experts=config.n_shared_experts, + first_k_dense_replace=dense_prefix, + n_group=int(metadata.get(f"{arch}.expert_group_count", 1)), + topk_group=int(metadata.get(f"{arch}.expert_group_used_count", 1)), + routed_scaling_factor=route_scale, + norm_topk_prob=bool(metadata.get(f"{arch}.expert_weights_norm")), + routing_weight_normalization_floor=( + 6.103515625e-5 if bool(metadata.get(f"{arch}.expert_weights_norm")) else None + ), + hidden_act="silu", + scoring_func="sigmoid", + topk_method="noaux_tc", + use_expert_bias=bool(bias_layers), + disable_qmoe=True, + rope_interleave=True, + indexer_rope_interleave=True, + indexer_types=indexer_types, + index_topk_freq=4 if config.max_position_embeddings >= 1_048_576 else 1, + index_skip_topk_offset=3 if config.max_position_embeddings >= 1_048_576 else 0, + partial_rotary_factor=1.0, + ) + return ArchitectureConfig(**fields) + + +def _glm_dsa_indexer_types( + config: ArchitectureConfig, + metadata: dict[str, Any], +) -> list[str]: + """Resolve the pinned bool/scalar indexer schedule without model-ID heuristics.""" + arch = "glm-dsa" + layers = config.num_hidden_layers + raw = metadata.get(f"{arch}.attention.indexer.types") + if raw is None: + if config.max_position_embeddings < 1_048_576: + result = ["full"] * layers + else: + if layers > 78: + raise ValueError( + "glm-dsa models with more than 78 trunk layers must serialize " + "attention.indexer.types" + ) + result = [ + "full" if index < 3 or (index - 2) % 4 == 0 else "shared" + for index in range(layers) + ] + elif isinstance(raw, (list, tuple, np.ndarray)): + if len(raw) != layers: + raise ValueError( + f"{arch}.attention.indexer.types has wrong array length; " + f"expected {layers}, got {len(raw)}" + ) + if any(value not in (0, 1, False, True) for value in raw): + raise ValueError(f"{arch}.attention.indexer.types entries must be bool or 0/1") + result = ["full" if bool(value) else "shared" for value in raw] + else: + if raw not in (0, 1, False, True): + raise ValueError(f"{arch}.attention.indexer.types scalar must be bool or 0/1") + result = ["full" if bool(raw) else "shared"] * layers + if not result or result[0] != "full": + raise ValueError("glm-dsa layer 0 must own a full indexer") + return result + + +def _minimax_m2_postprocess( + config: ArchitectureConfig, + metadata: dict[str, Any], + model: Any, +) -> ArchitectureConfig: + """Restore MiniMax-M2's full-vector Q/K norms and exact sigmoid router.""" + arch = model.architecture + _validate_closed_rope_scaling_metadata(metadata, arch) + gating = metadata[f"{arch}.expert_gating_func"] + if isinstance(gating, bool) or not isinstance(gating, (int, np.integer)): + raise TypeError(f"{arch}.expert_gating_func must be the integer SIGMOID enum") + if int(gating) != 2: + raise ValueError(f"{arch}.expert_gating_func must be SIGMOID (2), got {gating}") + + head_dim = int(metadata[f"{arch}.attention.key_length"]) + value_dim = int(metadata[f"{arch}.attention.value_length"]) + rope_dim = int(metadata[f"{arch}.rope.dimension_count"]) + intermediate = int(metadata[f"{arch}.feed_forward_length"]) + expert_intermediate = int(metadata[f"{arch}.expert_feed_forward_length"]) + if value_dim != head_dim: + raise ValueError(f"{arch}.attention.value_length must equal key_length") + if rope_dim <= 0 or rope_dim > head_dim or rope_dim % 2: + raise ValueError( + f"{arch}.rope.dimension_count must be positive, even, and <= head_dim" + ) + if expert_intermediate != intermediate: + raise ValueError( + f"{arch}.expert_feed_forward_length must equal feed_forward_length " + f"({intermediate}), got {expert_intermediate}" + ) + if int(metadata.get(f"{arch}.expert_shared_count", 0)): + raise ValueError("MiniMax-M2 does not support shared experts") + if int(metadata.get(f"{arch}.leading_dense_block_count", 0)): + raise ValueError("MiniMax-M2 uses routed experts in every layer") + if int(metadata.get(f"{arch}.attention.sliding_window", 0)): + raise ValueError("MiniMax-M2 does not use sliding-window attention") + if int(metadata.get(f"{arch}.expert_group_count", 0)) not in (0, 1): + raise ValueError("MiniMax-M2 does not use grouped expert selection") + if int(metadata.get(f"{arch}.nextn_predict_layers", 0)): + raise ValueError("MiniMax-M2 GGUF does not define an executable NextN graph") + route_scale = float(metadata.get(f"{arch}.expert_weights_scale", 0.0)) + if math.isclose(route_scale, 0.0): + route_scale = 1.0 + if not math.isfinite(route_scale) or route_scale <= 0: + raise ValueError("MiniMax-M2 expert_weights_scale must resolve to a positive value") + + return dataclasses.replace( + config, + model_type="minimax_m2_gguf", + head_dim=head_dim, + intermediate_size=intermediate, + moe_intermediate_size=intermediate, + hidden_act="silu", + tie_word_embeddings=False, + attn_qkv_bias=False, + attn_o_bias=False, + mlp_bias=False, + attn_qk_norm=True, + attn_qk_norm_full=True, + rope_type="default", + rope_interleave=False, + partial_rotary_factor=rope_dim / head_dim, + n_group=1, + topk_group=1, + norm_topk_prob=True, + routed_scaling_factor=route_scale, + routing_weight_normalization_floor=6.103515625e-5, + scoring_func="sigmoid", + topk_method="noaux_tc", + use_expert_bias=True, + disable_qmoe=True, + ) + + +def _mistral4_postprocess( + config: ArchitectureConfig, + metadata: dict[str, Any], + model: Any, +) -> ArchitectureConfig: + """Restore the pinned DeepSeek-V2 MLA/MoE contract serialized as Mistral4.""" + arch = model.architecture + _validate_conventional_moe_rope_scaling(metadata, arch) + q_lora_rank = int(metadata[f"{arch}.attention.q_lora_rank"]) + kv_lora_rank = int(metadata[f"{arch}.attention.kv_lora_rank"]) + qk_head_dim = int(metadata[f"{arch}.attention.key_length_mla"]) + rope_dim = int(metadata[f"{arch}.rope.dimension_count"]) + nope_dim = qk_head_dim - rope_dim + value_dim = int(metadata[f"{arch}.attention.value_length_mla"]) + serialized_key = int(metadata[f"{arch}.attention.key_length"]) + serialized_value = int(metadata[f"{arch}.attention.value_length"]) + if min(q_lora_rank, kv_lora_rank, nope_dim, rope_dim, value_dim) <= 0: + raise ValueError("Mistral4 requires positive Q/KV-LoRA and MLA head dimensions") + if serialized_key != kv_lora_rank + rope_dim or serialized_value != kv_lora_rank: + raise ValueError( + "Mistral4 compressed cache geometry must satisfy " + "key_length=kv_lora_rank+rope_dim and value_length=kv_lora_rank" + ) + if int(config.num_key_value_heads) != 1: + raise ValueError("Mistral4 GGUF must serialize attention.head_count_kv=1") + if int(metadata.get(f"{arch}.nextn_predict_layers", 0)): + raise ValueError("Mistral4's pinned graph does not execute a NextN sidecar") + temperature_scale = float(metadata.get(f"{arch}.attention.temperature_scale", 0.0)) + if not math.isclose(temperature_scale, 0.0): + raise ValueError( + "Mistral4 attention.temperature_scale is loader-optional but is not emitted " + "by the pinned converter; nonzero temperature scaling is outside this route" + ) + + layers = config.num_hidden_layers + dense_prefix = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) + if not 0 <= dense_prefix < layers: + raise ValueError("Mistral4 requires a valid dense prefix and at least one MoE layer") + experts = int(config.num_local_experts or 0) + top_k = int(config.num_experts_per_tok or 0) + shared = int(config.n_shared_experts or 0) + n_group = int(metadata.get(f"{arch}.expert_group_count", 1)) + topk_group = int(metadata.get(f"{arch}.expert_group_used_count", 1)) + if min(experts, top_k, shared) <= 0 or top_k > experts: + raise ValueError("Mistral4 requires routed experts, top-k, and shared experts") + if n_group != 1 or topk_group != 1: + raise ValueError("Mistral4's proven route requires one expert group") + + raw_gating = metadata.get(f"{arch}.expert_gating_func", 1) + if isinstance(raw_gating, bool) or not isinstance(raw_gating, (int, np.integer)): + raise TypeError(f"{arch}.expert_gating_func must be an integer enum") + gating = int(raw_gating) + if gating not in (1, 2): + raise ValueError(f"{arch}.expert_gating_func must be SOFTMAX (1) or SIGMOID (2)") + + routed_layers = set(range(dense_prefix, layers)) + names = set(model.tensor_names) + bias_layers = { + layer for layer in routed_layers if f"blk.{layer}.exp_probs_b.bias" in names + } + if bias_layers and bias_layers != routed_layers: + raise ValueError( + f"Mistral4 correction bias must be present in every routed layer or none; " + f"found {sorted(bias_layers)}, expected {sorted(routed_layers)}" + ) + route_scale = float(metadata.get(f"{arch}.expert_weights_scale", 0.0)) + if math.isclose(route_scale, 0.0): + route_scale = 1.0 + if not math.isfinite(route_scale) or route_scale <= 0: + raise ValueError("Mistral4 expert_weights_scale must resolve to a positive value") + norm_topk_prob = bool(metadata.get(f"{arch}.expert_weights_norm")) + + rope_scaling = None if config.rope_scaling is None else dict(config.rope_scaling) + yarn_log_multiplier = metadata.get(f"{arch}.rope.scaling.yarn_log_multiplier") + if yarn_log_multiplier is not None: + if rope_scaling is None or config.rope_type != "yarn": + raise ValueError("Mistral4 yarn_log_multiplier requires YaRN rope metadata") + rope_scaling["mscale"] = 1.0 + rope_scaling["mscale_all_dim"] = float(yarn_log_multiplier) / 0.1 + + return dataclasses.replace( + config, + model_type="mistral4_gguf", + head_dim=qk_head_dim, + num_key_value_heads=1, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=nope_dim, + qk_rope_head_dim=rope_dim, + v_head_dim=value_dim, + first_k_dense_replace=dense_prefix, + hidden_act="silu", + attn_qkv_bias=False, + attn_o_bias=False, + mlp_bias=False, + rope_scaling=rope_scaling, + rope_interleave=True, + partial_rotary_factor=None, + scoring_func="softmax" if gating == 1 else "sigmoid", + topk_method="greedy", + n_group=n_group, + topk_group=topk_group, + use_expert_bias=bool(bias_layers), + norm_topk_prob=norm_topk_prob, + routed_scaling_factor=route_scale, + routing_weight_normalization_floor=(6.103515625e-5 if norm_topk_prob else None), + disable_qmoe=True, + num_nextn_predict_layers=0, + ) + + def _kimi_k3_postprocess( config: ArchitectureConfig, metadata: dict[str, Any], @@ -4964,6 +5299,9 @@ def _qwen4exp_postprocess( "maincoder": _maincoder_postprocess, "t5": _t5_postprocess, "minimax": _minimax_postprocess, + "minimax_m2": _minimax_m2_postprocess, + "mistral4": _mistral4_postprocess, + "glm_dsa": _glm_dsa_postprocess, "kimi_linear": _kimi_linear_postprocess, "kimi_k3": _kimi_k3_postprocess, "minicpm": _minicpm_postprocess, @@ -5097,6 +5435,11 @@ def assert_glm_moe_dsa_resolvable( reasons.append(f"hidden_size must be > 0 (got {config.hidden_size!r})") # MoE expert stack. + if config.scoring_func != "sigmoid": + reasons.append( + "missing SIGMOID expert gate (GGUF '.expert_gating_func'=2); " + f"scoring_func={config.scoring_func!r}" + ) if not _positive(config.num_local_experts): reasons.append( "missing routed-expert count (GGUF '.expert_count'); " diff --git a/src/mobius/integrations/gguf/_conventional_decoder_test.py b/src/mobius/integrations/gguf/_conventional_decoder_test.py index ad2b5534a..cf518f442 100644 --- a/src/mobius/integrations/gguf/_conventional_decoder_test.py +++ b/src/mobius/integrations/gguf/_conventional_decoder_test.py @@ -393,7 +393,6 @@ def test_pangu_embedded_tied_output_has_single_embedding_owner() -> None: @pytest.mark.parametrize( ("architecture", "reason_fragments"), [ - ("mistral4", ("conditional dense/MoE", "overrides graph construction")), ("plamo3", ("fused QKV", "periodic full/sliding", "iSWA cache ABI")), ], ) diff --git a/src/mobius/integrations/gguf/_docs.py b/src/mobius/integrations/gguf/_docs.py index a5184895d..3d45a118b 100644 --- a/src/mobius/integrations/gguf/_docs.py +++ b/src/mobius/integrations/gguf/_docs.py @@ -21,6 +21,9 @@ _RUNTIME_VALIDATION_PENDING, iter_arch_specs, ) +from mobius.integrations.gguf._artifact_blocker_evidence import ( + iter_artifact_blocker_evidence, +) from mobius.integrations.gguf._mmproj_registry import ( MMPROJ_ARTIFACT_AVAILABILITY_PINS, MMPROJ_ARTIFACT_PINS, @@ -359,6 +362,16 @@ def _runtime_blocker_evidence_table() -> str: f"ORT {evidence.onnxruntime_version} / {evidence.execution_provider}; " f"{evidence.runtime} {evidence.runtime_version}; {blockers} | {withheld} |" ) + for artifact in iter_artifact_blocker_evidence(): + files = "
".join( + f"`{file.path}` {file.size:,} B `{file.lfs_sha256}`" for file in artifact.files + ) + rows.append( + f"| `{artifact.evidence_id}` | `{artifact.repository}@{artifact.revision}`
" + f"{files}
total {artifact.total_size:,} B | blocked; " + f"{artifact.blocker} | real-weight full-logit parity, runtime packaging, " + "deterministic generation |" + ) return "\n".join(rows) @@ -686,7 +699,6 @@ def render_document() -> str: {_runtime_evidence_table()} Runtime support above is independent from tokenizer materialization support below. - ### Fail-closed runtime evidence {_runtime_blocker_evidence_table()} diff --git a/src/mobius/integrations/gguf/_glm_moe_dsa_test.py b/src/mobius/integrations/gguf/_glm_moe_dsa_test.py index 6d6d598ca..123d8876a 100644 --- a/src/mobius/integrations/gguf/_glm_moe_dsa_test.py +++ b/src/mobius/integrations/gguf/_glm_moe_dsa_test.py @@ -18,6 +18,8 @@ from __future__ import annotations +import dataclasses + import pytest @@ -32,26 +34,43 @@ def get_metadata(self, key, default=None): @property def tensor_names(self) -> list[str]: - return ["output.weight", "blk.0.attn_q.weight"] + return [ + "output.weight", + "blk.0.attn_q.weight", + *(f"blk.{layer}.exp_probs_b.bias" for layer in range(3, 78)), + ] def _valid_glm_dsa_metadata() -> dict: return { - "glm-dsa.embedding_length": 5120, - "glm-dsa.block_count": 92, - "glm-dsa.attention.head_count": 96, - "glm-dsa.attention.head_count_kv": 96, + "glm-dsa.embedding_length": 6144, + "glm-dsa.context_length": 1048576, + "glm-dsa.block_count": 79, + "glm-dsa.nextn_predict_layers": 1, + "glm-dsa.attention.head_count": 64, + "glm-dsa.attention.head_count_kv": 1, + "glm-dsa.attention.layer_norm_rms_epsilon": 1e-5, "glm-dsa.feed_forward_length": 12288, - "glm-dsa.vocab_size": 151552, - "glm-dsa.expert_count": 160, + "glm-dsa.vocab_size": 154880, + "glm-dsa.expert_count": 256, "glm-dsa.expert_used_count": 8, - "glm-dsa.expert_feed_forward_length": 1536, + "glm-dsa.expert_feed_forward_length": 2048, "glm-dsa.expert_shared_count": 1, - "glm-dsa.attention.q_lora_rank": 1536, + "glm-dsa.expert_gating_func": 2, + "glm-dsa.expert_group_count": 1, + "glm-dsa.expert_group_used_count": 1, + "glm-dsa.expert_weights_norm": True, + "glm-dsa.expert_weights_scale": 2.5, + "glm-dsa.leading_dense_block_count": 3, + "glm-dsa.attention.q_lora_rank": 2048, "glm-dsa.attention.kv_lora_rank": 512, - "glm-dsa.attention.key_length": 128, + "glm-dsa.attention.key_length": 576, + "glm-dsa.attention.key_length_mla": 256, + "glm-dsa.attention.value_length": 512, + "glm-dsa.attention.value_length_mla": 256, "glm-dsa.rope.dimension_count": 64, - "glm-dsa.attention.indexer.head_count": 64, + "glm-dsa.rope.freq_base": 8_000_000.0, + "glm-dsa.attention.indexer.head_count": 32, "glm-dsa.attention.indexer.key_length": 128, "glm-dsa.attention.indexer.top_k": 2048, } @@ -85,6 +104,71 @@ def test_glm_dsa_config_resolves_to_glm_moe_dsa(): assert model_type == "glm_moe_dsa" +def test_glm_dsa_config_matches_official_checkpoint_geometry(): + from mobius.integrations.gguf._config_mapping import gguf_to_config + + config = gguf_to_config(_FakeGlmDsaModel(_valid_glm_dsa_metadata())) + + assert config.num_hidden_layers == 78 + assert config.num_attention_heads == 64 + assert config.num_key_value_heads == 64 + assert config.first_k_dense_replace == 3 + assert config.q_lora_rank == 2048 + assert config.kv_lora_rank == 512 + assert config.qk_nope_head_dim == 192 + assert config.qk_rope_head_dim == 64 + assert config.v_head_dim == 256 + assert config.scoring_func == "sigmoid" + assert config.topk_method == "noaux_tc" + assert config.use_expert_bias is True + assert config.index_topk_freq == 4 + assert config.index_skip_topk_offset == 3 + assert len(config.indexer_types) == 78 + assert config.indexer_types[:10] == [ + "full", + "full", + "full", + "shared", + "shared", + "shared", + "full", + "shared", + "shared", + "shared", + ] + + +@pytest.mark.parametrize( + ("gguf_name", "hf_name"), + [ + ("blk.4.attn_k_b.weight", "model.layers.4.self_attn.k_b_proj.weight"), + ("blk.4.attn_v_b.weight", "model.layers.4.self_attn.v_b_proj.weight"), + ("blk.4.indexer.attn_k.weight", "model.layers.4.self_attn.indexer.wk.weight"), + ("blk.4.indexer.attn_q_b.weight", "model.layers.4.self_attn.indexer.wq_b.weight"), + ( + "blk.4.indexer.proj.weight", + "model.layers.4.self_attn.indexer.weights_proj.weight", + ), + ( + "blk.4.ffn_gate_exps.weight", + "model.layers.4.mlp.experts.gate_proj.weight", + ), + ( + "blk.4.ffn_down_shexp.weight", + "model.layers.4.mlp.shared_experts.down_proj.weight", + ), + ( + "blk.4.exp_probs_b.bias", + "model.layers.4.mlp.gate.e_score_correction_bias", + ), + ], +) +def test_glm_dsa_tensor_mapping(gguf_name, hf_name): + from mobius.integrations.gguf._tensor_mapping import map_gguf_to_hf_names + + assert map_gguf_to_hf_names(gguf_name, "glm-dsa") == hf_name + + # --------------------------------------------------------------------------- # # assert_glm_moe_dsa_resolvable — valid / invalid # --------------------------------------------------------------------------- # @@ -101,48 +185,42 @@ def test_valid_glm_dsa_config_passes(): def test_missing_expert_count_rejected(): - from mobius.integrations.gguf._config_mapping import ( - GgufArchResolutionError, - assert_glm_moe_dsa_resolvable, - gguf_to_config, - ) + from mobius.integrations.gguf._config_mapping import gguf_to_config md = _valid_glm_dsa_metadata() del md["glm-dsa.expert_count"] + with pytest.raises(ValueError, match=r"(?i)expert"): + gguf_to_config(_FakeGlmDsaModel(md)) + + +def test_missing_sigmoid_gate_uses_pinned_loader_default(): + from mobius.integrations.gguf._config_mapping import gguf_to_config + + md = _valid_glm_dsa_metadata() + del md["glm-dsa.expert_gating_func"] config = gguf_to_config(_FakeGlmDsaModel(md)) - with pytest.raises(GgufArchResolutionError, match=r"(?i)expert"): - assert_glm_moe_dsa_resolvable(config, "glm-dsa", source="no_experts.gguf") + assert config.scoring_func == "sigmoid" def test_missing_mla_rank_rejected(): - from mobius.integrations.gguf._config_mapping import ( - GgufArchResolutionError, - assert_glm_moe_dsa_resolvable, - gguf_to_config, - ) + from mobius.integrations.gguf._config_mapping import gguf_to_config md = _valid_glm_dsa_metadata() del md["glm-dsa.attention.q_lora_rank"] del md["glm-dsa.attention.kv_lora_rank"] - config = gguf_to_config(_FakeGlmDsaModel(md)) - with pytest.raises(GgufArchResolutionError, match=r"(?i)MLA|latent|lora"): - assert_glm_moe_dsa_resolvable(config, "glm-dsa", source="no_mla.gguf") + with pytest.raises(ValueError, match=r"(?i)q_lora|kv_lora"): + gguf_to_config(_FakeGlmDsaModel(md)) def test_missing_dsa_indexer_rejected(): - from mobius.integrations.gguf._config_mapping import ( - GgufArchResolutionError, - assert_glm_moe_dsa_resolvable, - gguf_to_config, - ) + from mobius.integrations.gguf._config_mapping import gguf_to_config md = _valid_glm_dsa_metadata() del md["glm-dsa.attention.indexer.head_count"] del md["glm-dsa.attention.indexer.key_length"] del md["glm-dsa.attention.indexer.top_k"] - config = gguf_to_config(_FakeGlmDsaModel(md)) - with pytest.raises(GgufArchResolutionError, match=r"(?i)DSA|indexer"): - assert_glm_moe_dsa_resolvable(config, "glm-dsa", source="no_dsa.gguf") + with pytest.raises(ValueError, match=r"(?i)indexer"): + gguf_to_config(_FakeGlmDsaModel(md)) def test_rejection_lists_all_reasons(): @@ -154,15 +232,16 @@ def test_rejection_lists_all_reasons(): gguf_to_config, ) - md = { - "glm-dsa.embedding_length": 4096, - "glm-dsa.block_count": 32, - "glm-dsa.attention.head_count": 32, - "glm-dsa.attention.head_count_kv": 8, - "glm-dsa.feed_forward_length": 11008, - "glm-dsa.vocab_size": 128000, - } - config = gguf_to_config(_FakeGlmDsaModel(md)) + config = gguf_to_config(_FakeGlmDsaModel(_valid_glm_dsa_metadata())) + config = dataclasses.replace( + config, + num_local_experts=None, + q_lora_rank=None, + kv_lora_rank=None, + index_n_heads=None, + index_head_dim=None, + index_topk=None, + ) with pytest.raises(GgufArchResolutionError) as excinfo: assert_glm_moe_dsa_resolvable(config, "glm-dsa", source="bare_decoder.gguf") message = str(excinfo.value) diff --git a/src/mobius/integrations/gguf/_remaining_dense.py b/src/mobius/integrations/gguf/_remaining_dense.py new file mode 100644 index 000000000..37f4302ed --- /dev/null +++ b/src/mobius/integrations/gguf/_remaining_dense.py @@ -0,0 +1,492 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Strict tensor closure for MiniMax-M2, Mistral4, and GLM-DSA GGUF graphs.""" + +from __future__ import annotations + +import re + +from mobius._configs import ArchitectureConfig +from mobius.integrations.gguf._tensor_mapping import is_known_skip + + +def _tensor_shapes(model) -> dict[str, tuple[int, ...]]: + return { + name: tuple(int(dimension) for dimension in shape) + for name, _raw, _qtype, shape in model.tensor_items_raw() + } + + +def _validate_exact_closure( + architecture: str, + shapes: dict[str, tuple[int, ...]], + required: dict[str, tuple[int, ...]], + optional: dict[str, tuple[int, ...]], + *, + layers: int, +) -> None: + actual = set(shapes) + allowed = set(required) | set(optional) + unexpected = sorted(name for name in actual - allowed if not is_known_skip(name)) + out_of_range = sorted( + name + for name in actual + if (match := re.match(r"^blk\.(\d+)\.", name)) and int(match.group(1)) >= layers + ) + missing = sorted(set(required) - actual) + malformed = { + name: (expected, shapes.get(name)) + for name, expected in {**required, **optional}.items() + if name in shapes and shapes[name] != expected + } + if missing or unexpected or malformed or out_of_range: + raise ValueError( + f"Invalid {architecture} GGUF tensor closure: missing={missing}, " + f"unexpected={unexpected}, malformed={malformed}, " + f"out_of_range={out_of_range}" + ) + + +def _require_float_auxiliaries(model, names: set[str], *, architecture: str) -> None: + non_float = sorted( + name + for name, _raw, qtype, _shape in model.tensor_items_raw() + if name in names and getattr(qtype, "name", "") not in {"F32", "F16", "BF16"} + ) + if non_float: + raise ValueError( + f"{architecture} normalization/router sidecars must use float storage: {non_float}" + ) + + +def _validate_minimax_m2(model) -> None: + metadata = model.metadata + arch = "minimax-m2" + layers = int(metadata[f"{arch}.block_count"]) + hidden = int(metadata[f"{arch}.embedding_length"]) + heads = int(metadata[f"{arch}.attention.head_count"]) + kv_heads = int(metadata[f"{arch}.attention.head_count_kv"]) + head_dim = int(metadata[f"{arch}.attention.key_length"]) + value_dim = int(metadata[f"{arch}.attention.value_length"]) + intermediate = int(metadata[f"{arch}.feed_forward_length"]) + expert_width = int(metadata[f"{arch}.expert_feed_forward_length"]) + experts = int(metadata[f"{arch}.expert_count"]) + top_k = int(metadata[f"{arch}.expert_used_count"]) + vocab = int(metadata.get(f"{arch}.vocab_size", 0)) or len( + metadata.get("tokenizer.ggml.tokens", ()) + ) + if ( + min( + layers, + hidden, + heads, + kv_heads, + head_dim, + intermediate, + experts, + top_k, + vocab, + ) + <= 0 + or heads % kv_heads + or value_dim != head_dim + or expert_width != intermediate + or top_k > experts + ): + raise ValueError("MiniMax-M2 GGUF has inconsistent architecture geometry") + + actual = set(model.tensor_names) + fused = sorted(name for name in actual if ".attn_qkv." in name) + ignored_biases = sorted( + name for name in actual if name.endswith(("attn_q.bias", "attn_k.bias", "attn_v.bias")) + ) + if fused: + raise ValueError( + "MiniMax-M2 fused QKV is accepted by the pinned loader but its graph " + f"cannot execute it: {fused}" + ) + if ignored_biases: + raise ValueError( + "MiniMax-M2 Q/K/V biases are accepted by the pinned loader but ignored " + f"by its graph: {ignored_biases}" + ) + + q_width = heads * head_dim + kv_width = kv_heads * head_dim + required: dict[str, tuple[int, ...]] = { + "token_embd.weight": (vocab, hidden), + "output_norm.weight": (hidden,), + "output.weight": (vocab, hidden), + } + float_auxiliaries = {"output_norm.weight"} + for layer in range(layers): + prefix = f"blk.{layer}." + required.update( + { + prefix + "attn_norm.weight": (hidden,), + prefix + "attn_q.weight": (q_width, hidden), + prefix + "attn_k.weight": (kv_width, hidden), + prefix + "attn_v.weight": (kv_width, hidden), + prefix + "attn_q_norm.weight": (q_width,), + prefix + "attn_k_norm.weight": (kv_width,), + prefix + "attn_output.weight": (hidden, q_width), + prefix + "ffn_norm.weight": (hidden,), + prefix + "ffn_gate_inp.weight": (experts, hidden), + prefix + "ffn_gate_exps.weight": ( + experts, + intermediate, + hidden, + ), + prefix + "ffn_up_exps.weight": ( + experts, + intermediate, + hidden, + ), + prefix + "ffn_down_exps.weight": ( + experts, + hidden, + intermediate, + ), + prefix + "exp_probs_b.bias": (experts,), + } + ) + float_auxiliaries.update( + { + prefix + "attn_norm.weight", + prefix + "attn_q_norm.weight", + prefix + "attn_k_norm.weight", + prefix + "ffn_norm.weight", + prefix + "ffn_gate_inp.weight", + prefix + "exp_probs_b.bias", + } + ) + shapes = _tensor_shapes(model) + _validate_exact_closure(arch, shapes, required, {}, layers=layers) + _require_float_auxiliaries(model, float_auxiliaries, architecture=arch) + + +def _validate_mistral4(model) -> None: + metadata = model.metadata + arch = "mistral4" + nextn = int(metadata.get(f"{arch}.nextn_predict_layers", 0)) + if nextn: + raise ValueError("Mistral4's pinned graph does not execute NextN tensors") + layers = int(metadata[f"{arch}.block_count"]) + hidden = int(metadata[f"{arch}.embedding_length"]) + dense_width = int(metadata[f"{arch}.feed_forward_length"]) + heads = int(metadata[f"{arch}.attention.head_count"]) + kv_heads = int(metadata[f"{arch}.attention.head_count_kv"]) + q_lora = int(metadata[f"{arch}.attention.q_lora_rank"]) + kv_lora = int(metadata[f"{arch}.attention.kv_lora_rank"]) + qk_dim = int(metadata[f"{arch}.attention.key_length_mla"]) + rope_dim = int(metadata[f"{arch}.rope.dimension_count"]) + nope_dim = qk_dim - rope_dim + value_dim = int(metadata[f"{arch}.attention.value_length_mla"]) + dense_prefix = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) + experts = int(metadata[f"{arch}.expert_count"]) + top_k = int(metadata[f"{arch}.expert_used_count"]) + expert_width = int(metadata[f"{arch}.expert_feed_forward_length"]) + shared = int(metadata[f"{arch}.expert_shared_count"]) + vocab = int(metadata.get(f"{arch}.vocab_size", 0)) or len( + metadata.get("tokenizer.ggml.tokens", ()) + ) + if ( + min( + layers, + hidden, + dense_width, + heads, + q_lora, + kv_lora, + nope_dim, + rope_dim, + value_dim, + experts, + top_k, + expert_width, + shared, + vocab, + ) + <= 0 + or kv_heads != 1 + or int(metadata[f"{arch}.attention.key_length"]) != kv_lora + rope_dim + or int(metadata[f"{arch}.attention.value_length"]) != kv_lora + or top_k > experts + or not 0 <= dense_prefix < layers + ): + raise ValueError("Mistral4 GGUF has inconsistent MLA/MoE geometry") + + required: dict[str, tuple[int, ...]] = { + "token_embd.weight": (vocab, hidden), + "output_norm.weight": (hidden,), + } + optional: dict[str, tuple[int, ...]] = { + "output.weight": (vocab, hidden), + } + float_auxiliaries = {"output_norm.weight"} + actual = set(model.tensor_names) + for layer in range(layers): + prefix = f"blk.{layer}." + required.update( + { + prefix + "attn_norm.weight": (hidden,), + prefix + "attn_q_a_norm.weight": (q_lora,), + prefix + "attn_kv_a_norm.weight": (kv_lora,), + prefix + "attn_q_a.weight": (q_lora, hidden), + prefix + "attn_q_b.weight": (heads * qk_dim, q_lora), + prefix + "attn_kv_a_mqa.weight": ( + kv_lora + rope_dim, + hidden, + ), + prefix + "attn_k_b.weight": (heads, kv_lora, nope_dim), + prefix + "attn_v_b.weight": (heads, value_dim, kv_lora), + prefix + "attn_output.weight": (hidden, heads * value_dim), + prefix + "ffn_norm.weight": (hidden,), + } + ) + float_auxiliaries.update( + { + prefix + "attn_norm.weight", + prefix + "attn_q_a_norm.weight", + prefix + "attn_kv_a_norm.weight", + prefix + "ffn_norm.weight", + } + ) + if layer < dense_prefix: + required.update( + { + prefix + "ffn_gate.weight": (dense_width, hidden), + prefix + "ffn_up.weight": (dense_width, hidden), + prefix + "ffn_down.weight": (hidden, dense_width), + } + ) + continue + + required.update( + { + prefix + "ffn_gate_inp.weight": (experts, hidden), + prefix + "ffn_down_exps.weight": ( + experts, + hidden, + expert_width, + ), + prefix + "ffn_gate_shexp.weight": ( + expert_width * shared, + hidden, + ), + prefix + "ffn_up_shexp.weight": ( + expert_width * shared, + hidden, + ), + prefix + "ffn_down_shexp.weight": ( + hidden, + expert_width * shared, + ), + } + ) + fused = prefix + "ffn_gate_up_exps.weight" + gate = prefix + "ffn_gate_exps.weight" + up = prefix + "ffn_up_exps.weight" + has_fused = fused in actual + has_separate = gate in actual or up in actual + if has_fused == has_separate or (has_separate and not {gate, up} <= actual): + raise ValueError( + f"Mistral4 layer {layer} must contain exactly one fused or split " + "expert gate/up representation" + ) + if has_fused: + required[fused] = (experts, 2 * expert_width, hidden) + else: + required[gate] = (experts, expert_width, hidden) + required[up] = (experts, expert_width, hidden) + optional[prefix + "exp_probs_b.bias"] = (experts,) + float_auxiliaries.add(prefix + "ffn_gate_inp.weight") + if prefix + "exp_probs_b.bias" in actual: + float_auxiliaries.add(prefix + "exp_probs_b.bias") + + shapes = _tensor_shapes(model) + _validate_exact_closure(arch, shapes, required, optional, layers=layers) + _require_float_auxiliaries(model, float_auxiliaries, architecture=arch) + + +def _validate_glm_dsa(model) -> None: + metadata = model.metadata + arch = "glm-dsa" + nextn = int(metadata.get(f"{arch}.nextn_predict_layers", 0)) + total_layers = int(metadata[f"{arch}.block_count"]) + layers = total_layers - nextn + if layers <= 0: + raise ValueError("GLM-DSA block_count must exceed nextn_predict_layers") + hidden = int(metadata[f"{arch}.embedding_length"]) + dense_width = int(metadata[f"{arch}.feed_forward_length"]) + heads = int(metadata[f"{arch}.attention.head_count"]) + q_lora = int(metadata[f"{arch}.attention.q_lora_rank"]) + kv_lora = int(metadata[f"{arch}.attention.kv_lora_rank"]) + qk_dim = int(metadata[f"{arch}.attention.key_length_mla"]) + rope_dim = int(metadata[f"{arch}.rope.dimension_count"]) + nope_dim = qk_dim - rope_dim + value_dim = int(metadata[f"{arch}.attention.value_length_mla"]) + dense_prefix = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) + experts = int(metadata[f"{arch}.expert_count"]) + top_k = int(metadata[f"{arch}.expert_used_count"]) + expert_width = int(metadata[f"{arch}.expert_feed_forward_length"]) + shared = int(metadata[f"{arch}.expert_shared_count"]) + index_heads = int(metadata[f"{arch}.attention.indexer.head_count"]) + index_dim = int(metadata[f"{arch}.attention.indexer.key_length"]) + vocab = int(metadata.get(f"{arch}.vocab_size", 0)) or len( + metadata.get("tokenizer.ggml.tokens", ()) + ) + if ( + min( + hidden, + dense_width, + heads, + q_lora, + kv_lora, + nope_dim, + rope_dim, + value_dim, + experts, + top_k, + expert_width, + shared, + index_heads, + index_dim, + vocab, + ) + <= 0 + ): + raise ValueError("GLM-DSA GGUF has inconsistent MLA/DSA/MoE geometry") + if ( + int(metadata[f"{arch}.attention.head_count_kv"]) != 1 + or int(metadata[f"{arch}.attention.key_length"]) != kv_lora + rope_dim + or int(metadata[f"{arch}.attention.value_length"]) != kv_lora + or top_k > experts + or not 0 <= dense_prefix < layers + ): + raise ValueError("GLM-DSA GGUF has invalid expert or dense-prefix geometry") + + from mobius.integrations.gguf._config_mapping import _glm_dsa_indexer_types + + indexer_types = _glm_dsa_indexer_types( + ArchitectureConfig( + num_hidden_layers=layers, + max_position_embeddings=int(metadata[f"{arch}.context_length"]), + ), + metadata, + ) + required: dict[str, tuple[int, ...]] = { + "token_embd.weight": (vocab, hidden), + "output_norm.weight": (hidden,), + } + optional: dict[str, tuple[int, ...]] = { + "output.weight": (vocab, hidden), + } + actual = set(model.tensor_names) + float_auxiliaries = {"output_norm.weight"} + for layer in range(layers): + prefix = f"blk.{layer}." + required.update( + { + prefix + "attn_norm.weight": (hidden,), + prefix + "attn_q_a_norm.weight": (q_lora,), + prefix + "attn_kv_a_norm.weight": (kv_lora,), + prefix + "attn_q_a.weight": (q_lora, hidden), + prefix + "attn_q_b.weight": (heads * qk_dim, q_lora), + prefix + "attn_kv_a_mqa.weight": ( + kv_lora + rope_dim, + hidden, + ), + prefix + "attn_k_b.weight": (heads, kv_lora, nope_dim), + prefix + "attn_v_b.weight": (heads, value_dim, kv_lora), + prefix + "attn_output.weight": (hidden, heads * value_dim), + prefix + "ffn_norm.weight": (hidden,), + } + ) + float_auxiliaries.update( + { + prefix + "attn_norm.weight", + prefix + "attn_q_a_norm.weight", + prefix + "attn_kv_a_norm.weight", + prefix + "ffn_norm.weight", + } + ) + if indexer_types[layer] == "full": + required.update( + { + prefix + "indexer.k_norm.weight": (index_dim,), + prefix + "indexer.k_norm.bias": (index_dim,), + prefix + "indexer.proj.weight": (index_heads, hidden), + prefix + "indexer.attn_k.weight": (index_dim, hidden), + prefix + "indexer.attn_q_b.weight": ( + index_heads * index_dim, + q_lora, + ), + } + ) + float_auxiliaries.update( + { + prefix + "indexer.k_norm.weight", + prefix + "indexer.k_norm.bias", + } + ) + if layer < dense_prefix: + required.update( + { + prefix + "ffn_gate.weight": (dense_width, hidden), + prefix + "ffn_up.weight": (dense_width, hidden), + prefix + "ffn_down.weight": (hidden, dense_width), + } + ) + continue + required.update( + { + prefix + "ffn_gate_inp.weight": (experts, hidden), + prefix + "ffn_gate_exps.weight": ( + experts, + expert_width, + hidden, + ), + prefix + "ffn_up_exps.weight": ( + experts, + expert_width, + hidden, + ), + prefix + "ffn_down_exps.weight": ( + experts, + hidden, + expert_width, + ), + prefix + "ffn_gate_shexp.weight": ( + expert_width * shared, + hidden, + ), + prefix + "ffn_up_shexp.weight": ( + expert_width * shared, + hidden, + ), + prefix + "ffn_down_shexp.weight": ( + hidden, + expert_width * shared, + ), + } + ) + optional[prefix + "exp_probs_b.bias"] = (experts,) + float_auxiliaries.add(prefix + "ffn_gate_inp.weight") + if prefix + "exp_probs_b.bias" in actual: + float_auxiliaries.add(prefix + "exp_probs_b.bias") + + shapes = _tensor_shapes(model) + _validate_exact_closure(arch, shapes, required, optional, layers=layers) + _require_float_auxiliaries(model, float_auxiliaries, architecture=arch) + + +def validate_remaining_dense_tensor_contract(model) -> None: + """Validate the exact executable tensor subset for the three promoted routes.""" + if model.architecture == "minimax-m2": + _validate_minimax_m2(model) + elif model.architecture == "mistral4": + _validate_mistral4(model) + elif model.architecture == "glm-dsa": + _validate_glm_dsa(model) diff --git a/src/mobius/integrations/gguf/_remaining_dense_test.py b/src/mobius/integrations/gguf/_remaining_dense_test.py new file mode 100644 index 000000000..abed5d266 --- /dev/null +++ b/src/mobius/integrations/gguf/_remaining_dense_test.py @@ -0,0 +1,807 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Executable coverage for the remaining MiniMax-M2, Mistral4, and GLM-DSA routes.""" + +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import onnx_ir as ir +import pytest +import torch +from onnxscript import GraphBuilder + +from mobius._builder import build_from_module +from mobius._constants import OPSET_VERSION +from mobius._registry import registry +from mobius._testing import make_config +from mobius._testing.ort_inference import OnnxModelSession +from mobius.integrations.gguf import build_from_gguf +from mobius.integrations.gguf._arch_registry import get_arch_spec +from mobius.integrations.gguf._builder import _serialize_route_graph_config +from mobius.integrations.gguf._config_mapping import gguf_to_config +from mobius.integrations.gguf._remaining_dense import ( + validate_remaining_dense_tensor_contract, +) +from mobius.integrations.gguf._spec import Support +from mobius.integrations.gguf._tensor_mapping import map_gguf_to_hf_names +from mobius.integrations.gguf._upstream import upstream_architectures +from mobius.models.deepseek import DeepSeekV3CausalLMModel +from mobius.models.gguf_minimax_m2 import ( + MiniMaxM2Gate, + MiniMaxM2GGUFCausalLMModel, +) +from mobius.models.gguf_mistral4 import ( + Mistral4GGUFCausalLMModel, + Mistral4LatentAttention, +) +from mobius.tasks import CausalLMTask + + +@dataclass +class _FakeGGUF: + architecture: str + metadata: dict[str, object] + tensors: dict[str, np.ndarray] + + def __post_init__(self) -> None: + self.tensor_names = list(self.tensors) + self.qtypes = { + name: SimpleNamespace(value=0, name="F32") for name in self.tensor_names + } + + def get_metadata(self, key: str, default=None): + return self.metadata.get(key, default) + + def tensor_items_raw(self): + for name, tensor in self.tensors.items(): + yield name, None, self.qtypes[name], tensor.shape + + +def _values( + rng: np.random.Generator, + shape: tuple[int, ...], + *, + norm: bool = False, +) -> np.ndarray: + if norm: + return rng.uniform(0.7, 1.3, shape).astype(np.float32) + return (rng.standard_normal(shape) * 0.08).astype(np.float32) + + +def _m2_fixture() -> _FakeGGUF: + arch = "minimax-m2" + hidden, layers, heads, kv_heads, head_dim = 8, 1, 2, 1, 8 + intermediate, experts, top_k, vocab = 6, 3, 2, 16 + metadata: dict[str, object] = { + f"{arch}.context_length": 16, + f"{arch}.embedding_length": hidden, + f"{arch}.feed_forward_length": intermediate, + f"{arch}.block_count": layers, + f"{arch}.attention.head_count": heads, + f"{arch}.attention.head_count_kv": kv_heads, + f"{arch}.attention.key_length": head_dim, + f"{arch}.attention.value_length": head_dim, + f"{arch}.attention.layer_norm_rms_epsilon": 1e-6, + f"{arch}.rope.freq_base": 5_000_000.0, + f"{arch}.rope.dimension_count": 4, + f"{arch}.vocab_size": vocab, + f"{arch}.expert_count": experts, + f"{arch}.expert_used_count": top_k, + f"{arch}.expert_feed_forward_length": intermediate, + f"{arch}.expert_gating_func": 2, + } + rng = np.random.default_rng(11) + q_width = heads * head_dim + kv_width = kv_heads * head_dim + tensors = { + "token_embd.weight": _values(rng, (vocab, hidden)), + "output_norm.weight": _values(rng, (hidden,), norm=True), + "output.weight": _values(rng, (vocab, hidden)), + "blk.0.attn_norm.weight": _values(rng, (hidden,), norm=True), + "blk.0.attn_q.weight": _values(rng, (q_width, hidden)), + "blk.0.attn_k.weight": _values(rng, (kv_width, hidden)), + "blk.0.attn_v.weight": _values(rng, (kv_width, hidden)), + "blk.0.attn_q_norm.weight": np.linspace(0.3, 1.7, q_width, dtype=np.float32), + "blk.0.attn_k_norm.weight": np.linspace(0.6, 1.4, kv_width, dtype=np.float32), + "blk.0.attn_output.weight": _values(rng, (hidden, q_width)), + "blk.0.ffn_norm.weight": _values(rng, (hidden,), norm=True), + "blk.0.ffn_gate_inp.weight": _values(rng, (experts, hidden)), + "blk.0.ffn_gate_exps.weight": _values(rng, (experts, intermediate, hidden)), + "blk.0.ffn_up_exps.weight": _values(rng, (experts, intermediate, hidden)), + "blk.0.ffn_down_exps.weight": _values(rng, (experts, hidden, intermediate)), + "blk.0.exp_probs_b.bias": np.array([0.8, -0.5, 0.1], dtype=np.float32), + } + return _FakeGGUF(arch, metadata, tensors) + + +def _m4_fixture() -> _FakeGGUF: + arch = "mistral4" + hidden, layers, heads = 8, 2, 2 + dense_width, q_lora, kv_lora = 6, 4, 4 + nope, rope, value_dim = 2, 2, 2 + experts, top_k, expert_width, vocab = 2, 1, 3, 16 + metadata: dict[str, object] = { + f"{arch}.context_length": 16, + f"{arch}.embedding_length": hidden, + f"{arch}.feed_forward_length": dense_width, + f"{arch}.block_count": layers, + f"{arch}.attention.head_count": heads, + f"{arch}.attention.head_count_kv": 1, + f"{arch}.attention.key_length": kv_lora + rope, + f"{arch}.attention.value_length": kv_lora, + f"{arch}.attention.key_length_mla": nope + rope, + f"{arch}.attention.value_length_mla": value_dim, + f"{arch}.attention.q_lora_rank": q_lora, + f"{arch}.attention.kv_lora_rank": kv_lora, + f"{arch}.attention.layer_norm_rms_epsilon": 1e-5, + f"{arch}.rope.freq_base": 10_000.0, + f"{arch}.rope.dimension_count": rope, + f"{arch}.vocab_size": vocab, + f"{arch}.expert_count": experts, + f"{arch}.expert_used_count": top_k, + f"{arch}.expert_feed_forward_length": expert_width, + f"{arch}.expert_shared_count": 1, + f"{arch}.expert_weights_norm": True, + f"{arch}.leading_dense_block_count": 1, + } + rng = np.random.default_rng(17) + tensors: dict[str, np.ndarray] = { + "token_embd.weight": _values(rng, (vocab, hidden)), + "output_norm.weight": _values(rng, (hidden,), norm=True), + "output.weight": _values(rng, (vocab, hidden)), + } + for layer in range(layers): + prefix = f"blk.{layer}." + tensors.update( + { + prefix + "attn_norm.weight": _values(rng, (hidden,), norm=True), + prefix + "attn_q_a_norm.weight": _values(rng, (q_lora,), norm=True), + prefix + "attn_kv_a_norm.weight": _values(rng, (kv_lora,), norm=True), + prefix + "attn_q_a.weight": _values(rng, (q_lora, hidden)), + prefix + "attn_q_b.weight": _values(rng, (heads * (nope + rope), q_lora)), + prefix + "attn_kv_a_mqa.weight": _values(rng, (kv_lora + rope, hidden)), + # GGUF stores K-B as [heads, kv_lora, nope]. + prefix + "attn_k_b.weight": _values(rng, (heads, kv_lora, nope)), + prefix + "attn_v_b.weight": _values(rng, (heads, value_dim, kv_lora)), + prefix + "attn_output.weight": _values(rng, (hidden, heads * value_dim)), + prefix + "ffn_norm.weight": _values(rng, (hidden,), norm=True), + } + ) + tensors.update( + { + "blk.0.ffn_gate.weight": _values(rng, (dense_width, hidden)), + "blk.0.ffn_up.weight": _values(rng, (dense_width, hidden)), + "blk.0.ffn_down.weight": _values(rng, (hidden, dense_width)), + "blk.1.ffn_gate_inp.weight": _values(rng, (experts, hidden)), + "blk.1.ffn_gate_up_exps.weight": _values(rng, (experts, 2 * expert_width, hidden)), + "blk.1.ffn_down_exps.weight": _values(rng, (experts, hidden, expert_width)), + "blk.1.ffn_gate_shexp.weight": _values(rng, (expert_width, hidden)), + "blk.1.ffn_up_shexp.weight": _values(rng, (expert_width, hidden)), + "blk.1.ffn_down_shexp.weight": _values(rng, (hidden, expert_width)), + } + ) + return _FakeGGUF(arch, metadata, tensors) + + +def _glm_fixture() -> _FakeGGUF: + source = _m4_fixture() + arch = "glm-dsa" + source.architecture = arch + source.metadata = { + key.replace("mistral4.", f"{arch}."): value for key, value in source.metadata.items() + } + source.metadata.update( + { + f"{arch}.attention.indexer.head_count": 2, + f"{arch}.attention.indexer.key_length": 4, + f"{arch}.attention.indexer.top_k": 4, + f"{arch}.attention.indexer.types": [True, False], + f"{arch}.expert_gating_func": 2, + f"{arch}.expert_weights_scale": 2.5, + } + ) + rng = np.random.default_rng(23) + hidden = int(source.metadata[f"{arch}.embedding_length"]) + q_lora = int(source.metadata[f"{arch}.attention.q_lora_rank"]) + experts = int(source.metadata[f"{arch}.expert_count"]) + expert_width = int(source.metadata[f"{arch}.expert_feed_forward_length"]) + source.tensors.update( + { + "blk.0.indexer.k_norm.weight": _values(rng, (4,), norm=True), + "blk.0.indexer.k_norm.bias": _values(rng, (4,)), + "blk.0.indexer.proj.weight": _values(rng, (2, hidden)), + "blk.0.indexer.attn_k.weight": _values(rng, (4, hidden)), + "blk.0.indexer.attn_q_b.weight": _values(rng, (8, q_lora)), + "blk.1.exp_probs_b.bias": _values(rng, (experts,)), + "blk.1.ffn_gate_exps.weight": _values(rng, (experts, expert_width, hidden)), + "blk.1.ffn_up_exps.weight": _values(rng, (experts, expert_width, hidden)), + } + ) + source.tensors.pop("blk.1.ffn_gate_up_exps.weight") + source.tensor_names = list(source.tensors) + source.qtypes = { + name: SimpleNamespace(value=0, name="F32") for name in source.tensor_names + } + return source + + +def _write_gguf(path: Path, source: _FakeGGUF) -> None: + from gguf import GGUFWriter + + writer = GGUFWriter(str(path), source.architecture) + for key, value in source.metadata.items(): + if isinstance(value, bool): + writer.add_bool(key, value) + elif isinstance(value, (int, np.integer)): + writer.add_uint32(key, int(value)) + elif isinstance(value, (float, np.floating)): + writer.add_float32(key, float(value)) + else: + writer.add_array(key, value) + for name, tensor in source.tensors.items(): + writer.add_tensor(name, tensor) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + +def _run_component( + inputs: list[ir.Value], + output, + feeds: dict[str, np.ndarray], +) -> list[np.ndarray]: + graph = ir.Graph( + inputs=inputs, + outputs=[], + nodes=[], + name="remaining_dense_component", + opset_imports={"": OPSET_VERSION}, + ) + builder = GraphBuilder(graph) + realized = output(builder.op) + outputs = realized if isinstance(realized, tuple) else (realized,) + for index, value in enumerate(outputs): + value.name = f"output_{index}" + graph.outputs.append(value) + return list(OnnxModelSession(ir.Model(graph, ir_version=11)).run(feeds).values()) + + +def test_specs_promote_only_explicit_float_graph_routes() -> None: + expected = { + "minimax-m2": ("minimax_m2_gguf", MiniMaxM2GGUFCausalLMModel), + "mistral4": ("mistral4_gguf", Mistral4GGUFCausalLMModel), + "glm-dsa": ("glm_moe_dsa", registry.get("glm_moe_dsa")), + } + for architecture, (model_type, model_class) in expected.items(): + spec = get_arch_spec(architecture) + assert spec.is_importable + assert spec.runtime is Support.DEFERRED + assert spec.quantized_import is Support.REJECTED + assert spec.model_type == model_type + assert registry.get(spec.module_type or spec.model_type) is model_class + + +def test_routing_floor_fingerprint_is_isolated_to_consuming_routes() -> None: + config = make_config(routing_weight_normalization_floor=6.103515625e-5) + changed = dataclasses.replace( + config, + routing_weight_normalization_floor=0.25, + ) + for architecture in ("glm-dsa", "minimax-m2", "mistral4"): + assert _serialize_route_graph_config( + config, + architecture, + ) != _serialize_route_graph_config(changed, architecture) + assert _serialize_route_graph_config( + config, + "llama", + ) == _serialize_route_graph_config(changed, "llama") + + +def test_pinned_mistral4_inventory_names_the_deepseek2_loader() -> None: + record = upstream_architectures()["mistral4"] + assert record.loader_source == "src/models/deepseek2.cpp" + assert record.loader_helpers == () + assert "blk.{bid}.attn_kv_a_mqa.weight" in record.tensor_names + assert "blk.{bid}.ffn_gate_up_exps.weight" in record.tensor_names + assert "blk.{bid}.attn_qkv.weight" not in record.tensor_names + assert "rope_factors_long.weight" not in record.tensor_names + + +@pytest.mark.parametrize( + ("architecture", "gguf_name", "hf_name"), + [ + ( + "minimax-m2", + "blk.4.exp_probs_b.bias", + "model.layers.4.mlp.gate.e_score_correction_bias", + ), + ( + "minimax-m2", + "blk.4.ffn_gate_exps.weight", + "model.layers.4.mlp.experts.gate_proj.weight", + ), + ( + "mistral4", + "blk.4.attn_k_b.weight", + "model.layers.4.self_attn.k_b_proj.weight", + ), + ( + "mistral4", + "blk.4.ffn_gate_up_exps.weight", + "model.layers.4.mlp.experts.gate_up_proj.weight", + ), + ( + "glm-dsa", + "blk.4.indexer.proj.weight", + "model.layers.4.self_attn.indexer.weights_proj.weight", + ), + ], +) +def test_suffix_exact_tensor_mapping( + architecture: str, + gguf_name: str, + hf_name: str, +) -> None: + assert map_gguf_to_hf_names(gguf_name, architecture) == hf_name + + +def test_config_extraction_restores_architecture_discriminators() -> None: + m2 = gguf_to_config(_m2_fixture()) + assert m2.model_type == "minimax_m2_gguf" + assert m2.head_dim == 8 + assert m2.partial_rotary_factor == pytest.approx(0.5) + assert m2.attn_qk_norm and m2.attn_qk_norm_full + assert m2.scoring_func == "sigmoid" + assert m2.routing_weight_normalization_floor == pytest.approx(6.103515625e-5) + + m4 = gguf_to_config(_m4_fixture()) + assert m4.model_type == "mistral4_gguf" + assert m4.num_key_value_heads == 1 + assert (m4.qk_nope_head_dim, m4.qk_rope_head_dim, m4.v_head_dim) == (2, 2, 2) + assert m4.first_k_dense_replace == 1 + assert m4.scoring_func == "softmax" + + glm = gguf_to_config(_glm_fixture()) + assert glm.model_type == "glm_moe_dsa" + assert glm.indexer_types == ["full", "shared"] + assert glm.use_expert_bias + assert glm.scoring_func == "sigmoid" + + +def test_mistral4_restores_pinned_yarn_log_multiplier() -> None: + source = _m4_fixture() + source.metadata.update( + { + "mistral4.rope.scaling.type": "yarn", + "mistral4.rope.scaling.factor": 4.0, + "mistral4.rope.scaling.original_context_length": 8, + "mistral4.rope.scaling.yarn_beta_fast": 32.0, + "mistral4.rope.scaling.yarn_beta_slow": 1.0, + "mistral4.rope.scaling.yarn_log_multiplier": 0.0707, + } + ) + config = gguf_to_config(source) + assert config.rope_type == "yarn" + assert config.rope_scaling is not None + assert config.rope_scaling["mscale"] == pytest.approx(1.0) + assert config.rope_scaling["mscale_all_dim"] == pytest.approx(0.707) + expected_scale = (1.0 + 0.1 * np.log(4.0)) ** 2 / np.sqrt(4.0) + assert Mistral4LatentAttention(config).scaling == pytest.approx(expected_scale) + + +def test_glm_dsa_indexer_schedule_preserves_explicit_and_legacy_defaults() -> None: + source = _glm_fixture() + assert gguf_to_config(source).indexer_types == ["full", "shared"] + + source.metadata.pop("glm-dsa.attention.indexer.types") + assert gguf_to_config(source).indexer_types == ["full", "full"] + + source.metadata["glm-dsa.context_length"] = 1_048_576 + assert gguf_to_config(source).indexer_types == ["full", "full"] + + source.metadata["glm-dsa.attention.indexer.types"] = [True] + with pytest.raises(ValueError, match=r"expected 2, got 1"): + gguf_to_config(source) + + +def test_glm_dsa_rejects_unowned_rope_scaling() -> None: + source = _glm_fixture() + source.metadata["glm-dsa.rope.scaling.type"] = "yarn" + with pytest.raises(ValueError, match="unsupported RoPE scaling"): + gguf_to_config(source) + + +@pytest.mark.parametrize("factory", [_m2_fixture, _m4_fixture, _glm_fixture]) +def test_exact_tensor_closures_accept_synthetic_routes(factory) -> None: + validate_remaining_dense_tensor_contract(factory()) + + +@pytest.mark.parametrize( + ("factory", "mutation", "message"), + [ + (_m2_fixture, "bias", "ignored"), + (_m2_fixture, "fused", "cannot execute"), + (_m2_fixture, "missing", "missing="), + (_m4_fixture, "legacy_mla", "unexpected="), + (_m4_fixture, "partial_experts", "exactly one fused or split"), + (_glm_fixture, "shared_indexer", "unexpected="), + (_glm_fixture, "fused_experts", "unexpected="), + ], +) +def test_exact_tensor_closures_fail_closed( + factory, + mutation: str, + message: str, +) -> None: + source = factory() + if mutation == "bias": + source.tensors["blk.0.attn_q.bias"] = np.zeros(16, dtype=np.float32) + elif mutation == "fused": + source.tensors["blk.0.attn_qkv.weight"] = np.zeros((32, 8), dtype=np.float32) + for projection in ("q", "k", "v"): + source.tensors.pop(f"blk.0.attn_{projection}.weight") + elif mutation == "missing": + source.tensors.pop("blk.0.ffn_down_exps.weight") + elif mutation == "legacy_mla": + source.tensors["blk.0.attn_kv_b.weight"] = np.zeros((8, 4), dtype=np.float32) + elif mutation == "partial_experts": + source.tensors.pop("blk.1.ffn_gate_up_exps.weight") + source.tensors["blk.1.ffn_gate_exps.weight"] = np.zeros((2, 3, 8), dtype=np.float32) + elif mutation == "shared_indexer": + source.tensors["blk.1.indexer.k_norm.weight"] = np.ones(4, dtype=np.float32) + else: + source.tensors["blk.1.ffn_gate_up_exps.weight"] = np.zeros((2, 6, 8), dtype=np.float32) + source.tensor_names = list(source.tensors) + source.qtypes.update( + { + name: SimpleNamespace(value=0, name="F32") + for name in source.tensor_names + if name not in source.qtypes + } + ) + with pytest.raises(ValueError, match=message): + validate_remaining_dense_tensor_contract(source) + + +@pytest.mark.parametrize( + ("factory", "expected_model_type"), + [ + (_m2_fixture, "minimax_m2_gguf"), + (_m4_fixture, "mistral4_gguf"), + (_glm_fixture, "glm_moe_dsa"), + ], +) +def test_tiny_gguf_closes_graph_and_weight_loading( + factory, + expected_model_type: str, + tmp_path: Path, +) -> None: + source = factory() + path = tmp_path / f"{source.architecture}.gguf" + _write_gguf(path, source) + package = build_from_gguf(path, keep_quantized=False) + + assert package.config.model_type == expected_model_type + assert all( + initializer.const_value is not None + for initializer in package["model"].graph.initializers.values() + ) + if source.architecture == "glm-dsa": + assert ( + "shared-indexer-selection" in package["model"].metadata_props["mobius.cache_abi"] + ) + + +def test_minimax_m2_prefill_matches_cached_decode(tmp_path: Path) -> None: + source = _m2_fixture() + path = tmp_path / "minimax-m2.gguf" + _write_gguf(path, source) + model = build_from_gguf(path, keep_quantized=False)["model"] + session = OnnxModelSession(model) + tokens = np.array([[2, 5, 7]], dtype=np.int64) + empty = np.empty((1, 1, 0, 8), dtype=np.float32) + + def run( + input_ids: np.ndarray, + past_key: np.ndarray, + past_value: np.ndarray, + ): + total = past_key.shape[2] + input_ids.shape[1] + return session.run( + { + "input_ids": input_ids, + "attention_mask": np.ones((1, total), dtype=np.int64), + "position_ids": np.arange( + past_key.shape[2], + total, + dtype=np.int64, + ).reshape(1, -1), + "past_key_values.0.key": past_key, + "past_key_values.0.value": past_value, + } + ) + + full = run(tokens, empty, empty) + prefix = run(tokens[:, :2], empty, empty) + decoded = run( + tokens[:, 2:], + prefix["present.0.key"], + prefix["present.0.value"], + ) + np.testing.assert_allclose( + decoded["logits"], + full["logits"][:, -1:], + rtol=2e-4, + atol=2e-5, + ) + assert full["present.0.key"].shape == (1, 1, 3, 8) + q_norm = model.graph.initializers["model.layers.0.self_attn.q_norm.weight"] + assert tuple(q_norm.shape) == (16,) + q_norm_node = next( + node + for node in model.graph + if node.op_type == "RMSNormalization" + and node.inputs[1].name == "model.layers.0.self_attn.q_norm.weight" + ) + assert q_norm_node.inputs[0].producer().op_type == "MatMul" + rotary_nodes = [node for node in model.graph if node.op_type == "RotaryEmbedding"] + assert rotary_nodes + assert { + int(node.attributes["rotary_embedding_dim"].as_int()) for node in rotary_nodes + } == {4} + + +def test_minimax_m2_router_matches_selection_only_bias_and_floor() -> None: + config = make_config( + hidden_size=2, + num_local_experts=3, + num_experts_per_tok=2, + scoring_func="sigmoid", + norm_topk_prob=True, + routed_scaling_factor=1.0, + routing_weight_normalization_floor=6.103515625e-5, + ) + gate = MiniMaxM2Gate(config) + gate.weight.const_value = ir.tensor( + np.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]], dtype=np.float32) + ) + gate.e_score_correction_bias.const_value = ir.tensor( + np.array([-2.0, 1.0, 2.0], dtype=np.float32) + ) + hidden = ir.Value( + name="hidden", + shape=ir.Shape([1, 1, 2]), + type=ir.TensorType(ir.DataType.FLOAT), + ) + hidden_values = np.array([[[2.0, 1.0]]], dtype=np.float32) + actual_weights, actual_experts = _run_component( + [hidden], + lambda op: gate(op, hidden), + {"hidden": hidden_values}, + ) + + logits = hidden_values @ gate.weight.const_value.numpy().T + probabilities = 1.0 / (1.0 + np.exp(-logits)) + bias = gate.e_score_correction_bias.const_value.numpy() + expected_experts = np.argsort(-(probabilities + bias), axis=-1)[..., :2] + assert set(actual_experts.reshape(-1)) == set(expected_experts.reshape(-1)) + expected_weights = np.take_along_axis(probabilities, actual_experts, axis=-1) + expected_weights /= np.maximum( + expected_weights.sum(axis=-1, keepdims=True), + 6.103515625e-5, + ) + np.testing.assert_allclose(actual_weights, expected_weights, rtol=1e-6) + + +def test_minimax_m2_static_cache_uses_standard_kv_contract(tmp_path: Path) -> None: + source = _m2_fixture() + path = tmp_path / "minimax-m2-static.gguf" + _write_gguf(path, source) + graph = build_from_gguf( + path, + keep_quantized=False, + static_cache=True, + max_seq_len=8, + )["model"].graph + inputs = {value.name: tuple(value.shape) for value in graph.inputs} + assert inputs["key_cache.0"] == ("batch", 8, 8) + assert inputs["value_cache.0"] == ("batch", 8, 8) + assert any(node.op_type == "TensorScatter" for node in graph) + + +def test_mistral4_latent_cache_prefill_matches_decode(tmp_path: Path) -> None: + source = _m4_fixture() + path = tmp_path / "mistral4.gguf" + _write_gguf(path, source) + model = build_from_gguf(path, keep_quantized=False)["model"] + session = OnnxModelSession(model) + tokens = np.array([[1, 4, 9]], dtype=np.int64) + empty = np.empty((1, 1, 0, 6), dtype=np.float32) + + def run(input_ids: np.ndarray, past: list[np.ndarray]): + total = past[0].shape[2] + input_ids.shape[1] + feeds = { + "input_ids": input_ids, + "attention_mask": np.ones((1, total), dtype=np.int64), + "position_ids": np.arange( + past[0].shape[2], + total, + dtype=np.int64, + ).reshape(1, -1), + } + feeds.update( + {f"past_key_values.{layer}.key": value for layer, value in enumerate(past)} + ) + return session.run(feeds) + + full = run(tokens, [empty, empty]) + prefix = run(tokens[:, :2], [empty, empty]) + decoded = run( + tokens[:, 2:], + [prefix["present.0.key"], prefix["present.1.key"]], + ) + np.testing.assert_allclose( + decoded["logits"], + full["logits"][:, -1:], + rtol=2e-4, + atol=2e-5, + ) + assert full["present.0.key"].shape == (1, 1, 3, 6) + assert "present.0.value" not in full + assert "no-value-cache" in model.metadata_props["mobius.cache_abi"] + + +def test_mistral4_latent_cache_matches_expanded_mla_reference( + tmp_path: Path, +) -> None: + source = _m4_fixture() + path = tmp_path / "mistral4-reference.gguf" + _write_gguf(path, source) + package = build_from_gguf(path, keep_quantized=False) + latent_model = package["model"] + config = dataclasses.replace( + package.config, + partial_rotary_factor=1.0, + num_key_value_heads=package.config.num_attention_heads, + ) + expanded_model = build_from_module( + DeepSeekV3CausalLMModel(config), + config, + CausalLMTask(), + )["model"] + + latent_initializers = latent_model.graph.initializers + for name, initializer in expanded_model.graph.initializers.items(): + if initializer.const_value is not None: + continue + if name.endswith(".self_attn.kv_b_proj.weight"): + prefix = name[: -len("kv_b_proj.weight")] + key = latent_initializers[prefix + "k_b_proj.weight_t"].const_value + value = latent_initializers[prefix + "v_b_proj.weight_t"].const_value + assert key is not None and value is not None + key_array = key.numpy().T.reshape( + config.num_attention_heads, + config.qk_nope_head_dim, + config.kv_lora_rank, + ) + value_array = value.numpy().T.reshape( + config.num_attention_heads, + config.v_head_dim, + config.kv_lora_rank, + ) + fused = np.concatenate((key_array, value_array), axis=1).reshape( + -1, + config.kv_lora_rank, + ) + initializer.const_value = ir.tensor(fused, name=name) + continue + source_initializer = latent_initializers.get(name) + transpose = False + if source_initializer is None and name.endswith(".weight"): + source_initializer = latent_initializers.get(name + "_t") + transpose = source_initializer is not None + assert source_initializer is not None, name + assert source_initializer.const_value is not None, name + value = source_initializer.const_value.numpy() + if transpose: + value = value.T + initializer.const_value = ir.tensor( + value, + name=name, + ) + + tokens = np.array([[1, 4, 9]], dtype=np.int64) + attention_mask = np.ones_like(tokens) + position_ids = np.arange(tokens.shape[1], dtype=np.int64).reshape(1, -1) + latent = OnnxModelSession(latent_model).run( + { + "input_ids": tokens, + "attention_mask": attention_mask, + "position_ids": position_ids, + "past_key_values.0.key": np.empty((1, 1, 0, 6), dtype=np.float32), + "past_key_values.1.key": np.empty((1, 1, 0, 6), dtype=np.float32), + } + ) + expanded = OnnxModelSession(expanded_model).run( + { + "input_ids": tokens, + "attention_mask": attention_mask, + "position_ids": position_ids, + "past_key_values.0.key": np.empty((1, 2, 0, 4), dtype=np.float32), + "past_key_values.0.value": np.empty((1, 2, 0, 2), dtype=np.float32), + "past_key_values.1.key": np.empty((1, 2, 0, 4), dtype=np.float32), + "past_key_values.1.value": np.empty((1, 2, 0, 2), dtype=np.float32), + } + ) + np.testing.assert_allclose( + latent["logits"], + expanded["logits"], + rtol=2e-4, + atol=2e-5, + ) + + +def test_glm_dsa_explicit_indexer_schedule_and_dense_kv_reassembly() -> None: + from mobius.integrations.gguf._tensor_processors import process_tensors + + source = _glm_fixture() + config = gguf_to_config(source) + assert config.indexer_types == ["full", "shared"] + + dense_config = SimpleNamespace( + model_type="glm_moe_dsa", + _gguf_arch="glm-dsa", + use_dsa=False, + ) + state = { + "model.layers.0.self_attn.k_b_proj.weight": torch.from_numpy( + source.tensors["blk.0.attn_k_b.weight"] + ), + "model.layers.0.self_attn.v_b_proj.weight": torch.from_numpy( + source.tensors["blk.0.attn_v_b.weight"] + ), + } + processed = process_tensors(state, dense_config) + assert set(processed) == {"model.layers.0.self_attn.kv_b_proj.weight"} + assert processed["model.layers.0.self_attn.kv_b_proj.weight"].shape == (8, 4) + + +@pytest.mark.parametrize("task_name", ["glm-moe-dsa", "mistral4-gguf-text-generation"]) +def test_specialized_cache_tasks_reject_static_cache(task_name: str) -> None: + from mobius.tasks import get_task + + task_class = type(get_task(task_name)) + with pytest.raises(ValueError, match="static cache"): + task_class(static_cache=True) + + +@pytest.mark.parametrize( + ("factory", "message"), + [ + (_glm_fixture, "dsa_kv_cache_specs"), + (_m4_fixture, "latent K-only cache"), + ], +) +def test_specialized_gguf_routes_reject_static_cache_before_graph_build( + factory, + message: str, + tmp_path: Path, +) -> None: + source = factory() + path = tmp_path / f"{source.architecture}-static.gguf" + _write_gguf(path, source) + with pytest.raises(ValueError, match=message): + build_from_gguf( + path, + keep_quantized=False, + static_cache=True, + max_seq_len=8, + ) diff --git a/src/mobius/integrations/gguf/_route_census_test.py b/src/mobius/integrations/gguf/_route_census_test.py index 27f5be7e3..8dc075ba3 100644 --- a/src/mobius/integrations/gguf/_route_census_test.py +++ b/src/mobius/integrations/gguf/_route_census_test.py @@ -94,6 +94,9 @@ def test_route_reasons_are_sourced_from_authoritative_records() -> None: def test_known_route_boundaries_are_not_collapsed() -> None: by_id = {item.route_id: item for item in iter_remaining_route_work()} assert by_id["architecture:bitnet"].category == "evidence-only" + assert by_id["architecture:glm-dsa"].category == "evidence-only" + assert by_id["architecture:minimax-m2"].category == "evidence-only" + assert by_id["architecture:mistral4"].category == "evidence-only" assert by_id["architecture:deepseek4"].category == "dependency-or-runtime-abi-blocked" assert by_id["architecture:rwkv6"].category == "dependency-or-runtime-abi-blocked" assert by_id["architecture:bailingmoe2"].category == "intentionally-rejected" diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index 4e021194a..d5ecca9b1 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -1268,6 +1268,25 @@ "blk.{bid}.ffn_down_exps": "model.layers.{bid}.mlp.experts.down_proj", } +_MINIMAX_M2_MAPPING: dict[str, str] = { + "token_embd": "model.embed_tokens", + "output": "lm_head", + "output_norm": "model.norm", + "blk.{bid}.attn_norm": "model.layers.{bid}.input_layernorm", + "blk.{bid}.attn_q": "model.layers.{bid}.self_attn.q_proj", + "blk.{bid}.attn_k": "model.layers.{bid}.self_attn.k_proj", + "blk.{bid}.attn_v": "model.layers.{bid}.self_attn.v_proj", + "blk.{bid}.attn_q_norm": "model.layers.{bid}.self_attn.q_norm", + "blk.{bid}.attn_k_norm": "model.layers.{bid}.self_attn.k_norm", + "blk.{bid}.attn_output": "model.layers.{bid}.self_attn.o_proj", + "blk.{bid}.ffn_norm": "model.layers.{bid}.post_attention_layernorm", + "blk.{bid}.ffn_gate_inp": "model.layers.{bid}.mlp.gate", + "blk.{bid}.exp_probs_b": ("model.layers.{bid}.mlp.gate.e_score_correction_bias@"), + "blk.{bid}.ffn_gate_exps": "model.layers.{bid}.mlp.experts.gate_proj", + "blk.{bid}.ffn_up_exps": "model.layers.{bid}.mlp.experts.up_proj", + "blk.{bid}.ffn_down_exps": "model.layers.{bid}.mlp.experts.down_proj", +} + _KIMI_LINEAR_MAPPING: dict[str, str] = { "token_embd": "model.embed_tokens", "output": "lm_head", @@ -1326,6 +1345,65 @@ del _KIMI_K3_MAPPING["blk.{bid}.ssm_g_a"] del _KIMI_K3_MAPPING["blk.{bid}.ssm_g_b"] +_GLM_DSA_MAPPING: dict[str, str] = { + "token_embd": "model.embed_tokens", + "output": "lm_head", + "output_norm": "model.norm", + "blk.{bid}.attn_norm": "model.layers.{bid}.input_layernorm", + "blk.{bid}.attn_q_a": "model.layers.{bid}.self_attn.q_a_proj", + "blk.{bid}.attn_q_a_norm": "model.layers.{bid}.self_attn.q_a_layernorm", + "blk.{bid}.attn_q_b": "model.layers.{bid}.self_attn.q_b_proj", + "blk.{bid}.attn_kv_a_mqa": "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", + "blk.{bid}.attn_kv_a_norm": "model.layers.{bid}.self_attn.kv_a_layernorm", + "blk.{bid}.attn_k_b": "model.layers.{bid}.self_attn.k_b_proj", + "blk.{bid}.attn_v_b": "model.layers.{bid}.self_attn.v_b_proj", + "blk.{bid}.attn_output": "model.layers.{bid}.self_attn.o_proj", + "blk.{bid}.ffn_norm": "model.layers.{bid}.post_attention_layernorm", + "blk.{bid}.ffn_gate": "model.layers.{bid}.mlp.gate_proj", + "blk.{bid}.ffn_up": "model.layers.{bid}.mlp.up_proj", + "blk.{bid}.ffn_down": "model.layers.{bid}.mlp.down_proj", + "blk.{bid}.ffn_gate_inp": "model.layers.{bid}.mlp.gate", + "blk.{bid}.exp_probs_b": "model.layers.{bid}.mlp.gate.e_score_correction_bias@", + "blk.{bid}.ffn_gate_exps": "model.layers.{bid}.mlp.experts.gate_proj", + "blk.{bid}.ffn_up_exps": "model.layers.{bid}.mlp.experts.up_proj", + "blk.{bid}.ffn_down_exps": "model.layers.{bid}.mlp.experts.down_proj", + "blk.{bid}.ffn_gate_shexp": "model.layers.{bid}.mlp.shared_experts.gate_proj", + "blk.{bid}.ffn_up_shexp": "model.layers.{bid}.mlp.shared_experts.up_proj", + "blk.{bid}.ffn_down_shexp": "model.layers.{bid}.mlp.shared_experts.down_proj", + "blk.{bid}.indexer.attn_k": "model.layers.{bid}.self_attn.indexer.wk", + "blk.{bid}.indexer.attn_q_b": "model.layers.{bid}.self_attn.indexer.wq_b", + "blk.{bid}.indexer.k_norm": "model.layers.{bid}.self_attn.indexer.k_norm", + "blk.{bid}.indexer.proj": "model.layers.{bid}.self_attn.indexer.weights_proj", +} + +_MISTRAL4_MAPPING: dict[str, str] = { + "token_embd": "model.embed_tokens", + "output": "lm_head", + "output_norm": "model.norm", + "blk.{bid}.attn_norm": "model.layers.{bid}.input_layernorm", + "blk.{bid}.attn_q_a": "model.layers.{bid}.self_attn.q_a_proj", + "blk.{bid}.attn_q_a_norm": "model.layers.{bid}.self_attn.q_a_layernorm", + "blk.{bid}.attn_q_b": "model.layers.{bid}.self_attn.q_b_proj", + "blk.{bid}.attn_kv_a_mqa": "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", + "blk.{bid}.attn_kv_a_norm": "model.layers.{bid}.self_attn.kv_a_layernorm", + "blk.{bid}.attn_k_b": "model.layers.{bid}.self_attn.k_b_proj", + "blk.{bid}.attn_v_b": "model.layers.{bid}.self_attn.v_b_proj", + "blk.{bid}.attn_output": "model.layers.{bid}.self_attn.o_proj", + "blk.{bid}.ffn_norm": "model.layers.{bid}.post_attention_layernorm", + "blk.{bid}.ffn_gate": "model.layers.{bid}.mlp.gate_proj", + "blk.{bid}.ffn_up": "model.layers.{bid}.mlp.up_proj", + "blk.{bid}.ffn_down": "model.layers.{bid}.mlp.down_proj", + "blk.{bid}.ffn_gate_inp": "model.layers.{bid}.mlp.gate", + "blk.{bid}.exp_probs_b": ("model.layers.{bid}.mlp.gate.e_score_correction_bias@"), + "blk.{bid}.ffn_gate_exps": "model.layers.{bid}.mlp.experts.gate_proj", + "blk.{bid}.ffn_up_exps": "model.layers.{bid}.mlp.experts.up_proj", + "blk.{bid}.ffn_down_exps": "model.layers.{bid}.mlp.experts.down_proj", + "blk.{bid}.ffn_gate_up_exps": "model.layers.{bid}.mlp.experts.gate_up_proj", + "blk.{bid}.ffn_gate_shexp": "model.layers.{bid}.mlp.shared_experts.gate_proj", + "blk.{bid}.ffn_up_shexp": "model.layers.{bid}.mlp.shared_experts.up_proj", + "blk.{bid}.ffn_down_shexp": "model.layers.{bid}.mlp.shared_experts.down_proj", +} + _MINICPM3_MAPPING: dict[str, str] = { "token_embd": "model.embed_tokens", "output": "lm_head", @@ -1457,8 +1535,11 @@ "hunyuan_extras": _HUNYUAN_EXTRAS, "muse_glimmer_extras": _MUSE_GLIMMER_EXTRAS, "minimax": _MINIMAX_MAPPING, + "minimax_m2": _MINIMAX_M2_MAPPING, + "mistral4": _MISTRAL4_MAPPING, "kimi_linear": _KIMI_LINEAR_MAPPING, "kimi_k3": _KIMI_K3_MAPPING, + "glm_dsa": _GLM_DSA_MAPPING, "minicpm3": _MINICPM3_MAPPING, } ) diff --git a/src/mobius/integrations/gguf/_tensor_processors.py b/src/mobius/integrations/gguf/_tensor_processors.py index 98f5ce0b1..dcdd7a344 100644 --- a/src/mobius/integrations/gguf/_tensor_processors.py +++ b/src/mobius/integrations/gguf/_tensor_processors.py @@ -486,6 +486,84 @@ def _process_kimi_linear( return state_dict +def _process_split_mla_kv( + state_dict: dict[str, torch.Tensor], + *, + fuse_for_expanded_graph: bool, +) -> dict[str, torch.Tensor]: + """Restore split MLA K/V-B matrices to flat Linear row order.""" + k_suffix = ".k_b_proj.weight" + v_suffix = ".v_b_proj.weight" + k_by_prefix = { + name[: -len(k_suffix)]: tensor + for name, tensor in state_dict.items() + if name.endswith(k_suffix) + } + v_by_prefix = { + name[: -len(v_suffix)]: tensor + for name, tensor in state_dict.items() + if name.endswith(v_suffix) + } + if set(k_by_prefix) != set(v_by_prefix): + missing_k = sorted(set(v_by_prefix) - set(k_by_prefix)) + missing_v = sorted(set(k_by_prefix) - set(v_by_prefix)) + raise ValueError( + "GLM-5.2 split K/V-B projections must be paired per layer: " + f"missing_k={missing_k}, missing_v={missing_v}" + ) + + for prefix in sorted(k_by_prefix): + key = k_by_prefix[prefix] + value = v_by_prefix[prefix] + if key.dim() != 3 or value.dim() != 3: + raise ValueError( + "GLM-5.2 split K/V-B tensors must both be rank 3, got " + f"{prefix + k_suffix!r} rank {key.dim()} and " + f"{prefix + v_suffix!r} rank {value.dim()}" + ) + key = key.transpose(1, 2) + if not fuse_for_expanded_graph: + state_dict[prefix + k_suffix] = key.reshape( + key.shape[0] * key.shape[1], key.shape[2] + ) + state_dict[prefix + v_suffix] = value.reshape( + value.shape[0] * value.shape[1], value.shape[2] + ) + continue + + if key.shape[0] != value.shape[0] or key.shape[2] != value.shape[2]: + raise ValueError( + f"GLM-5.2 split K/V-B tensors for {prefix!r} have incompatible shapes " + f"{tuple(key.shape)} and {tuple(value.shape)}" + ) + del state_dict[prefix + k_suffix] + del state_dict[prefix + v_suffix] + state_dict[prefix + ".kv_b_proj.weight"] = torch.cat((key, value), dim=1).reshape( + -1, key.shape[2] + ) + return state_dict + + +def _process_glm_dsa( + state_dict: dict[str, torch.Tensor], + config: Any, +) -> dict[str, torch.Tensor]: + """Restore GLM-DSA split MLA weights for sparse or dense graph selection.""" + return _process_split_mla_kv( + state_dict, + fuse_for_expanded_graph=not getattr(config, "use_dsa", True), + ) + + +def _process_mistral4( + state_dict: dict[str, torch.Tensor], + config: Any, +) -> dict[str, torch.Tensor]: + """Restore Mistral4 K/V-B matrices while retaining its split latent graph.""" + del config + return _process_split_mla_kv(state_dict, fuse_for_expanded_graph=False) + + def _process_kimi_k3( state_dict: dict[str, torch.Tensor], config: Any, @@ -678,6 +756,8 @@ def _process_plamo( "plamo": _process_plamo, "plamo2": _process_plamo2, "granitehybrid": _process_granitehybrid, + "glm_dsa": _process_glm_dsa, + "mistral4": _process_mistral4, "kimi_linear": _process_kimi_linear, "kimi_k3": _process_kimi_k3, "bloom": _process_bloom, diff --git a/src/mobius/integrations/gguf/_tensor_processors_test.py b/src/mobius/integrations/gguf/_tensor_processors_test.py index 0b3725c18..1ef4af3fa 100644 --- a/src/mobius/integrations/gguf/_tensor_processors_test.py +++ b/src/mobius/integrations/gguf/_tensor_processors_test.py @@ -61,6 +61,33 @@ def test_rejects_non_negative_decay(self) -> None: ) +class TestProcessTensorsGlmDsa: + def test_restores_split_mla_projection_layouts(self) -> None: + heads, rank, nope, value_dim = 2, 3, 4, 5 + key_hf = torch.arange(heads * nope * rank, dtype=torch.float32).reshape( + heads, nope, rank + ) + value_hf = torch.arange(heads * value_dim * rank, dtype=torch.float32).reshape( + heads, value_dim, rank + ) + state = { + "model.layers.0.self_attn.k_b_proj.weight": key_hf.transpose(1, 2), + "model.layers.0.self_attn.v_b_proj.weight": value_hf, + } + config = SimpleNamespace(model_type="glm_moe_dsa", _gguf_arch="glm-dsa") + + result = process_tensors(state, config) + + torch.testing.assert_close( + result["model.layers.0.self_attn.k_b_proj.weight"], + key_hf.reshape(heads * nope, rank), + ) + torch.testing.assert_close( + result["model.layers.0.self_attn.v_b_proj.weight"], + value_hf.reshape(heads * value_dim, rank), + ) + + class TestProcessTensorsKimiK3: def test_restores_collapsed_conv_and_decay(self) -> None: config = SimpleNamespace(model_type="kimi_k3", _gguf_arch="kimi-k3") diff --git a/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json b/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json index 997adac60..82e12b23f 100644 --- a/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json +++ b/src/mobius/integrations/gguf/_upstream_data/llamacpp_pin.json @@ -112,7 +112,7 @@ "minimax-m2": {"cohort": "C02-moe-dense-attn", "cpp_loader": true, "dual_moe": false, "tensor_families": ["token_embd", "output_norm", "output", "blk.{bid}.attn_norm", "blk.{bid}.attn_q", "blk.{bid}.attn_q_norm", "blk.{bid}.attn_k", "blk.{bid}.attn_k_norm", "blk.{bid}.attn_v", "blk.{bid}.attn_output", "blk.{bid}.ffn_norm", "blk.{bid}.ffn_gate_inp", "blk.{bid}.ffn_gate_exps", "blk.{bid}.ffn_down_exps", "blk.{bid}.ffn_up_exps", "blk.{bid}.exp_probs_b"], "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_q_norm.weight", "blk.{bid}.attn_k_norm.weight", "blk.{bid}.ffn_norm.weight", "blk.{bid}.ffn_gate_inp.weight", "blk.{bid}.ffn_gate_exps.weight", "blk.{bid}.ffn_down_exps.weight", "blk.{bid}.ffn_up_exps.weight", "blk.{bid}.exp_probs_b.bias", "blk.{bid}.attn_qkv.weight", "blk.{bid}.attn_qkv.bias", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_q.bias", "blk.{bid}.attn_k.weight", "blk.{bid}.attn_k.bias", "blk.{bid}.attn_v.weight", "blk.{bid}.attn_v.bias"], "tensor_closure_status": "audited-direct-loader-conditional-union", "converter_inventory_status": "exact-pinned-MODEL_TENSORS-family-inventory", "loader_source": "src/models/minimax-m2.cpp", "required_metadata": ["attention.layer_norm_rms_epsilon", "expert_feed_forward_length"], "optional_metadata": ["expert_gating_func"], "loader_helpers": ["create_tensor_qkv"]}, "minimax-m3": {"cohort": "C02-moe-dense-attn", "cpp_loader": true, "dual_moe": true, "tensor_families": ["token_embd", "output_norm", "output", "blk.{bid}.attn_norm", "blk.{bid}.attn_q", "blk.{bid}.attn_q_norm", "blk.{bid}.attn_k", "blk.{bid}.attn_k_norm", "blk.{bid}.attn_v", "blk.{bid}.attn_output", "blk.{bid}.ffn_norm", "blk.{bid}.ffn_gate_inp", "blk.{bid}.exp_probs_b", "blk.{bid}.ffn_gate_exps", "blk.{bid}.ffn_down_exps", "blk.{bid}.ffn_up_exps", "blk.{bid}.ffn_gate_shexp", "blk.{bid}.ffn_down_shexp", "blk.{bid}.ffn_up_shexp", "blk.{bid}.ffn_gate", "blk.{bid}.ffn_down", "blk.{bid}.ffn_up", "blk.{bid}.indexer.q_proj", "blk.{bid}.indexer.k_proj", "blk.{bid}.indexer.q_norm", "blk.{bid}.indexer.k_norm"], "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_q_norm.weight", "blk.{bid}.attn_k_norm.weight", "blk.{bid}.ffn_norm.weight", "blk.{bid}.ffn_gate.weight", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_gate_inp.weight", "blk.{bid}.exp_probs_b.bias", "blk.{bid}.ffn_gate_exps.weight", "blk.{bid}.ffn_down_exps.weight", "blk.{bid}.ffn_up_exps.weight", "blk.{bid}.ffn_gate_shexp.weight", "blk.{bid}.ffn_down_shexp.weight", "blk.{bid}.ffn_up_shexp.weight", "blk.{bid}.indexer.q_proj.weight", "blk.{bid}.indexer.k_proj.weight", "blk.{bid}.indexer.q_norm.weight", "blk.{bid}.indexer.k_norm.weight", "blk.{bid}.attn_qkv.weight", "blk.{bid}.attn_qkv.bias", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_q.bias", "blk.{bid}.attn_k.weight", "blk.{bid}.attn_k.bias", "blk.{bid}.attn_v.weight", "blk.{bid}.attn_v.bias"], "tensor_closure_status": "audited-direct-loader-conditional-union", "converter_inventory_status": "exact-pinned-MODEL_TENSORS-family-inventory", "loader_source": "src/models/minimax-m3.cpp", "required_metadata": ["attention.layer_norm_rms_epsilon", "expert_feed_forward_length", "expert_shared_count", "expert_gating_func", "attention.indexer.head_count", "attention.indexer.key_length", "attention.indexer.top_k", "attention.indexer.block_size", "attention.indexer.local_blocks"], "optional_metadata": ["leading_dense_block_count", "expert_weights_scale", "expert_weights_norm"], "loader_helpers": ["create_tensor_qkv"]}, "mistral3": {"cohort": "C02-moe-dense-attn", "cpp_loader": true, "dual_moe": true, "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_qkv.weight", "blk.{bid}.attn_qkv.bias", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_q.bias", "blk.{bid}.attn_k.weight", "blk.{bid}.attn_k.bias", "blk.{bid}.attn_v.weight", "blk.{bid}.attn_v.bias", "blk.{bid}.attn_output.weight", "blk.{bid}.attn_output.bias", "blk.{bid}.ffn_norm.weight", "rope_freqs.weight", "rope_factors_long.weight", "rope_factors_short.weight", "blk.{bid}.ffn_gate.weight", "blk.{bid}.ffn_gate.bias", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_down.bias", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_up.bias", "blk.{bid}.ffn_gate_inp.weight", "blk.{bid}.ffn_gate_exps.weight", "blk.{bid}.ffn_down_exps.weight", "blk.{bid}.ffn_up_exps.weight", "blk.{bid}.ffn_gate_shexp.weight", "blk.{bid}.ffn_down_shexp.weight", "blk.{bid}.ffn_up_shexp.weight"], "expert_tensor_suffixes": ["weight", "scale", "input_scale"], "tensor_suffixes": {"output": ["weight", "scale", "input_scale"], "blk.{bid}.attn_qkv": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.attn_q": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.attn_k": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.attn_v": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.attn_output": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.ffn_gate": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.ffn_down": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.ffn_up": ["weight", "bias", "scale", "input_scale"], "blk.{bid}.ffn_gate_exps": ["weight", "scale", "input_scale"], "blk.{bid}.ffn_down_exps": ["weight", "scale", "input_scale"], "blk.{bid}.ffn_up_exps": ["weight", "scale", "input_scale"], "blk.{bid}.ffn_gate_shexp": ["weight", "scale", "input_scale"], "blk.{bid}.ffn_down_shexp": ["weight", "scale", "input_scale"], "blk.{bid}.ffn_up_shexp": ["weight", "scale", "input_scale"]}, "tensor_closure_status": "exact-direct-loader-conditional-union", "converter_inventory_status": "unresolved-inherited-conditional-converter-hooks"}, - "mistral4": {"cohort": "C02-moe-dense-attn", "cpp_loader": true, "dual_moe": true, "tensor_families": ["token_embd", "output_norm", "output", "rope_freqs", "blk.{bid}.attn_norm", "blk.{bid}.attn_q", "blk.{bid}.attn_q_a", "blk.{bid}.attn_q_b", "blk.{bid}.attn_kv_a_mqa", "blk.{bid}.attn_kv_b", "blk.{bid}.attn_k_b", "blk.{bid}.attn_v_b", "blk.{bid}.attn_q_a_norm", "blk.{bid}.attn_kv_a_norm", "blk.{bid}.attn_output", "blk.{bid}.attn_rot_embd", "blk.{bid}.ffn_gate_inp", "blk.{bid}.ffn_norm", "blk.{bid}.ffn_gate", "blk.{bid}.ffn_down", "blk.{bid}.ffn_up", "blk.{bid}.ffn_gate_exps", "blk.{bid}.ffn_down_exps", "blk.{bid}.ffn_up_exps", "blk.{bid}.ffn_gate_up_exps", "blk.{bid}.ffn_gate_shexp", "blk.{bid}.ffn_down_shexp", "blk.{bid}.ffn_up_shexp", "blk.{bid}.exp_probs_b"], "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_output.weight", "blk.{bid}.attn_output.bias", "blk.{bid}.ffn_norm.weight", "rope_factors_long.weight", "rope_factors_short.weight", "rope_freqs.weight", "blk.{bid}.ffn_gate.weight", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_gate.bias", "blk.{bid}.ffn_down.bias", "blk.{bid}.ffn_up.bias", "blk.{bid}.ffn_gate_inp.weight", "blk.{bid}.ffn_gate_exps.weight", "blk.{bid}.ffn_down_exps.weight", "blk.{bid}.ffn_up_exps.weight", "blk.{bid}.ffn_gate_shexp.weight", "blk.{bid}.ffn_up_shexp.weight", "blk.{bid}.ffn_down_shexp.weight", "blk.{bid}.attn_qkv.weight", "blk.{bid}.attn_qkv.bias", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_q.bias", "blk.{bid}.attn_k.weight", "blk.{bid}.attn_k.bias", "blk.{bid}.attn_v.weight", "blk.{bid}.attn_v.bias"], "tensor_closure_status": "audited-direct-loader-conditional-union", "converter_inventory_status": "exact-pinned-MODEL_TENSORS-family-inventory", "loader_source": "src/models/mistral3.cpp", "required_metadata": ["attention.layer_norm_rms_epsilon"], "optional_metadata": ["attention.temperature_scale", "rope.scaling.yarn_beta_fast", "rope.scaling.yarn_beta_slow", "rope.scaling.yarn_log_multiplier"], "loader_helpers": ["create_tensor_qkv"]}, + "mistral4": {"cohort": "C02-moe-dense-attn", "cpp_loader": true, "dual_moe": true, "tensor_families": ["token_embd", "output_norm", "output", "rope_freqs", "blk.{bid}.attn_norm", "blk.{bid}.attn_q", "blk.{bid}.attn_q_a", "blk.{bid}.attn_q_b", "blk.{bid}.attn_kv_a_mqa", "blk.{bid}.attn_kv_b", "blk.{bid}.attn_k_b", "blk.{bid}.attn_v_b", "blk.{bid}.attn_q_a_norm", "blk.{bid}.attn_kv_a_norm", "blk.{bid}.attn_output", "blk.{bid}.attn_rot_embd", "blk.{bid}.ffn_gate_inp", "blk.{bid}.ffn_norm", "blk.{bid}.ffn_gate", "blk.{bid}.ffn_down", "blk.{bid}.ffn_up", "blk.{bid}.ffn_gate_exps", "blk.{bid}.ffn_down_exps", "blk.{bid}.ffn_up_exps", "blk.{bid}.ffn_gate_up_exps", "blk.{bid}.ffn_gate_shexp", "blk.{bid}.ffn_down_shexp", "blk.{bid}.ffn_up_shexp", "blk.{bid}.exp_probs_b"], "tensor_names": ["token_embd.weight", "output_norm.weight", "output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_q_a_norm.weight", "blk.{bid}.attn_kv_a_norm.weight", "blk.{bid}.attn_q_a.weight", "blk.{bid}.attn_q_b.weight", "blk.{bid}.attn_q.weight", "blk.{bid}.attn_kv_a_mqa.weight", "blk.{bid}.attn_kv_b.weight", "blk.{bid}.attn_k_b.weight", "blk.{bid}.attn_v_b.weight", "blk.{bid}.attn_output.weight", "blk.{bid}.ffn_norm.weight", "blk.{bid}.ffn_gate.weight", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_gate_inp.weight", "blk.{bid}.exp_probs_b.bias", "blk.{bid}.ffn_gate_exps.weight", "blk.{bid}.ffn_down_exps.weight", "blk.{bid}.ffn_up_exps.weight", "blk.{bid}.ffn_gate_up_exps.weight", "blk.{bid}.ffn_gate_shexp.weight", "blk.{bid}.ffn_down_shexp.weight", "blk.{bid}.ffn_up_shexp.weight", "blk.{bid}.nextn.eh_proj.weight", "blk.{bid}.nextn.enorm.weight", "blk.{bid}.nextn.hnorm.weight", "blk.{bid}.nextn.embed_tokens.weight", "blk.{bid}.nextn.shared_head_head.weight", "blk.{bid}.nextn.shared_head_norm.weight"], "tensor_closure_status": "audited-direct-loader-conditional-union", "converter_inventory_status": "exact-pinned-MODEL_TENSORS-family-inventory", "loader_source": "src/models/deepseek2.cpp", "required_metadata": ["attention.layer_norm_rms_epsilon", "attention.q_lora_rank", "attention.kv_lora_rank", "expert_feed_forward_length", "expert_shared_count"], "optional_metadata": ["leading_dense_block_count", "attention.key_length_mla", "attention.value_length_mla", "expert_weights_scale", "expert_weights_norm", "expert_gating_func", "expert_group_count", "expert_group_used_count", "rope.scaling.yarn_beta_fast", "rope.scaling.yarn_beta_slow", "rope.scaling.yarn_log_multiplier", "nextn_predict_layers", "attention.temperature_scale", "attention.temperature_length"], "loader_helpers": []}, "modern-bert": {"cohort": "C03-encoder-embed-rerank", "cpp_loader": true, "dual_moe": false, "tensor_families": ["token_embd", "token_embd_norm", "output_norm", "cls", "cls_norm", "cls_out", "blk.{bid}.attn_norm", "blk.{bid}.attn_qkv", "blk.{bid}.attn_output", "blk.{bid}.ffn_up", "blk.{bid}.ffn_down", "blk.{bid}.ffn_norm"], "tensor_names": ["token_embd.weight", "token_embd_norm.weight", "output_norm.weight", "cls.weight", "cls_norm.weight", "cls_out.weight", "cls_out.bias", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_qkv.weight", "blk.{bid}.attn_output.weight", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_norm.weight"]}, "mpt": {"cohort": "C01-dense-transformer", "cpp_loader": true, "dual_moe": false, "tensor_families": ["token_embd", "position_embd", "output_norm", "output", "blk.{bid}.attn_norm", "blk.{bid}.attn_qkv", "blk.{bid}.attn_q_norm", "blk.{bid}.attn_k_norm", "blk.{bid}.attn_output", "blk.{bid}.ffn_norm", "blk.{bid}.ffn_up", "blk.{bid}.ffn_down", "blk.{bid}.ffn_act"], "tensor_names": ["token_embd.weight", "position_embd.weight", "output_norm.weight", "output_norm.bias", "output.weight", "blk.{bid}.attn_norm.weight", "blk.{bid}.attn_norm.bias", "blk.{bid}.attn_qkv.weight", "blk.{bid}.attn_qkv.bias", "blk.{bid}.attn_q_norm.weight", "blk.{bid}.attn_q_norm.bias", "blk.{bid}.attn_k_norm.weight", "blk.{bid}.attn_k_norm.bias", "blk.{bid}.attn_output.weight", "blk.{bid}.attn_output.bias", "blk.{bid}.ffn_norm.weight", "blk.{bid}.ffn_norm.bias", "blk.{bid}.ffn_up.weight", "blk.{bid}.ffn_up.bias", "blk.{bid}.ffn_down.weight", "blk.{bid}.ffn_down.bias", "blk.{bid}.ffn_act.scales"], "tensor_closure_status": "audited-direct-loader-conditional-union", "converter_inventory_status": "exact-pinned-MODEL_TENSORS-family-inventory", "loader_source": "src/models/mpt.cpp", "required_metadata": ["attention.layer_norm_epsilon"], "optional_metadata": ["attention.clamp_kqv", "attention.max_alibi_bias"], "loader_helpers": []}, "muse-glimmer": {"cohort": "C01-dense-transformer", "cpp_loader": true, "dual_moe": false}, diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 31d0c5caa..6769da63c 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -71,6 +71,8 @@ "Gemma4UnifiedModel", "GemmaCausalLMModel", "GemmaEmbeddingGGUFModel", + "MiniMaxM2GGUFCausalLMModel", + "Mistral4GGUFCausalLMModel", "Glm4CausalLMModel", "GlmAsrForConditionalGeneration", "Glm4MoECausalLMModel", @@ -286,6 +288,8 @@ NomicBertMoEGGUFModel, ) from mobius.models.gguf_legacy_decoders import ExactLegacyGGUFCausalLMModel +from mobius.models.gguf_minimax_m2 import MiniMaxM2GGUFCausalLMModel +from mobius.models.gguf_mistral4 import Mistral4GGUFCausalLMModel from mobius.models.gguf_plamo import PlamoGGUFCausalLMModel from mobius.models.glm import Glm4CausalLMModel, GlmCausalLMModel from mobius.models.glm_asr import GlmAsrForConditionalGeneration diff --git a/src/mobius/models/gguf_minimax_m2.py b/src/mobius/models/gguf_minimax_m2.py new file mode 100644 index 000000000..a67494ee7 --- /dev/null +++ b/src/mobius/models/gguf_minimax_m2.py @@ -0,0 +1,104 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Exact explicit-float decoder for ``general.architecture=minimax-m2`` GGUF.""" + +from __future__ import annotations + +import math + +import onnx_ir as ir +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig +from mobius.components import RMSNorm +from mobius.models.base import CausalLMModel, TextModel +from mobius.models.moe import MoEDecoderLayer + +_F16_NORMAL_MIN = 6.103515625e-5 + + +class MiniMaxM2Gate(nn.Module): + """F32 sigmoid router with selection-only correction bias.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + if config.num_local_experts is None or config.num_experts_per_tok is None: + raise ValueError("MiniMax-M2 requires expert count and top-k") + if config.scoring_func != "sigmoid": + raise ValueError("MiniMax-M2 requires sigmoid expert routing") + if not config.norm_topk_prob: + raise ValueError("MiniMax-M2 requires normalized selected-expert weights") + floor = config.routing_weight_normalization_floor + if floor is None or not math.isclose(floor, _F16_NORMAL_MIN): + raise ValueError("MiniMax-M2 requires the llama.cpp F16-normal routing floor") + + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + self.normalization_floor = floor + self.output_scale = config.routed_scaling_factor + self.weight = nn.Parameter( + [self.num_experts, config.hidden_size], + dtype=ir.DataType.FLOAT, + ) + self.e_score_correction_bias = nn.Parameter( + [self.num_experts], + dtype=ir.DataType.FLOAT, + ) + self.weight._keep_float32 = True # type: ignore[attr-defined] + self.e_score_correction_bias._keep_float32 = True # type: ignore[attr-defined] + + def forward(self, op: OpBuilder, hidden_states: ir.Value): + router_input = op.Cast(hidden_states, to=ir.DataType.FLOAT) + logits = op.MatMul(router_input, op.Transpose(self.weight, perm=[1, 0])) + probabilities = op.Sigmoid(logits) + selection_scores = op.Add(probabilities, self.e_score_correction_bias) + _, selected_experts = op.TopK( + selection_scores, + op.Constant(value_ints=[self.top_k]), + axis=-1, + largest=1, + sorted=0, + _outputs=2, + ) + routing_weights = op.GatherElements(probabilities, selected_experts, axis=-1) + denominator = op.ReduceSum(routing_weights, [-1], keepdims=True) + denominator = op.Max(denominator, float(self.normalization_floor)) + routing_weights = op.Div(routing_weights, denominator) + if not math.isclose(self.output_scale, 1.0): + routing_weights = op.Mul(routing_weights, float(self.output_scale)) + return op.CastLike(routing_weights, hidden_states), selected_experts + + +class MiniMaxM2DecoderLayer(MoEDecoderLayer): + """Standard pre-norm attention block with MiniMax-M2 routed experts.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__(config, gate=MiniMaxM2Gate(config), norm_class=RMSNorm) + + +class MiniMaxM2TextModel(TextModel): + """MiniMax-M2 backbone with full-vector Q/K norms and partial NeoX RoPE.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + self.layers = nn.ModuleList( + [MiniMaxM2DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + + +class MiniMaxM2GGUFCausalLMModel(CausalLMModel): + """MiniMax-M2 GGUF decoder with full attention and all-layer routed MoE.""" + + default_task: str = "text-generation" + category: str = "Mixture of Experts" + + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + self._replace_text_model(MiniMaxM2TextModel(config)) + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + return super().preprocess_weights(state_dict) diff --git a/src/mobius/models/gguf_mistral4.py b/src/mobius/models/gguf_mistral4.py new file mode 100644 index 000000000..02270267f --- /dev/null +++ b/src/mobius/models/gguf_mistral4.py @@ -0,0 +1,327 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Exact explicit-float Mistral4 GGUF decoder with a latent K-only cache.""" + +from __future__ import annotations + +import dataclasses +import math + +import onnx_ir as ir +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig +from mobius.components import ( + MLP, + Embedding, + Linear, + RMSNorm, + create_attention_bias, + initialize_rope, +) +from mobius.components._rotary_embedding import apply_rotary_pos_emb +from mobius.models.base import CausalLMModel +from mobius.models.deepseek import DeepSeekMoEGate, DeepSeekV3CausalLMModel, _DeepSeekMoEFFN + + +class Mistral4LatentAttention(nn.Module): + """DeepSeek-V2 MLA with the graph-visible ``[latent | RoPE]`` K cache.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + if ( + config.q_lora_rank is None + or config.q_lora_rank <= 0 + or config.kv_lora_rank is None + or config.kv_lora_rank <= 0 + or config.qk_nope_head_dim is None + or config.qk_nope_head_dim <= 0 + or config.qk_rope_head_dim is None + or config.qk_rope_head_dim <= 0 + or config.v_head_dim is None + or config.v_head_dim <= 0 + ): + raise ValueError("Mistral4 requires positive Q/KV-LoRA and MLA head dimensions") + + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self._rope_interleave = config.rope_interleave + self.scaling = _mistral4_attention_scale(config, self.qk_head_dim) + + self.q_a_proj = Linear(config.hidden_size, self.q_lora_rank, bias=False) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = Linear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + ) + self.kv_a_proj_with_mqa = Linear( + config.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.k_b_proj = Linear( + self.kv_lora_rank, + self.num_heads * self.qk_nope_head_dim, + bias=False, + ) + self.v_b_proj = Linear( + self.kv_lora_rank, + self.num_heads * self.v_head_dim, + bias=False, + ) + self.o_proj = Linear( + self.num_heads * self.v_head_dim, + config.hidden_size, + bias=False, + ) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple[ir.Value, ir.Value], + past_key_value: ir.Value | None = None, + ): + # Q: (B, S, H*(D_nope + D_rope)) -> (B, S, H, D_qk). + query = self.q_b_proj( + op, + self.q_a_layernorm(op, self.q_a_proj(op, hidden_states)), + ) + query = op.Reshape(query, [0, 0, self.num_heads, self.qk_head_dim]) + query_nope, query_rope = op.Split( + query, + [self.qk_nope_head_dim, self.qk_rope_head_dim], + axis=-1, + _outputs=2, + ) + query_rope = apply_rotary_pos_emb( + op, + op.Reshape(query_rope, [0, 0, -1]), + position_embeddings, + num_heads=self.num_heads, + rotary_embedding_dim=0, + interleaved=self._rope_interleave, + ) + query = op.Concat( + query_nope, + op.Reshape( + query_rope, + [0, 0, self.num_heads, self.qk_rope_head_dim], + ), + axis=-1, + ) + + # Cache exactly the normalized latent content followed by the rotated + # single-head RoPE key: (B, 1, T, L + D_rope). + compressed = self.kv_a_proj_with_mqa(op, hidden_states) + latent, key_rope = op.Split( + compressed, + [self.kv_lora_rank, self.qk_rope_head_dim], + axis=-1, + _outputs=2, + ) + latent = self.kv_a_layernorm(op, latent) + key_rope = apply_rotary_pos_emb( + op, + key_rope, + position_embeddings, + num_heads=1, + rotary_embedding_dim=0, + interleaved=self._rope_interleave, + ) + current = op.Unsqueeze(op.Concat(latent, key_rope, axis=-1), [1]) + present = ( + current if past_key_value is None else op.Concat(past_key_value, current, axis=2) + ) + + # Re-expand the graph-visible latent cache for portable ONNX Attention. + # This is algebraically equivalent to llama.cpp's absorbed Q/value path. + cached = op.Squeeze(present, [1]) + cached_latent, cached_rope = op.Split( + cached, + [self.kv_lora_rank, self.qk_rope_head_dim], + axis=-1, + _outputs=2, + ) + key_nope = self.k_b_proj(op, cached_latent) + key_nope = op.Reshape( + key_nope, + [0, 0, self.num_heads, self.qk_nope_head_dim], + ) + cached_rope = op.Reshape( + cached_rope, + [0, 0, 1, self.qk_rope_head_dim], + ) + cached_rope = op.Expand(cached_rope, [1, 1, self.num_heads, 1]) + key = op.Concat(key_nope, cached_rope, axis=-1) + value = self.v_b_proj(op, cached_latent) + + output = op.Attention( + op.Reshape(query, [0, 0, -1]), + op.Reshape(key, [0, 0, -1]), + value, + attention_bias, + q_num_heads=self.num_heads, + kv_num_heads=self.num_heads, + scale=self.scaling, + ) + return self.o_proj(op, output), present + + +class Mistral4DecoderLayer(nn.Module): + """One pre-norm latent-attention block with dense or shared/routed MoE FFN.""" + + def __init__(self, config: ArchitectureConfig, *, is_moe: bool): + super().__init__() + self.self_attn = Mistral4LatentAttention(config) + self.mlp = _DeepSeekMoEFFN(config, DeepSeekMoEGate(config)) if is_moe else MLP(config) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + ) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple[ir.Value, ir.Value], + past_key_value: ir.Value | None = None, + ): + residual = hidden_states + attention_output, present = self.self_attn( + op, + self.input_layernorm(op, hidden_states), + attention_bias, + position_embeddings, + past_key_value, + ) + hidden_states = op.Add(residual, attention_output) + + residual = hidden_states + hidden_states = self.mlp( + op, + self.post_attention_layernorm(op, hidden_states), + ) + return op.Add(residual, hidden_states), present + + +class Mistral4TextModel(nn.Module): + """Mistral4 trunk with a per-layer compressed latent K cache.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + if config.first_k_dense_replace >= config.num_hidden_layers: + raise ValueError("Mistral4 requires at least one routed-expert layer") + self._dtype = config.dtype + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + config.pad_token_id, + ) + self.layers = nn.ModuleList( + [ + Mistral4DecoderLayer( + config, + is_moe=layer >= config.first_k_dense_replace, + ) + for layer in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + assert config.qk_rope_head_dim is not None + rope_config = dataclasses.replace( + config, + head_dim=config.qk_rope_head_dim, + partial_rotary_factor=1.0, + ) + rotary_emb = initialize_rope(rope_config) + if rotary_emb is None: + raise ValueError("Mistral4 requires RoPE metadata") + self.rotary_emb = rotary_emb + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + inputs_embeds: ir.Value | None = None, + ): + hidden_states = ( + inputs_embeds if inputs_embeds is not None else self.embed_tokens(op, input_ids) + ) + position_embeddings = self.rotary_emb(op, position_ids) + attention_bias = create_attention_bias( + op, + input_ids=hidden_states if input_ids is None else input_ids, + attention_mask=attention_mask, + dtype=self._dtype, + ) + present_key_values = [] + past_values = past_key_values or [None] * len(self.layers) + for layer, past_key_value in zip(self.layers, past_values): + hidden_states, present = layer( + op, + hidden_states, + attention_bias, + position_embeddings, + past_key_value, + ) + present_key_values.append(present) + return self.norm(op, hidden_states), present_key_values + + +class Mistral4GGUFCausalLMModel(DeepSeekV3CausalLMModel): + """Mistral4 GGUF decoder with MLA, mandatory shared/routed MoE, and latent cache.""" + + default_task: str = "mistral4-gguf-text-generation" + category: str = "Mixture of Experts" + + def __init__(self, config: ArchitectureConfig): + base_config = ( + config + if config.partial_rotary_factor is not None + else dataclasses.replace(config, partial_rotary_factor=1.0) + ) + CausalLMModel.__init__(self, base_config) + self.config = config + self._replace_text_model(Mistral4TextModel(config)) + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + return DeepSeekV3CausalLMModel.preprocess_weights(self, state_dict) + + def latent_cache_width(self) -> int: + """Return the serialized ``[latent | RoPE]`` width for every layer.""" + assert self.config.kv_lora_rank is not None + assert self.config.qk_rope_head_dim is not None + return self.config.kv_lora_rank + self.config.qk_rope_head_dim + + +def _mistral4_attention_scale( + config: ArchitectureConfig, + qk_head_dim: int, +) -> float: + """Return pinned llama.cpp's Mistral4 MLA softmax scale.""" + scale = qk_head_dim**-0.5 + if config.rope_type != "yarn" or not config.rope_scaling: + return scale + factor = float(config.rope_scaling.get("factor", 1.0)) + if factor <= 1.0: + return scale + mscale = 1.0 + 0.1 * math.log(factor) + return scale * mscale * mscale diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index ad45ba8d8..585e225c7 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -131,6 +131,8 @@ def __init__(self, config: ArchitectureConfig, linear_class: type | None = None) self.softmax_scale = self.head_dim**-0.5 self.wq_b = linear_class(config.q_lora_rank, self.n_heads * self.head_dim, bias=False) self.wk = linear_class(config.hidden_size, self.head_dim, bias=False) + # Follow the authoritative Transformers implementation's hard-coded + # epsilon. The pinned llama.cpp loader leaves f_norm_eps at zero. self.k_norm = LayerNorm(self.head_dim, eps=1e-6) self.weights_proj = linear_class(config.hidden_size, self.n_heads, bias=False) @@ -274,7 +276,7 @@ def __init__( indexer_type: str, linear_class: type | None = None, ): - super().__init__(config, linear_class=linear_class) + super().__init__(config, linear_class=linear_class, split_kv_b=True) self.indexer_type = indexer_type self.dtype = config.dtype self.main_key_dim = self.num_heads * self.qk_head_dim @@ -457,10 +459,13 @@ def forward( compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], axis=-1, _outputs=2 ) kv_latent = self.kv_a_layernorm(op, kv_latent) - kv = self.kv_b_proj(op, kv_latent) - kv = op.Reshape(kv, [0, 0, self.num_heads, self.qk_nope_head_dim + self.v_head_dim]) - k_nope, value = op.Split( - kv, [self.qk_nope_head_dim, self.v_head_dim], axis=-1, _outputs=2 + k_nope = op.Reshape( + self.k_b_proj(op, kv_latent), + [0, 0, self.num_heads, self.qk_nope_head_dim], + ) + value = op.Reshape( + self.v_b_proj(op, kv_latent), + [0, 0, self.num_heads, self.v_head_dim], ) value = op.Reshape(value, [0, 0, -1]) k_rope = apply_rotary_pos_emb( @@ -667,6 +672,41 @@ def __init__(self, config: ArchitectureConfig): def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: + state_dict = dict(state_dict) + kv_b_suffix = ".self_attn.kv_b_proj.weight" + for key in tuple(state_dict): + if not self.config.use_dsa or not key.endswith(kv_b_suffix): + continue + tensor = state_dict.pop(key) + expected_rows = self.config.num_attention_heads * ( + self.config.qk_nope_head_dim + self.config.v_head_dim + ) + if tensor.dim() != 2 or tensor.shape != ( + expected_rows, + self.config.kv_lora_rank, + ): + raise ValueError( + f"GLM-5.2 fused KV-B tensor {key!r} must have shape " + f"({expected_rows}, {self.config.kv_lora_rank}), got " + f"{tuple(tensor.shape)}" + ) + per_head = tensor.reshape( + self.config.num_attention_heads, + self.config.qk_nope_head_dim + self.config.v_head_dim, + self.config.kv_lora_rank, + ) + key_rows, value_rows = per_head.split( + [self.config.qk_nope_head_dim, self.config.v_head_dim], + dim=1, + ) + prefix = key[: -len("kv_b_proj.weight")] + state_dict[f"{prefix}k_b_proj.weight"] = key_rows.reshape( + -1, self.config.kv_lora_rank + ) + state_dict[f"{prefix}v_b_proj.weight"] = value_rows.reshape( + -1, self.config.kv_lora_rank + ) + mtp_keys = [] for key in state_dict: match = _LAYER_RE.match(key) diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py index 69d84d459..6db67946f 100644 --- a/src/mobius/models/glm_moe_dsa_test.py +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -374,17 +374,10 @@ def test_moe_layers_fuse_to_qmoe_when_quantized(self): class TestPreprocessWeights: def _state_dict(self, config: ArchitectureConfig) -> dict[str, torch.Tensor]: model = GlmMoeDsaCausalLMModel(config) - state = { + return { name: torch.zeros(tuple(int(d) for d in p.shape)) for name, p in model.named_parameters() } - # Convert Mobius's own attribute-path names into the HF-style - # ``model.layers.N.*`` names preprocess_weights expects on input. - renamed = { - f"model.{name}" if not name.startswith("lm_head") else name: v - for name, v in state.items() - } - return renamed def test_drops_mtp_layer_weights_with_warning(self, caplog): config = _glm_config(num_hidden_layers=4) @@ -420,6 +413,70 @@ def test_dsa_path_keeps_indexer_weights(self): out = model.preprocess_weights(state) assert any(".self_attn.indexer." in k for k in out) + @pytest.mark.parametrize("use_dsa", [True, False], ids=["dsa", "dense"]) + def test_preprocessed_weights_cover_every_graph_parameter(self, use_dsa: bool): + config = _glm_config( + num_hidden_layers=1, + indexer_types=["full"], + use_dsa=use_dsa, + ) + model = GlmMoeDsaCausalLMModel(config) + graph = build_from_module(model, config, task="glm-moe-dsa")["model"].graph + state = { + name: torch.zeros(tuple(int(d) for d in initializer.shape)) + for name, initializer in graph.initializers.items() + if initializer.const_value is None + } + expected = set(state) + + if use_dsa: + prefix = "model.layers.0.self_attn." + key_rows = state.pop(f"{prefix}k_b_proj.weight").reshape( + config.num_attention_heads, + config.qk_nope_head_dim, + config.kv_lora_rank, + ) + value_rows = state.pop(f"{prefix}v_b_proj.weight").reshape( + config.num_attention_heads, + config.v_head_dim, + config.kv_lora_rank, + ) + state[f"{prefix}kv_b_proj.weight"] = torch.cat( + (key_rows, value_rows), dim=1 + ).reshape(-1, config.kv_lora_rank) + + out = model.preprocess_weights(state) + + assert expected + assert set(out) == expected + + def test_splits_hf_fused_kv_b_projection(self): + config = _glm_config(num_hidden_layers=1, indexer_types=["full"]) + model = GlmMoeDsaCausalLMModel(config) + rows = config.num_attention_heads * (config.qk_nope_head_dim + config.v_head_dim) + fused = torch.arange(rows * config.kv_lora_rank, dtype=torch.float32).reshape( + rows, config.kv_lora_rank + ) + + out = model.preprocess_weights({"model.layers.0.self_attn.kv_b_proj.weight": fused}) + + per_head = fused.reshape( + config.num_attention_heads, + config.qk_nope_head_dim + config.v_head_dim, + config.kv_lora_rank, + ) + expected_k, expected_v = per_head.split( + [config.qk_nope_head_dim, config.v_head_dim], dim=1 + ) + torch.testing.assert_close( + out["model.layers.0.self_attn.k_b_proj.weight"], + expected_k.reshape(-1, config.kv_lora_rank), + ) + torch.testing.assert_close( + out["model.layers.0.self_attn.v_b_proj.weight"], + expected_v.reshape(-1, config.kv_lora_rank), + ) + # -------------------------------------------------------------------------- # Numeric parity: GlmMoeDsaIndexer vs the real transformers reference diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 2a680db21..4decfea6d 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -63,6 +63,7 @@ "ModelTask", "MllamaVisionLanguageTask", "MageVLTask", + "Mistral4GGUFCausalLMTask", "MiniCPMVLTask", "MuseGlimmerVLTask", "MaskedDiffusionTask", @@ -157,6 +158,7 @@ MiniMaxMusic3RVQTask, MiniMaxMusic3VocoderTask, ) +from mobius.tasks._mistral4_gguf import Mistral4GGUFCausalLMTask from mobius.tasks._moshi import MoshiDepformerTask, MoshiTemporalTask from mobius.tasks._multimodal import MultiModalTask from mobius.tasks._object_detection import ObjectDetectionTask @@ -270,6 +272,7 @@ "gemma4-unified": Gemma4UnifiedTask, "gemma4-assistant": Gemma4AssistantTask, "glm-moe-dsa": GlmMoeDsaTask, + "mistral4-gguf-text-generation": Mistral4GGUFCausalLMTask, "hunyuan-vl-mot": HunYuanVLMoTTask, "sensenova-u1": SenseNovaU1Task, "multimodal": MultiModalTask, diff --git a/src/mobius/tasks/_glm_moe_dsa.py b/src/mobius/tasks/_glm_moe_dsa.py index 30ece3047..6f9f43f91 100644 --- a/src/mobius/tasks/_glm_moe_dsa.py +++ b/src/mobius/tasks/_glm_moe_dsa.py @@ -36,6 +36,13 @@ class GlmMoeDsaTask(ModelTask): input/output wiring. """ + def __init__(self, *, static_cache: bool = False): + if static_cache: + raise ValueError( + "GLM-DSA static cache is not implemented: dsa_kv_cache_specs() " + "describes a per-layer-varying packed dynamic cache" + ) + def build(self, module, config: ArchitectureConfig) -> ModelPackage: if not config.use_dsa: return CausalLMTask().build(module, config) @@ -104,4 +111,12 @@ def _build_dsa(self, module, config: ArchitectureConfig) -> ModelPackage: builder.add_output(present_key, f"present.{i}.key") builder.add_output(present_value, f"present.{i}.value") - return ModelPackage({"model": _make_model(graph)}, config=config) + model = _make_model(graph) + model.metadata_props["mobius.cache_abi"] = ( + "glm-dsa:per-layer=expanded-key[+indexer-key-on-full-layers]," + "expanded-value;dynamic-concat;shared-indexer-selection" + ) + model.metadata_props["mobius.runtime_support"] = ( + "deferred: MTP, packed expert, and real-weight package parity are not evidenced" + ) + return ModelPackage({"model": model}, config=config) diff --git a/src/mobius/tasks/_mistral4_gguf.py b/src/mobius/tasks/_mistral4_gguf.py new file mode 100644 index 000000000..b8855b24a --- /dev/null +++ b/src/mobius/tasks/_mistral4_gguf.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Graph task for Mistral4 GGUF's per-layer latent K-only cache.""" + +from __future__ import annotations + +import onnx_ir as ir + +from mobius._configs import ArchitectureConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + + +class Mistral4GGUFCausalLMTask(ModelTask): + """Build Mistral4 with one ``[latent | RoPE]`` cache tensor per layer.""" + + def __init__(self, *, static_cache: bool = False): + if static_cache: + raise ValueError( + "Mistral4 GGUF static cache is not implemented; the dedicated task " + "currently owns only the dynamic latent K-only cache contract" + ) + + def build(self, module, config: ArchitectureConfig) -> ModelPackage: + width_fn = getattr(module, "latent_cache_width", None) + if not callable(width_fn): + raise TypeError( + f"{type(module).__name__} must implement latent_cache_width() " + "for Mistral4GGUFCausalLMTask" + ) + cache_width = int(width_fn()) + if cache_width <= 0: + raise ValueError("Mistral4 latent cache width must be positive") + + graph, builder = _make_graph("mistral4_gguf") + batch = "batch" + sequence = "sequence_len" + past_sequence = "past_sequence_len" + input_ids = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, sequence], + ) + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_sequence_len + sequence_len"], + ) + position_ids = builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, sequence], + ) + past_key_values = [ + builder.input( + f"past_key_values.{layer}.key", + dtype=config.dtype, + shape=[batch, 1, past_sequence, cache_width], + ) + for layer in range(config.num_hidden_layers) + ] + + logits, present_key_values = module( + builder.op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + builder.add_output(logits, "logits") + total_sequence = "past_sequence_len + sequence_len" + for layer, present in enumerate(present_key_values): + present.type = ir.TensorType(config.dtype) + present.shape = ir.Shape([batch, 1, total_sequence, cache_width]) + builder.add_output(present, f"present.{layer}.key") + + model = _make_model(graph) + model.metadata_props["mobius.cache_abi"] = ( + "mistral4:per-layer=latent_key[normalized_kv|rotated_rope];" + "dynamic-concat;no-value-cache" + ) + model.metadata_props["mobius.runtime_support"] = ( + "deferred: real-weight parity and runtime packaging are not evidenced" + ) + return ModelPackage({"model": model}, config=config) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index b881813e8..7d8e3d07b 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -600,6 +600,25 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + ( + "minimax_m2_gguf", + { + "hidden_act": "silu", + "head_dim": 16, + "partial_rotary_factor": 0.5, + "attn_qk_norm": True, + "attn_qk_norm_full": True, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 32, + "scoring_func": "sigmoid", + "norm_topk_prob": True, + "routing_weight_normalization_floor": 6.103515625e-5, + "use_expert_bias": True, + "disable_qmoe": True, + }, + True, + ), # === Mixture of Experts === ( "phimoe", diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 2686c30f6..ee4b312f2 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -6072,6 +6072,8 @@ def test_jamba_preprocesses_fused_experts_in_numeric_order(self): _SPECIALIZED_TEST_MODEL_TYPES: set[str] = { # Internal GGUF-only graph covered by _exact_legacy_decoder_test.py. "gguf_legacy", + # Mistral4 owns a K-only latent cache, covered by _remaining_dense_test.py. + "mistral4_gguf", # T5 encoder-only hidden-state contract (co-located models/t5_test.py). "t5encoder", # LLaDA masked-diffusion LM (co-located src/mobius/models/llada_test.py): diff --git a/tests/gguf_remaining_dense_artifact_test.py b/tests/gguf_remaining_dense_artifact_test.py new file mode 100644 index 000000000..31eef6ef2 --- /dev/null +++ b/tests/gguf_remaining_dense_artifact_test.py @@ -0,0 +1,50 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Metadata-only immutable artifact evidence for the remaining dense GGUF cohort.""" + +from __future__ import annotations + +import pytest +from huggingface_hub import HfApi + +from mobius.integrations.gguf._artifact_blocker_evidence import ( + MAX_BOUNDED_ARTIFACT_BYTES, + GGUFArtifactBlockerEvidence, + iter_artifact_blocker_evidence, +) + + +@pytest.mark.integration +@pytest.mark.parametrize( + "artifact", + iter_artifact_blocker_evidence(), + ids=lambda item: item.architecture, +) +def test_real_candidate_identity_and_budget_blocker( + artifact: GGUFArtifactBlockerEvidence, +) -> None: + """Verify exact LFS identities without downloading any GGUF payload.""" + records = HfApi().get_paths_info( + artifact.repository, + [file.path for file in artifact.files], + revision=artifact.revision, + ) + observed = {record.path: record for record in records} + assert set(observed) == {file.path for file in artifact.files} + for expected in artifact.files: + record = observed[expected.path] + assert record.size == expected.size + assert record.lfs is not None + assert record.lfs.sha256 == expected.lfs_sha256 + assert artifact.total_size > MAX_BOUNDED_ARTIFACT_BYTES + + +def test_no_real_candidate_can_enter_the_bounded_payload_set() -> None: + artifacts = iter_artifact_blocker_evidence() + assert {artifact.architecture for artifact in artifacts} == { + "glm-dsa", + "minimax-m2", + "mistral4", + } + assert all(artifact.total_size > MAX_BOUNDED_ARTIFACT_BYTES for artifact in artifacts) diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 5209fce6a..47f2b303e 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -211,6 +211,12 @@ def _all_registered_with_test_id() -> dict[str, str]: "smallthinker_gguf": "GGUF-only SmallThinker route whose routing and SWA/NoPE schedules " "are restored by GGUF metadata postprocessing, not the native HF config path; dedicated " "config/graph/synthetic execution tests cover it while real-weight L4/L5 is deferred.", + "minimax_m2_gguf": "GGUF-only MiniMax-M2 route with dedicated config, tensor closure, " + "full-vector Q/K norm, partial-RoPE, routing, and cache execution tests; the smallest " + "immutable public GGUF is 46,514,882,176 bytes, above the 16 GiB evidence budget.", + "mistral4_gguf": "GGUF-only Mistral4 latent-cache route with dedicated config, tensor " + "closure, and prefill/cached-decode execution tests; the smallest immutable public GGUF " + "is 32,306,941,632 bytes, above the 16 GiB evidence budget.", "arctic_gguf": "GGUF-only Arctic route whose residual-MoE tensor layout is reconstructed " "from GGUF metadata; dedicated exact closure, graph, and synthetic execution tests cover " "it while real-weight L2/L4/L5 remains deferred.", diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index ba74f3341..f1ba6f267 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -266,6 +266,8 @@ "grok_gguf", "grovemoe_gguf", "hunyuan_moe_gguf", + "minimax_m2_gguf", + "mistral4_gguf", # gemma4_unified_text: mobius-internal alias for the gemma-4-12B text # backbone (reuses Gemma4CausalLMModel). No matching HF model_type is # registered with AutoModelForCausalLM, so a reference model cannot be From 1e9e85176f6eeb4779d478c50f20e06083c4ea5c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 27 Aug 2026 23:32:12 -0700 Subject: [PATCH 2/4] Fix dense GGUF review findings Canonicalize GLM-DSA alias dispatch while retaining raw metadata prefixes, force the dedicated task, and add alias-spelled weight/execution closure. Restore Mistral4 YaRN scaling from serialized metadata and add an independent mutation discriminator. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0ee8718a-33b9-4015-9e20-3745bad54686 Signed-off-by: Justin Chu --- src/mobius/integrations/gguf/_builder.py | 6 +- .../integrations/gguf/_config_mapping.py | 7 +- .../integrations/gguf/_remaining_dense.py | 12 ++- .../gguf/_remaining_dense_test.py | 89 ++++++++++++++++++- src/mobius/models/gguf_mistral4.py | 29 ++---- 5 files changed, 112 insertions(+), 31 deletions(-) diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 114c193c4..f860fb314 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -7118,7 +7118,7 @@ def build_from_gguf( "smallthinker GGUF only supports the dedicated " "'smallthinker-gguf-text-generation' task" ) - if gguf_arch == "glm-dsa": + if spec.gguf_arch == "glm-dsa": from mobius.tasks import GlmMoeDsaTask if static_cache: @@ -7128,6 +7128,8 @@ def build_from_gguf( ) if task is not None and task != "glm-moe-dsa" and not isinstance(task, GlmMoeDsaTask): raise ValueError("glm-dsa GGUF only supports the dedicated 'glm-moe-dsa' task") + if task is None: + task = GlmMoeDsaTask() if gguf_arch == "mistral4": from mobius.tasks import Mistral4GGUFCausalLMTask @@ -7348,7 +7350,7 @@ def build_from_gguf( # ``_gguf_arch`` matters most: it is the key ``process_tensors`` dispatches # on, so losing it here would silently demote every non-float32 and every # quantized import to the ``model_type`` fallback. - config._gguf_arch = gguf_arch + config._gguf_arch = spec.gguf_arch config._gguf_model_type = model_type config._gguf_nextn_predict_layers = mtp_predict_layers config._gguf_mtp_block_indices = mtp_block_indices diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 0a58f0d06..a457ec2d7 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -4188,7 +4188,7 @@ def _glm_dsa_postprocess( f"{config.num_hidden_layers}], got {dense_prefix}" ) - indexer_types = _glm_dsa_indexer_types(config, metadata) + indexer_types = _glm_dsa_indexer_types(config, metadata, arch=arch) names = set(model.tensor_names) routed_layers = set(range(dense_prefix, config.num_hidden_layers)) bias_layers = { @@ -4248,9 +4248,10 @@ def _glm_dsa_postprocess( def _glm_dsa_indexer_types( config: ArchitectureConfig, metadata: dict[str, Any], + *, + arch: str = "glm-dsa", ) -> list[str]: - """Resolve the pinned bool/scalar indexer schedule without model-ID heuristics.""" - arch = "glm-dsa" + """Resolve the raw-prefix bool/scalar schedule without model-ID heuristics.""" layers = config.num_hidden_layers raw = metadata.get(f"{arch}.attention.indexer.types") if raw is None: diff --git a/src/mobius/integrations/gguf/_remaining_dense.py b/src/mobius/integrations/gguf/_remaining_dense.py index 37f4302ed..ae5ccb4f4 100644 --- a/src/mobius/integrations/gguf/_remaining_dense.py +++ b/src/mobius/integrations/gguf/_remaining_dense.py @@ -8,6 +8,7 @@ import re from mobius._configs import ArchitectureConfig +from mobius.integrations.gguf._arch_registry import try_get_arch_spec from mobius.integrations.gguf._tensor_mapping import is_known_skip @@ -312,7 +313,7 @@ def _validate_mistral4(model) -> None: def _validate_glm_dsa(model) -> None: metadata = model.metadata - arch = "glm-dsa" + arch = model.architecture nextn = int(metadata.get(f"{arch}.nextn_predict_layers", 0)) total_layers = int(metadata[f"{arch}.block_count"]) layers = total_layers - nextn @@ -375,6 +376,7 @@ def _validate_glm_dsa(model) -> None: max_position_embeddings=int(metadata[f"{arch}.context_length"]), ), metadata, + arch=arch, ) required: dict[str, tuple[int, ...]] = { "token_embd.weight": (vocab, hidden), @@ -484,9 +486,11 @@ def _validate_glm_dsa(model) -> None: def validate_remaining_dense_tensor_contract(model) -> None: """Validate the exact executable tensor subset for the three promoted routes.""" - if model.architecture == "minimax-m2": + spec = try_get_arch_spec(model.architecture) + architecture = spec.gguf_arch if spec is not None else model.architecture + if architecture == "minimax-m2": _validate_minimax_m2(model) - elif model.architecture == "mistral4": + elif architecture == "mistral4": _validate_mistral4(model) - elif model.architecture == "glm-dsa": + elif architecture == "glm-dsa": _validate_glm_dsa(model) diff --git a/src/mobius/integrations/gguf/_remaining_dense_test.py b/src/mobius/integrations/gguf/_remaining_dense_test.py index abed5d266..12bf0d7b9 100644 --- a/src/mobius/integrations/gguf/_remaining_dense_test.py +++ b/src/mobius/integrations/gguf/_remaining_dense_test.py @@ -6,6 +6,7 @@ from __future__ import annotations import dataclasses +import json from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -231,6 +232,15 @@ def _glm_fixture() -> _FakeGGUF: return source +def _glm_alias_fixture() -> _FakeGGUF: + source = _glm_fixture() + source.architecture = "glm_dsa" + source.metadata = { + key.replace("glm-dsa.", "glm_dsa."): value for key, value in source.metadata.items() + } + return source + + def _write_gguf(path: Path, source: _FakeGGUF) -> None: from gguf import GGUFWriter @@ -393,8 +403,20 @@ def test_mistral4_restores_pinned_yarn_log_multiplier() -> None: assert config.rope_scaling is not None assert config.rope_scaling["mscale"] == pytest.approx(1.0) assert config.rope_scaling["mscale_all_dim"] == pytest.approx(0.707) - expected_scale = (1.0 + 0.1 * np.log(4.0)) ** 2 / np.sqrt(4.0) - assert Mistral4LatentAttention(config).scaling == pytest.approx(expected_scale) + factor = float(source.metadata["mistral4.rope.scaling.factor"]) + serialized_multiplier = float(source.metadata["mistral4.rope.scaling.yarn_log_multiplier"]) + qk_head_dim = int(source.metadata["mistral4.attention.key_length_mla"]) + expected_scale = (1.0 + serialized_multiplier * np.log(factor)) ** 2 / np.sqrt(qk_head_dim) + actual_scale = Mistral4LatentAttention(config).scaling + assert actual_scale == pytest.approx(expected_scale) + assert actual_scale == pytest.approx(0.6028141, rel=1e-6) + + source.metadata["mistral4.rope.scaling.yarn_log_multiplier"] = 0.1 + mutated_config = gguf_to_config(source) + mutated_scale = Mistral4LatentAttention(mutated_config).scaling + mutated_expected = (1.0 + 0.1 * np.log(factor)) ** 2 / np.sqrt(qk_head_dim) + assert mutated_scale == pytest.approx(mutated_expected) + assert mutated_scale != pytest.approx(actual_scale) def test_glm_dsa_indexer_schedule_preserves_explicit_and_legacy_defaults() -> None: @@ -412,6 +434,69 @@ def test_glm_dsa_indexer_schedule_preserves_explicit_and_legacy_defaults() -> No gguf_to_config(source) +def test_glm_dsa_alias_closes_schedule_weights_and_dense_execution( + tmp_path: Path, +) -> None: + from mobius.integrations.gguf._builder import _normalize_gguf_weights + from mobius.integrations.gguf._tensor_processors import process_tensors + from mobius.models.glm_moe_dsa import GlmMoeDsaCausalLMModel + + source = _glm_alias_fixture() + config = gguf_to_config(source) + assert config._gguf_arch == "glm-dsa" + assert config.indexer_types == ["full", "shared"] + validate_remaining_dense_tensor_contract(source) + + path = tmp_path / "glm_dsa-alias.gguf" + _write_gguf(path, source) + package = build_from_gguf(path, keep_quantized=False) + assert package.gguf_architecture == "glm-dsa" + assert package.config._gguf_arch == "glm-dsa" + route = json.loads(package.gguf_import_route) + assert route["task"]["class"].endswith(".GlmMoeDsaTask") + initializers = package["model"].graph.initializers + assert all(initializer.const_value is not None for initializer in initializers.values()) + assert any(name.startswith("model.layers.0.self_attn.indexer.") for name in initializers) + assert not any( + name.startswith("model.layers.1.self_attn.indexer.") for name in initializers + ) + + dense_config = dataclasses.replace(config, use_dsa=False) + state_dict = { + mapped: torch.from_numpy(np.array(value, copy=True)) + for name, value in source.tensors.items() + if (mapped := map_gguf_to_hf_names(name, source.architecture)) is not None + } + state_dict = process_tensors(state_dict, dense_config) + state_dict = _normalize_gguf_weights( + state_dict, + "glm-dsa", + dense_config, + ) + module = GlmMoeDsaCausalLMModel(dense_config) + state_dict = module.preprocess_weights(state_dict) + dense_package = build_from_module(module, dense_config, task="glm-moe-dsa") + dense_package.apply_weights(state_dict) + session = OnnxModelSession(dense_package["model"]) + tokens = np.array([[1, 3]], dtype=np.int64) + feeds: dict[str, np.ndarray] = { + "input_ids": tokens, + "attention_mask": np.ones_like(tokens), + "position_ids": np.arange(tokens.shape[1], dtype=np.int64).reshape(1, -1), + } + for layer in range(dense_config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = np.empty( + (1, dense_config.num_attention_heads, 0, 4), + dtype=np.float32, + ) + feeds[f"past_key_values.{layer}.value"] = np.empty( + (1, dense_config.num_attention_heads, 0, 2), + dtype=np.float32, + ) + result = session.run(feeds) + assert np.isfinite(result["logits"]).all() + + def test_glm_dsa_rejects_unowned_rope_scaling() -> None: source = _glm_fixture() source.metadata["glm-dsa.rope.scaling.type"] = "yarn" diff --git a/src/mobius/models/gguf_mistral4.py b/src/mobius/models/gguf_mistral4.py index 02270267f..1ec710125 100644 --- a/src/mobius/models/gguf_mistral4.py +++ b/src/mobius/models/gguf_mistral4.py @@ -6,7 +6,6 @@ from __future__ import annotations import dataclasses -import math import onnx_ir as ir import torch @@ -21,8 +20,7 @@ create_attention_bias, initialize_rope, ) -from mobius.components._rotary_embedding import apply_rotary_pos_emb -from mobius.models.base import CausalLMModel +from mobius.components._rotary_embedding import apply_rotary_pos_emb, yarn_apply_mscale from mobius.models.deepseek import DeepSeekMoEGate, DeepSeekV3CausalLMModel, _DeepSeekMoEFFN @@ -53,7 +51,13 @@ def __init__(self, config: ArchitectureConfig): self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim self.v_head_dim = config.v_head_dim self._rope_interleave = config.rope_interleave - self.scaling = _mistral4_attention_scale(config, self.qk_head_dim) + # The GGUF postprocessor restores yarn_log_multiplier as + # mscale_all_dim; keep the softmax scale metadata-owned. + self.scaling = yarn_apply_mscale( + config.rope_type, + config.rope_scaling, + self.qk_head_dim**-0.5, + ) self.q_a_proj = Linear(config.hidden_size, self.q_lora_rank, bias=False) self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) @@ -296,7 +300,7 @@ def __init__(self, config: ArchitectureConfig): if config.partial_rotary_factor is not None else dataclasses.replace(config, partial_rotary_factor=1.0) ) - CausalLMModel.__init__(self, base_config) + super().__init__(base_config) self.config = config self._replace_text_model(Mistral4TextModel(config)) @@ -310,18 +314,3 @@ def latent_cache_width(self) -> int: assert self.config.kv_lora_rank is not None assert self.config.qk_rope_head_dim is not None return self.config.kv_lora_rank + self.config.qk_rope_head_dim - - -def _mistral4_attention_scale( - config: ArchitectureConfig, - qk_head_dim: int, -) -> float: - """Return pinned llama.cpp's Mistral4 MLA softmax scale.""" - scale = qk_head_dim**-0.5 - if config.rope_type != "yarn" or not config.rope_scaling: - return scale - factor = float(config.rope_scaling.get("factor", 1.0)) - if factor <= 1.0: - return scale - mscale = 1.0 + 0.1 * math.log(factor) - return scale * mscale * mscale From ae9a6c56787ecc768df3d9368b93f26c9a97d834 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 28 Aug 2026 00:35:17 -0700 Subject: [PATCH 3/4] Reject invalid Mistral4 YaRN multipliers Validate the serialized float32 yarn_log_multiplier before graph construction, accept the exact encoded [0.0, 0.1] endpoints, and reject non-finite or out-of-range values with real GGUF round-trip coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0ee8718a-33b9-4015-9e20-3745bad54686 Signed-off-by: Justin Chu --- .../integrations/gguf/_config_mapping.py | 19 ++++- .../gguf/_remaining_dense_test.py | 71 ++++++++++++++++++- 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index a457ec2d7..4c42593c5 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -61,6 +61,8 @@ logger = logging.getLogger(__name__) +_MISTRAL4_YARN_LOG_MULTIPLIER_UNIT = float(np.float32(0.1)) + # Map GGUF architecture names → our registry model_type strings. # @@ -4436,8 +4438,23 @@ def _mistral4_postprocess( if yarn_log_multiplier is not None: if rope_scaling is None or config.rope_type != "yarn": raise ValueError("Mistral4 yarn_log_multiplier requires YaRN rope metadata") + if isinstance(yarn_log_multiplier, bool) or not isinstance( + yarn_log_multiplier, + (int, float, np.integer, np.floating), + ): + raise TypeError("Mistral4 yarn_log_multiplier must be a numeric scalar") + yarn_log_multiplier = float(yarn_log_multiplier) + if ( + not math.isfinite(yarn_log_multiplier) + or not 0.0 <= yarn_log_multiplier <= _MISTRAL4_YARN_LOG_MULTIPLIER_UNIT + ): + raise ValueError( + "Mistral4 yarn_log_multiplier must be finite and within [0.0, 0.1]" + ) rope_scaling["mscale"] = 1.0 - rope_scaling["mscale_all_dim"] = float(yarn_log_multiplier) / 0.1 + rope_scaling["mscale_all_dim"] = ( + yarn_log_multiplier / _MISTRAL4_YARN_LOG_MULTIPLIER_UNIT + ) return dataclasses.replace( config, diff --git a/src/mobius/integrations/gguf/_remaining_dense_test.py b/src/mobius/integrations/gguf/_remaining_dense_test.py index 12bf0d7b9..603a00baa 100644 --- a/src/mobius/integrations/gguf/_remaining_dense_test.py +++ b/src/mobius/integrations/gguf/_remaining_dense_test.py @@ -246,7 +246,9 @@ def _write_gguf(path: Path, source: _FakeGGUF) -> None: writer = GGUFWriter(str(path), source.architecture) for key, value in source.metadata.items(): - if isinstance(value, bool): + if isinstance(value, str): + writer.add_string(key, value) + elif isinstance(value, bool): writer.add_bool(key, value) elif isinstance(value, (int, np.integer)): writer.add_uint32(key, int(value)) @@ -419,6 +421,73 @@ def test_mistral4_restores_pinned_yarn_log_multiplier() -> None: assert mutated_scale != pytest.approx(actual_scale) +@pytest.mark.parametrize( + "yarn_log_multiplier", + [ + float("nan"), + float("inf"), + float("-inf"), + -1e-6, + float(np.nextafter(np.float32(0.1), np.float32(np.inf))), + ], + ids=["nan", "positive-inf", "negative-inf", "below-range", "above-range"], +) +def test_mistral4_rejects_invalid_yarn_log_multiplier( + yarn_log_multiplier: float, +) -> None: + source = _m4_fixture() + source.metadata.update( + { + "mistral4.rope.scaling.type": "yarn", + "mistral4.rope.scaling.factor": 4.0, + "mistral4.rope.scaling.original_context_length": 8, + "mistral4.rope.scaling.yarn_beta_fast": 32.0, + "mistral4.rope.scaling.yarn_beta_slow": 1.0, + "mistral4.rope.scaling.yarn_log_multiplier": yarn_log_multiplier, + } + ) + with pytest.raises( + ValueError, + match=r"yarn_log_multiplier must be finite and within \[0\.0, 0\.1\]", + ): + gguf_to_config(source) + + +def test_mistral4_accepts_finite_yarn_log_multiplier_boundaries( + tmp_path: Path, +) -> None: + from mobius.integrations.gguf._reader import GGUFModel + + scales = [] + for yarn_log_multiplier in (0.0, 0.1): + source = _m4_fixture() + source.metadata.update( + { + "mistral4.rope.scaling.type": "yarn", + "mistral4.rope.scaling.factor": 4.0, + "mistral4.rope.scaling.original_context_length": 8, + "mistral4.rope.scaling.yarn_beta_fast": 32.0, + "mistral4.rope.scaling.yarn_beta_slow": 1.0, + "mistral4.rope.scaling.yarn_log_multiplier": yarn_log_multiplier, + } + ) + path = tmp_path / f"mistral4-yarn-{yarn_log_multiplier}.gguf" + _write_gguf(path, source) + roundtripped = GGUFModel(path) + config = gguf_to_config(roundtripped) + assert config.rope_scaling is not None + assert config.rope_scaling["mscale_all_dim"] == pytest.approx( + yarn_log_multiplier / 0.1 + ) + attention = Mistral4LatentAttention(config) + model = Mistral4GGUFCausalLMModel(config) + scales.append(attention.scaling) + assert np.isfinite(attention.scaling) + assert np.isfinite(model.model.rotary_emb.cos_cache.const_value.numpy()).all() + assert np.isfinite(model.model.rotary_emb.sin_cache.const_value.numpy()).all() + assert scales[0] != pytest.approx(scales[1]) + + def test_glm_dsa_indexer_schedule_preserves_explicit_and_legacy_defaults() -> None: source = _glm_fixture() assert gguf_to_config(source).indexer_types == ["full", "shared"] From 208f2cb3c7c5215ea9bcf8547a54109acf768bee Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 28 Aug 2026 00:41:30 -0700 Subject: [PATCH 4/4] Reconcile GGUF cohorts after rebase Pin the combined supported-architecture and remaining-route census after integrating the Grok, GroveMoE, and Hunyuan-MoE cohort from main with the MiniMax-M2, Mistral4, and GLM-DSA cohort. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0ee8718a-33b9-4015-9e20-3745bad54686 Signed-off-by: Justin Chu --- src/mobius/integrations/gguf/_arch_registry_test.py | 2 +- src/mobius/integrations/gguf/_route_census_test.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mobius/integrations/gguf/_arch_registry_test.py b/src/mobius/integrations/gguf/_arch_registry_test.py index 6edb68ce5..b94d07ec3 100644 --- a/src/mobius/integrations/gguf/_arch_registry_test.py +++ b/src/mobius/integrations/gguf/_arch_registry_test.py @@ -60,7 +60,7 @@ #: Number of importable architectures. Pinned so that adding support is a #: deliberate act that also updates the documented support matrix, and so that #: accidentally losing an architecture is a failure rather than a silence. -_EXPECTED_SUPPORTED_COUNT = 102 +_EXPECTED_SUPPORTED_COUNT = 105 _PROMOTED_CONVENTIONAL_DECODERS = frozenset( { "bitnet", diff --git a/src/mobius/integrations/gguf/_route_census_test.py b/src/mobius/integrations/gguf/_route_census_test.py index 8dc075ba3..665e27925 100644 --- a/src/mobius/integrations/gguf/_route_census_test.py +++ b/src/mobius/integrations/gguf/_route_census_test.py @@ -64,8 +64,8 @@ def test_every_route_has_one_actionable_classification() -> None: } assert Counter(item.category for item in items) == { "dependency-or-runtime-abi-blocked": 107, - "evidence-only": 97, - "immediately-implementable": 50, + "evidence-only": 100, + "immediately-implementable": 47, "intentionally-rejected": 19, "artifact-unavailable": 5, }