diff --git a/docs/api/build_from_gguf.md b/docs/api/build_from_gguf.md index 685edd67d..b9b6a74ae 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': 46, 'rejected': 2, 'supported': 100}; importable: 99; quantized import: {'rejected': 41, 'supported': 107}; runtime: {'deferred': 142, 'rejected': 2, 'supported': 4} | +| Architectures | 148 | graph verdicts: {'deferred': 46, 'rejected': 2, 'supported': 100}; importable: 100; quantized import: {'rejected': 41, 'supported': 107}; runtime: {'deferred': 142, 'rejected': 2, 'supported': 4} | | 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': 11, 'validated-pinned-source': 31} | @@ -98,13 +98,12 @@ remain machine-readable in `_route_census.py`; this table groups only shared nex | `dependency-or-runtime-abi-blocked` | `mtp-specialized-abi` | `mtp:bailingmoe3`, `mtp:cohere2moe`, `mtp:deepseek2`, `mtp:deepseek32`, `mtp:deepseek4`, `mtp:glm-dsa`, `mtp:mimo2`, `mtp:nemotron_h_moe`, `mtp:qwen35moe`, `mtp:qwen3next`, `mtp:step35` | specialized sidecar graph; routed/cache state ABI | | `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 | -| `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:gpt2`, `architecture:gptneox`, `architecture:granite`, `architecture:granitemoe`, `architecture:hunyuan-dense`, `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:mpt`, `architecture:muse-glimmer`, `architecture:nemotron`, `architecture:nemotron_h`, `architecture:neo-bert`, `architecture:nomic-bert`, `architecture:nomic-bert-moe`, `architecture:olmo`, `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:starcoder`, `architecture:starcoder2`, `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:gpt2`, `architecture:gptneox`, `architecture:granite`, `architecture:granitemoe`, `architecture:hunyuan-dense`, `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:mpt`, `architecture:muse-glimmer`, `architecture:nemotron`, `architecture:nemotron_h`, `architecture:neo-bert`, `architecture:nomic-bert`, `architecture:nomic-bert-moe`, `architecture:olmo`, `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:starcoder`, `architecture:starcoder2`, `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 | | `evidence-only` | `tokenizer-artifact-evidence` | `tokenizer:bailingmoe`, `tokenizer:bailingmoe2`, `tokenizer:chatglm-bpe`, `tokenizer:cohere2moe`, `tokenizer:glm4`, `tokenizer:llada-moe`, `tokenizer:tiny_aya` | immutable GGUF/source pair; ordered vocabulary and encoding parity | | `immediately-implementable` | `architecture-implementation` | `architecture:grok`, `architecture:grovemoe`, `architecture:hunyuan-moe`, `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 | @@ -199,7 +198,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=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. | | `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. | 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 e260af984..4c538e9fc 100644 --- a/src/mobius/integrations/gguf/_arch_registry.py +++ b/src/mobius/integrations/gguf/_arch_registry.py @@ -2079,13 +2079,11 @@ model_type="glm_moe_dsa", aliases=frozenset({"glm_dsa"}), config_key_map="glm_dsa", - tensor_map=Support.DEFERRED, - 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 - ), + config_postprocessor="glm_dsa", + tensor_map_recipe=("glm_dsa",), + tensor_processor="glm_dsa", + rope_interleave=True, + reason=_RUNTIME_VALIDATION_PENDING, ), GGUFArchitectureSpec( gguf_arch="apertus", diff --git a/src/mobius/integrations/gguf/_arch_registry_test.py b/src/mobius/integrations/gguf/_arch_registry_test.py index cbbaf2cca..a4dd0b3c2 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 = 99 +_EXPECTED_SUPPORTED_COUNT = 100 _PROMOTED_CONVENTIONAL_DECODERS = frozenset( { "bitnet", @@ -190,6 +190,7 @@ "gemma3", "gemma4", "granite", + "glm-dsa", "granitemoe", "hunyuan-dense", "jamba", diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 937c6cd1c..456c1d86a 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -7760,7 +7760,7 @@ def _replace_native_block_linears( hf_name = name_mapper(gguf_name, gguf_arch) if hf_name is None: continue - if gguf_arch in {"kimi-linear", "kimi-k3"} and hf_name.endswith( + if gguf_arch in {"glm-dsa", "kimi-linear", "kimi-k3"} and hf_name.endswith( (".k_b_proj.weight", ".v_b_proj.weight") ): continue @@ -8409,11 +8409,12 @@ 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 { + "glm-dsa", "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 @@ -8422,7 +8423,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) @@ -9301,11 +9302,12 @@ 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 { + "glm-dsa", "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 = [] @@ -9360,7 +9362,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}: " @@ -9627,7 +9629,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 da0aea79d..43fafd318 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -136,10 +136,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", @@ -147,6 +148,9 @@ "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", @@ -3920,6 +3924,76 @@ 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 + raw_gating = metadata.get(f"{arch}.expert_gating_func") + gating = int(raw_gating) if raw_gating is not None else None + if gating not in {None, 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}" + ) + + 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=float(metadata.get(f"{arch}.expert_weights_scale", 1.0)), + norm_topk_prob=bool(metadata.get(f"{arch}.expert_weights_norm", True)), + hidden_act="silu", + scoring_func="sigmoid" if gating == 2 else config.scoring_func, + topk_method="noaux_tc", + use_expert_bias=True, + disable_qmoe=True, + rope_interleave=True, + indexer_rope_interleave=True, + index_topk_freq=4, + index_skip_topk_offset=3, + partial_rotary_factor=None, + ) + return ArchitectureConfig(**fields) + + def _kimi_k3_postprocess( config: ArchitectureConfig, metadata: dict[str, Any], @@ -4745,6 +4819,7 @@ def _qwen4exp_postprocess( "maincoder": _maincoder_postprocess, "t5": _t5_postprocess, "minimax": _minimax_postprocess, + "glm_dsa": _glm_dsa_postprocess, "kimi_linear": _kimi_linear_postprocess, "kimi_k3": _kimi_k3_postprocess, "minicpm": _minicpm_postprocess, @@ -4876,6 +4951,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/_glm_moe_dsa_test.py b/src/mobius/integrations/gguf/_glm_moe_dsa_test.py index 6d6d598ca..5f323853d 100644 --- a/src/mobius/integrations/gguf/_glm_moe_dsa_test.py +++ b/src/mobius/integrations/gguf/_glm_moe_dsa_test.py @@ -37,21 +37,31 @@ def tensor_names(self) -> list[str]: 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.block_count": 79, + "glm-dsa.nextn_predict_layers": 1, + "glm-dsa.attention.head_count": 64, + "glm-dsa.attention.head_count_kv": 1, "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.attention.indexer.head_count": 32, "glm-dsa.attention.indexer.key_length": 128, "glm-dsa.attention.indexer.top_k": 2048, } @@ -85,6 +95,58 @@ 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 + + +@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 # --------------------------------------------------------------------------- # @@ -114,6 +176,20 @@ def test_missing_expert_count_rejected(): assert_glm_moe_dsa_resolvable(config, "glm-dsa", source="no_experts.gguf") +def test_missing_sigmoid_gate_rejected(): + from mobius.integrations.gguf._config_mapping import ( + GgufArchResolutionError, + assert_glm_moe_dsa_resolvable, + 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)SIGMOID|gating"): + assert_glm_moe_dsa_resolvable(config, "glm-dsa", source="no_gating.gguf") + + def test_missing_mla_rank_rejected(): from mobius.integrations.gguf._config_mapping import ( GgufArchResolutionError, @@ -159,6 +235,10 @@ def test_rejection_lists_all_reasons(): "glm-dsa.block_count": 32, "glm-dsa.attention.head_count": 32, "glm-dsa.attention.head_count_kv": 8, + "glm-dsa.attention.key_length": 192, + "glm-dsa.attention.key_length_mla": 128, + "glm-dsa.attention.value_length_mla": 128, + "glm-dsa.rope.dimension_count": 64, "glm-dsa.feed_forward_length": 11008, "glm-dsa.vocab_size": 128000, } diff --git a/src/mobius/integrations/gguf/_route_census_test.py b/src/mobius/integrations/gguf/_route_census_test.py index bb65297c4..e373fe184 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": 100, - "evidence-only": 107, - "immediately-implementable": 53, + "evidence-only": 108, + "immediately-implementable": 52, "intentionally-rejected": 19, "artifact-unavailable": 5, } @@ -94,6 +94,7 @@ 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: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 04fc7421b..4ef24958c 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -1192,6 +1192,37 @@ 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", +} + _MINICPM3_MAPPING: dict[str, str] = { "token_embd": "model.embed_tokens", "output": "lm_head", @@ -1322,6 +1353,7 @@ "minimax": _MINIMAX_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..7b4669d63 100644 --- a/src/mobius/integrations/gguf/_tensor_processors.py +++ b/src/mobius/integrations/gguf/_tensor_processors.py @@ -486,6 +486,33 @@ def _process_kimi_linear( return state_dict +def _process_glm_dsa( + state_dict: dict[str, torch.Tensor], + config: Any, +) -> dict[str, torch.Tensor]: + """Restore GLM-5.2's split MLA K/V-B matrices to Linear row order.""" + del config + for name in tuple(state_dict): + tensor = state_dict[name] + if name.endswith(".k_b_proj.weight"): + if tensor.dim() != 3: + raise ValueError( + f"GLM-5.2 K-B tensor {name!r} must be rank 3, got {tensor.dim()}" + ) + state_dict[name] = tensor.transpose(1, 2).reshape( + tensor.shape[0] * tensor.shape[2], tensor.shape[1] + ) + elif name.endswith(".v_b_proj.weight"): + if tensor.dim() != 3: + raise ValueError( + f"GLM-5.2 V-B tensor {name!r} must be rank 3, got {tensor.dim()}" + ) + state_dict[name] = tensor.reshape( + tensor.shape[0] * tensor.shape[1], tensor.shape[2] + ) + return state_dict + + def _process_kimi_k3( state_dict: dict[str, torch.Tensor], config: Any, @@ -678,6 +705,7 @@ def _process_plamo( "plamo": _process_plamo, "plamo2": _process_plamo2, "granitehybrid": _process_granitehybrid, + "glm_dsa": _process_glm_dsa, "kimi_linear": _process_kimi_linear, "kimi_k3": _process_kimi_k3, "bloom": _process_bloom, @@ -693,6 +721,7 @@ def _process_plamo( { "kimi_k3", "kimi_linear", + "glm_dsa", "llama", "mamba", "muse_glimmer", 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/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index ad45ba8d8..add790041 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -274,7 +274,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 +457,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 +670,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