From 887070a77d60d4291b499bf80cd32749839c1d6e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 16 Jul 2026 14:34:45 +0000 Subject: [PATCH 01/23] Add GLM-5.2 MoE export support Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/research/glm52-export.md | 131 +++++++++++++++ src/mobius/_configs/_base.py | 22 ++- src/mobius/_registry.py | 4 + src/mobius/components/_deepseek_mla.py | 15 +- src/mobius/components/_moe.py | 11 +- src/mobius/integrations/gguf/_builder.py | 153 +++++++++++++++++- src/mobius/integrations/gguf/_builder_test.py | 42 +++-- .../integrations/gguf/_config_mapping.py | 96 +++++++++++ src/mobius/integrations/gguf/_reader_test.py | 72 +++++++++ src/mobius/integrations/gguf/_repacker.py | 27 +++- .../integrations/gguf/_repacker_test.py | 17 ++ .../integrations/gguf/_tensor_mapping.py | 41 ++++- .../integrations/gguf/_tensor_mapping_test.py | 16 ++ src/mobius/models/__init__.py | 2 + src/mobius/models/deepseek.py | 95 ++++++++--- src/mobius/models/glm_moe_dsa.py | 48 ++++++ src/mobius/models/glm_moe_dsa_test.py | 126 +++++++++++++++ tests/cli_test.py | 29 ++++ 18 files changed, 900 insertions(+), 47 deletions(-) create mode 100644 docs/research/glm52-export.md create mode 100644 src/mobius/models/glm_moe_dsa.py create mode 100644 src/mobius/models/glm_moe_dsa_test.py diff --git a/docs/research/glm52-export.md b/docs/research/glm52-export.md new file mode 100644 index 00000000..e85de313 --- /dev/null +++ b/docs/research/glm52-export.md @@ -0,0 +1,131 @@ +# GLM-5.2 export design + +## Verified architecture + +Sources: + +- [`zai-org/GLM-5.2` config](https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json) +- Transformers `GlmMoeDsaForCausalLM` implementation +- llama.cpp `glm-dsa` GGUF metadata and tensor names +- [`unsloth/GLM-5.2-GGUF`](https://huggingface.co/unsloth/GLM-5.2-GGUF) +- [IndexShare paper, arXiv:2603.12201](https://arxiv.org/abs/2603.12201) +- [GLM-5 report, arXiv:2602.15763](https://arxiv.org/abs/2602.15763) + +The Hugging Face model type is `glm_moe_dsa`, with +`GlmMoeDsaForCausalLM` as its architecture. llama.cpp writes the exact +`general.architecture` value `glm-dsa`. + +| Field | GLM-5.2 value | +| --- | --- | +| Backbone layers | 78 | +| MTP layers | 1 additional NextN layer | +| Hidden / dense FFN size | 6144 / 12288 | +| Attention heads / KV heads | 64 / 64 | +| QK / NoPE / RoPE / V head dimensions | 256 / 192 / 64 / 256 | +| Q-LoRA / KV-LoRA rank | 2048 / 512 | +| Routed experts / experts per token | 256 / 8 | +| Routed / shared-expert FFN size | 2048 / 2048 | +| Shared experts | 1 | +| Dense prefix | first 3 layers | +| Router | float32 sigmoid, `noaux_tc`, normalized top-k, scale 2.5 | +| Norm / activation | pre-RMSNorm, epsilon 1e-5 / SiLU | +| RoPE | interleaved, theta 8,000,000 | +| Maximum context | 1,048,576 | +| Indexer | 32 heads, head dimension 128, top-k 2048 | + +GGUF reports `block_count=79` and `nextn_predict_layers=1`; therefore the +causal-LM backbone has 78 layers. GGUF also splits the MLA KV decompression +weight into `attn_k_b` and `attn_v_b`, while Transformers exposes one +`kv_b_proj`. + +## Implemented scope + +This increment registers `glm_moe_dsa` and maps `glm-dsa` GGUF files to it. +It exports the 78-layer, two-norm MLA backbone, the three dense FFN layers, +the 75 routed-MoE layers, the shared expert, sigmoid/no-aux router, interleaved +RoPE, KV cache, and untied output head. It uses the existing portable ONNX +Attention and MoE components, so the graph can be packaged for both +ONNX Runtime GenAI runtime spellings (`onnx-genai` and `ort-genai`). + +The importer: + +- derives all MLA, MoE, IndexShare, RoPE, and NextN fields from GGUF metadata; +- removes the GGUF-only MTP block from the backbone layer count; +- fuses split `attn_k_b` and `attn_v_b` tensors into `kv_b_proj`; +- streams stacked routed-expert tensors one expert at a time instead of + materializing all 256 experts as one float tensor; and +- preserves supported 4-bit and 8-bit tensors as MatMulNBits, requantizing + mixed source types to the selected graph-wide layout when necessary. + +The generic portable MoE implementation evaluates every expert and masks the +outputs. It is correct for routing, but a fused sparse-MoE runtime kernel will +be required for practical GLM-5.2 execution. + +## Deliberately deferred + +### IndexShare sparse attention + +Layers 0-2 own full indexers. Starting at layer 6, one full indexer is used +every four layers and the following three layers reuse its selected token +indices. The current export omits indexer tensors and evaluates full causal +MLA instead. This is a coherent backbone/exporter increment, but it does **not** +preserve DSA numerics or its long-context cost. A follow-up should add an +indexer component, cross-layer top-k state, sparse prefill/decode attention, +and cache-aware position handling as one tested feature rather than partially +wiring index weights into dense Attention. + +### Improved MTP + +The final NextN block and `index_share_for_mtp_iteration` behavior are omitted. +A follow-up should model the GLM-specific MTP inputs, shared index selection, +cache contract, outputs, and ORT GenAI speculative-decoding configuration. + +## Unsloth dynamic quantization + +The six `UD-IQ1_M` GGUF shard headers were inspected. The dominant routed +expert tensors are below the currently targeted 4-bit floor: + +- routed gate/up projections: 76 `IQ1_M` and 74 `IQ2_XXS` tensors; +- routed down projections: 71 `IQ3_XXS` and 4 `IQ4_XS` tensors; +- shared experts: mainly `Q5_K` / `Q6_K`; +- MLA projections: mainly `Q8_0` / `Q5_K`; +- norms: F32; embedding and output: Q4_K. + +This increment can dequantize IQ sources and requantize them to a supported +Q4 MatMulNBits layout. That enables export but loses the intended sub-4-bit +size and may alter accuracy. Native preservation should use one or both of the +following paths. + +### Runtime custom-op requirement + +Add an ONNX Runtime GenAI execution-provider capability for a fused sparse-MoE +operator that: + +1. consumes top-k expert indices and routing weights without evaluating all + 256 experts; +2. supports the llama.cpp IQ1_M, IQ2_XXS, and IQ3_XXS block/codebook layouts + (plus IQ4_XS), including their scales/minimums; +3. fuses gate/up activation and down projection where profitable; +4. loads only selected expert blocks and supports CPU and target accelerator + implementations; and +5. advertises supported bit layouts through Mobius EP capabilities so export + selects native packed initializers, otherwise falling back to Q4 + MatMulNBits requantization. + +### Draft ONNX Runtime issue + +**Title:** Support sub-4-bit GGUF/IQ weights for MatMulNBits and sparse MoE + +**Body:** + +> GLM-5.2 dynamic GGUF distributions use IQ1_M and IQ2_XXS for most routed +> gate/up expert matrices and IQ3_XXS for most down matrices. Exporting these +> models to ONNX currently requires expansion or requantization to 4-bit, +> losing the distribution's size advantage and making a 256-expert model +> impractical. Please define supported 1/2/3-bit MatMulNBits representations, +> or provide an extensible packed/codebook-weight contract for the llama.cpp +> IQ formats, with CPU and execution-provider kernels. For MoE, an operator +> that accepts expert weights plus per-token top-k indices/weights should avoid +> evaluating or unpacking every expert. Required formats for the first target +> are IQ1_M, IQ2_XXS, IQ3_XXS, and IQ4_XS. Please also document zero-point, +> scale/codebook packing, block-size, alignment, and external-data constraints. diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 33dd17aa..098d6cc7 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -414,6 +414,7 @@ class ArchitectureConfig(BaseModelConfig): topk_method: str = "greedy" first_k_dense_replace: int = 0 n_shared_experts: int | None = None + mlp_layer_types: list[str] | None = None # Multi-head Latent Attention (MLA) config — DeepSeek-V2/V3 q_lora_rank: int | None = None @@ -423,6 +424,15 @@ class ArchitectureConfig(BaseModelConfig): v_head_dim: int | None = None rope_interleave: bool = False + # Deep Sparse Attention / IndexShare config — GLM-5.2. + # The dense-attention fallback preserves these fields for future DSA export. + index_topk: int | None = None + index_head_dim: int | None = None + index_n_heads: int | None = None + indexer_types: list[str] | None = None + index_share_for_mtp_iteration: bool = False + num_nextn_predict_layers: int = 0 + # Vision shared fields (accessed as top-level config.X by tasks) mm_tokens_per_image: int | None = None image_token_id: int | None = None @@ -580,7 +590,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: config, "rope_interleave", (getattr(config, "qk_rope_head_dim", None) or 0) > 0 - or model_type in ("glm", "glm4", "glm4_moe", "chatglm"), + or model_type in ("glm", "glm4", "glm4_moe", "glm_moe_dsa", "chatglm"), ) if rope_config is not None: rope_config = dataclasses.replace(rope_config, rope_interleave=rope_interleave) @@ -800,12 +810,22 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: topk_method=getattr(config, "topk_method", "greedy"), first_k_dense_replace=getattr(config, "first_k_dense_replace", 0), n_shared_experts=getattr(config, "n_shared_experts", None), + mlp_layer_types=getattr(config, "mlp_layer_types", None), # Multi-head Latent Attention (MLA) q_lora_rank=getattr(config, "q_lora_rank", None), kv_lora_rank=getattr(config, "kv_lora_rank", None), qk_nope_head_dim=getattr(config, "qk_nope_head_dim", None), qk_rope_head_dim=getattr(config, "qk_rope_head_dim", None), v_head_dim=getattr(config, "v_head_dim", None), + # Deep Sparse Attention / IndexShare + index_topk=getattr(config, "index_topk", None), + index_head_dim=getattr(config, "index_head_dim", None), + index_n_heads=getattr(config, "index_n_heads", None), + indexer_types=getattr(config, "indexer_types", None), + index_share_for_mtp_iteration=getattr( + config, "index_share_for_mtp_iteration", False + ), + num_nextn_predict_layers=getattr(config, "num_nextn_predict_layers", 0), # Encoder-specific type_vocab_size=getattr(config, "type_vocab_size", 0), # Encoder-decoder diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index c7ca586d..11d1dcf8 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -56,6 +56,7 @@ Glm4CausalLMModel, Glm4MoECausalLMModel, GlmCausalLMModel, + GlmMoeDsaCausalLMModel, GPTOSSCausalLMModel, GraniteCausalLMModel, GraniteMoECausalLMModel, @@ -523,6 +524,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "ernie4_5_moe": ModelRegistration(Ernie45MoECausalLMModel), "flex_olmo": ModelRegistration(MoECausalLMModel), "glm4_moe": ModelRegistration(Glm4MoECausalLMModel), + "glm_moe_dsa": ModelRegistration(GlmMoeDsaCausalLMModel), "granitemoe": ModelRegistration(GraniteMoECausalLMModel), "granitemoehybrid": ModelRegistration(GraniteMoeHybridCausalLMModel), "granitemoeshared": ModelRegistration(GraniteMoECausalLMModel), @@ -899,6 +901,7 @@ def _create_default_registry() -> ModelRegistry: "ernie4_5_moe": "baidu/ERNIE-4.5-21B-A3B-PT", "flex_olmo": "allenai/Flex-reddit-2x7B-1T", "glm4_moe": "zai-org/GLM-4.5-Air", + "glm_moe_dsa": "zai-org/GLM-5.2", "granitemoehybrid": "ibm-granite/granite-4.0-tiny-preview", "granitemoeshared": "ibm-research/moe-7b-1b-active-shared-experts", "qwen3_omni_moe": "Qwen/Qwen3-Omni-30B-A3B-Instruct", @@ -1222,6 +1225,7 @@ def _create_default_registry() -> ModelRegistry: "deepseek_v2": "mla", "deepseek_v2_moe": "mla+moe", "deepseek_v3": "mla+moe", + "glm_moe_dsa": "mla+moe-full-attention", "phi3small": "blocksparse", "falcon_h1": "hybrid-ssm", "mamba": "ssm", diff --git a/src/mobius/components/_deepseek_mla.py b/src/mobius/components/_deepseek_mla.py index 8434abf6..a16cb741 100644 --- a/src/mobius/components/_deepseek_mla.py +++ b/src/mobius/components/_deepseek_mla.py @@ -43,8 +43,11 @@ def __init__( self, config: ArchitectureConfig, scale: float | None = None, + linear_class: type | None = None, ): super().__init__() + if linear_class is None: + linear_class = Linear self.num_heads = config.num_attention_heads self.q_lora_rank = config.q_lora_rank self.kv_lora_rank = config.kv_lora_rank @@ -55,16 +58,16 @@ def __init__( # Q path: LoRA compression → norm → decompression if self.q_lora_rank is not None and self.q_lora_rank > 0: - self.q_a_proj = Linear(config.hidden_size, self.q_lora_rank, bias=False) + self.q_a_proj = linear_class(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_b_proj = linear_class( self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, ) else: # No LoRA — direct projection - self.q_proj = Linear( + self.q_proj = linear_class( config.hidden_size, self.num_heads * self.qk_head_dim, bias=False, @@ -72,20 +75,20 @@ def __init__( # KV path: joint projection for latent KV + RoPE key # Output: (kv_lora_rank) for latent KV + (qk_rope_head_dim) for k_rope - self.kv_a_proj_with_mqa = Linear( + self.kv_a_proj_with_mqa = linear_class( 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) # Decompresses latent KV into per-head k_nope + v - self.kv_b_proj = Linear( + 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( + self.o_proj = linear_class( self.num_heads * self.v_head_dim, config.hidden_size, bias=False, diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index d0bfa8dc..87858dd4 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -190,7 +190,12 @@ class MoELayer(nn.Module): then results are masked and weighted. """ - def __init__(self, config: ArchitectureConfig, gate: nn.Module | None = None): + def __init__( + self, + config: ArchitectureConfig, + gate: nn.Module | None = None, + linear_class: type | None = None, + ): super().__init__() assert config.num_local_experts is not None assert config.num_experts_per_tok is not None @@ -206,7 +211,9 @@ def __init__(self, config: ArchitectureConfig, gate: nn.Module | None = None): if config.moe_intermediate_size is not None else config ) - self.experts = nn.ModuleList([MLP(expert_config) for _ in range(self.num_experts)]) + self.experts = nn.ModuleList( + [MLP(expert_config, linear_class=linear_class) for _ in range(self.num_experts)] + ) def forward(self, op: OpBuilder, hidden_states: ir.Value): routing_weights, selected_experts = self.gate(op, hidden_states) diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index d969e1d7..849eaa2a 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -322,7 +322,19 @@ def _normalize_gguf_weights( import torch result: dict[str, torch.Tensor] = {} + kv_b_parts: dict[str, dict[str, torch.Tensor]] = {} for key, value in state_dict.items(): + for part in ("k", "v"): + suffix = f".self_attn.kv_b_proj.{part}_proj.weight" + if key.endswith(suffix): + prefix = key[: -len(suffix)] + kv_b_parts.setdefault(prefix, {})[part] = value + break + else: + part = None + if part is not None: + continue + # Stacked expert weights [num_experts, out, in] → per-expert unpacked = False for proj in ("gate_proj", "up_proj", "down_proj"): @@ -351,8 +363,31 @@ def _normalize_gguf_weights( result[key[: -len(".bias")]] = value continue + if key.endswith(".gate.e_score_correction_bias.bias"): + result[key[: -len(".bias")]] = value + continue + result[key] = value + for prefix, parts in kv_b_parts.items(): + if set(parts) != {"k", "v"}: + missing = {"k", "v"} - set(parts) + raise ValueError( + f"Incomplete split kv_b_proj for {prefix}: missing {sorted(missing)}" + ) + k_proj = parts["k"].transpose(1, 2) + v_proj = parts["v"] + if k_proj.shape[0] != v_proj.shape[0] or k_proj.shape[2] != v_proj.shape[2]: + raise ValueError( + f"Incompatible split kv_b_proj shapes for {prefix}: " + f"k={tuple(parts['k'].shape)}, v={tuple(v_proj.shape)}" + ) + fused = torch.cat((k_proj, v_proj), dim=1).reshape( + -1, + k_proj.shape[-1], + ) + result[f"{prefix}.self_attn.kv_b_proj.weight"] = fused + return result @@ -521,12 +556,11 @@ def _require_supported_requantization( block_size: int, tensor_name: str, ) -> None: - if bits != 4 or block_size != 32: + if bits not in (4, 8) or block_size != 32: raise ValueError( "keep_quantized MatMulNBits requantization currently supports only " - f"4-bit/block-32 targets; got bits={bits} block={block_size} " - f"for tensor {tensor_name}. Use keep_quantized=False or a " - "4-bit/block-32 target." + f"4-bit or 8-bit block-32 targets; got bits={bits} block={block_size} " + f"for tensor {tensor_name}. Use keep_quantized=False or a supported target." ) @@ -649,6 +683,7 @@ def _load_quantized_state_dict( can_repack, repack_dequantized_tensor, repack_gguf_tensor, + repack_quant_params, ) from mobius.integrations.gguf._tencent_q1_0 import ( is_tencent_q1_0_layout, @@ -694,6 +729,7 @@ def _load_quantized_state_dict( target_bits = config.quantization.bits target_block_size = config.quantization.group_size target_symmetric = config.quantization.sym + split_kv_b: dict[str, dict[str, np.ndarray]] = {} for gguf_name, raw, qtype, np_shape in tqdm.tqdm( gguf_model.tensor_items_raw(), @@ -708,6 +744,112 @@ def _load_quantized_state_dict( # Determine the int value of the quant type for can_repack qtype_val = qtype.value if hasattr(qtype, "value") else qtype + # GGUF stores every routed expert projection in one 3D tensor. Repack + # one expert at a time to avoid materializing a multi-gigabyte float + # tensor for GLM-5.2's 256 experts. + expert_proj = next( + ( + proj + for proj in ("gate_proj", "up_proj", "down_proj") + if hf_name.endswith(f".mlp.experts.{proj}.weight") + ), + None, + ) + if expert_proj is not None and len(np_shape) == 3: + num_experts, n_out, k_in = (int(dim) for dim in np_shape) + raw_flat = raw.reshape(-1) + if raw_flat.size % num_experts != 0: + raise ValueError( + f"Cannot split GGUF expert tensor {gguf_name}: " + f"{raw_flat.size} blocks across {num_experts} experts" + ) + blocks_per_expert = raw_flat.size // num_experts + source_params = repack_quant_params(qtype_val) + for expert_idx in range(num_experts): + raw_expert = raw_flat[ + expert_idx * blocks_per_expert : (expert_idx + 1) * blocks_per_expert + ] + if source_params == (target_bits, target_block_size): + repacked = repack_gguf_tensor( + raw_expert.ravel().view(np.uint8), + qtype_val, + (n_out, k_in), + ) + else: + _require_supported_requantization( + bits=target_bits, + block_size=target_block_size, + tensor_name=hf_name, + ) + values = gguf_model.dequantize_raw_tensor( + raw_expert, + qtype, + (n_out, k_in), + ) + repacked = repack_dequantized_tensor( + values, + bits=target_bits, + block_size=target_block_size, + symmetric=target_symmetric, + ) + n_requantized += 1 + + expert_name = hf_name.replace( + f".mlp.experts.{expert_proj}.weight", + f".mlp.experts.{expert_idx}.{expert_proj}.weight", + ) + expert_stem = expert_name[: -len(".weight")] + state_dict[expert_name] = torch.from_numpy(repacked.weight) + state_dict[f"{expert_stem}.scales"] = torch.from_numpy(repacked.scales) + if repacked.zero_points is not None: + state_dict[f"{expert_stem}.zero_points"] = torch.from_numpy( + repacked.zero_points + ) + n_repacked += 1 + continue + + # GLM-5.2 GGUF splits the MLA kv_b projection into per-head K and V + # tensors. Fuse them in float and immediately requantize the resulting + # 2D projection so only one layer's temporary arrays are resident. + kv_b_part = None + for part in ("k", "v"): + suffix = f".self_attn.kv_b_proj.{part}_proj.weight" + if hf_name.endswith(suffix): + kv_b_part = part + kv_b_prefix = hf_name[: -len(suffix)] + break + if kv_b_part is not None: + values = gguf_model.dequantize_raw_tensor(raw, qtype, np_shape) + parts = split_kv_b.setdefault(kv_b_prefix, {}) + parts[kv_b_part] = values + if set(parts) == {"k", "v"}: + k_proj = parts["k"].transpose(0, 2, 1) + v_proj = parts["v"] + fused = np.concatenate((k_proj, v_proj), axis=1).reshape( + -1, + k_proj.shape[-1], + ) + _require_supported_requantization( + bits=target_bits, + block_size=target_block_size, + tensor_name=f"{kv_b_prefix}.self_attn.kv_b_proj.weight", + ) + repacked = repack_dequantized_tensor( + fused, + bits=target_bits, + block_size=target_block_size, + symmetric=target_symmetric, + ) + stem = f"{kv_b_prefix}.self_attn.kv_b_proj" + state_dict[f"{stem}.weight"] = torch.from_numpy(repacked.weight) + state_dict[f"{stem}.scales"] = torch.from_numpy(repacked.scales) + if repacked.zero_points is not None: + state_dict[f"{stem}.zero_points"] = torch.from_numpy(repacked.zero_points) + del split_kv_b[kv_b_prefix] + n_repacked += 1 + n_requantized += 1 + continue + # Repack every target QuantizedLinear weight. Mixed GGUF presets # otherwise leave unsupported source types as full float matrices, # which cannot fit the graph's packed MatMulNBits initializer shape. @@ -801,6 +943,9 @@ def _load_quantized_state_dict( arr = dequantize(raw, qtype).reshape(np_shape) state_dict[hf_name] = torch.from_numpy(arr) + if split_kv_b: + raise ValueError(f"Incomplete split kv_b_proj tensors: {sorted(split_kv_b)}") + logger.info( "Loaded %d state_dict entries (%d GGUF tensors repacked for quantized ops, " "%d requantized from mixed source types)", diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index 1fda82d6..d01bf93e 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -269,20 +269,18 @@ def test_untied_quantized_head_uses_q4_matmulnbits( assert list(model.graph.initializers["lm_head.scales"].shape) == [256, 2] assert "lm_head.weight_t" not in model.graph.initializers - def test_unsupported_requantization_target_has_clear_error( - self, q8_0_projection_q4_head_gguf: Path - ): - """Mixed targets outside 4-bit/block-32 fail before the Q4 repacker.""" + def test_q8_target_requantizes_mixed_q4_head(self, q8_0_projection_q4_head_gguf: Path): + """A mixed Q8 target requantizes its Q4 output head to block-32 Q8.""" from mobius.integrations.gguf import build_from_gguf - with pytest.raises( - ValueError, - match=( - "keep_quantized MatMulNBits requantization currently supports only " - r"4-bit/block-32 targets; got bits=8 block=32 for tensor lm_head\.weight" - ), - ): - build_from_gguf(q8_0_projection_q4_head_gguf, keep_quantized=True) + model = build_from_gguf(q8_0_projection_q4_head_gguf, keep_quantized=True)["model"] + head = next( + node + for node in model.graph + if node.op_type == "MatMulNBits" and node.outputs[0].name == "logits" + ) + assert head.attributes["bits"].value == 8 + assert head.attributes["block_size"].value == 32 def test_norms_are_float(self, q4_0_gguf: Path): """Norm weights remain float, not quantized.""" @@ -427,3 +425,23 @@ def test_dequantize_raw_tensor_matches_get_tensor(self, q4_0_gguf: Path): expected = model.get_tensor(name) actual = model.dequantize_raw_tensor(raw, qtype, shape) np.testing.assert_array_equal(actual, expected) + + +def test_normalize_glm_dsa_split_kv_b_and_router_bias(): + import torch + + from mobius.integrations.gguf._builder import _normalize_gguf_weights + + k_proj = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4) + v_proj = torch.arange(2 * 5 * 3, dtype=torch.float32).reshape(2, 5, 3) + result = _normalize_gguf_weights( + { + "model.layers.0.self_attn.kv_b_proj.k_proj.weight": k_proj, + "model.layers.0.self_attn.kv_b_proj.v_proj.weight": v_proj, + "model.layers.0.mlp.gate.e_score_correction_bias.bias": torch.zeros(2), + } + ) + + expected = torch.cat((k_proj.transpose(1, 2), v_proj), dim=1).reshape(18, 3) + assert torch.equal(result["model.layers.0.self_attn.kv_b_proj.weight"], expected) + assert "model.layers.0.mlp.gate.e_score_correction_bias" in result diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 29e4c41e..29a390c0 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -43,6 +43,7 @@ "qwen3": "qwen3", "qwen3_moe": "qwen3_moe", "qwen35moe": "qwen3_5_moe", + "glm-dsa": "glm_moe_dsa", "gemma2": "gemma2", "gemma3": "gemma3_text", # Gemma 4 GGUF contains the text backbone only — no vision or audio encoder. @@ -89,6 +90,26 @@ "ssm.conv_kernel": "linear_conv_kernel_dim", } +_ARCH_KEY_MAPS: dict[str, dict[str, str]] = { + "glm-dsa": { + "leading_dense_block_count": "first_k_dense_replace", + "expert_shared_count": "n_shared_experts", + "expert_group_count": "n_group", + "expert_group_used_count": "topk_group", + "expert_weights_scale": "routed_scaling_factor", + "expert_weights_norm": "norm_topk_prob", + "expert_gating_func": "expert_gating_func", + "attention.q_lora_rank": "q_lora_rank", + "attention.kv_lora_rank": "kv_lora_rank", + "attention.key_length_mla": "qk_head_dim", + "attention.value_length_mla": "v_head_dim", + "attention.indexer.head_count": "index_n_heads", + "attention.indexer.key_length": "index_head_dim", + "attention.indexer.top_k": "index_topk", + "nextn_predict_layers": "num_nextn_predict_layers", + }, +} + # GGUF hidden_act values → HuggingFace activation function names _ACTIVATION_MAP: dict[str, str] = { @@ -154,6 +175,11 @@ def _extract_config_fields( len(hf_fields), ) + for gguf_suffix, hf_key in _ARCH_KEY_MAPS.get(gguf_arch, {}).items(): + full_key = f"{gguf_arch}.{gguf_suffix}" + if full_key in metadata: + hf_fields[hf_key] = metadata[full_key] + # Extract vocab_size from tokenizer token list if not in metadata if "vocab_size" not in hf_fields: tokens = metadata.get("tokenizer.ggml.tokens") @@ -346,6 +372,25 @@ def gguf_to_config( num_experts_per_tok=hf_fields.get("num_experts_per_tok"), moe_intermediate_size=hf_fields.get("moe_intermediate_size"), shared_expert_intermediate_size=hf_fields.get("shared_expert_intermediate_size"), + norm_topk_prob=hf_fields.get("norm_topk_prob", True), + n_group=hf_fields.get("n_group", 1), + topk_group=hf_fields.get("topk_group", 1), + routed_scaling_factor=hf_fields.get("routed_scaling_factor", 1.0), + scoring_func=hf_fields.get("scoring_func", "softmax"), + topk_method=hf_fields.get("topk_method", "greedy"), + first_k_dense_replace=hf_fields.get("first_k_dense_replace", 0), + n_shared_experts=hf_fields.get("n_shared_experts"), + # Multi-head Latent Attention fields + q_lora_rank=hf_fields.get("q_lora_rank"), + kv_lora_rank=hf_fields.get("kv_lora_rank"), + qk_nope_head_dim=hf_fields.get("qk_nope_head_dim"), + qk_rope_head_dim=hf_fields.get("qk_rope_head_dim"), + v_head_dim=hf_fields.get("v_head_dim"), + # DSA / MTP metadata + index_topk=hf_fields.get("index_topk"), + index_head_dim=hf_fields.get("index_head_dim"), + index_n_heads=hf_fields.get("index_n_heads"), + num_nextn_predict_layers=hf_fields.get("num_nextn_predict_layers", 0), # Hybrid architecture fields layer_types=layer_types, full_attention_interval=full_attention_interval, @@ -639,12 +684,63 @@ def _gemma3_postprocess( return config +def _glm_moe_dsa_postprocess( + config: ArchitectureConfig, + metadata: dict[str, Any], +) -> ArchitectureConfig: + """Restore GLM-5.2 fields that GGUF represents with MLA-specific keys.""" + arch = "glm-dsa" + n_mtp = int(metadata.get(f"{arch}.nextn_predict_layers", 0)) + num_hidden_layers = config.num_hidden_layers - n_mtp + qk_head_dim = int(metadata[f"{arch}.attention.key_length_mla"]) + qk_rope_head_dim = int(metadata[f"{arch}.rope.dimension_count"]) + qk_nope_head_dim = qk_head_dim - qk_rope_head_dim + n_shared_experts = int(metadata.get(f"{arch}.expert_shared_count", 0)) + first_k_dense_replace = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) + + # GLM-5.2 shares one full indexer across each four-layer DSA group. + indexer_types = [ + "full" if i < first_k_dense_replace or (i - 2) % 4 == 0 else "shared" + for i in range(num_hidden_layers) + ] + mlp_layer_types = [ + "dense" if i < first_k_dense_replace else "sparse" for i in range(num_hidden_layers) + ] + + return dataclasses.replace( + config, + num_hidden_layers=num_hidden_layers, + num_key_value_heads=config.num_attention_heads, + head_dim=qk_rope_head_dim, + partial_rotary_factor=1.0, + rope_interleave=True, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=int(metadata[f"{arch}.attention.value_length_mla"]), + q_lora_rank=int(metadata[f"{arch}.attention.q_lora_rank"]), + kv_lora_rank=int(metadata[f"{arch}.attention.kv_lora_rank"]), + n_shared_experts=n_shared_experts, + shared_expert_intermediate_size=( + config.moe_intermediate_size * n_shared_experts + if config.moe_intermediate_size is not None + else None + ), + first_k_dense_replace=first_k_dense_replace, + mlp_layer_types=mlp_layer_types, + scoring_func="sigmoid", + topk_method="noaux_tc", + indexer_types=indexer_types, + num_nextn_predict_layers=n_mtp, + ) + + # Architecture-specific config postprocessors. # Each takes a base ArchitectureConfig + raw metadata and returns # an architecture-specific config subclass. _CONFIG_POSTPROCESSORS: dict[str, Any] = { "gemma3_text": _gemma3_postprocess, "gemma4_text": _gemma4_postprocess, + "glm_moe_dsa": _glm_moe_dsa_postprocess, } diff --git a/src/mobius/integrations/gguf/_reader_test.py b/src/mobius/integrations/gguf/_reader_test.py index 4a2088a4..d705f084 100644 --- a/src/mobius/integrations/gguf/_reader_test.py +++ b/src/mobius/integrations/gguf/_reader_test.py @@ -10,6 +10,7 @@ from __future__ import annotations from pathlib import Path +from typing import ClassVar import numpy as np import pytest @@ -409,6 +410,77 @@ def test_arch_to_model_type_mapping(self): assert GGUF_ARCH_TO_MODEL_TYPE["mistral"] == "llama" assert GGUF_ARCH_TO_MODEL_TYPE["qwen2"] == "qwen2" assert GGUF_ARCH_TO_MODEL_TYPE["phi3"] == "phi3" + assert GGUF_ARCH_TO_MODEL_TYPE["glm-dsa"] == "glm_moe_dsa" + + def test_glm_dsa_config_extraction(self): + class _FakeGlmDsa: + architecture = "glm-dsa" + tensor_names: ClassVar[list[str]] = ["token_embd.weight", "output.weight"] + metadata: ClassVar[dict] = { + "glm-dsa.embedding_length": 6144, + "glm-dsa.feed_forward_length": 12288, + "glm-dsa.block_count": 79, + "glm-dsa.attention.head_count": 64, + "glm-dsa.attention.head_count_kv": 1, + "glm-dsa.attention.layer_norm_rms_epsilon": 1e-5, + "glm-dsa.attention.q_lora_rank": 2048, + "glm-dsa.attention.kv_lora_rank": 512, + "glm-dsa.attention.key_length_mla": 256, + "glm-dsa.attention.value_length_mla": 256, + "glm-dsa.attention.indexer.head_count": 32, + "glm-dsa.attention.indexer.key_length": 128, + "glm-dsa.attention.indexer.top_k": 2048, + "glm-dsa.expert_count": 256, + "glm-dsa.expert_used_count": 8, + "glm-dsa.expert_feed_forward_length": 2048, + "glm-dsa.expert_shared_count": 1, + "glm-dsa.expert_group_count": 1, + "glm-dsa.expert_group_used_count": 1, + "glm-dsa.expert_weights_scale": 2.5, + "glm-dsa.expert_weights_norm": True, + "glm-dsa.expert_gating_func": 2, + "glm-dsa.leading_dense_block_count": 3, + "glm-dsa.nextn_predict_layers": 1, + "glm-dsa.rope.dimension_count": 64, + "glm-dsa.rope.freq_base": 8_000_000.0, + "glm-dsa.context_length": 1_048_576, + "glm-dsa.vocab_size": 154880, + } + + def get_metadata(self, key, default=None): + return self.metadata.get(key, default) + + config = gguf_to_config(_FakeGlmDsa()) + + assert config.model_type == "glm_moe_dsa" + assert config.num_hidden_layers == 78 + assert config.hidden_size == 6144 + assert config.num_attention_heads == 64 + assert config.num_key_value_heads == 64 + assert config.head_dim == 64 + 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.num_local_experts == 256 + assert config.num_experts_per_tok == 8 + assert config.moe_intermediate_size == 2048 + assert config.shared_expert_intermediate_size == 2048 + assert config.first_k_dense_replace == 3 + assert config.routed_scaling_factor == pytest.approx(2.5) + assert config.scoring_func == "sigmoid" + assert config.rope_interleave is True + assert config.max_position_embeddings == 1_048_576 + assert config.indexer_types[:7] == [ + "full", + "full", + "full", + "shared", + "shared", + "shared", + "full", + ] def test_tie_embeddings_detected(self, tied_gguf: Path): model = GGUFModel(tied_gguf) diff --git a/src/mobius/integrations/gguf/_repacker.py b/src/mobius/integrations/gguf/_repacker.py index d089b251..7f369c9b 100644 --- a/src/mobius/integrations/gguf/_repacker.py +++ b/src/mobius/integrations/gguf/_repacker.py @@ -451,10 +451,10 @@ def repack_dequantized_tensor( """ if values.ndim != 2: raise ValueError(f"Expected 2D values (N, K), got shape {values.shape}") - if bits != 4 or block_size != _BLOCK_SIZE: + if bits not in (4, 8) or block_size != _BLOCK_SIZE: raise ValueError( "Float requantization currently supports only " - f"bits=4, block_size={_BLOCK_SIZE}; got bits={bits}, " + f"bits=4/8, block_size={_BLOCK_SIZE}; got bits={bits}, " f"block_size={block_size}" ) @@ -465,6 +465,29 @@ def repack_dequantized_tensor( padded[:, :k_in] = values.astype(np.float32, copy=False) blocks = padded.reshape(n_out, n_blocks, block_size) + if bits == 8: + block_min = np.minimum(blocks.min(axis=-1), 0.0) + block_max = np.maximum(blocks.max(axis=-1), 0.0) + if symmetric: + scales = np.maximum(-block_min / 128.0, block_max / 127.0) + zero_points = np.full_like(scales, 128, dtype=np.uint8) + else: + scales = (block_max - block_min) / 255.0 + safe_scales = np.where(scales != 0, scales, 1.0) + zero_points = np.clip(np.rint(-block_min / safe_scales), 0, 255).astype(np.uint8) + safe_scales = np.where(scales != 0, scales, 1.0) + quants = np.rint(blocks / safe_scales[:, :, None]) + quants += zero_points[:, :, None] + weight = np.clip(quants, 0, 255).astype(np.uint8) + weight = np.where(scales[:, :, None] != 0, weight, 0).astype(np.uint8) + return RepackedTensor( + weight=weight, + scales=scales.astype(np.float32), + zero_points=zero_points, + block_size=block_size, + bits=bits, + ) + if symmetric: # MatMulNBits' implicit uint4 zero-point is 8. Choose a scale that # covers the asymmetric signed code range [-8, 7]. diff --git a/src/mobius/integrations/gguf/_repacker_test.py b/src/mobius/integrations/gguf/_repacker_test.py index f9899f0a..7cbc5531 100644 --- a/src/mobius/integrations/gguf/_repacker_test.py +++ b/src/mobius/integrations/gguf/_repacker_test.py @@ -665,6 +665,23 @@ def test_asymmetric_q4_round_trip_bound(self): assert result.weight.shape == (3, 3, 16) assert result.zero_points.shape == (3, 2) + def test_asymmetric_q8_round_trip_bound(self): + rng = np.random.default_rng(1) + values = rng.normal(size=(3, 70)).astype(np.float32) + result = repack_dequantized_tensor(values, bits=8, block_size=32) + + n_blocks = result.scales.shape[1] + dequantized = ( + result.weight.astype(np.float32) + - result.zero_points[:, :, None].astype(np.float32) + ) * result.scales[:, :, None] + dequantized = dequantized.reshape(3, n_blocks * 32)[:, : values.shape[1]] + + max_abs_diff = float(np.max(np.abs(dequantized - values))) + assert max_abs_diff <= float(result.scales.max()) * 0.51 + 1e-6 + assert result.weight.shape == (3, 3, 32) + assert result.zero_points.shape == (3, 3) + # ---- Q1_0 tests ---- diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index 473cf6cb..dac71e96 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -187,6 +187,36 @@ "blk.{bid}.ffn_down_shexp": ("model.layers.{bid}.mlp.shared_expert.down_proj"), } +# GLM-5.2 uses MLA with split K/V decompression weights in GGUF. +# The temporary ``kv_b_proj.{k,v}_proj`` names are fused by +# ``_normalize_gguf_weights`` before model-specific preprocessing. +_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_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_k_b": "model.layers.{bid}.self_attn.kv_b_proj.k_proj", + "blk.{bid}.attn_v_b": "model.layers.{bid}.self_attn.kv_b_proj.v_proj", + "blk.{bid}.attn_q_a_norm": "model.layers.{bid}.self_attn.q_a_layernorm", + "blk.{bid}.attn_kv_a_norm": "model.layers.{bid}.self_attn.kv_a_layernorm", + "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", +} + # Qwen3.5-MoE hybrid extensions: DeltaNet (SSM) + full-attention + MoE. # DeltaNet layers use linear_attn.* naming; full-attention layers add # q_norm/k_norm under self_attn; both use post_attention_layernorm. @@ -254,7 +284,12 @@ def is_known_skip(gguf_name: str) -> bool: """ if gguf_name.startswith("tokenizer."): return True - if "rope_freqs" in gguf_name or "attn_rot_embd" in gguf_name: + if ( + "rope_freqs" in gguf_name + or "attn_rot_embd" in gguf_name + or ".indexer." in gguf_name + or ".nextn." in gguf_name + ): return True return False @@ -305,13 +340,15 @@ def _build_mapping( result.update(_MOE_EXTRAS) if arch == "qwen35moe": result.update(_QWEN35MOE_EXTRAS) + elif arch == "glm-dsa": + result = dict(_GLM_DSA_MAPPING) else: supported = sorted( _LLAMA_FAMILY | _GEMMA_FAMILY | _MOE_FAMILY | _HUNYUAN_FAMILY - | {"gemma4", "phi3", "falcon", "gpt2", "mamba"} + | {"gemma4", "glm-dsa", "phi3", "falcon", "gpt2", "mamba"} ) raise ValueError( f"Unsupported GGUF architecture: {architecture!r}. " diff --git a/src/mobius/integrations/gguf/_tensor_mapping_test.py b/src/mobius/integrations/gguf/_tensor_mapping_test.py index 39f77d4d..b81e57c8 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping_test.py +++ b/src/mobius/integrations/gguf/_tensor_mapping_test.py @@ -289,6 +289,22 @@ def test_moe_extras(self) -> None: "model.layers.0.mlp.experts.gate_proj.weight" ) + def test_glm_dsa_mapping(self) -> None: + assert map_gguf_to_hf_names("blk.3.attn_q_a.weight", "glm-dsa") == ( + "model.layers.3.self_attn.q_a_proj.weight" + ) + assert map_gguf_to_hf_names("blk.3.attn_k_b.weight", "glm-dsa") == ( + "model.layers.3.self_attn.kv_b_proj.k_proj.weight" + ) + assert map_gguf_to_hf_names("blk.3.ffn_gate_exps.weight", "glm-dsa") == ( + "model.layers.3.mlp.experts.gate_proj.weight" + ) + assert map_gguf_to_hf_names("blk.3.exp_probs_b.bias", "glm-dsa") == ( + "model.layers.3.mlp.gate.e_score_correction_bias.bias" + ) + assert map_gguf_to_hf_names("blk.3.indexer.attn_k.weight", "glm-dsa") is None + assert map_gguf_to_hf_names("blk.78.nextn.eh_proj.weight", "glm-dsa") is None + # ---- Unsupported architecture ---- def test_unsupported_raises(self) -> None: diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 82ba1682..0aacdead 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -55,6 +55,7 @@ "Glm4CausalLMModel", "Glm4MoECausalLMModel", "GlmCausalLMModel", + "GlmMoeDsaCausalLMModel", "GraniteCausalLMModel", "GraniteMoECausalLMModel", "GraniteMoeHybridCausalLMModel", @@ -186,6 +187,7 @@ ) from mobius.models.gemma4_assistant import Gemma4AssistantCausalLMModel from mobius.models.glm import Glm4CausalLMModel, GlmCausalLMModel +from mobius.models.glm_moe_dsa import GlmMoeDsaCausalLMModel from mobius.models.gpt2 import GPT2CausalLMModel from mobius.models.gpt_neox import GPTNeoXCausalLMModel, GPTNeoXJapaneseCausalLMModel from mobius.models.gptj_codegen import CodeGenCausalLMModel, GPTJCausalLMModel diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index ed835417..4d2d47c9 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -21,9 +21,11 @@ Embedding, Linear, MoELayer, + QuantizedEmbedding, RMSNorm, create_attention_bias, initialize_rope, + make_quantized_linear_factory, ) from mobius.components._deepseek_mla import DeepSeekMLA from mobius.models.base import CausalLMModel @@ -32,6 +34,21 @@ import onnx_ir as ir +def _linear_class(config: ArchitectureConfig) -> type: + qc = config.quantization + if qc is None or qc.quant_method == "none": + return Linear + import onnx_ir as ir + + zero_point_dtype = config.dtype if qc.float_zero_point else ir.DataType.UINT8 + return make_quantized_linear_factory( + bits=qc.bits, + block_size=qc.group_size, + has_zero_point=not qc.sym, + zero_point_dtype=zero_point_dtype, + ) + + class DeepSeekMoEGate(nn.Module): """Expert routing gate for DeepSeek-V2/V3 MoE. @@ -150,14 +167,19 @@ class DeepSeekMLADecoderLayer(nn.Module): Forward signature matches DecoderLayer for compatibility with TextModel. """ - def __init__(self, config: ArchitectureConfig, is_moe: bool = False): + def __init__( + self, + config: ArchitectureConfig, + is_moe: bool = False, + linear_class: type | None = None, + ): super().__init__() - self.self_attn = DeepSeekMLA(config) + self.self_attn = DeepSeekMLA(config, linear_class=linear_class) if is_moe: gate = DeepSeekMoEGate(config) - self.mlp = _DeepSeekMoEFFN(config, gate) + self.mlp = _DeepSeekMoEFFN(config, gate, linear_class=linear_class) else: - self.mlp = MLP(config) + self.mlp = MLP(config, linear_class=linear_class) 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) @@ -197,14 +219,19 @@ class _DeepSeekStandardDecoderLayer(nn.Module): Forward signature matches DeepSeekMLADecoderLayer for compatibility. """ - def __init__(self, config: ArchitectureConfig, is_moe: bool = False): + def __init__( + self, + config: ArchitectureConfig, + is_moe: bool = False, + linear_class: type | None = None, + ): super().__init__() - self.self_attn = Attention(config) + self.self_attn = Attention(config, linear_class=linear_class) if is_moe: gate = DeepSeekMoEGate(config) - self.mlp = _DeepSeekMoEFFN(config, gate) + self.mlp = _DeepSeekMoEFFN(config, gate, linear_class=linear_class) else: - self.mlp = MLP(config) + self.mlp = MLP(config, linear_class=linear_class) 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) @@ -244,15 +271,24 @@ class _DeepSeekMoEFFN(nn.Module): Output = moe_routed_output + shared_expert_output. """ - def __init__(self, config: ArchitectureConfig, gate: nn.Module): + def __init__( + self, + config: ArchitectureConfig, + gate: nn.Module, + linear_class: type | None = None, + ): super().__init__() assert config.num_local_experts is not None assert config.moe_intermediate_size is not None - self.moe = MoELayer(config, gate=gate) + self.moe = MoELayer(config, gate=gate, linear_class=linear_class) # Shared expert uses moe_intermediate_size * n_shared_experts n_shared = config.n_shared_experts or 1 shared_intermediate = config.moe_intermediate_size * n_shared - self.shared_experts = _SharedExpertMLP(config, shared_intermediate) + self.shared_experts = _SharedExpertMLP( + config, + shared_intermediate, + linear_class=linear_class, + ) def forward(self, op: OpBuilder, hidden_states: ir.Value): moe_output = self.moe(op, hidden_states) @@ -263,11 +299,18 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): class _SharedExpertMLP(nn.Module): """Shared expert MLP (same architecture as gate/up/down SiLU MLP).""" - def __init__(self, config: ArchitectureConfig, intermediate_size: int): + def __init__( + self, + config: ArchitectureConfig, + intermediate_size: int, + linear_class: type | None = None, + ): super().__init__() - self.gate_proj = Linear(config.hidden_size, intermediate_size, bias=False) - self.up_proj = Linear(config.hidden_size, intermediate_size, bias=False) - self.down_proj = Linear(intermediate_size, config.hidden_size, bias=False) + if linear_class is None: + linear_class = Linear + self.gate_proj = linear_class(config.hidden_size, intermediate_size, bias=False) + self.up_proj = linear_class(config.hidden_size, intermediate_size, bias=False) + self.down_proj = linear_class(intermediate_size, config.hidden_size, bias=False) def forward(self, op: OpBuilder, hidden_states: ir.Value): gate_out = self.gate_proj(op, hidden_states) @@ -290,7 +333,19 @@ class DeepSeekV3TextModel(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() - self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + linear_class = _linear_class(config) + qc = config.quantization + if qc is not None and qc.quantize_embeddings: + self.embed_tokens = QuantizedEmbedding( + config.vocab_size, + config.hidden_size, + bits=qc.bits, + block_size=qc.group_size, + has_zero_point=not qc.sym, + padding_idx=config.pad_token_id, + ) + else: + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) self._dtype = config.dtype # Detect MLA vs standard attention @@ -304,7 +359,7 @@ def __init__(self, config: ArchitectureConfig): first_k = config.num_hidden_layers self.layers = nn.ModuleList( [ - LayerClass(config, is_moe=(i >= first_k)) + LayerClass(config, is_moe=(i >= first_k), linear_class=linear_class) for i in range(config.num_hidden_layers) ] ) @@ -369,7 +424,11 @@ def __init__(self, config: ArchitectureConfig): nn.Module.__init__(self) self.config = config self.model = DeepSeekV3TextModel(config) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + qc = config.quantization + lm_head_class = ( + _linear_class(config) if qc is not None and qc.quantize_lm_head else Linear + ) + self.lm_head = lm_head_class(config.hidden_size, config.vocab_size, bias=False) def preprocess_weights( self, state_dict: dict[str, torch.Tensor] diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py new file mode 100644 index 00000000..650f44aa --- /dev/null +++ b/src/mobius/models/glm_moe_dsa.py @@ -0,0 +1,48 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""GLM-5/5.2 MoE backbone with MLA and a full-attention fallback. + +GLM-5.2 uses Deep Sparse Attention (DSA) with IndexShare. This first export +increment preserves the exact MLA, MoE, shared-expert, routing, RoPE, and norm +backbone while evaluating every attention layer densely. IndexShare and the +checkpoint's final MTP layer are intentionally not exported yet. +""" + +from __future__ import annotations + +import re + +from mobius._configs import ArchitectureConfig +from mobius.models.deepseek import DeepSeekV3CausalLMModel, DeepSeekV3TextModel + + +class GlmMoeDsaTextModel(DeepSeekV3TextModel): + """GLM-MoE-DSA backbone using full MLA attention instead of sparse DSA.""" + + +class GlmMoeDsaCausalLMModel(DeepSeekV3CausalLMModel): + """GLM-5/5.2 causal LM with MLA, hybrid dense/MoE FFNs, and shared experts.""" + + category: str = "Mixture of Experts" + + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + self.model = GlmMoeDsaTextModel(config) + + def preprocess_weights(self, state_dict): + """Map HF/GGUF GLM-MoE names to the shared DeepSeek-style modules.""" + filtered = {} + for key, value in state_dict.items(): + match = re.search(r"\.layers\.(\d+)\.", key) + if match is not None and int(match.group(1)) >= self.config.num_hidden_layers: + continue + if ".indexer." in key or ".nextn." in key: + continue + filtered[key] = value + + processed = super().preprocess_weights(filtered) + return { + key.replace(".mlp.experts.", ".mlp.moe.experts."): value + for key, value in processed.items() + } diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py new file mode 100644 index 00000000..7645733c --- /dev/null +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -0,0 +1,126 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import torch +from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig + +from mobius._builder import build_from_module +from mobius._configs import ArchitectureConfig, QuantizationConfig +from mobius._registry import registry +from mobius._testing import count_op_type, make_config +from mobius.models.glm_moe_dsa import GlmMoeDsaCausalLMModel + + +def _tiny_glm_moe_dsa_config(**overrides): + defaults = dict( + model_type="glm_moe_dsa", + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=4, + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=8, + head_dim=4, + rope_interleave=True, + num_local_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + n_shared_experts=1, + first_k_dense_replace=1, + norm_topk_prob=True, + routed_scaling_factor=2.5, + scoring_func="sigmoid", + topk_method="noaux_tc", + index_topk=8, + index_head_dim=8, + index_n_heads=2, + indexer_types=["full", "shared", "shared", "shared"], + num_nextn_predict_layers=1, + ) + defaults.update(overrides) + return make_config(**defaults) + + +class TestGlmMoeDsaExport: + def test_hf_config_extraction(self): + hf_config = GlmMoeDsaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=4, + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=8, + n_routed_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=32, + n_shared_experts=1, + first_k_dense_replace=1, + index_topk=8, + index_head_dim=8, + index_n_heads=2, + ) + + config = ArchitectureConfig.from_transformers(hf_config) + + assert config.model_type == "glm_moe_dsa" + assert config.num_local_experts == 2 + assert config.shared_expert_intermediate_size == 32 + assert config.qk_nope_head_dim == 12 + assert config.qk_rope_head_dim == 4 + assert config.index_topk == 8 + assert config.rope_interleave is True + + def test_registry(self): + assert registry.get("glm_moe_dsa") is GlmMoeDsaCausalLMModel + + def test_builds_full_attention_mla_moe_graph(self): + config = _tiny_glm_moe_dsa_config() + graph = build_from_module(GlmMoeDsaCausalLMModel(config), config)["model"].graph + + assert count_op_type(graph, "Attention") == config.num_hidden_layers + assert count_op_type(graph, "TopK") == config.num_hidden_layers - 1 + assert count_op_type(graph, "Sigmoid") >= config.num_hidden_layers - 1 + assert not any("indexer" in name for name in graph.initializers) + assert {value.name for value in graph.outputs} >= { + "logits", + "present.0.key", + "present.0.value", + } + + def test_quantized_graph_uses_matmul_nbits(self): + config = _tiny_glm_moe_dsa_config( + quantization=QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=False, + ) + ) + graph = build_from_module(GlmMoeDsaCausalLMModel(config), config)["model"].graph + assert count_op_type(graph, "MatMulNBits") > 0 + + def test_preprocess_drops_indexer_and_mtp_layer(self): + config = _tiny_glm_moe_dsa_config() + model = GlmMoeDsaCausalLMModel(config) + state_dict = { + "model.layers.1.mlp.gate.weight": torch.zeros(2, 64), + "model.layers.1.mlp.experts.0.gate_proj.weight": torch.zeros(32, 64), + "model.layers.1.indexer.wk.weight": torch.zeros(8, 64), + "model.layers.4.self_attn.q_a_proj.weight": torch.zeros(16, 64), + } + + result = model.preprocess_weights(state_dict) + + assert "model.layers.1.mlp.moe.gate.weight" in result + assert "model.layers.1.mlp.moe.experts.0.gate_proj.weight" in result + assert not any("indexer" in key for key in result) + assert not any(".layers.4." in key for key in result) diff --git a/tests/cli_test.py b/tests/cli_test.py index a2401c63..3df44cb3 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -253,6 +253,35 @@ def test_runtime_ort_genai_calls_write_ort_genai_config(self): call_kwargs = mock_export.call_args assert call_kwargs.kwargs.get("hf_model_id") == "Qwen/Qwen2.5-0.5B" + def test_runtime_onnx_genai_calls_write_inference_metadata(self): + """--runtime onnx-genai writes inference_metadata (not genai_config).""" + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius.integrations.onnx_genai.write_inference_metadata", + return_value="inference_metadata.yaml", + ) as mock_meta, + mock.patch( + "mobius.integrations.ort_genai.write_ort_genai_config", + return_value={}, + ) as mock_ort, + mock.patch("mobius._model_package.ModelPackage.save"), + ): + main( + [ + "build", + "--model", + "Qwen/Qwen2.5-0.5B", + tmpdir, + "--no-weights", + "--runtime", + "onnx-genai", + ] + ) + + mock_meta.assert_called_once() + mock_ort.assert_not_called() + def test_no_runtime_does_not_call_write_ort_genai_config(self): """Omitting --runtime does NOT call write_ort_genai_config().""" with ( From 77e2b0c58e670b0e4b47b934098a86978f6b9d2a Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 16 Jul 2026 15:24:51 +0000 Subject: [PATCH 02/23] Implement GLM-5.2 IndexShare DSA and MTP export Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/research/glm52-export.md | 62 +- src/mobius/__main__.py | 16 + src/mobius/_builder.py | 5 + src/mobius/_configs/_base.py | 9 +- src/mobius/_registry.py | 4 +- .../integrations/gguf/_config_mapping.py | 11 +- .../integrations/gguf/_tensor_mapping.py | 13 +- .../integrations/gguf/_tensor_mapping_test.py | 14 +- .../integrations/ort_genai/auto_export.py | 43 +- .../ort_genai/auto_export_test.py | 19 + src/mobius/models/glm_moe_dsa.py | 537 +++++++++++++++++- src/mobius/models/glm_moe_dsa_test.py | 66 ++- src/mobius/tasks/__init__.py | 3 + src/mobius/tasks/_glm_moe_dsa.py | 138 +++++ tests/_test_configs.py | 28 + tests/model_coverage_test.py | 1 + 16 files changed, 901 insertions(+), 68 deletions(-) create mode 100644 src/mobius/tasks/_glm_moe_dsa.py diff --git a/docs/research/glm52-export.md b/docs/research/glm52-export.md index e85de313..088a6827 100644 --- a/docs/research/glm52-export.md +++ b/docs/research/glm52-export.md @@ -40,18 +40,36 @@ weight into `attn_k_b` and `attn_v_b`, while Transformers exposes one ## Implemented scope -This increment registers `glm_moe_dsa` and maps `glm-dsa` GGUF files to it. -It exports the 78-layer, two-norm MLA backbone, the three dense FFN layers, -the 75 routed-MoE layers, the shared expert, sigmoid/no-aux router, interleaved -RoPE, KV cache, and untied output head. It uses the existing portable ONNX -Attention and MoE components, so the graph can be packaged for both -ONNX Runtime GenAI runtime spellings (`onnx-genai` and `ort-genai`). +The exporter now preserves the 78-layer backbone, IndexShare DSA, and the +additional improved-MTP layer. Full indexers are emitted for layers 0, 1, 2, +6, 10, ... 74; each full layer's INT32 top-k result is reused by the following +three shared layers. The indexer implements the checkpoint equations in fp32: +Q-LoRA residual projection, LayerNorm+interleaved-RoPE key projection, ReLU +scores, learned signed head weights, causal bias, and dynamic +`min(2048, key_length)` TopK. + +Each full-indexer layer packs its 128-value index key beside the decompressed +MLA key in the standard `past_key_values.N.key` tensor. Shared layers retain +the normal MLA cache. This keeps the ORT GenAI key/value naming contract +without duplicating the index key across 64 attention heads. + +The MTP component exports: + +1. `enorm(next-token embedding)` and `hnorm(target hidden state)`; +2. concatenation and the `2H -> H` `eh_proj`; +3. the complete layer-78 DSA+MoE decoder block; and +4. `shared_head.norm`, producing `mtp_hidden` for the shared target LM head. + +Multi-component packages contain `model/model.onnx`, `mtp/model.onnx`, and an +`mtp_config.json` sidecar. Both runtime aliases (`onnx-genai` and `ort-genai`) +emit these artifacts. ORT GenAI does not yet orchestrate the sidecar itself. The importer: - derives all MLA, MoE, IndexShare, RoPE, and NextN fields from GGUF metadata; -- removes the GGUF-only MTP block from the backbone layer count; +- separates the GGUF MTP block from the 78-layer backbone and exports it as `mtp`; - fuses split `attn_k_b` and `attn_v_b` tensors into `kv_b_proj`; +- maps `blk.N.indexer.*` and supported `blk.78.nextn.*` tensors; - streams stacked routed-expert tensors one expert at a time instead of materializing all 256 experts as one float tensor; and - preserves supported 4-bit and 8-bit tensors as MatMulNBits, requantizing @@ -61,26 +79,24 @@ The generic portable MoE implementation evaluates every expert and masks the outputs. It is correct for routing, but a fused sparse-MoE runtime kernel will be required for practical GLM-5.2 execution. -## Deliberately deferred +## Runtime follow-ups -### IndexShare sparse attention +Portable ONNX preserves DSA selection numerics by converting the selected +indices to an additive mask for the standard `Attention` operator. No custom op +is required for correctness. Practical million-token performance still needs +a selected-token sparse-attention kernel so the runtime avoids evaluating +masked K/V positions. -Layers 0-2 own full indexers. Starting at layer 6, one full indexer is used -every four layers and the following three layers reuse its selected token -indices. The current export omits indexer tensors and evaluates full causal -MLA instead. This is a coherent backbone/exporter increment, but it does **not** -preserve DSA numerics or its long-context cost. A follow-up should add an -indexer component, cross-layer top-k state, sparse prefill/decode attention, -and cache-aware position handling as one tested feature rather than partially -wiring index weights into dense Attention. +`index_share_for_mtp_iteration` also needs runtime orchestration. MTP step 0 can +return `topk_indices`, but accepted speculative positions make the MTP MLA and +indexer caches advance at different logical lengths. ORT GenAI currently has +no separate indexer-cache/control-state contract, so steps 1+ recompute the +index selection rather than silently applying an incorrect packed-cache reuse. -### Improved MTP +Use `--glm-full-attention` to retain the prior dense MLA fallback. The fallback +disables the DSA-dependent MTP artifact. -The final NextN block and `index_share_for_mtp_iteration` behavior are omitted. -A follow-up should model the GLM-specific MTP inputs, shared index selection, -cache contract, outputs, and ORT GenAI speculative-decoding configuration. - -## Unsloth dynamic quantization +## Separate workstream: Unsloth dynamic quantization The six `UD-IQ1_M` GGUF shard headers were inspected. The dominant routed expert tensors are below the currently targeted 4-bit floor: diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 54494b3d..eb220e37 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -218,6 +218,12 @@ def _cmd_build(args: argparse.Namespace) -> None: config = _config_from_hf(hf_config, parent_config=parent_config) if dtype_override is not None: config = dataclasses.replace(config, dtype=dtype_override) + if args.glm_full_attention: + config = dataclasses.replace( + config, + use_dsa=False, + num_nextn_predict_layers=0, + ) if task is None: task = _default_task_for_model(model_type) module_class = registry.get(model_type) @@ -241,6 +247,11 @@ def _cmd_build(args: argparse.Namespace) -> None: trust_remote_code=trust_remote_code, execution_provider=execution_provider, text_only=args.text_only, + config_overrides=( + {"use_dsa": False, "num_nextn_predict_layers": 0} + if args.glm_full_attention + else None + ), ) _save_package(pkg, output_dir, args, optimize, component_filter) @@ -563,6 +574,11 @@ def main(argv: list[str] | None = None) -> None: action="store_true", help="Do not include weights in the output ONNX model.", ) + build_parser.add_argument( + "--glm-full-attention", + action="store_true", + help="Disable GLM-5.2 IndexShare DSA and export the dense MLA fallback.", + ) build_parser.add_argument( "--trust-remote-code", action="store_true", diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index a080f36a..38ff8e23 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -354,6 +354,7 @@ def build( execution_provider: str = "default", trace_optimization: bool = False, text_only: bool = False, + config_overrides: dict[str, object] | None = None, ) -> ModelPackage: """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID. @@ -417,6 +418,8 @@ def build( Raises :class:`ValueError` if the resolved ``model_type`` has no text-only sibling. Currently supported for ``gemma4_unified`` (``google/gemma-4-12B``). + config_overrides: Optional dataclass field overrides applied after + HuggingFace config extraction and dtype resolution. Returns: A :class:`ModelPackage` containing the built model(s). @@ -568,6 +571,8 @@ def build( if dtype is not None: dtype = resolve_dtype(dtype) config = dataclasses.replace(config, dtype=dtype) + if config_overrides: + config = dataclasses.replace(config, **config_overrides) if output_layer_indices is not None: # Opt-in: emit additional `hidden_states.{k}` ONNX outputs for each diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 098d6cc7..f5627c9c 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -425,10 +425,13 @@ class ArchitectureConfig(BaseModelConfig): rope_interleave: bool = False # Deep Sparse Attention / IndexShare config — GLM-5.2. - # The dense-attention fallback preserves these fields for future DSA export. + use_dsa: bool = True index_topk: int | None = None index_head_dim: int | None = None index_n_heads: int | None = None + index_topk_freq: int = 4 + index_skip_topk_offset: int = 3 + indexer_rope_interleave: bool = True indexer_types: list[str] | None = None index_share_for_mtp_iteration: bool = False num_nextn_predict_layers: int = 0 @@ -818,9 +821,13 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: qk_rope_head_dim=getattr(config, "qk_rope_head_dim", None), v_head_dim=getattr(config, "v_head_dim", None), # Deep Sparse Attention / IndexShare + use_dsa=getattr(config, "use_dsa", True), index_topk=getattr(config, "index_topk", None), index_head_dim=getattr(config, "index_head_dim", None), index_n_heads=getattr(config, "index_n_heads", None), + index_topk_freq=getattr(config, "index_topk_freq", 4), + index_skip_topk_offset=getattr(config, "index_skip_topk_offset", 3), + indexer_rope_interleave=getattr(config, "indexer_rope_interleave", True), indexer_types=getattr(config, "indexer_types", None), index_share_for_mtp_iteration=getattr( config, "index_share_for_mtp_iteration", False diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 11d1dcf8..fa14926f 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -524,7 +524,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "ernie4_5_moe": ModelRegistration(Ernie45MoECausalLMModel), "flex_olmo": ModelRegistration(MoECausalLMModel), "glm4_moe": ModelRegistration(Glm4MoECausalLMModel), - "glm_moe_dsa": ModelRegistration(GlmMoeDsaCausalLMModel), + "glm_moe_dsa": ModelRegistration(GlmMoeDsaCausalLMModel, task="glm-moe-dsa"), "granitemoe": ModelRegistration(GraniteMoECausalLMModel), "granitemoehybrid": ModelRegistration(GraniteMoeHybridCausalLMModel), "granitemoeshared": ModelRegistration(GraniteMoECausalLMModel), @@ -1225,7 +1225,7 @@ def _create_default_registry() -> ModelRegistry: "deepseek_v2": "mla", "deepseek_v2_moe": "mla+moe", "deepseek_v3": "mla+moe", - "glm_moe_dsa": "mla+moe-full-attention", + "glm_moe_dsa": "mla+moe-indexshare-dsa+mtp", "phi3small": "blocksparse", "falcon_h1": "hybrid-ssm", "mamba": "ssm", diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 29a390c0..4eeb181b 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -698,9 +698,11 @@ def _glm_moe_dsa_postprocess( n_shared_experts = int(metadata.get(f"{arch}.expert_shared_count", 0)) first_k_dense_replace = int(metadata.get(f"{arch}.leading_dense_block_count", 0)) - # GLM-5.2 shares one full indexer across each four-layer DSA group. + # Official GLM-5.2 schedule: full indexers in layers 0, 1, 2, then + # 6, 10, ...; each later full layer shares its selection with the next + # three layers. indexer_types = [ - "full" if i < first_k_dense_replace or (i - 2) % 4 == 0 else "shared" + "full" if i <= 2 or (i >= 6 and (i - 6) % 4 == 0) else "shared" for i in range(num_hidden_layers) ] mlp_layer_types = [ @@ -729,7 +731,12 @@ def _glm_moe_dsa_postprocess( mlp_layer_types=mlp_layer_types, scoring_func="sigmoid", topk_method="noaux_tc", + use_dsa=True, + index_topk_freq=4, + index_skip_topk_offset=3, + indexer_rope_interleave=True, indexer_types=indexer_types, + index_share_for_mtp_iteration=n_mtp > 0, num_nextn_predict_layers=n_mtp, ) diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index dac71e96..7e5415ee 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -209,6 +209,17 @@ "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", + # GLM-5.2 IndexShare DSA. + "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", + # GLM-5.2 improved MTP. The shared embedding/head tensors are tied to + # token_embd/output and therefore intentionally remain unmapped. + "blk.{bid}.nextn.eh_proj": "model.layers.{bid}.eh_proj", + "blk.{bid}.nextn.enorm": "model.layers.{bid}.enorm", + "blk.{bid}.nextn.hnorm": "model.layers.{bid}.hnorm", + "blk.{bid}.nextn.shared_head_norm": "model.layers.{bid}.shared_head.norm", "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", @@ -287,8 +298,6 @@ def is_known_skip(gguf_name: str) -> bool: if ( "rope_freqs" in gguf_name or "attn_rot_embd" in gguf_name - or ".indexer." in gguf_name - or ".nextn." in gguf_name ): return True return False diff --git a/src/mobius/integrations/gguf/_tensor_mapping_test.py b/src/mobius/integrations/gguf/_tensor_mapping_test.py index b81e57c8..7364af5c 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping_test.py +++ b/src/mobius/integrations/gguf/_tensor_mapping_test.py @@ -302,8 +302,18 @@ def test_glm_dsa_mapping(self) -> None: assert map_gguf_to_hf_names("blk.3.exp_probs_b.bias", "glm-dsa") == ( "model.layers.3.mlp.gate.e_score_correction_bias.bias" ) - assert map_gguf_to_hf_names("blk.3.indexer.attn_k.weight", "glm-dsa") is None - assert map_gguf_to_hf_names("blk.78.nextn.eh_proj.weight", "glm-dsa") is None + assert map_gguf_to_hf_names("blk.3.indexer.attn_k.weight", "glm-dsa") == ( + "model.layers.3.self_attn.indexer.wk.weight" + ) + assert map_gguf_to_hf_names("blk.3.indexer.k_norm.bias", "glm-dsa") == ( + "model.layers.3.self_attn.indexer.k_norm.bias" + ) + assert map_gguf_to_hf_names("blk.78.nextn.eh_proj.weight", "glm-dsa") == ( + "model.layers.78.eh_proj.weight" + ) + assert map_gguf_to_hf_names( + "blk.78.nextn.shared_head_norm.weight", "glm-dsa" + ) == "model.layers.78.shared_head.norm.weight" # ---- Unsupported architecture ---- diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 1337785c..81518750 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -86,6 +86,7 @@ # HunYuan-V1 dense / Hy-MT1.5 — generic decoder LLM type accepted by # ORT GenAI (see onnxruntime-genai/src/models/model_type.h LLM list). "hunyuan_v1_dense": "decoder", + "glm_moe_dsa": "decoder", # Qwen VL models all use the same GenAI pipeline as qwen2_5_vl "qwen2_vl": "qwen2_5_vl", "qwen3_vl": "qwen2_5_vl", @@ -718,7 +719,11 @@ def _write_genai_config( decoder_inputs["past_value_names"] = "past_key_values.%d.value" # Derive decoder filename from the actual package key - decoder_filename = f"{decoder_key}/model.onnx" if decoder_key != "model" else "model.onnx" + decoder_filename = ( + f"{decoder_key}/model.onnx" + if len(pkg) > 1 or decoder_key != "model" + else "model.onnx" + ) # ORT GenAI's ``past_present_share_buffer`` mode requires the decoder # graph to write the KV cache in place. Only ``com.microsoft. @@ -985,6 +990,42 @@ def write_ort_genai_config( result: dict[str, str] = {"genai_config": genai_path} + if "mtp" in pkg: + mtp_path = os.path.join(directory, "mtp_config.json") + with open(mtp_path, "w") as f: + json.dump( + { + "model": {"filename": "mtp/model.onnx"}, + "inputs": [ + "inputs_embeds", + "hidden_states", + "attention_mask", + "position_ids", + "past_key_values.0.key", + "past_key_values.0.value", + ], + "outputs": [ + "mtp_hidden", + "present.0.key", + "present.0.value", + "topk_indices", + ], + "num_nextn_predict_layers": getattr( + config, "num_nextn_predict_layers", 0 + ), + "index_share_for_mtp_iteration": getattr( + config, "index_share_for_mtp_iteration", False + ), + "shared_embedding": "model.embed_tokens", + "shared_lm_head": "lm_head", + "runtime_orchestration": "external", + }, + f, + indent=2, + ) + f.write("\n") + result["mtp_config"] = mtp_path + # Copy tokenizer files — HF Hub takes precedence; local dir is the fallback # for --config mode where no HF model ID is available. if hf_model_id is not None: diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 40c86a0f..bd83b477 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -75,6 +75,9 @@ def test_hunyuan_v1_dense_maps_to_decoder(self): # decoder-only causal LM not in its built-in registry. assert _resolve_ort_genai_model_type("hunyuan_v1_dense") == "decoder" + def test_glm_moe_dsa_maps_to_decoder(self): + assert _resolve_ort_genai_model_type("glm_moe_dsa") == "decoder" + def test_unknown_model_type_passthrough(self): assert _resolve_ort_genai_model_type("my_custom") == "my_custom" @@ -548,6 +551,22 @@ def test_genai_config_json_is_written(self, tmp_path): assert "model" in data assert data["model"]["type"] == "qwen2" + def test_mtp_component_writes_sidecar_and_nested_decoder_path(self, tmp_path): + pkg = self._make_pkg() + pkg["mtp"] = mock.MagicMock() + + result = write_ort_genai_config(pkg, str(tmp_path)) + + assert "mtp_config" in result + with open(result["genai_config"]) as f: + genai = json.load(f) + assert genai["model"]["decoder"]["filename"] == "model/model.onnx" + with open(result["mtp_config"]) as f: + mtp = json.load(f) + assert mtp["model"]["filename"] == "mtp/model.onnx" + assert mtp["outputs"][-1] == "topk_indices" + assert mtp["runtime_orchestration"] == "external" + def test_processor_config_written_with_vision(self, tmp_path): """image_processor.json is written when pkg.config.vision is set.""" import dataclasses diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index 650f44aa..6bac24f3 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -1,48 +1,531 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""GLM-5/5.2 MoE backbone with MLA and a full-attention fallback. - -GLM-5.2 uses Deep Sparse Attention (DSA) with IndexShare. This first export -increment preserves the exact MLA, MoE, shared-expert, routing, RoPE, and norm -backbone while evaluating every attention layer densely. IndexShare and the -checkpoint's final MTP layer are intentionally not exported yet. -""" +"""GLM-5.2 MLA+MoE with IndexShare deep sparse attention and MTP export.""" from __future__ import annotations import re +from typing import TYPE_CHECKING + +import onnx_ir as ir +from onnxscript import OpBuilder, nn from mobius._configs import ArchitectureConfig -from mobius.models.deepseek import DeepSeekV3CausalLMModel, DeepSeekV3TextModel +from mobius.components import ( + MLP, + Embedding, + Linear, + RMSNorm, + create_attention_bias, + initialize_rope, +) +from mobius.components._common import LayerNorm +from mobius.components._deepseek_mla import DeepSeekMLA +from mobius.components._rotary_embedding import apply_rotary_pos_emb +from mobius.models.deepseek import ( + DeepSeekMoEGate, + DeepSeekV3CausalLMModel, + DeepSeekV3TextModel, + _DeepSeekMoEFFN, + _linear_class, +) + +if TYPE_CHECKING: + import torch + + +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def _indexer_types(config: ArchitectureConfig) -> list[str]: + """Return the authoritative full/shared IndexShare schedule.""" + if config.indexer_types is not None: + return list(config.indexer_types) + full_layers = {0, 1, 2} + full_layers.update( + range( + config.index_skip_topk_offset * 2, + config.num_hidden_layers, + config.index_topk_freq, + ) + ) + return ["full" if i in full_layers else "shared" for i in range(config.num_hidden_layers)] + + +class GlmMoeDsaIndexer(nn.Module): + """GLM-5.2 fp32 token indexer used by full IndexShare layers.""" + + def __init__(self, config: ArchitectureConfig, linear_class: type): + super().__init__() + assert config.q_lora_rank is not None + assert config.index_head_dim is not None + assert config.index_n_heads is not None + self.index_head_dim = config.index_head_dim + self.index_n_heads = config.index_n_heads + self.index_topk = int(config.index_topk or 2048) + self.rope_interleave = config.indexer_rope_interleave + self.wq_b = linear_class( + config.q_lora_rank, + self.index_n_heads * self.index_head_dim, + bias=False, + ) + self.wk = linear_class(config.hidden_size, self.index_head_dim, bias=False) + self.k_norm = LayerNorm(self.index_head_dim, eps=1e-6) + self.weights_proj = linear_class(config.hidden_size, self.index_n_heads, bias=False) + + def project_key( + self, + op: OpBuilder, + hidden_states: ir.Value, + position_embeddings: tuple, + ) -> ir.Value: + key = self.k_norm(op, self.wk(op, hidden_states)) + return apply_rotary_pos_emb( + op, + key, + position_embeddings, + num_heads=1, + rotary_embedding_dim=self.index_head_dim // 2, + interleaved=self.rope_interleave, + ) + + def select( + self, + op: OpBuilder, + hidden_states: ir.Value, + q_resid: ir.Value, + all_index_keys: ir.Value, + attention_bias: ir.Value, + ) -> ir.Value: + query = self.wq_b(op, q_resid) + query = op.Reshape(query, [0, 0, self.index_n_heads, self.index_head_dim]) + query = op.Cast(query, to=ir.DataType.FLOAT) + keys = op.Cast(all_index_keys, to=ir.DataType.FLOAT) + scores = op.MatMul(query, op.Transpose(keys, perm=[0, 2, 1])) + scores = op.Relu(scores) + + weights = self.weights_proj(op, hidden_states) + weights = op.Cast(weights, to=ir.DataType.FLOAT) + weights = op.Mul(weights, self.index_n_heads**-0.5) + scores = op.ReduceSum(op.Mul(scores, op.Unsqueeze(weights, [3])), [2], keepdims=False) + scores = op.Add(scores, op.Cast(op.Squeeze(attention_bias, [1]), to=ir.DataType.FLOAT)) + + key_length = op.Shape(scores, start=2, end=3) + k = op.Min(key_length, op.Constant(value_ints=[self.index_topk])) + _, indices = op.TopK(scores, k, axis=-1, largest=1, sorted=0, _outputs=2) + return op.Cast(indices, to=ir.DataType.INT32) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + q_resid: ir.Value, + position_embeddings: tuple, + past_index_key: ir.Value | None, + attention_bias: ir.Value, + ): + current_index_key = self.project_key(op, hidden_states, position_embeddings) + all_index_keys = ( + current_index_key + if past_index_key is None + else op.Concat(past_index_key, current_index_key, axis=1) + ) + indices = self.select( + op, + hidden_states, + q_resid, + all_index_keys, + attention_bias, + ) + return all_index_keys, indices + + +class GlmMoeDsaAttention(DeepSeekMLA): + """MLA attention with a packed indexer-key cache and sparse additive mask.""" + + def __init__( + self, + config: ArchitectureConfig, + indexer_type: str, + linear_class: type, + ): + super().__init__(config, linear_class=linear_class) + self.indexer_type = indexer_type + self.dtype = config.dtype + self.main_key_dim = self.num_heads * self.qk_head_dim + self.main_value_dim = self.num_heads * self.v_head_dim + self.index_head_dim = int(config.index_head_dim or 0) + if indexer_type == "full": + self.indexer = GlmMoeDsaIndexer(config, linear_class) + + def _unpack_past(self, op: OpBuilder, past_key_value: tuple): + packed_key, packed_value = past_key_value + key_tokens = op.Squeeze(packed_key, [1]) + main_key = op.Slice(key_tokens, [0], [self.main_key_dim], [2]) + main_key = op.Transpose( + op.Reshape(main_key, [0, 0, self.num_heads, self.qk_head_dim]), + perm=[0, 2, 1, 3], + ) + value_tokens = op.Squeeze(packed_value, [1]) + main_value = op.Transpose( + op.Reshape(value_tokens, [0, 0, self.num_heads, self.v_head_dim]), + perm=[0, 2, 1, 3], + ) + index_key = None + if self.indexer_type == "full": + index_key = op.Slice( + key_tokens, + [self.main_key_dim], + [self.main_key_dim + self.index_head_dim], + [2], + ) + return (main_key, main_value), index_key + + def _pack_present( + self, + op: OpBuilder, + present: tuple, + index_keys: ir.Value | None, + ) -> tuple: + key = op.Reshape(op.Transpose(present[0], perm=[0, 2, 1, 3]), [0, 0, self.main_key_dim]) + if index_keys is not None: + key = op.Concat(key, index_keys, axis=-1) + value = op.Reshape( + op.Transpose(present[1], perm=[0, 2, 1, 3]), + [0, 0, self.main_value_dim], + ) + return op.Unsqueeze(key, [1]), op.Unsqueeze(value, [1]) + + def _sparse_bias( + self, + op: OpBuilder, + indices: ir.Value, + attention_bias: ir.Value, + ) -> ir.Value: + shape = op.Shape(op.Squeeze(attention_bias, [1])) + masked = op.Expand(op.Constant(value_float=float(ir.DataType.FLOAT.min)), shape) + updates = op.Mul(op.Cast(indices, to=ir.DataType.FLOAT), 0.0) + sparse = op.ScatterElements(masked, indices, updates, axis=-1, reduction="none") + sparse = op.Cast(op.Unsqueeze(sparse, [1]), to=self.dtype) + return op.Add(attention_bias, sparse) + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple, + past_key_value: tuple | None = None, + shared_topk_indices: ir.Value | None = None, + ): + q_resid = self.q_a_layernorm(op, self.q_a_proj(op, hidden_states)) + q = self.q_b_proj(op, q_resid) + q = op.Reshape(q, [0, 0, self.num_heads, self.qk_head_dim]) + q_nope, q_rope = op.Split( + q, + [self.qk_nope_head_dim, self.qk_rope_head_dim], + axis=-1, + _outputs=2, + ) + q_rope = op.Reshape(q_rope, [0, 0, -1]) + q_rope = apply_rotary_pos_emb( + op, + q_rope, + position_embeddings, + num_heads=self.num_heads, + interleaved=self._rope_interleave, + ) + q = op.Concat( + op.Reshape(q_nope, [0, 0, self.num_heads, self.qk_nope_head_dim]), + op.Reshape(q_rope, [0, 0, self.num_heads, self.qk_rope_head_dim]), + axis=-1, + ) + q = op.Reshape(q, [0, 0, self.num_heads * self.qk_head_dim]) -class GlmMoeDsaTextModel(DeepSeekV3TextModel): - """GLM-MoE-DSA backbone using full MLA attention instead of sparse DSA.""" + compressed_kv = self.kv_a_proj_with_mqa(op, hidden_states) + kv_latent, k_rope = op.Split( + 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, + ) + value = op.Reshape(value, [0, 0, -1]) + k_rope = apply_rotary_pos_emb( + op, + k_rope, + position_embeddings, + num_heads=1, + interleaved=self._rope_interleave, + ) + k_rope = op.Tile(k_rope, [1, 1, self.num_heads]) + key = op.Concat( + op.Reshape(k_nope, [0, 0, self.num_heads, self.qk_nope_head_dim]), + op.Reshape(k_rope, [0, 0, self.num_heads, self.qk_rope_head_dim]), + axis=-1, + ) + key = op.Reshape(key, [0, 0, self.num_heads * self.qk_head_dim]) + + main_past = None + past_index_key = None + if past_key_value is not None: + main_past, past_index_key = self._unpack_past(op, past_key_value) + + all_index_keys = None + topk_indices = shared_topk_indices + if self.indexer_type == "full": + all_index_keys, topk_indices = self.indexer( + op, + hidden_states, + q_resid, + position_embeddings, + past_index_key, + attention_bias, + ) + if topk_indices is None: + raise ValueError("Shared GLM DSA layers require top-k indices from a preceding full layer") + + sparse_bias = self._sparse_bias(op, topk_indices, attention_bias) + attn_output, present_key, present_value = op.Attention( + q, + key, + value, + attn_mask=sparse_bias, + past_key=main_past[0] if main_past is not None else None, + past_value=main_past[1] if main_past is not None else None, + q_num_heads=self.num_heads, + kv_num_heads=self.num_heads, + scale=self.scaling, + _outputs=3, + ) + attn_output = self.o_proj(op, attn_output) + present = self._pack_present( + op, + (present_key, present_value), + all_index_keys, + ) + return attn_output, present, topk_indices + + +class GlmMoeDsaDecoderLayer(nn.Module): + def __init__( + self, + config: ArchitectureConfig, + indexer_type: str, + is_moe: bool, + linear_class: type, + ): + super().__init__() + self.self_attn = GlmMoeDsaAttention(config, indexer_type, linear_class) + if is_moe: + self.mlp = _DeepSeekMoEFFN( + config, + DeepSeekMoEGate(config), + linear_class=linear_class, + ) + else: + self.mlp = MLP(config, linear_class=linear_class) + 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, + past_key_value: tuple | None, + shared_topk_indices: ir.Value | None, + ): + residual = hidden_states + hidden_states = self.input_layernorm(op, hidden_states) + hidden_states, present, topk_indices = self.self_attn( + op, + hidden_states, + attention_bias, + position_embeddings, + past_key_value, + shared_topk_indices, + ) + hidden_states = op.Add(residual, hidden_states) + residual = hidden_states + hidden_states = self.post_attention_layernorm(op, hidden_states) + hidden_states = self.mlp(op, hidden_states) + return op.Add(residual, hidden_states), present, topk_indices + + +class GlmMoeDsaTextModel(nn.Module): + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + self.rotary_emb = initialize_rope(config) + linear_class = _linear_class(config) + types = _indexer_types(config) + self.layers = nn.ModuleList( + [ + GlmMoeDsaDecoderLayer( + config, + types[i], + is_moe=i >= config.first_k_dense_replace, + linear_class=linear_class, + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None, + ): + hidden_states = self.embed_tokens(op, input_ids) + position_embeddings = self.rotary_emb(op, position_ids) + attention_bias = create_attention_bias( + op, + input_ids, + attention_mask, + dtype=self.config.dtype, + ) + presents = [] + shared_topk_indices = None + for i, layer in enumerate(self.layers): + past = past_key_values[i] if past_key_values is not None else None + hidden_states, present, shared_topk_indices = layer( + op, + hidden_states, + attention_bias, + position_embeddings, + past, + shared_topk_indices, + ) + presents.append(present) + return self.norm(op, hidden_states), presents + + +class _SharedHead(nn.Module): + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class GlmMoeDsaMtp(nn.Module): + """Improved GLM-5.2 MTP head: embed/hidden fusion plus a full DSA MoE block.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + linear_class = _linear_class(config) + self.config = config + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = linear_class(config.hidden_size * 2, config.hidden_size, bias=False) + self.layer = GlmMoeDsaDecoderLayer( + config, + "full", + is_moe=True, + linear_class=linear_class, + ) + self.shared_head = _SharedHead(config) + self.rotary_emb = initialize_rope(config) + + def forward( + self, + op: OpBuilder, + inputs_embeds: ir.Value, + hidden_states: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_value: tuple, + ): + fused = self.eh_proj( + op, + op.Concat(self.enorm(op, inputs_embeds), self.hnorm(op, hidden_states), axis=-1), + ) + position_embeddings = self.rotary_emb(op, position_ids) + attention_bias = create_attention_bias( + op, + position_ids, + attention_mask, + dtype=self.config.dtype, + ) + output, present, topk_indices = self.layer( + op, + fused, + attention_bias, + position_embeddings, + past_key_value, + None, + ) + return self.shared_head.norm(op, output), present, topk_indices class GlmMoeDsaCausalLMModel(DeepSeekV3CausalLMModel): - """GLM-5/5.2 causal LM with MLA, hybrid dense/MoE FFNs, and shared experts.""" + """GLM-5.2 target model with portable IndexShare and an exported MTP graph.""" - category: str = "Mixture of Experts" + default_task = "glm-moe-dsa" def __init__(self, config: ArchitectureConfig): - super().__init__(config) - self.model = GlmMoeDsaTextModel(config) + if config.indexer_types is None: + config.indexer_types = _indexer_types(config) + nn.Module.__init__(self) + self.config = config + self.model = ( + GlmMoeDsaTextModel(config) if config.use_dsa else DeepSeekV3TextModel(config) + ) + qc = config.quantization + lm_head_class = ( + _linear_class(config) if qc is not None and qc.quantize_lm_head else Linear + ) + self.lm_head = lm_head_class(config.hidden_size, config.vocab_size, bias=False) + if config.num_nextn_predict_layers > 0: + if config.num_nextn_predict_layers != 1: + raise ValueError("GLM MTP export currently supports exactly one NextN layer") + if not config.use_dsa: + raise ValueError( + "GLM-5.2 MTP export currently requires DSA; set num_nextn_predict_layers=0 " + "when using the full-attention fallback" + ) + self.mtp = GlmMoeDsaMtp(config) + else: + self.mtp = None - def preprocess_weights(self, state_dict): - """Map HF/GGUF GLM-MoE names to the shared DeepSeek-style modules.""" - filtered = {} + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + state_dict = DeepSeekV3CausalLMModel.preprocess_weights(self, state_dict) + result: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - match = re.search(r"\.layers\.(\d+)\.", key) + match = _LAYER_RE.match(key) if match is not None and int(match.group(1)) >= self.config.num_hidden_layers: + layer_idx = int(match.group(1)) + if ( + self.mtp is None + or layer_idx + >= self.config.num_hidden_layers + self.config.num_nextn_predict_layers + ): + continue + suffix = match.group(2) + if suffix.startswith(("enorm.", "hnorm.", "eh_proj.", "shared_head.")): + key = f"mtp.{suffix}" + else: + key = f"mtp.layer.{suffix}" + elif not self.config.use_dsa and ".indexer." in key: continue - if ".indexer." in key or ".nextn." in key: - continue - filtered[key] = value - - processed = super().preprocess_weights(filtered) - return { - key.replace(".mlp.experts.", ".mlp.moe.experts."): value - for key, value in processed.items() - } + result[key.replace(".mlp.experts.", ".mlp.moe.experts.")] = value + return result diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py index 7645733c..db86256e 100644 --- a/src/mobius/models/glm_moe_dsa_test.py +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -7,6 +7,7 @@ from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig from mobius._builder import build_from_module +from mobius._config_resolver import _default_task_for_model from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius._registry import registry from mobius._testing import count_op_type, make_config @@ -67,6 +68,9 @@ def test_hf_config_extraction(self): index_topk=8, index_head_dim=8, index_n_heads=2, + index_topk_freq=4, + index_skip_topk_offset=3, + index_share_for_mtp_iteration=True, ) config = ArchitectureConfig.from_transformers(hf_config) @@ -77,24 +81,45 @@ def test_hf_config_extraction(self): assert config.qk_nope_head_dim == 12 assert config.qk_rope_head_dim == 4 assert config.index_topk == 8 + assert config.index_topk_freq == 4 + assert config.index_skip_topk_offset == 3 + assert config.index_share_for_mtp_iteration is True assert config.rope_interleave is True def test_registry(self): assert registry.get("glm_moe_dsa") is GlmMoeDsaCausalLMModel + assert _default_task_for_model("glm_moe_dsa") == "glm-moe-dsa" def test_builds_full_attention_mla_moe_graph(self): config = _tiny_glm_moe_dsa_config() - graph = build_from_module(GlmMoeDsaCausalLMModel(config), config)["model"].graph + package = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + ) + graph = package["model"].graph assert count_op_type(graph, "Attention") == config.num_hidden_layers - assert count_op_type(graph, "TopK") == config.num_hidden_layers - 1 + assert count_op_type(graph, "ScatterElements") >= 1 assert count_op_type(graph, "Sigmoid") >= config.num_hidden_layers - 1 - assert not any("indexer" in name for name in graph.initializers) + assert any("layers.0.self_attn.indexer.wk.weight" in name for name in graph.initializers) + assert not any("layers.1.self_attn.indexer" in name for name in graph.initializers) assert {value.name for value in graph.outputs} >= { "logits", "present.0.key", "present.0.value", } + assert set(package) == {"model", "mtp"} + + mtp_graph = package["mtp"].graph + assert count_op_type(mtp_graph, "Attention") == 1 + assert count_op_type(mtp_graph, "ScatterElements") == 1 + assert {value.name for value in mtp_graph.outputs} == { + "mtp_hidden", + "present.0.key", + "present.0.value", + "topk_indices", + } def test_quantized_graph_uses_matmul_nbits(self): config = _tiny_glm_moe_dsa_config( @@ -105,22 +130,47 @@ def test_quantized_graph_uses_matmul_nbits(self): sym=False, ) ) - graph = build_from_module(GlmMoeDsaCausalLMModel(config), config)["model"].graph + graph = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + )["model"].graph assert count_op_type(graph, "MatMulNBits") > 0 - def test_preprocess_drops_indexer_and_mtp_layer(self): + def test_preprocess_maps_indexer_and_mtp_layer(self): config = _tiny_glm_moe_dsa_config() model = GlmMoeDsaCausalLMModel(config) state_dict = { "model.layers.1.mlp.gate.weight": torch.zeros(2, 64), "model.layers.1.mlp.experts.0.gate_proj.weight": torch.zeros(32, 64), - "model.layers.1.indexer.wk.weight": torch.zeros(8, 64), + "model.layers.0.self_attn.indexer.wk.weight": torch.zeros(8, 64), "model.layers.4.self_attn.q_a_proj.weight": torch.zeros(16, 64), + "model.layers.4.eh_proj.weight": torch.zeros(64, 128), + "model.layers.4.shared_head.norm.weight": torch.zeros(64), } result = model.preprocess_weights(state_dict) assert "model.layers.1.mlp.moe.gate.weight" in result assert "model.layers.1.mlp.moe.experts.0.gate_proj.weight" in result - assert not any("indexer" in key for key in result) - assert not any(".layers.4." in key for key in result) + assert "model.layers.0.self_attn.indexer.wk.weight" in result + assert "mtp.layer.self_attn.q_a_proj.weight" in result + assert "mtp.eh_proj.weight" in result + assert "mtp.shared_head.norm.weight" in result + + def test_full_attention_fallback(self): + config = _tiny_glm_moe_dsa_config( + use_dsa=False, + num_nextn_predict_layers=0, + ) + package = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + ) + graph = package["model"].graph + + assert set(package) == {"model"} + assert count_op_type(graph, "Attention") == config.num_hidden_layers + assert count_op_type(graph, "ScatterElements") == 0 + assert not any("indexer" in name for name in graph.initializers) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 4ca4fc14..dc7c8dbb 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -37,6 +37,7 @@ "Gemma4AssistantTask", "Gemma4Task", "Gemma4UnifiedTask", + "GlmMoeDsaTask", "Gemma4TextCausalLMTask", "HybridCausalLMTask", "HybridQwenVLTask", @@ -97,6 +98,7 @@ Gemma4UnifiedTask, ) from mobius.tasks._gemma4_assistant import Gemma4AssistantTask +from mobius.tasks._glm_moe_dsa import GlmMoeDsaTask from mobius.tasks._hunyuan_vl_mot import HunYuanVLMoTTask from mobius.tasks._image_classification import ImageClassificationTask from mobius.tasks._masked_diffusion import MaskedDiffusionTask @@ -159,6 +161,7 @@ "gemma4-text-generation": Gemma4TextCausalLMTask, "gemma4-unified": Gemma4UnifiedTask, "gemma4-assistant": Gemma4AssistantTask, + "glm-moe-dsa": GlmMoeDsaTask, "hunyuan-vl-mot": HunYuanVLMoTTask, "multimodal": MultiModalTask, "phi4mm-multimodal": Phi4MMMultiModalTask, diff --git a/src/mobius/tasks/_glm_moe_dsa.py b/src/mobius/tasks/_glm_moe_dsa.py new file mode 100644 index 00000000..094a3043 --- /dev/null +++ b/src/mobius/tasks/_glm_moe_dsa.py @@ -0,0 +1,138 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Graph tasks for GLM-5.2 IndexShare DSA and its MTP head.""" + +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 GlmMoeDsaTask(ModelTask): + """Build the target decoder and optional GLM-5.2 MTP component.""" + + model_roles = {"model": "decoder", "mtp": "decoder"} + + @staticmethod + def _cache_dims(config: ArchitectureConfig, layer_idx: int) -> tuple[int, int]: + qk_dim = config.num_attention_heads * ( + int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0) + ) + value_dim = config.num_attention_heads * int(config.v_head_dim or 0) + if config.use_dsa and config.indexer_types: + if config.indexer_types[layer_idx] == "full": + qk_dim += int(config.index_head_dim or 0) + return qk_dim, value_dim + + def _build_target(self, module, config: ArchitectureConfig): + graph, builder = _make_graph("glm_moe_dsa") + input_ids = builder.input("input_ids", ir.DataType.INT64, ["batch_size", "sequence_length"]) + attention_mask = builder.input( + "attention_mask", ir.DataType.INT64, ["batch_size", "total_sequence_length"] + ) + position_ids = builder.input( + "position_ids", ir.DataType.INT64, ["batch_size", "sequence_length"] + ) + + past_key_values = [] + for i in range(config.num_hidden_layers): + key_dim, value_dim = self._cache_dims(config, i) + if config.use_dsa: + key_shape = ["batch_size", 1, "past_sequence_length", key_dim] + value_shape = ["batch_size", 1, "past_sequence_length", value_dim] + else: + key_shape = [ + "batch_size", + config.num_attention_heads, + "past_sequence_length", + int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0), + ] + value_shape = [ + "batch_size", + config.num_attention_heads, + "past_sequence_length", + int(config.v_head_dim or 0), + ] + key = builder.input( + f"past_key_values.{i}.key", + config.dtype, + key_shape, + ) + value = builder.input( + f"past_key_values.{i}.value", + config.dtype, + value_shape, + ) + past_key_values.append((key, value)) + + logits, presents = module( + builder.op, + input_ids, + attention_mask, + position_ids, + past_key_values, + ) + builder.add_output(logits, "logits") + for i, (key, value) in enumerate(presents): + builder.add_output(key, f"present.{i}.key") + builder.add_output(value, f"present.{i}.value") + return _make_model(graph) + + def _build_mtp(self, module, config: ArchitectureConfig): + graph, builder = _make_graph("glm_moe_dsa_mtp") + inputs_embeds = builder.input( + "inputs_embeds", + config.dtype, + ["batch_size", "sequence_length", config.hidden_size], + ) + hidden_states = builder.input( + "hidden_states", + config.dtype, + ["batch_size", "sequence_length", config.hidden_size], + ) + attention_mask = builder.input( + "attention_mask", ir.DataType.INT64, ["batch_size", "total_sequence_length"] + ) + position_ids = builder.input( + "position_ids", ir.DataType.INT64, ["batch_size", "sequence_length"] + ) + key_dim = ( + config.num_attention_heads + * (int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0)) + + int(config.index_head_dim or 0) + ) + value_dim = config.num_attention_heads * int(config.v_head_dim or 0) + past_key = builder.input( + "past_key_values.0.key", + config.dtype, + ["batch_size", 1, "past_sequence_length", key_dim], + ) + past_value = builder.input( + "past_key_values.0.value", + config.dtype, + ["batch_size", 1, "past_sequence_length", value_dim], + ) + + mtp_hidden, present, topk_indices = module.mtp( + builder.op, + inputs_embeds, + hidden_states, + attention_mask, + position_ids, + (past_key, past_value), + ) + builder.add_output(mtp_hidden, "mtp_hidden") + builder.add_output(present[0], "present.0.key") + builder.add_output(present[1], "present.0.value") + builder.add_output(topk_indices, "topk_indices") + return _make_model(graph) + + def build(self, module, config: ArchitectureConfig) -> ModelPackage: + models = {"model": self._build_target(module, config)} + if getattr(module, "mtp", None) is not None: + models["mtp"] = self._build_mtp(module, config) + return ModelPackage(models, config=config) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 9fb16b00..8b21702c 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -587,6 +587,34 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + ( + "glm_moe_dsa", + { + "num_key_value_heads": TINY_HEADS, + "q_lora_rank": 32, + "kv_lora_rank": 16, + "qk_nope_head_dim": 16, + "qk_rope_head_dim": 8, + "v_head_dim": 16, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 32, + "n_group": 1, + "topk_group": 1, + "routed_scaling_factor": 2.5, + "scoring_func": "sigmoid", + "topk_method": "noaux_tc", + "first_k_dense_replace": 1, + "n_shared_experts": 1, + "rope_interleave": True, + "index_topk": 8, + "index_head_dim": 8, + "index_n_heads": 2, + "indexer_types": ["full", "full"], + "num_nextn_predict_layers": 0, + }, + True, + ), ( "deepseek_v2", { diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index c52cd1c0..7912e789 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -208,6 +208,7 @@ def _all_registered_with_test_id() -> dict[str, str]: "arctic": "Very large MoE (480B) — no small public checkpoint", "dbrx": "Large MoE (132B) — no small public checkpoint", "deepseek_v3": "Very large MoE (671B) — no small public checkpoint", + "glm_moe_dsa": "Very large MoE — no small public checkpoint or golden data", "llama4_text": "Very large MoE (109B) — no small public checkpoint", "qwen3_5_moe": "Large MoE (22B) — no small public checkpoint", # --- Models without test_model_id --- From cff5bf2c8fddf06d4bde056b3fec39e224855b4a Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 19 Jul 2026 16:37:48 +0000 Subject: [PATCH 03/23] fix(glm_moe_dsa): indexer RoPE must rotate full index_head_dim The IndexShare DSA indexer shares the model's rotary_emb cos/sin cache, which is sized qk_rope_head_dim/2. GlmMoeDsaIndexer.project_key passed rotary_embedding_dim=index_head_dim//2 to the opset-24 RotaryEmbedding op, which then expects a cos cache of (index_head_dim//2)/2 and fails shape validation at runtime: Input 'cos_cache' dimension 2 should be same as head_size / 2 or rotary_embedding_dim / 2, got 4 The indexer key spans the full index_head_dim and must be fully rotated, matching the main MLA q_rope/k_rope calls (rotary_embedding_dim=0). Fixes end-to-end inference of glm_moe_dsa models through the onnx-genai engine. Adds a regression test asserting indexer RotaryEmbedding nodes use full rotation (rotary_embedding_dim=0). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/glm_moe_dsa.py | 8 ++++++- src/mobius/models/glm_moe_dsa_test.py | 34 +++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index 6bac24f3..7ec85bda 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -81,12 +81,18 @@ def project_key( position_embeddings: tuple, ) -> ir.Value: key = self.k_norm(op, self.wk(op, hidden_states)) + # The indexer key shares the model's rotary_emb cos/sin cache (sized + # qk_rope_head_dim / 2). RoPE rotates the *entire* index_head_dim, so + # rotary_embedding_dim must be 0 (full rotation) — matching the main + # MLA q_rope/k_rope calls. Passing index_head_dim // 2 here makes the + # opset-24 RotaryEmbedding op expect a cos cache of that_value / 2, + # which mismatches the shared cache and fails shape validation. return apply_rotary_pos_emb( op, key, position_embeddings, num_heads=1, - rotary_embedding_dim=self.index_head_dim // 2, + rotary_embedding_dim=0, interleaved=self.rope_interleave, ) diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py index db86256e..545bc3dc 100644 --- a/src/mobius/models/glm_moe_dsa_test.py +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -121,6 +121,40 @@ def test_builds_full_attention_mla_moe_graph(self): "topk_indices", } + def test_indexer_rotary_uses_full_rotation(self): + """Regression: indexer RoPE must rotate the full index_head_dim. + + The indexer key shares the model's rotary_emb cos/sin cache (sized + qk_rope_head_dim / 2). The opset-24 RotaryEmbedding op validates that + the cos cache last dim equals head_size / 2 or rotary_embedding_dim / 2. + Passing ``rotary_embedding_dim = index_head_dim // 2`` made the op + expect a cache of that_value / 2, which mismatched the shared cache and + raised "cos_cache dimension 2 should be same as head_size / 2 or + rotary_embedding_dim / 2" at runtime. Full rotation (0) matches the + shared cache like the main MLA q_rope/k_rope calls do. + """ + config = _tiny_glm_moe_dsa_config() + package = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + ) + graph = package["model"].graph + + indexer_rope_nodes = [ + node + for node in graph + if node.op_type == "RotaryEmbedding" and "indexer" in (node.name or "") + ] + assert indexer_rope_nodes, "expected at least one indexer RotaryEmbedding node" + for node in indexer_rope_nodes: + attr = node.attributes.get("rotary_embedding_dim") + dim = attr.value if attr is not None else 0 + assert dim == 0, ( + f"indexer RotaryEmbedding {node.name!r} must use full rotation " + f"(rotary_embedding_dim=0) to match the shared cos cache, got {dim}" + ) + def test_quantized_graph_uses_matmul_nbits(self): config = _tiny_glm_moe_dsa_config( quantization=QuantizationConfig( From 5d9fd61289b912217232d640a359c301154227e2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 19 Jul 2026 16:50:30 +0000 Subject: [PATCH 04/23] test(export): add tiny DeepSeek-V2 ONNX helper Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- export_deepseek_v2_tiny.py | 64 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 export_deepseek_v2_tiny.py diff --git a/export_deepseek_v2_tiny.py b/export_deepseek_v2_tiny.py new file mode 100644 index 00000000..fe24047a --- /dev/null +++ b/export_deepseek_v2_tiny.py @@ -0,0 +1,64 @@ +"""Export the tiny synthetic DeepSeek-V2 MLA + MoE model for onnx-genai E2E.""" + +from __future__ import annotations + +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests")) + +import numpy as np +import onnx_ir as ir + +from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config +from mobius._config_resolver import _default_task_for_model +from mobius._registry import registry +from mobius.integrations.onnx_genai import write_inference_metadata +from mobius.tasks import get_task + +OUT = "/home/justinchu/ds-e2e-artifacts/deepseek-v2-tiny" + + +def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: + for init in model.graph.initializers.values(): + if init.const_value is not None: + continue + shape = tuple(int(d) for d in init.shape) + if not shape: + continue + if init.dtype == ir.DataType.FLOAT: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + elif init.dtype == ir.DataType.FLOAT16: + data = (rng.standard_normal(shape) * 0.02).astype(np.float16) + elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32): + npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32 + data = rng.integers(0, 10, size=shape).astype(npd) + else: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + init.const_value = ir.Tensor(data) + + +def main() -> None: + overrides = dict( + next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "deepseek_v2") + ) + config = _base_config(**overrides) + config.dtype = ir.DataType.FLOAT + + module = registry.get("deepseek_v2")(config) + task = get_task(_default_task_for_model("deepseek_v2")) + pkg = task.build(module, config) + + rng = np.random.default_rng(0) + for model in pkg.values(): + _fill_random_weights(model, rng) + + os.makedirs(OUT, exist_ok=True) + pkg.save(OUT, external_data="onnx", check_weights=False) + print("inference_metadata:", write_inference_metadata(pkg, OUT)) + print("Saved to", OUT) + print("files:", sorted(os.listdir(OUT))) + + +if __name__ == "__main__": + main() From baff68c64a2ddc82a79f2e534dc40b2400fbd751 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 19 Jul 2026 16:53:33 +0000 Subject: [PATCH 05/23] tools: add quantized tiny glm_moe_dsa export helper for onnx-genai E2E Sibling of export_glm_tiny.py. Attaches an int4 block-32 asymmetric QuantizationConfig so the tiny glm_moe_dsa linear projections and per-expert MoE MLPs are emitted as com.microsoft::MatMulNBits, producing an onnx-genai-loadable artifact for the quantized E2E smoke test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- export_glm_tiny_quant.py | 104 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 export_glm_tiny_quant.py diff --git a/export_glm_tiny_quant.py b/export_glm_tiny_quant.py new file mode 100644 index 00000000..0d91c900 --- /dev/null +++ b/export_glm_tiny_quant.py @@ -0,0 +1,104 @@ +"""Export a QUANTIZED tiny synthetic glm_moe_dsa model for onnx-genai E2E. + +Same tiny `glm_moe_dsa` config as ``export_glm_tiny.py`, but with a +``QuantizationConfig`` attached so the linear projections (and per-expert +MoE MLPs) are emitted as ``com.microsoft::MatMulNBits`` (int4, block-32). + +The GLM/DeepSeek MoE path decomposes experts into per-expert MLPs, so each +expert's gate/up/down projections also become MatMulNBits when quantized. +mobius has no fused ``com.microsoft::QMoE`` emitter for any model (only an +unquantized ``com.microsoft::MoE`` path in gemma4), so QMoE is NOT emitted +here — see the decision note for the scoping gap. + +Writes an onnx-genai-loadable artifact directory: + /model.onnx (+ external data), inference_metadata.yaml, tokenizer.json +""" + +from __future__ import annotations + +import os +import sys +import shutil + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests")) + +import numpy as np +import onnx_ir as ir + +from _test_configs import _base_config, ALL_CAUSAL_LM_CONFIGS +from mobius._config_resolver import _default_task_for_model +from mobius._configs import QuantizationConfig +from mobius._registry import registry +from mobius.tasks import get_task +from mobius.integrations.onnx_genai import write_inference_metadata + +OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-q4" +FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny" + + +def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: + for init in model.graph.initializers.values(): + if init.const_value is not None: + continue + shape = tuple(int(d) for d in init.shape) + if not shape: + continue + if init.dtype == ir.DataType.FLOAT: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + elif init.dtype == ir.DataType.FLOAT16: + data = (rng.standard_normal(shape) * 0.02).astype(np.float16) + elif init.dtype == ir.DataType.UINT8: + # Packed int4/int8 quantized weights or bit-packed zero points. + data = rng.integers(0, 256, size=shape).astype(np.uint8) + elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32): + npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32 + data = rng.integers(0, 10, size=shape).astype(npd) + else: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + init.const_value = ir.Tensor(data) + + +def main() -> None: + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "glm_moe_dsa")) + config = _base_config(**overrides) + config.dtype = ir.DataType.FLOAT + # int4, block-32, asymmetric (GGUF/AWQ-style) -> MatMulNBits with uint8 + # zero points. This matches the validated glm_moe_dsa quant unit test. + config.quantization = QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=False, + ) + + model_type = "glm_moe_dsa" + module = registry.get(model_type)(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + + rng = np.random.default_rng(0) + for model in pkg.values(): + _fill_random_weights(model, rng) + + os.makedirs(OUT, exist_ok=True) + pkg.save(OUT, external_data="onnx", check_weights=False) + path = write_inference_metadata(pkg, OUT) + print("inference_metadata:", path) + + # Reuse tokenizer from the fp32 artifacts. + tok_src = os.path.join(FP32_OUT, "tokenizer.json") + if os.path.exists(tok_src): + shutil.copy(tok_src, os.path.join(OUT, "tokenizer.json")) + + for attr in ("num_hidden_layers", "num_attention_heads", "num_key_value_heads", + "head_dim", "hidden_size", "vocab_size", "max_position_embeddings", + "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim", "kv_lora_rank", + "q_lora_rank", "num_local_experts", "moe_intermediate_size"): + print(f" {attr} =", getattr(config, attr, None)) + + print("Saved to", OUT) + print("files:", sorted(os.listdir(OUT))) + + +if __name__ == "__main__": + main() From bc93ccf3336e203fab59782fbde11305dde06226 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 19 Jul 2026 17:37:18 +0000 Subject: [PATCH 06/23] feat(moe): fused com.microsoft::QMoE emitter for GLM/DeepSeek routed experts Add FusedQuantizedMoE component that emits a single int4/block-32 expert-major com.microsoft::QMoE node (swiglu_fusion=1) for the shared GLM/DeepSeek routed expert stack, gated behind ArchitectureConfig.fused_quantized_moe. GLM sigmoid + noaux_tc routing is preserved exactly: scores_for_choice feeds QMoE router_probs for top-k selection, and the normalized+scaled combine weights are scattered into the optional router_weights input (normalize_routing_weights=0). Adds DeepSeekMoEGate.route_for_qmoe, export_glm_tiny_qmoe.py, and a fused-path unit test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- export_glm_tiny_qmoe.py | 107 ++++++++++++++++ src/mobius/_configs/_base.py | 4 + src/mobius/components/__init__.py | 2 + src/mobius/components/_moe.py | 168 +++++++++++++++++++++++++- src/mobius/models/deepseek.py | 52 +++++++- src/mobius/models/glm_moe_dsa_test.py | 28 ++++- 6 files changed, 357 insertions(+), 4 deletions(-) create mode 100644 export_glm_tiny_qmoe.py diff --git a/export_glm_tiny_qmoe.py b/export_glm_tiny_qmoe.py new file mode 100644 index 00000000..057a37d0 --- /dev/null +++ b/export_glm_tiny_qmoe.py @@ -0,0 +1,107 @@ +"""Export a tiny synthetic glm_moe_dsa model with a FUSED ``com.microsoft::QMoE``. + +Same tiny `glm_moe_dsa` config as ``export_glm_tiny_quant.py`, but sets +``config.fused_quantized_moe = True`` so the routed MoE experts are emitted as a +single fused ``com.microsoft::QMoE`` node (int4, block-32, expert-major layout) +instead of a per-expert unroll of ``com.microsoft::MatMulNBits``. + +Routing (GLM sigmoid + noaux_tc, n_group=1) is preserved exactly: the QMoE +kernel re-derives top-k selection from ``router_probs`` (the gate's +``scores_for_choice``), and the normalized+scaled combine weights are scattered +into the optional ``router_weights`` input with ``normalize_routing_weights=0``. + +Writes an onnx-genai-loadable artifact directory: + /model.onnx (+ external data), inference_metadata.yaml, tokenizer.json +""" + +from __future__ import annotations + +import os +import shutil +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests")) + +import numpy as np +import onnx_ir as ir +from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config + +from mobius._config_resolver import _default_task_for_model +from mobius._configs import QuantizationConfig +from mobius._registry import registry +from mobius.integrations.onnx_genai import write_inference_metadata +from mobius.tasks import get_task + +OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-qmoe" +FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny" + + +def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: + for init in model.graph.initializers.values(): + if init.const_value is not None: + continue + shape = tuple(int(d) for d in init.shape) + if not shape: + continue + if init.dtype == ir.DataType.FLOAT: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + elif init.dtype == ir.DataType.FLOAT16: + data = (rng.standard_normal(shape) * 0.02).astype(np.float16) + elif init.dtype == ir.DataType.UINT8: + # Packed int4/int8 quantized weights or bit-packed zero points. + data = rng.integers(0, 256, size=shape).astype(np.uint8) + elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32): + npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32 + data = rng.integers(0, 10, size=shape).astype(npd) + else: + data = rng.standard_normal(shape).astype(np.float32) * 0.02 + init.const_value = ir.Tensor(data) + + +def main() -> None: + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "glm_moe_dsa")) + config = _base_config(**overrides) + config.dtype = ir.DataType.FLOAT + # int4, block-32, symmetric -> fused QMoE (no zero points; the kernel + # defaults per-block zero-point to 1 << (bits - 1)). + config.quantization = QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=True, + ) + # Emit routed experts as a single fused com.microsoft::QMoE node. + config.fused_quantized_moe = True + + model_type = "glm_moe_dsa" + module = registry.get(model_type)(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + + rng = np.random.default_rng(0) + for model in pkg.values(): + _fill_random_weights(model, rng) + + os.makedirs(OUT, exist_ok=True) + pkg.save(OUT, external_data="onnx", check_weights=False) + path = write_inference_metadata(pkg, OUT) + print("inference_metadata:", path) + + # Reuse tokenizer from the fp32 artifacts if present. + tok_src = os.path.join(FP32_OUT, "tokenizer.json") + if os.path.exists(tok_src): + shutil.copy(tok_src, os.path.join(OUT, "tokenizer.json")) + + for attr in ("num_hidden_layers", "num_attention_heads", "num_key_value_heads", + "head_dim", "hidden_size", "vocab_size", "max_position_embeddings", + "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim", "kv_lora_rank", + "q_lora_rank", "num_local_experts", "num_experts_per_tok", + "moe_intermediate_size"): + print(f" {attr} =", getattr(config, attr, None)) + + print("Saved to", OUT) + print("files:", sorted(os.listdir(OUT))) + + +if __name__ == "__main__": + main() diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index f5627c9c..1431bf35 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -415,6 +415,10 @@ class ArchitectureConfig(BaseModelConfig): first_k_dense_replace: int = 0 n_shared_experts: int | None = None mlp_layer_types: list[str] | None = None + # When True (and quantization is active), routed MoE experts are emitted as a + # single fused ``com.microsoft::QMoE`` op instead of a per-expert unroll of + # ``MatMulNBits``. Only wired for the GLM/DeepSeek MoE path today. + fused_quantized_moe: bool = False # Multi-head Latent Attention (MLA) config — DeepSeek-V2/V3 q_lora_rank: int | None = None diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 4ef07186..aa258915 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -47,6 +47,7 @@ "MLP", "MLPMultiModalProjector", "MoELayer", + "FusedQuantizedMoE", "OffsetRMSNorm", "PatchEmbed", "PatchEmbedding", @@ -162,6 +163,7 @@ from mobius.components._mamba_block import MambaBlock as MambaBlock from mobius.components._mlp import FCMLP, MLP, FusedGateUpMLP, GatedMLP from mobius.components._moe import ( + FusedQuantizedMoE, MoELayer, SigmoidTopKGate, SoftmaxTopKGate, diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 87858dd4..010426a2 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -8,13 +8,14 @@ import dataclasses from typing import TYPE_CHECKING +import onnx_ir as ir from onnxscript import OpBuilder, nn from mobius._configs import ArchitectureConfig from mobius.components._mlp import MLP if TYPE_CHECKING: - import onnx_ir as ir + pass class TopKGate(nn.Module): @@ -233,3 +234,168 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): result = op.Add(result, contribution) return result + + +class FusedQuantizedMoE(nn.Module): + """Routed MoE experts emitted as a single fused ``com.microsoft::QMoE`` op. + + Replaces the per-expert ``MatMulNBits`` unroll (:class:`MoELayer`) for + weight-only int-quantized MoE. All routed experts are packed into + expert-major integer weight tensors laid out exactly as the ORT contrib + ``QMoE`` kernel expects:: + + fc1_experts_weights: [E, 2*inter, hidden // pack_size] uint8 + fc1_scales: [E, 2*inter, hidden // block_size] float32 + fc2_experts_weights: [E, hidden, inter // pack_size] uint8 + fc2_scales: [E, hidden, inter // block_size] float32 + + where ``pack_size = 8 // bits``. ``fc1`` fuses the SwiGLU gate/up + projections in the **interleaved** layout ``[g_0, u_0, g_1, u_1, ...]`` and is + consumed with ``swiglu_fusion=1`` (the only SwiGLU layout the ORT CPU QMoE + kernel supports), so the kernel computes ``silu(gate) * up`` — matching + GLM/DeepSeek's SiLU-gated experts (``activation_alpha=1``, ``activation_beta=0``, + ``swiglu_limit=inf``, all ORT defaults). + + Routing is delegated to ``gate.route_for_qmoe`` (GLM sigmoid + noaux_tc / + DeepSeek softmax group-limited). The kernel re-derives top-k selection from + ``router_probs`` (the gate's ``scores_for_choice``); the exact combine + weights are scattered into a dense ``[rows, E]`` aggregation tensor and fed + through the optional ``router_weights`` input with + ``normalize_routing_weights=0``, so the fused op reproduces the per-expert + path's routing bit-for-bit. + + Symmetric int quantization is used (no zero-points): the kernel defaults the + per-block zero-point to ``1 << (bits - 1)``. + """ + + _MICROSOFT_DOMAIN = "com.microsoft" + + def __init__( + self, + config: ArchitectureConfig, + gate: nn.Module, + ): + super().__init__() + assert config.num_local_experts is not None + assert config.num_experts_per_tok is not None + assert config.moe_intermediate_size is not None + assert config.quantization is not None + if not hasattr(gate, "route_for_qmoe"): + raise TypeError( + f"gate {type(gate).__name__} does not support the fused QMoE path " + "(missing route_for_qmoe); use MoELayer instead" + ) + + qc = config.quantization + bits = qc.bits + block_size = qc.group_size + if bits not in (1, 2, 4, 8): + raise ValueError(f"QMoE expert_weight_bits must be 1/2/4/8, got {bits}") + if block_size < 16 or (block_size & (block_size - 1)): + raise ValueError(f"QMoE block_size must be a power of 2 >= 16, got {block_size}") + + self.gate = gate + self._num_experts = config.num_local_experts + self._top_k = config.num_experts_per_tok + self._hidden = config.hidden_size + self._inter = config.moe_intermediate_size + self._bits = bits + self._block_size = block_size + + if self._hidden % block_size != 0: + raise ValueError( + f"hidden_size {self._hidden} must be divisible by block_size {block_size}" + ) + if self._inter % block_size != 0: + raise ValueError( + f"moe_intermediate_size {self._inter} must be divisible by " + f"block_size {block_size}" + ) + + pack_size = 8 // bits + e = self._num_experts + fc1_out = 2 * self._inter # interleaved [g_0, u_0, ...] for swiglu_fusion=1 + # fc1: [E, 2*inter, hidden] quantized along hidden (K) + self.fc1_experts_weights = nn.Parameter( + [e, fc1_out, self._hidden // pack_size], + dtype=ir.DataType.UINT8, + ) + self.fc1_scales = nn.Parameter( + [e, fc1_out, self._hidden // block_size], + dtype=ir.DataType.FLOAT, + ) + # fc2: [E, hidden, inter] quantized along inter (K) + self.fc2_experts_weights = nn.Parameter( + [e, self._hidden, self._inter // pack_size], + dtype=ir.DataType.UINT8, + ) + self.fc2_scales = nn.Parameter( + [e, self._hidden, self._inter // block_size], + dtype=ir.DataType.FLOAT, + ) + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + hidden = self._hidden + + # QMoE requires 2-D router_probs, so flatten [B, S, H] -> [rows, H]. + orig_shape = op.Shape(hidden_states) + flat = op.Reshape(hidden_states, op.Constant(value_ints=[-1, hidden])) + # QMoE input/router_probs must be float32. + flat_f32 = op.Cast(flat, to=1) + + scores_for_choice, routing_weights, selected_experts = self._route(op, flat_f32) + + # Dense [rows, E] aggregation weights: combine weight at each selected + # expert position, 0 elsewhere. QMoE reads these at its own top-k picks. + zeros = op.Mul(scores_for_choice, 0.0) + aggregation = op.ScatterElements( + zeros, selected_experts, routing_weights, axis=-1 + ) + + moe_out = op.QMoE( + flat_f32, # 0: input + scores_for_choice, # 1: router_probs (selection logits) + self.fc1_experts_weights, # 2 + self.fc1_scales, # 3 + None, # 4: fc1_experts_bias + self.fc2_experts_weights, # 5 + self.fc2_scales, # 6 + None, # 7: fc2_experts_bias + None, # 8: fc3_experts_weights + None, # 9: fc3_scales + None, # 10: fc3_experts_bias + None, # 11: fc1_zero_points + None, # 12: fc2_zero_points + None, # 13: fc3_zero_points + aggregation, # 14: router_weights (explicit combine weights) + activation_type="swiglu", + k=self._top_k, + normalize_routing_weights=0, + swiglu_fusion=1, + expert_weight_bits=self._bits, + block_size=self._block_size, + quant_type="int", + _domain=self._MICROSOFT_DOMAIN, + ) + moe_out = op.CastLike(moe_out, hidden_states) + return op.Reshape(moe_out, orig_shape) + + def _route(self, op: OpBuilder, hidden_states: ir.Value): + """Invoke ``gate.route_for_qmoe`` under the gate's module scope. + + ``route_for_qmoe`` is a plain method (not ``forward``), so it never goes + through :meth:`nn.Module.__call__`. We replicate the parameter-realization + step here so the gate's ``weight`` / ``e_score_correction_bias`` are + registered as graph initializers under the ``...moe.gate`` scope (matching + the per-expert path) instead of dangling as unqualified names. + """ + builder = op.builder + module_name = self.gate._name or "gate" + class_name = type(self.gate).__qualname__ + builder.push_module(module_name, class_name) + try: + for param in self.gate._parameters.values(): + param._realize(builder) + return self.gate.route_for_qmoe(op, hidden_states) + finally: + builder.pop_module() diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index 4d2d47c9..125c2fd5 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -19,6 +19,7 @@ MLP, Attention, Embedding, + FusedQuantizedMoE, Linear, MoELayer, QuantizedEmbedding, @@ -120,7 +121,45 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return routing_weights, selected_experts - def _group_topk_selection(self, op, scores_for_choice): + def route_for_qmoe(self, op: OpBuilder, hidden_states: ir.Value): + """Routing outputs shaped for the fused ``com.microsoft::QMoE`` op. + + Returns a triple ``(scores_for_choice, routing_weights, selected_experts)`` + computed identically to :meth:`forward`, but also exposes + ``scores_for_choice`` — the per-expert selection scores (sigmoid + bias + for V3, softmax for V2, with group masking already applied). The QMoE + kernel performs its own top-k over ``scores_for_choice`` (passed as + ``router_probs``), so feeding these masked scores reproduces GLM's + group-limited / noaux_tc selection exactly. ``routing_weights`` and + ``selected_experts`` are scattered into a dense aggregation tensor by the + caller and consumed via the optional ``router_weights`` input, so the + combine weights (normalized sigmoid * routed_scaling_factor) match the + per-expert path bit-for-bit. + """ + weight_t = op.Transpose(self.weight, perm=[1, 0]) + router_logits = op.MatMul(op.Cast(hidden_states, to=1), weight_t) + + if self.scoring_func == "sigmoid": + scores = op.Sigmoid(router_logits) + scores_for_choice = op.Add(scores, self.e_score_correction_bias) + else: + scores = op.Softmax(router_logits, axis=-1) + scores_for_choice = scores + + if self.n_group > 1 and self.topk_method != "greedy": + scores_for_choice = self._group_topk_selection(op, scores_for_choice) + + k_val = op.Constant(value_ints=[self.top_k]) + _, selected_experts = op.TopK(scores_for_choice, k_val, axis=-1, _outputs=2) + + routing_weights = op.GatherElements(scores, selected_experts, axis=-1) + if self.norm_topk_prob: + weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True) + eps = 1e-20 + routing_weights = op.Div(routing_weights, op.Add(weight_sum, eps)) + routing_weights = op.Mul(routing_weights, float(self.routed_scaling_factor)) + + return scores_for_choice, routing_weights, selected_experts """Group-based expert selection: pick topk_group groups first.""" experts_per_group = self.num_experts // self.n_group @@ -280,7 +319,16 @@ def __init__( super().__init__() assert config.num_local_experts is not None assert config.moe_intermediate_size is not None - self.moe = MoELayer(config, gate=gate, linear_class=linear_class) + qc = config.quantization + use_fused_qmoe = ( + config.fused_quantized_moe + and qc is not None + and qc.quant_method != "none" + ) + if use_fused_qmoe: + self.moe = FusedQuantizedMoE(config, gate=gate) + else: + self.moe = MoELayer(config, gate=gate, linear_class=linear_class) # Shared expert uses moe_intermediate_size * n_shared_experts n_shared = config.n_shared_experts or 1 shared_intermediate = config.moe_intermediate_size * n_shared diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py index 545bc3dc..fff21ca8 100644 --- a/src/mobius/models/glm_moe_dsa_test.py +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -171,7 +171,33 @@ def test_quantized_graph_uses_matmul_nbits(self): )["model"].graph assert count_op_type(graph, "MatMulNBits") > 0 - def test_preprocess_maps_indexer_and_mtp_layer(self): + def test_fused_quantized_moe_emits_qmoe(self): + config = _tiny_glm_moe_dsa_config( + quantization=QuantizationConfig( + bits=4, + group_size=32, + quant_method="gguf", + sym=True, + ), + fused_quantized_moe=True, + ) + graph = build_from_module( + GlmMoeDsaCausalLMModel(config), + config, + task="glm-moe-dsa", + )["model"].graph + # Routed experts collapse into a single fused QMoE node per MoE layer. + n_moe_layers = config.num_hidden_layers - config.first_k_dense_replace + assert count_op_type(graph, "QMoE") == n_moe_layers + # Routed experts no longer emit per-expert MatMulNBits: the only + # MatMulNBits left come from attention, the dense layer, shared experts, + # and lm_head — none named under the routed-expert ``moe.experts`` scope. + routed = [ + node.name + for node in graph + if node.op_type == "MatMulNBits" and "moe.experts" in (node.name or "") + ] + assert routed == [], f"routed experts still emit MatMulNBits: {routed}" config = _tiny_glm_moe_dsa_config() model = GlmMoeDsaCausalLMModel(config) state_dict = { From 87056ad61c394f58d03ef5da1bb1e79cbd96de29 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 19 Jul 2026 17:49:27 +0000 Subject: [PATCH 07/23] fix(deepseek): restore group top-k selection Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/deepseek.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index 125c2fd5..7579ef6e 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -160,6 +160,8 @@ def route_for_qmoe(self, op: OpBuilder, hidden_states: ir.Value): routing_weights = op.Mul(routing_weights, float(self.routed_scaling_factor)) return scores_for_choice, routing_weights, selected_experts + + def _group_topk_selection(self, op, scores_for_choice): """Group-based expert selection: pick topk_group groups first.""" experts_per_group = self.num_experts // self.n_group From 8867ef3d6b06d80209245a7cb4ec96a8976a6ede Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jul 2026 22:50:54 +0000 Subject: [PATCH 08/23] style: fix GLM export lint Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- export_deepseek_v2_tiny.py | 6 ++--- export_glm_tiny_qmoe.py | 22 ++++++++++++---- export_glm_tiny_quant.py | 26 ++++++++++++++----- src/mobius/components/_moe.py | 4 +-- .../integrations/gguf/_tensor_mapping.py | 5 +--- .../integrations/gguf/_tensor_mapping_test.py | 7 ++--- .../integrations/ort_genai/auto_export.py | 8 ++---- src/mobius/models/deepseek.py | 4 +-- src/mobius/models/glm_moe_dsa.py | 8 ++++-- src/mobius/models/glm_moe_dsa_test.py | 4 ++- src/mobius/tasks/_glm_moe_dsa.py | 16 +++++++----- 11 files changed, 65 insertions(+), 45 deletions(-) diff --git a/export_deepseek_v2_tiny.py b/export_deepseek_v2_tiny.py index fe24047a..a994a392 100644 --- a/export_deepseek_v2_tiny.py +++ b/export_deepseek_v2_tiny.py @@ -9,8 +9,8 @@ import numpy as np import onnx_ir as ir - from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config + from mobius._config_resolver import _default_task_for_model from mobius._registry import registry from mobius.integrations.onnx_genai import write_inference_metadata @@ -39,9 +39,7 @@ def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: def main() -> None: - overrides = dict( - next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "deepseek_v2") - ) + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "deepseek_v2")) config = _base_config(**overrides) config.dtype = ir.DataType.FLOAT diff --git a/export_glm_tiny_qmoe.py b/export_glm_tiny_qmoe.py index 057a37d0..4f173fba 100644 --- a/export_glm_tiny_qmoe.py +++ b/export_glm_tiny_qmoe.py @@ -92,11 +92,23 @@ def main() -> None: if os.path.exists(tok_src): shutil.copy(tok_src, os.path.join(OUT, "tokenizer.json")) - for attr in ("num_hidden_layers", "num_attention_heads", "num_key_value_heads", - "head_dim", "hidden_size", "vocab_size", "max_position_embeddings", - "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim", "kv_lora_rank", - "q_lora_rank", "num_local_experts", "num_experts_per_tok", - "moe_intermediate_size"): + for attr in ( + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "hidden_size", + "vocab_size", + "max_position_embeddings", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "kv_lora_rank", + "q_lora_rank", + "num_local_experts", + "num_experts_per_tok", + "moe_intermediate_size", + ): print(f" {attr} =", getattr(config, attr, None)) print("Saved to", OUT) diff --git a/export_glm_tiny_quant.py b/export_glm_tiny_quant.py index 0d91c900..0b2a7f5d 100644 --- a/export_glm_tiny_quant.py +++ b/export_glm_tiny_quant.py @@ -17,20 +17,20 @@ from __future__ import annotations import os -import sys import shutil +import sys sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests")) import numpy as np import onnx_ir as ir +from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config -from _test_configs import _base_config, ALL_CAUSAL_LM_CONFIGS from mobius._config_resolver import _default_task_for_model from mobius._configs import QuantizationConfig from mobius._registry import registry -from mobius.tasks import get_task from mobius.integrations.onnx_genai import write_inference_metadata +from mobius.tasks import get_task OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-q4" FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny" @@ -90,10 +90,22 @@ def main() -> None: if os.path.exists(tok_src): shutil.copy(tok_src, os.path.join(OUT, "tokenizer.json")) - for attr in ("num_hidden_layers", "num_attention_heads", "num_key_value_heads", - "head_dim", "hidden_size", "vocab_size", "max_position_embeddings", - "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim", "kv_lora_rank", - "q_lora_rank", "num_local_experts", "moe_intermediate_size"): + for attr in ( + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "hidden_size", + "vocab_size", + "max_position_embeddings", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "kv_lora_rank", + "q_lora_rank", + "num_local_experts", + "moe_intermediate_size", + ): print(f" {attr} =", getattr(config, attr, None)) print("Saved to", OUT) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 010426a2..12bd9d7c 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -348,9 +348,7 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: # Dense [rows, E] aggregation weights: combine weight at each selected # expert position, 0 elsewhere. QMoE reads these at its own top-k picks. zeros = op.Mul(scores_for_choice, 0.0) - aggregation = op.ScatterElements( - zeros, selected_experts, routing_weights, axis=-1 - ) + aggregation = op.ScatterElements(zeros, selected_experts, routing_weights, axis=-1) moe_out = op.QMoE( flat_f32, # 0: input diff --git a/src/mobius/integrations/gguf/_tensor_mapping.py b/src/mobius/integrations/gguf/_tensor_mapping.py index 7e5415ee..7c8ac3fa 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping.py +++ b/src/mobius/integrations/gguf/_tensor_mapping.py @@ -295,10 +295,7 @@ def is_known_skip(gguf_name: str) -> bool: """ if gguf_name.startswith("tokenizer."): return True - if ( - "rope_freqs" in gguf_name - or "attn_rot_embd" in gguf_name - ): + if "rope_freqs" in gguf_name or "attn_rot_embd" in gguf_name: return True return False diff --git a/src/mobius/integrations/gguf/_tensor_mapping_test.py b/src/mobius/integrations/gguf/_tensor_mapping_test.py index 7364af5c..30c94e29 100644 --- a/src/mobius/integrations/gguf/_tensor_mapping_test.py +++ b/src/mobius/integrations/gguf/_tensor_mapping_test.py @@ -311,9 +311,10 @@ def test_glm_dsa_mapping(self) -> None: assert map_gguf_to_hf_names("blk.78.nextn.eh_proj.weight", "glm-dsa") == ( "model.layers.78.eh_proj.weight" ) - assert map_gguf_to_hf_names( - "blk.78.nextn.shared_head_norm.weight", "glm-dsa" - ) == "model.layers.78.shared_head.norm.weight" + assert ( + map_gguf_to_hf_names("blk.78.nextn.shared_head_norm.weight", "glm-dsa") + == "model.layers.78.shared_head.norm.weight" + ) # ---- Unsupported architecture ---- diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 81518750..d6a68c8c 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -720,9 +720,7 @@ def _write_genai_config( # Derive decoder filename from the actual package key decoder_filename = ( - f"{decoder_key}/model.onnx" - if len(pkg) > 1 or decoder_key != "model" - else "model.onnx" + f"{decoder_key}/model.onnx" if len(pkg) > 1 or decoder_key != "model" else "model.onnx" ) # ORT GenAI's ``past_present_share_buffer`` mode requires the decoder @@ -1010,9 +1008,7 @@ def write_ort_genai_config( "present.0.value", "topk_indices", ], - "num_nextn_predict_layers": getattr( - config, "num_nextn_predict_layers", 0 - ), + "num_nextn_predict_layers": getattr(config, "num_nextn_predict_layers", 0), "index_share_for_mtp_iteration": getattr( config, "index_share_for_mtp_iteration", False ), diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index 7579ef6e..0a3f400e 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -323,9 +323,7 @@ def __init__( assert config.moe_intermediate_size is not None qc = config.quantization use_fused_qmoe = ( - config.fused_quantized_moe - and qc is not None - and qc.quant_method != "none" + config.fused_quantized_moe and qc is not None and qc.quant_method != "none" ) if use_fused_qmoe: self.moe = FusedQuantizedMoE(config, gate=gate) diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index 7ec85bda..94a5d59d 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -194,7 +194,9 @@ def _pack_present( present: tuple, index_keys: ir.Value | None, ) -> tuple: - key = op.Reshape(op.Transpose(present[0], perm=[0, 2, 1, 3]), [0, 0, self.main_key_dim]) + key = op.Reshape( + op.Transpose(present[0], perm=[0, 2, 1, 3]), [0, 0, self.main_key_dim] + ) if index_keys is not None: key = op.Concat(key, index_keys, axis=-1) value = op.Reshape( @@ -301,7 +303,9 @@ def forward( attention_bias, ) if topk_indices is None: - raise ValueError("Shared GLM DSA layers require top-k indices from a preceding full layer") + raise ValueError( + "Shared GLM DSA layers require top-k indices from a preceding full layer" + ) sparse_bias = self._sparse_bias(op, topk_indices, attention_bias) attn_output, present_key, present_value = op.Attention( diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py index fff21ca8..f0f17e8f 100644 --- a/src/mobius/models/glm_moe_dsa_test.py +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -102,7 +102,9 @@ def test_builds_full_attention_mla_moe_graph(self): assert count_op_type(graph, "Attention") == config.num_hidden_layers assert count_op_type(graph, "ScatterElements") >= 1 assert count_op_type(graph, "Sigmoid") >= config.num_hidden_layers - 1 - assert any("layers.0.self_attn.indexer.wk.weight" in name for name in graph.initializers) + assert any( + "layers.0.self_attn.indexer.wk.weight" in name for name in graph.initializers + ) assert not any("layers.1.self_attn.indexer" in name for name in graph.initializers) assert {value.name for value in graph.outputs} >= { "logits", diff --git a/src/mobius/tasks/_glm_moe_dsa.py b/src/mobius/tasks/_glm_moe_dsa.py index 094a3043..484acd20 100644 --- a/src/mobius/tasks/_glm_moe_dsa.py +++ b/src/mobius/tasks/_glm_moe_dsa.py @@ -5,6 +5,8 @@ from __future__ import annotations +from typing import ClassVar + import onnx_ir as ir from mobius._configs import ArchitectureConfig @@ -15,7 +17,7 @@ class GlmMoeDsaTask(ModelTask): """Build the target decoder and optional GLM-5.2 MTP component.""" - model_roles = {"model": "decoder", "mtp": "decoder"} + model_roles: ClassVar[dict[str, str]] = {"model": "decoder", "mtp": "decoder"} @staticmethod def _cache_dims(config: ArchitectureConfig, layer_idx: int) -> tuple[int, int]: @@ -30,7 +32,9 @@ def _cache_dims(config: ArchitectureConfig, layer_idx: int) -> tuple[int, int]: def _build_target(self, module, config: ArchitectureConfig): graph, builder = _make_graph("glm_moe_dsa") - input_ids = builder.input("input_ids", ir.DataType.INT64, ["batch_size", "sequence_length"]) + input_ids = builder.input( + "input_ids", ir.DataType.INT64, ["batch_size", "sequence_length"] + ) attention_mask = builder.input( "attention_mask", ir.DataType.INT64, ["batch_size", "total_sequence_length"] ) @@ -100,11 +104,9 @@ def _build_mtp(self, module, config: ArchitectureConfig): position_ids = builder.input( "position_ids", ir.DataType.INT64, ["batch_size", "sequence_length"] ) - key_dim = ( - config.num_attention_heads - * (int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0)) - + int(config.index_head_dim or 0) - ) + key_dim = config.num_attention_heads * ( + int(config.qk_nope_head_dim or 0) + int(config.qk_rope_head_dim or 0) + ) + int(config.index_head_dim or 0) value_dim = config.num_attention_heads * int(config.v_head_dim or 0) past_key = builder.input( "past_key_values.0.key", From ab78ba004e84c2db64dfac41d333242d34187b52 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jul 2026 22:56:16 +0000 Subject: [PATCH 09/23] test: update onnx-genai runtime writer mock Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/cli_test.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/cli_test.py b/tests/cli_test.py index 3df44cb3..bc69057c 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -253,14 +253,14 @@ def test_runtime_ort_genai_calls_write_ort_genai_config(self): call_kwargs = mock_export.call_args assert call_kwargs.kwargs.get("hf_model_id") == "Qwen/Qwen2.5-0.5B" - def test_runtime_onnx_genai_calls_write_inference_metadata(self): - """--runtime onnx-genai writes inference_metadata (not genai_config).""" + def test_runtime_onnx_genai_calls_write_onnx_genai_config(self): + """--runtime onnx-genai calls the unified config writer.""" with ( tempfile.TemporaryDirectory() as tmpdir, mock.patch( - "mobius.integrations.onnx_genai.write_inference_metadata", - return_value="inference_metadata.yaml", - ) as mock_meta, + "mobius.integrations.onnx_genai.write_onnx_genai_config", + return_value={}, + ) as mock_export, mock.patch( "mobius.integrations.ort_genai.write_ort_genai_config", return_value={}, @@ -279,7 +279,7 @@ def test_runtime_onnx_genai_calls_write_inference_metadata(self): ] ) - mock_meta.assert_called_once() + mock_export.assert_called_once() mock_ort.assert_not_called() def test_no_runtime_does_not_call_write_ort_genai_config(self): From 2c4fdc92327dfb399301d9b7b0f1c79423676b52 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jul 2026 23:15:08 +0000 Subject: [PATCH 10/23] test: cover streamed GLM GGUF quantization Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/integrations/gguf/_builder_test.py | 107 ++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index d01bf93e..98b0a02e 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -445,3 +445,110 @@ def test_normalize_glm_dsa_split_kv_b_and_router_bias(): expected = torch.cat((k_proj.transpose(1, 2), v_proj), dim=1).reshape(18, 3) assert torch.equal(result["model.layers.0.self_attn.kv_b_proj.weight"], expected) assert "model.layers.0.mlp.gate.e_score_correction_bias" in result + + +def test_load_quantized_glm_experts_repacked_individually(monkeypatch): + from types import SimpleNamespace + + from gguf import GGMLQuantizationType + + from mobius.integrations.gguf import _repacker, _tencent_q1_0, _tensor_mapping + from mobius.integrations.gguf._builder import _load_quantized_state_dict + + repacked_inputs = [] + + def _repack(raw, qtype, shape): + repacked_inputs.append((raw.copy(), qtype, shape)) + return SimpleNamespace( + weight=np.full((2, 1, 16), len(repacked_inputs), dtype=np.uint8), + scales=np.ones((2, 1), dtype=np.float32), + zero_points=None, + ) + + monkeypatch.setattr(_tencent_q1_0, "is_tencent_q1_0_layout", lambda _model: False) + monkeypatch.setattr( + _tensor_mapping, + "map_gguf_to_hf_names", + lambda _name, _arch: "model.layers.0.mlp.experts.gate_proj.weight", + ) + monkeypatch.setattr(_repacker, "repack_quant_params", lambda _qtype: (4, 32)) + monkeypatch.setattr(_repacker, "repack_gguf_tensor", _repack) + + raw = np.arange(8, dtype=np.uint8) + model = SimpleNamespace( + _tensor_index={"experts": object()}, + tensor_items_raw=lambda: [("experts", raw, GGMLQuantizationType.Q4_0, (2, 2, 2))], + ) + module = SimpleNamespace(named_modules=list) + config = SimpleNamespace( + quantization=SimpleNamespace(bits=4, group_size=32, sym=True), + num_attention_heads=1, + num_key_value_heads=1, + model_type="glm_moe_dsa", + ) + + result = _load_quantized_state_dict(model, "glm-dsa", module, config) + + assert [call[2] for call in repacked_inputs] == [(2, 2), (2, 2)] + assert np.array_equal(repacked_inputs[0][0], raw[:4]) + assert np.array_equal(repacked_inputs[1][0], raw[4:]) + assert "model.layers.0.mlp.experts.0.gate_proj.weight" in result + assert "model.layers.0.mlp.experts.1.gate_proj.weight" in result + + +def test_load_quantized_glm_fuses_split_kv_b(monkeypatch): + from types import SimpleNamespace + + from gguf import GGMLQuantizationType + + from mobius.integrations.gguf import _repacker, _tencent_q1_0, _tensor_mapping + from mobius.integrations.gguf._builder import _load_quantized_state_dict + + fused_values = [] + + def _repack(values, **_kwargs): + fused_values.append(values.copy()) + return SimpleNamespace( + weight=np.zeros((18, 1, 16), dtype=np.uint8), + scales=np.ones((18, 1), dtype=np.float32), + zero_points=np.zeros((18, 1), dtype=np.uint8), + ) + + names = { + "k": "model.layers.0.self_attn.kv_b_proj.k_proj.weight", + "v": "model.layers.0.self_attn.kv_b_proj.v_proj.weight", + } + monkeypatch.setattr(_tencent_q1_0, "is_tencent_q1_0_layout", lambda _model: False) + monkeypatch.setattr( + _tensor_mapping, + "map_gguf_to_hf_names", + lambda name, _arch: names[name], + ) + monkeypatch.setattr(_repacker, "repack_dequantized_tensor", _repack) + + k_proj = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) + v_proj = np.arange(2 * 5 * 3, dtype=np.float32).reshape(2, 5, 3) + tensors = [ + ("k", k_proj, GGMLQuantizationType.F32, k_proj.shape), + ("v", v_proj, GGMLQuantizationType.F32, v_proj.shape), + ] + model = SimpleNamespace( + _tensor_index={"k": object(), "v": object()}, + tensor_items_raw=lambda: tensors, + dequantize_raw_tensor=lambda raw, _qtype, _shape: raw, + ) + module = SimpleNamespace(named_modules=list) + config = SimpleNamespace( + quantization=SimpleNamespace(bits=4, group_size=32, sym=False), + num_attention_heads=1, + num_key_value_heads=1, + model_type="glm_moe_dsa", + ) + + result = _load_quantized_state_dict(model, "glm-dsa", module, config) + + expected = np.concatenate((k_proj.transpose(0, 2, 1), v_proj), axis=1).reshape(18, 3) + np.testing.assert_array_equal(fused_values[0], expected) + assert "model.layers.0.self_attn.kv_b_proj.weight" in result + assert "model.layers.0.self_attn.kv_b_proj.scales" in result + assert "model.layers.0.self_attn.kv_b_proj.zero_points" in result From a5f4c4961435207b3fdb07b8ac5216af91c61de7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jul 2026 23:24:26 +0000 Subject: [PATCH 11/23] ci: keep GPU validation on CUDA 12 ORT Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/workflows/gpu_l4_golden_parity.yml | 3 ++- .github/workflows/gpu_l5_generation_e2e.yml | 3 ++- .github/workflows/main.yml | 3 ++- .github/workflows/validation_examples_gpu.yml | 3 ++- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 383fd239..71668543 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -59,7 +59,8 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. + pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda - name: Check for golden test data id: check_golden diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index f976b29d..3c4ffb7d 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -59,7 +59,8 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. + pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda - name: Run L5 generation E2E tests env: diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 92ec133e..fc3ca7c3 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -283,7 +283,8 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. + pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda - name: Run fast integration tests env: HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/.github/workflows/validation_examples_gpu.yml b/.github/workflows/validation_examples_gpu.yml index 2b754bb0..313ceaaf 100644 --- a/.github/workflows/validation_examples_gpu.yml +++ b/.github/workflows/validation_examples_gpu.yml @@ -83,7 +83,8 @@ jobs: pip install -r requirements/ci/requirements.txt pip install -e '.[testing,transformers]' pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8. + pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda - name: Run ${{ matrix.name }} env: From 671fc2dfc3ff069482967d9a0106e96bb3e7ff92 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jul 2026 23:26:51 +0000 Subject: [PATCH 12/23] test: handle keyed recurrent cache states Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/integration_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/integration_test.py b/tests/integration_test.py index 843594ac..83107c7f 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -3703,7 +3703,10 @@ def test_qwen35_deltanet_single_layer_parity(): hidden_states=torch.from_numpy(hidden_np).float(), cache_params=cache, ).numpy() - hf_rec = cache.layers[0].recurrent_states.numpy() + hf_recurrent_states = cache.layers[0].recurrent_states + if isinstance(hf_recurrent_states, dict): + hf_recurrent_states = hf_recurrent_states[0] + hf_rec = hf_recurrent_states.numpy() # ONNX forward sess = _make_session(onnx_model) From 2f7c414bc164079fe568a11ea396e640adfc096d Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Jul 2026 23:43:18 +0000 Subject: [PATCH 13/23] test: close ONNX temp files before Windows saves Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/llada_test.py | 8 +++++--- src/mobius/models/unet_parity_test.py | 29 ++++++++++++++++----------- 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/src/mobius/models/llada_test.py b/src/mobius/models/llada_test.py index c8025c27..ff0872ac 100644 --- a/src/mobius/models/llada_test.py +++ b/src/mobius/models/llada_test.py @@ -20,6 +20,7 @@ from __future__ import annotations import tempfile +from pathlib import Path import numpy as np import pytest @@ -167,11 +168,12 @@ def _build_onnx_session(config: ArchitectureConfig, state: dict[str, torch.Tenso model = MaskedDiffusionTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(state)) - with tempfile.NamedTemporaryFile(suffix=".onnx") as handle: - onnx_ir.save(model, handle.name) + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "model.onnx" + onnx_ir.save(model, model_path) # onnxruntime reads the model fully at construction, so the temp file # can be removed as soon as the session exists. - return ort.InferenceSession(handle.name) + return ort.InferenceSession(model_path) def test_llada_matches_torch_reference(): diff --git a/src/mobius/models/unet_parity_test.py b/src/mobius/models/unet_parity_test.py index fbdf2b1a..5dd7b348 100644 --- a/src/mobius/models/unet_parity_test.py +++ b/src/mobius/models/unet_parity_test.py @@ -12,6 +12,7 @@ import re import tempfile +from pathlib import Path import numpy as np import pytest @@ -63,9 +64,10 @@ def test_cross_attention_block_matches_diffusers(): model = _make_model(graph) apply_weights(model, _remap_transformer(hf.state_dict())) - with tempfile.NamedTemporaryFile(suffix=".onnx") as handle: - onnx_ir.save(model, handle.name) - session = ort.InferenceSession(handle.name) + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "model.onnx" + onnx_ir.save(model, model_path) + session = ort.InferenceSession(model_path) actual = session.run(None, {"hidden": hidden.numpy(), "context": context.numpy()})[0] assert np.abs(actual - expected).max() < 1e-4 @@ -115,9 +117,10 @@ def test_unet_matches_diffusers(): model = DenoisingTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(dict(hf.state_dict()))) - with tempfile.NamedTemporaryFile(suffix=".onnx") as handle: - onnx_ir.save(model, handle.name) - session = ort.InferenceSession(handle.name) + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "model.onnx" + onnx_ir.save(model, model_path) + session = ort.InferenceSession(model_path) actual = session.run( None, { @@ -185,9 +188,10 @@ def test_unet_sd1x_mixed_block_types_matches_diffusers(): model = DenoisingTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(dict(hf.state_dict()))) - with tempfile.NamedTemporaryFile(suffix=".onnx") as handle: - onnx_ir.save(model, handle.name) - session = ort.InferenceSession(handle.name) + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "model.onnx" + onnx_ir.save(model, model_path) + session = ort.InferenceSession(model_path) actual = session.run( None, { @@ -271,9 +275,10 @@ def test_unet_lora_gate_parity(): weights.update(remap_diffusers_unet_lora(lora_state, "test")) apply_weights(model, weights) - with tempfile.NamedTemporaryFile(suffix=".onnx") as handle: - onnx_ir.save(model, handle.name) - session = ort.InferenceSession(handle.name) + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "model.onnx" + onnx_ir.save(model, model_path) + session = ort.InferenceSession(model_path) feed = { "sample": sample.numpy(), "timestep": timestep.numpy().astype(np.int64), From e98e6984ddc39c6ee120df90eda2afa592f34871 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 14:10:44 +0000 Subject: [PATCH 14/23] fix: address GLM4 GPTQ review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_configs/_base.py | 4 +- src/mobius/_configs_test.py | 17 ++ src/mobius/models/chatglm.py | 71 ++++++-- src/mobius/models/chatglm_test.py | 151 ++++++++++++++++++ .../rewrite_rules/_group_query_attention.py | 27 +++- .../_group_query_attention_test.py | 47 ++++++ 6 files changed, 297 insertions(+), 20 deletions(-) create mode 100644 src/mobius/models/chatglm_test.py diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 1431bf35..09ac2a45 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -51,7 +51,7 @@ def _resolve_hidden_act(config, model_type: str) -> str | None: dense_act_fn — some BERT variants activation — generic fallback afn — older BERT configs - "silu" (qwen) — Qwen v1 hardcodes silu; no activation attr + "silu" (qwen, chatglm) — Qwen v1 and ChatGLM hardcode silu; no activation attr "gelu" (XLM) — gelu_activation=True is a boolean flag "relu" (ctrl) — CTRL hardcodes relu; no hidden_act attr """ @@ -65,7 +65,7 @@ def _resolve_hidden_act(config, model_type: str) -> str | None: or getattr(config, "afn", None) # LLaDA/OLMo expose the activation as ``activation_type`` (e.g. "silu"). or getattr(config, "activation_type", None) - or ("silu" if model_type in ("qwen",) else None) + or ("silu" if model_type in ("qwen", "chatglm") else None) # gelu_activation is a boolean (XLM) — must be after all string # attrs so it cannot override an explicit hidden_act. or ("gelu" if getattr(config, "gelu_activation", False) else None) diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index 219355a1..75e7a2e1 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -1052,6 +1052,23 @@ def test_boa_token_id_none_when_absent(self): class TestActivationFallbacks: """Tests for hidden_act extraction fallbacks (ff_activation, gelu_activation).""" + def test_chatglm_silu_fallback(self): + """ChatGLM hardcodes SiLU without exposing an activation field.""" + + class FakeConfig: + model_type = "chatglm" + num_attention_heads = 8 + num_key_value_heads = 8 + num_hidden_layers = 2 + vocab_size = 1000 + hidden_size = 256 + intermediate_size = 512 + max_position_embeddings = 1024 + head_dim = 32 + + config = ArchitectureConfig.from_transformers(FakeConfig()) + assert config.hidden_act == "silu" + def test_ff_activation_fallback(self): """ff_activation is used when hidden_act is absent (XLNet pattern).""" diff --git a/src/mobius/models/chatglm.py b/src/mobius/models/chatglm.py index ff6a4b66..bf47963a 100644 --- a/src/mobius/models/chatglm.py +++ b/src/mobius/models/chatglm.py @@ -5,29 +5,68 @@ import torch -from mobius._weight_utils import rename_weight_keys +from mobius._weight_utils import split_fused_qkv, split_gate_up_proj from mobius.models.base import CausalLMModel class ChatGLMCausalLMModel(CausalLMModel): - """ChatGLM model with partial rotary (0.5 factor) and MLP name remapping. + """ChatGLM model with partial rotary and fused projection imports. - ChatGLM uses interleaved RoPE and has a different MLP attribute naming - convention (dense_4h_to_h instead of down_proj). + GLM-4's original ChatGLM checkpoint layout nests decoder weights below + ``transformer.encoder`` and stores QKV and gate/up projections fused. + Canonicalize those names before generic GPTQ preprocessing, then split the + converted QKV and gate/up tensors along their output dimension. Splitting + the repacked gate/up tensors is lossless and avoids a runtime activation + Split between captured CUDA graph segments. """ def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: - state_dict = super().preprocess_weights(state_dict) - # ChatGLM uses HF-specific attribute names: - # dense_h_to_4h / dense_4h_to_h → up_proj / down_proj (FCMLP naming) - # self_attention → self_attn - return rename_weight_keys( - state_dict, - [ - ("dense_4h_to_h", "down_proj"), - ("dense_h_to_4h", "up_proj"), - ("self_attention", "self_attn"), - ], - ) + renamed: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + key = key.replace("transformer.embedding.word_embeddings.", "model.embed_tokens.") + key = key.replace("transformer.encoder.final_layernorm.", "model.norm.") + key = key.replace("transformer.output_layer.", "lm_head.") + key = key.replace("transformer.encoder.layers.", "model.layers.") + key = key.replace(".self_attention.", ".self_attn.") + key = key.replace(".self_attn.dense.", ".self_attn.o_proj.") + key = key.replace(".mlp.dense_h_to_4h.", ".mlp.gate_up_proj.") + key = key.replace(".mlp.dense_4h_to_h.", ".mlp.down_proj.") + key = key.replace(".self_attn.query_key_value.", ".self_attn.qkv_proj.") + renamed[key] = value + + state_dict = super().preprocess_weights(renamed) + + quantization = getattr(self.config, "quantization", None) + if quantization is not None and quantization.sym: + state_dict = { + key: value + for key, value in state_dict.items() + if not key.endswith(".zero_points") + } + + for key in list(state_dict): + if ".self_attn.qkv_proj." not in key: + continue + q, k, v = split_fused_qkv( + state_dict.pop(key), + self.config.num_attention_heads, + self.config.num_key_value_heads, + self.config.head_dim, + ) + state_dict[key.replace("qkv_proj", "q_proj")] = q + state_dict[key.replace("qkv_proj", "k_proj")] = k + state_dict[key.replace("qkv_proj", "v_proj")] = v + + for key in list(state_dict): + if ".mlp.gate_up_proj." not in key: + continue + gate, up = split_gate_up_proj( + state_dict.pop(key), + self.config.intermediate_size, + ) + state_dict[key.replace("gate_up_proj", "gate_proj")] = gate + state_dict[key.replace("gate_up_proj", "up_proj")] = up + + return state_dict diff --git a/src/mobius/models/chatglm_test.py b/src/mobius/models/chatglm_test.py new file mode 100644 index 00000000..be95d13d --- /dev/null +++ b/src/mobius/models/chatglm_test.py @@ -0,0 +1,151 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import torch + +from mobius._configs import ArchitectureConfig, QuantizationConfig +from mobius._weight_utils import preprocess_gptq_weights +from mobius.components import MLP +from mobius.models.chatglm import ChatGLMCausalLMModel + + +def _gptq_tensor_shapes(k: int, n: int, group_size: int = 16): + return { + "qweight": torch.zeros(k // 8, n, dtype=torch.int32), + "qzeros": torch.zeros(1, n, dtype=torch.int32), + "scales": torch.ones(k // group_size, n, dtype=torch.float16), + "g_idx": torch.arange(k, dtype=torch.int32) // group_size, + } + + +def _add_gptq(state_dict: dict[str, torch.Tensor], stem: str, k: int, n: int) -> None: + for suffix, value in _gptq_tensor_shapes(k, n).items(): + state_dict[f"{stem}.{suffix}"] = value + + +def test_glm4_gptq_checkpoint_remap_and_fused_projection_split(): + config = ArchitectureConfig( + model_type="chatglm", + vocab_size=64, + hidden_size=32, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + hidden_act="silu", + max_position_embeddings=128, + rope_type="default", + partial_rotary_factor=0.5, + rope_interleave=True, + attn_qkv_bias=True, + quantization=QuantizationConfig(bits=4, group_size=16, quant_method="gptq", sym=True), + ) + model = ChatGLMCausalLMModel(config) + assert isinstance(model.model.layers[0].mlp, MLP) + state_dict = { + "transformer.embedding.word_embeddings.weight": torch.zeros(64, 32), + "transformer.encoder.final_layernorm.weight": torch.ones(32), + "transformer.output_layer.weight": torch.zeros(64, 32), + "transformer.encoder.layers.0.input_layernorm.weight": torch.ones(32), + "transformer.encoder.layers.0.post_attention_layernorm.weight": torch.ones(32), + "transformer.encoder.layers.0.self_attention.query_key_value.bias": torch.zeros(64), + "transformer.encoder.layers.0.mlp.dense_h_to_4h.bias": torch.arange( + 32, dtype=torch.float16 + ), + } + _add_gptq( + state_dict, + "transformer.encoder.layers.0.self_attention.query_key_value", + 32, + 64, + ) + _add_gptq( + state_dict, + "transformer.encoder.layers.0.self_attention.dense", + 32, + 32, + ) + _add_gptq( + state_dict, + "transformer.encoder.layers.0.mlp.dense_h_to_4h", + 32, + 32, + ) + _add_gptq( + state_dict, + "transformer.encoder.layers.0.mlp.dense_4h_to_h", + 16, + 32, + ) + + result = model.preprocess_weights(state_dict) + + assert result["model.layers.0.self_attn.q_proj.weight"].shape == (32, 2, 8) + assert result["model.layers.0.self_attn.k_proj.weight"].shape == (16, 2, 8) + assert result["model.layers.0.self_attn.v_proj.weight"].shape == (16, 2, 8) + assert result["model.layers.0.self_attn.q_proj.scales"].shape == (32, 2) + assert result["model.layers.0.self_attn.k_proj.bias"].shape == (16,) + assert not any(key.endswith(".zero_points") for key in result) + assert "model.layers.0.self_attn.qkv_proj.weight" not in result + assert "model.layers.0.mlp.gate_up_proj.weight" not in result + assert result["model.layers.0.mlp.gate_proj.weight"].shape == (16, 2, 8) + assert result["model.layers.0.mlp.up_proj.weight"].shape == (16, 2, 8) + assert torch.equal( + result["model.layers.0.mlp.gate_proj.bias"], + torch.arange(16, dtype=torch.float16), + ) + assert torch.equal( + result["model.layers.0.mlp.up_proj.bias"], + torch.arange(16, 32, dtype=torch.float16), + ) + assert "model.layers.0.mlp.down_proj.weight" in result + assert "model.embed_tokens.weight" in result + assert "model.norm.weight" in result + assert "lm_head.weight" in result + + +def test_glm4_gptq_gate_up_split_preserves_packed_tensors(): + config = ArchitectureConfig( + model_type="chatglm", + vocab_size=64, + hidden_size=32, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + hidden_act="silu", + max_position_embeddings=128, + rope_type="default", + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gptq", + sym=False, + ), + ) + model = ChatGLMCausalLMModel(config) + source: dict[str, torch.Tensor] = {} + stem = "transformer.encoder.layers.0.mlp.dense_h_to_4h" + _add_gptq(source, stem, 32, 32) + source[f"{stem}.qweight"][:, 16:] = 0x12345678 + source[f"{stem}.scales"][:, 16:] = 2 + source[f"{stem}.qzeros"][:, 16:] = 0x11111111 + + expected = preprocess_gptq_weights( + { + key.replace(stem, "model.layers.0.mlp.gate_up_proj"): value + for key, value in source.items() + }, + bits=4, + group_size=16, + ) + result = model.preprocess_weights(source) + + for suffix in ("weight", "scales", "zero_points"): + fused = expected[f"model.layers.0.mlp.gate_up_proj.{suffix}"] + assert torch.equal(result[f"model.layers.0.mlp.gate_proj.{suffix}"], fused[:16]) + assert torch.equal(result[f"model.layers.0.mlp.up_proj.{suffix}"], fused[16:]) diff --git a/src/mobius/rewrite_rules/_group_query_attention.py b/src/mobius/rewrite_rules/_group_query_attention.py index df19cb06..f90b772f 100644 --- a/src/mobius/rewrite_rules/_group_query_attention.py +++ b/src/mobius/rewrite_rules/_group_query_attention.py @@ -50,6 +50,23 @@ ) +def _has_unequal_kv_head_dimensions(k, v, past_key, past_value) -> bool: + """Return whether static K/V shapes prove incompatible GQA head dimensions.""" + + def _static_last_dim(value): + if value is None or value.shape is None or len(value.shape) == 0: + return None + dim = value.shape[-1] + return dim if isinstance(dim, int) else None + + for key, value in ((k, v), (past_key, past_value)): + key_dim = _static_last_dim(key) + value_dim = _static_last_dim(value) + if key_dim is not None and value_dim is not None and key_dim != value_dim: + return True + return False + + class RotaryAttentionToGQA(RewriteRuleClassBase): """Replace RotaryEmbedding + Attention with GroupQueryAttention. @@ -119,7 +136,7 @@ def pattern(self, op, q_pre, k_pre, v, attention_bias, past_key, past_value, cos # ------------------------------------------------------------------ check - def check(self, context, attn_out, cos, sin, past_key, past_value, **_): + def check(self, context, attn_out, k_pre, v, cos, sin, past_key, past_value, **_): result = MatchResult() attn = attn_out.producer() @@ -147,6 +164,9 @@ def check(self, context, attn_out, cos, sin, past_key, past_value, **_): if past_value.producer() is not None: return result.fail("past_value is not a graph input") + if _has_unequal_kv_head_dimensions(k_pre, v, past_key, past_value): + return result.fail("K and V head dimensions differ; retain standard Attention") + return result # ------------------------------------------------------------------ rewrite @@ -547,7 +567,7 @@ def pattern(self, op, q, k, v, attention_bias, past_key, past_value): # ------------------------------------------------------------------ check - def check(self, context, attn_out, past_key, past_value, **_): + def check(self, context, attn_out, k, v, past_key, past_value, **_): result = MatchResult() attn = attn_out.producer() @@ -571,6 +591,9 @@ def check(self, context, attn_out, past_key, past_value, **_): if not any(gi.name == "attention_mask" for gi in graph.inputs): return result.fail("No attention_mask graph input — cannot build seqlens_k") + if _has_unequal_kv_head_dimensions(k, v, past_key, past_value): + return result.fail("K and V head dimensions differ; retain standard Attention") + return result # ------------------------------------------------------------------ rewrite diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index b67db66a..2cb5b957 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -113,6 +113,31 @@ ) +def _build_tiny_mla_graph(v_head_dim: int): + """Build a CPU-only MLA graph with configurable K/V head dimensions.""" + config = ArchitectureConfig( + model_type="deepseek_v2", + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + num_hidden_layers=1, + vocab_size=256, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=v_head_dim, + ) + return build_from_module(registry.get("deepseek_v2")(config), config)["model"] + + class TestGroupQueryAttentionRules: def test_rules_returns_rule_set(self): rules = group_query_attention_rules() @@ -209,6 +234,28 @@ def test_preserves_non_matching_model(self): assert vision_counts.get("Attention", 0) == vision_attn_before assert vision_counts.get("GroupQueryAttention", 0) == 0 + def test_retains_attention_for_unequal_mla_kv_head_dimensions(self): + """MLA's unequal K/V dimensions are incompatible with GQA.""" + model = _build_tiny_mla_graph(v_head_dim=8) + assert count_ops(model).get("Attention", 0) == 1 + + rewrite(model, pattern_rewrite_rules=group_query_attention_rules()) + + counts = count_ops(model) + assert counts.get("Attention", 0) == 1 + assert counts.get("GroupQueryAttention", 0) == 0 + + def test_fuses_attention_for_equal_mla_kv_head_dimensions(self): + """The unequal-dimension guard does not reject compatible MLA graphs.""" + model = _build_tiny_mla_graph(v_head_dim=16) + assert count_ops(model).get("Attention", 0) == 1 + + rewrite(model, pattern_rewrite_rules=group_query_attention_rules()) + + counts = count_ops(model) + assert counts.get("Attention", 0) == 0 + assert counts.get("GroupQueryAttention", 0) == 1 + def test_fallback_attention_to_gqa_no_rope(self): """AttentionToGQA fallback fires when applied in isolation (do_rotary=0). From 958f2eb49bcf8365d8f36e6ee25a45b52536bffc Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 14:11:36 +0000 Subject: [PATCH 15/23] Address GLM export review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/research/glm52-export.md | 2 +- export_deepseek_v2_tiny.py | 26 ++++++--- export_glm_tiny_qmoe.py | 36 ++++++++---- export_glm_tiny_quant.py | 47 +++++++++------ .../integrations/onnx_genai/auto_export.py | 58 +++++++++++++++++-- .../onnx_genai/auto_export_test.py | 17 ++++++ src/mobius/models/glm_moe_dsa.py | 16 +++-- src/mobius/models/glm_moe_dsa_test.py | 9 ++- tests/cli_test.py | 8 ++- 9 files changed, 171 insertions(+), 48 deletions(-) diff --git a/docs/research/glm52-export.md b/docs/research/glm52-export.md index 088a6827..c225d1f3 100644 --- a/docs/research/glm52-export.md +++ b/docs/research/glm52-export.md @@ -42,7 +42,7 @@ weight into `attn_k_b` and `attn_v_b`, while Transformers exposes one The exporter now preserves the 78-layer backbone, IndexShare DSA, and the additional improved-MTP layer. Full indexers are emitted for layers 0, 1, 2, -6, 10, ... 74; each full layer's INT32 top-k result is reused by the following +6, 10, ... 74; each full layer's INT64 top-k result is reused by the following three shared layers. The indexer implements the checkpoint equations in fp32: Q-LoRA residual projection, LayerNorm+interleaved-RoPE key projection, ReLU scores, learned signed head weights, causal bias, and dynamic diff --git a/export_deepseek_v2_tiny.py b/export_deepseek_v2_tiny.py index a994a392..68a27a7a 100644 --- a/export_deepseek_v2_tiny.py +++ b/export_deepseek_v2_tiny.py @@ -2,6 +2,7 @@ from __future__ import annotations +import argparse import os import sys @@ -13,10 +14,12 @@ from mobius._config_resolver import _default_task_for_model from mobius._registry import registry -from mobius.integrations.onnx_genai import write_inference_metadata +from mobius.integrations.onnx_genai import write_onnx_genai_config from mobius.tasks import get_task -OUT = "/home/justinchu/ds-e2e-artifacts/deepseek-v2-tiny" +ARTIFACTS_DIR = os.environ.get( + "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts") +) def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: @@ -39,6 +42,13 @@ def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + default=os.path.join(ARTIFACTS_DIR, "deepseek-v2-tiny"), + ) + args = parser.parse_args() + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "deepseek_v2")) config = _base_config(**overrides) config.dtype = ir.DataType.FLOAT @@ -51,11 +61,13 @@ def main() -> None: for model in pkg.values(): _fill_random_weights(model, rng) - os.makedirs(OUT, exist_ok=True) - pkg.save(OUT, external_data="onnx", check_weights=False) - print("inference_metadata:", write_inference_metadata(pkg, OUT)) - print("Saved to", OUT) - print("files:", sorted(os.listdir(OUT))) + os.makedirs(args.output_dir, exist_ok=True) + pkg.save(args.output_dir, external_data="onnx", check_weights=False) + artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config) + for name, path in artifacts.items(): + print(f"{name}:", path) + print("Saved to", args.output_dir) + print("files:", sorted(os.listdir(args.output_dir))) if __name__ == "__main__": diff --git a/export_glm_tiny_qmoe.py b/export_glm_tiny_qmoe.py index 4f173fba..3847f666 100644 --- a/export_glm_tiny_qmoe.py +++ b/export_glm_tiny_qmoe.py @@ -16,6 +16,7 @@ from __future__ import annotations +import argparse import os import shutil import sys @@ -29,11 +30,12 @@ from mobius._config_resolver import _default_task_for_model from mobius._configs import QuantizationConfig from mobius._registry import registry -from mobius.integrations.onnx_genai import write_inference_metadata +from mobius.integrations.onnx_genai import write_onnx_genai_config from mobius.tasks import get_task -OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-qmoe" -FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny" +ARTIFACTS_DIR = os.environ.get( + "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts") +) def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: @@ -59,6 +61,17 @@ def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny-qmoe"), + ) + parser.add_argument( + "--fp32-output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny"), + ) + args = parser.parse_args() + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "glm_moe_dsa")) config = _base_config(**overrides) config.dtype = ir.DataType.FLOAT @@ -82,15 +95,16 @@ def main() -> None: for model in pkg.values(): _fill_random_weights(model, rng) - os.makedirs(OUT, exist_ok=True) - pkg.save(OUT, external_data="onnx", check_weights=False) - path = write_inference_metadata(pkg, OUT) - print("inference_metadata:", path) + os.makedirs(args.output_dir, exist_ok=True) + pkg.save(args.output_dir, external_data="onnx", check_weights=False) + artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config) + for name, path in artifacts.items(): + print(f"{name}:", path) # Reuse tokenizer from the fp32 artifacts if present. - tok_src = os.path.join(FP32_OUT, "tokenizer.json") + tok_src = os.path.join(args.fp32_output_dir, "tokenizer.json") if os.path.exists(tok_src): - shutil.copy(tok_src, os.path.join(OUT, "tokenizer.json")) + shutil.copy(tok_src, os.path.join(args.output_dir, "tokenizer.json")) for attr in ( "num_hidden_layers", @@ -111,8 +125,8 @@ def main() -> None: ): print(f" {attr} =", getattr(config, attr, None)) - print("Saved to", OUT) - print("files:", sorted(os.listdir(OUT))) + print("Saved to", args.output_dir) + print("files:", sorted(os.listdir(args.output_dir))) if __name__ == "__main__": diff --git a/export_glm_tiny_quant.py b/export_glm_tiny_quant.py index 0b2a7f5d..227eb641 100644 --- a/export_glm_tiny_quant.py +++ b/export_glm_tiny_quant.py @@ -1,14 +1,13 @@ -"""Export a QUANTIZED tiny synthetic glm_moe_dsa model for onnx-genai E2E. +"""Export a quantized tiny synthetic glm_moe_dsa model for onnx-genai E2E. Same tiny `glm_moe_dsa` config as ``export_glm_tiny.py`, but with a ``QuantizationConfig`` attached so the linear projections (and per-expert MoE MLPs) are emitted as ``com.microsoft::MatMulNBits`` (int4, block-32). -The GLM/DeepSeek MoE path decomposes experts into per-expert MLPs, so each -expert's gate/up/down projections also become MatMulNBits when quantized. -mobius has no fused ``com.microsoft::QMoE`` emitter for any model (only an -unquantized ``com.microsoft::MoE`` path in gemma4), so QMoE is NOT emitted -here — see the decision note for the scoping gap. +This script leaves ``fused_quantized_moe`` disabled, so the GLM/DeepSeek MoE +path decomposes experts into per-expert MLPs and each expert projection becomes +MatMulNBits. For the fused ``com.microsoft::QMoE`` path, use +``export_glm_tiny_qmoe.py``. Writes an onnx-genai-loadable artifact directory: /model.onnx (+ external data), inference_metadata.yaml, tokenizer.json @@ -16,6 +15,7 @@ from __future__ import annotations +import argparse import os import shutil import sys @@ -29,11 +29,12 @@ from mobius._config_resolver import _default_task_for_model from mobius._configs import QuantizationConfig from mobius._registry import registry -from mobius.integrations.onnx_genai import write_inference_metadata +from mobius.integrations.onnx_genai import write_onnx_genai_config from mobius.tasks import get_task -OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-q4" -FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny" +ARTIFACTS_DIR = os.environ.get( + "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts") +) def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: @@ -59,6 +60,17 @@ def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny-q4"), + ) + parser.add_argument( + "--fp32-output-dir", + default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny"), + ) + args = parser.parse_args() + overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "glm_moe_dsa")) config = _base_config(**overrides) config.dtype = ir.DataType.FLOAT @@ -80,15 +92,16 @@ def main() -> None: for model in pkg.values(): _fill_random_weights(model, rng) - os.makedirs(OUT, exist_ok=True) - pkg.save(OUT, external_data="onnx", check_weights=False) - path = write_inference_metadata(pkg, OUT) - print("inference_metadata:", path) + os.makedirs(args.output_dir, exist_ok=True) + pkg.save(args.output_dir, external_data="onnx", check_weights=False) + artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config) + for name, path in artifacts.items(): + print(f"{name}:", path) # Reuse tokenizer from the fp32 artifacts. - tok_src = os.path.join(FP32_OUT, "tokenizer.json") + tok_src = os.path.join(args.fp32_output_dir, "tokenizer.json") if os.path.exists(tok_src): - shutil.copy(tok_src, os.path.join(OUT, "tokenizer.json")) + shutil.copy(tok_src, os.path.join(args.output_dir, "tokenizer.json")) for attr in ( "num_hidden_layers", @@ -108,8 +121,8 @@ def main() -> None: ): print(f" {attr} =", getattr(config, attr, None)) - print("Saved to", OUT) - print("files:", sorted(os.listdir(OUT))) + print("Saved to", args.output_dir) + print("files:", sorted(os.listdir(args.output_dir))) if __name__ == "__main__": diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index b5fef98a..ae466cb4 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -12,6 +12,7 @@ from __future__ import annotations +import json import logging import os from typing import Any @@ -35,6 +36,41 @@ _DENOISER_KEYS = ("denoiser", "transformer", "unet") +def _write_mtp_config(output_dir: str, config: Any) -> str: + path = os.path.join(output_dir, "mtp_config.json") + with open(path, "w", encoding="utf-8") as handle: + json.dump( + { + "model": {"filename": "mtp/model.onnx"}, + "inputs": [ + "inputs_embeds", + "hidden_states", + "attention_mask", + "position_ids", + "past_key_values.0.key", + "past_key_values.0.value", + ], + "outputs": [ + "mtp_hidden", + "present.0.key", + "present.0.value", + "topk_indices", + ], + "num_nextn_predict_layers": getattr(config, "num_nextn_predict_layers", 0), + "index_share_for_mtp_iteration": getattr( + config, "index_share_for_mtp_iteration", False + ), + "shared_embedding": "model.embed_tokens", + "shared_lm_head": "lm_head", + "runtime_orchestration": "external", + }, + handle, + indent=2, + ) + handle.write("\n") + return path + + def _write_clip_tokenizer(output_dir: str, source: str | None) -> str | None: """Emit ``tokenizer.json`` for a text-conditioned diffusion package. @@ -501,14 +537,28 @@ def write_onnx_genai_config( artifacts["tokenizer"] = tokenizer_path return artifacts - # Fallback: a single-component decoder language model. A multi-component - # package that matched none of the composite shapes above would be silently - # mis-emitted as a bare decoder — fail loudly instead so an unsupported shape - # is obvious rather than producing wrong metadata. try: component_names = sorted(pkg.keys()) except (AttributeError, TypeError): component_names = [] + + if component_names == ["model", "mtp"]: + path = write_decoder_metadata( + output_dir, config=resolved_config, kv_native_dtype=kv_native_dtype + ) + artifacts = { + "inference_metadata": path, + "mtp_config": _write_mtp_config(output_dir, resolved_config), + } + tokenizer_path = _write_hf_tokenizer(output_dir, source) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path + return artifacts + + # Fallback: a single-component decoder language model. A multi-component + # package that matched none of the composite shapes above would be silently + # mis-emitted as a bare decoder — fail loudly instead so an unsupported shape + # is obvious rather than producing wrong metadata. if len(component_names) > 1: raise ValueError( "onnx-genai config emission does not recognize this multi-component " diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index a647518d..fc871b00 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -6,6 +6,7 @@ from __future__ import annotations import dataclasses +import json import pytest import yaml @@ -22,6 +23,8 @@ class _Cfg: max_position_embeddings: int = 8192 sliding_window: int | None = None model_type: str = "qwen" + num_nextn_predict_layers: int = 1 + index_share_for_mtp_iteration: bool = True class _DiffusionPkg(dict): @@ -442,6 +445,20 @@ def test_unrecognized_multi_component_package_fails_loudly(tmp_path): write_onnx_genai_config(pkg, str(tmp_path)) +def test_mtp_package_emits_decoder_metadata_and_sidecar(tmp_path): + pkg = _MultimodalPkg({"model": object(), "mtp": object()}) + + artifacts = write_onnx_genai_config(pkg, str(tmp_path)) + + assert set(artifacts) == {"inference_metadata", "mtp_config"} + with open(artifacts["mtp_config"]) as handle: + mtp = json.load(handle) + assert mtp["model"]["filename"] == "mtp/model.onnx" + assert mtp["num_nextn_predict_layers"] == 1 + assert mtp["index_share_for_mtp_iteration"] is True + assert mtp["runtime_orchestration"] == "external" + + def test_decoder_emits_tokenizer_from_source(tmp_path): # A text-producing package emits tokenizer.json from its HF source so the # onnx-genai package is self-contained. diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index 94a5d59d..c9fb1b53 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -14,14 +14,14 @@ from mobius._configs import ArchitectureConfig from mobius.components import ( MLP, + DeepSeekMLA, Embedding, + LayerNorm, Linear, RMSNorm, create_attention_bias, initialize_rope, ) -from mobius.components._common import LayerNorm -from mobius.components._deepseek_mla import DeepSeekMLA from mobius.components._rotary_embedding import apply_rotary_pos_emb from mobius.models.deepseek import ( DeepSeekMoEGate, @@ -41,7 +41,13 @@ def _indexer_types(config: ArchitectureConfig) -> list[str]: """Return the authoritative full/shared IndexShare schedule.""" if config.indexer_types is not None: - return list(config.indexer_types) + indexer_types = list(config.indexer_types) + if len(indexer_types) != config.num_hidden_layers: + raise ValueError( + "indexer_types must contain exactly one entry per hidden layer " + f"(expected {config.num_hidden_layers}, got {len(indexer_types)})" + ) + return indexer_types full_layers = {0, 1, 2} full_layers.update( range( @@ -120,7 +126,7 @@ def select( key_length = op.Shape(scores, start=2, end=3) k = op.Min(key_length, op.Constant(value_ints=[self.index_topk])) _, indices = op.TopK(scores, k, axis=-1, largest=1, sorted=0, _outputs=2) - return op.Cast(indices, to=ir.DataType.INT32) + return indices def forward( self, @@ -212,7 +218,7 @@ def _sparse_bias( attention_bias: ir.Value, ) -> ir.Value: shape = op.Shape(op.Squeeze(attention_bias, [1])) - masked = op.Expand(op.Constant(value_float=float(ir.DataType.FLOAT.min)), shape) + masked = op.Expand(op.Constant(value_float=float(self.dtype.min)), shape) updates = op.Mul(op.Cast(indices, to=ir.DataType.FLOAT), 0.0) sparse = op.ScatterElements(masked, indices, updates, axis=-1, reduction="none") sparse = op.Cast(op.Unsqueeze(sparse, [1]), to=self.dtype) diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py index f0f17e8f..990fc3e1 100644 --- a/src/mobius/models/glm_moe_dsa_test.py +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -3,6 +3,7 @@ from __future__ import annotations +import pytest import torch from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig @@ -11,7 +12,7 @@ from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius._registry import registry from mobius._testing import count_op_type, make_config -from mobius.models.glm_moe_dsa import GlmMoeDsaCausalLMModel +from mobius.models.glm_moe_dsa import GlmMoeDsaCausalLMModel, _indexer_types def _tiny_glm_moe_dsa_config(**overrides): @@ -90,6 +91,12 @@ def test_registry(self): assert registry.get("glm_moe_dsa") is GlmMoeDsaCausalLMModel assert _default_task_for_model("glm_moe_dsa") == "glm-moe-dsa" + def test_indexer_types_length_must_match_layers(self): + config = _tiny_glm_moe_dsa_config(indexer_types=["full"]) + + with pytest.raises(ValueError, match="expected 4, got 1"): + _indexer_types(config) + def test_builds_full_attention_mla_moe_graph(self): config = _tiny_glm_moe_dsa_config() package = build_from_module( diff --git a/tests/cli_test.py b/tests/cli_test.py index bc69057c..29a4d707 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -253,13 +253,16 @@ def test_runtime_ort_genai_calls_write_ort_genai_config(self): call_kwargs = mock_export.call_args assert call_kwargs.kwargs.get("hf_model_id") == "Qwen/Qwen2.5-0.5B" - def test_runtime_onnx_genai_calls_write_onnx_genai_config(self): + def test_runtime_onnx_genai_calls_write_onnx_genai_config(self, capsys): """--runtime onnx-genai calls the unified config writer.""" with ( tempfile.TemporaryDirectory() as tmpdir, mock.patch( "mobius.integrations.onnx_genai.write_onnx_genai_config", - return_value={}, + return_value={ + "inference_metadata": "inference_metadata.yaml", + "mtp_config": "mtp_config.json", + }, ) as mock_export, mock.patch( "mobius.integrations.ort_genai.write_ort_genai_config", @@ -281,6 +284,7 @@ def test_runtime_onnx_genai_calls_write_onnx_genai_config(self): mock_export.assert_called_once() mock_ort.assert_not_called() + assert "mtp_config: mtp_config.json" in capsys.readouterr().out def test_no_runtime_does_not_call_write_ort_genai_config(self): """Omitting --runtime does NOT call write_ort_genai_config().""" From 751645b7b1b57f689a9c3ca119376fdd5414b36a Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 14:24:03 +0000 Subject: [PATCH 16/23] feat(moe): pack routed experts for fused QMoE Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/components/_moe.py | 216 +++++++++++++- src/mobius/components/_moe_test.py | 266 +++++++++++++++++- src/mobius/integrations/gguf/_builder.py | 8 +- src/mobius/integrations/gguf/_builder_test.py | 59 ++++ src/mobius/models/deepseek.py | 23 +- 5 files changed, 559 insertions(+), 13 deletions(-) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 12bd9d7c..6bc5b2aa 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -6,12 +6,16 @@ from __future__ import annotations import dataclasses +import math +import re from typing import TYPE_CHECKING import onnx_ir as ir +import torch from onnxscript import OpBuilder, nn from mobius._configs import ArchitectureConfig +from mobius._weight_utils import preprocess_awq_weights, preprocess_gptq_weights from mobius.components._mlp import MLP if TYPE_CHECKING: @@ -106,7 +110,7 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): # Renormalize selected weights to sum to 1 (prevents vanishing gradients) weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True) routing_weights = op.Div(routing_weights, op.Add(weight_sum, 1e-9)) - if self.routed_scaling_factor != 1.0: # noqa: RUF069 + if self.routed_scaling_factor != 1.0: # ruff:ignore[float-equality-comparison] routing_weights = op.Mul(routing_weights, self.routed_scaling_factor) return routing_weights, selected_experts @@ -264,8 +268,8 @@ class FusedQuantizedMoE(nn.Module): ``normalize_routing_weights=0``, so the fused op reproduces the per-expert path's routing bit-for-bit. - Symmetric int quantization is used (no zero-points): the kernel defaults the - per-block zero-point to ``1 << (bits - 1)``. + Symmetric quantization omits zero-point inputs and uses the kernel's implicit + midpoint. Asymmetric quantization supplies packed per-block zero-points. """ _MICROSOFT_DOMAIN = "com.microsoft" @@ -333,6 +337,26 @@ def __init__( [e, self._hidden, self._inter // block_size], dtype=ir.DataType.FLOAT, ) + if qc.sym: + self.fc1_experts_zero_points = None + self.fc2_experts_zero_points = None + else: + self.fc1_experts_zero_points = nn.Parameter( + [ + e, + fc1_out, + math.ceil((self._hidden // block_size) * bits / 8), + ], + dtype=ir.DataType.UINT8, + ) + self.fc2_experts_zero_points = nn.Parameter( + [ + e, + self._hidden, + math.ceil((self._inter // block_size) * bits / 8), + ], + dtype=ir.DataType.UINT8, + ) def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: hidden = self._hidden @@ -354,16 +378,16 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: flat_f32, # 0: input scores_for_choice, # 1: router_probs (selection logits) self.fc1_experts_weights, # 2 - self.fc1_scales, # 3 + op.Cast(self.fc1_scales, to=1), # 3 None, # 4: fc1_experts_bias self.fc2_experts_weights, # 5 - self.fc2_scales, # 6 + op.Cast(self.fc2_scales, to=1), # 6 None, # 7: fc2_experts_bias None, # 8: fc3_experts_weights None, # 9: fc3_scales None, # 10: fc3_experts_bias - None, # 11: fc1_zero_points - None, # 12: fc2_zero_points + self.fc1_experts_zero_points, # 11 + self.fc2_experts_zero_points, # 12 None, # 13: fc3_zero_points aggregation, # 14: router_weights (explicit combine weights) activation_type="swiglu", @@ -397,3 +421,181 @@ def _route(self, op: OpBuilder, hidden_states: ir.Value): return self.gate.route_for_qmoe(op, hidden_states) finally: builder.pop_module() + + +_PER_EXPERT_RE = re.compile( + r"^(?P.*\.mlp)\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight|scales|zero_points)$" +) +_FUSED_EXPERT_RE = re.compile( + r"^(?P.*\.mlp)\.experts\." + r"(?Pgate_up_proj|down_proj)\." + r"(?Pqweight|qzeros|weight|scales|zero_points)$" +) + + +def pack_fused_quantized_moe_weights( + state_dict: dict[str, torch.Tensor], + config: ArchitectureConfig, +) -> dict[str, torch.Tensor]: + """Pack routed expert tensors into the expert-major QMoE initializer ABI. + + Accepts either fused 3-D GPTQ/AWQ checkpoint tensors or already-repacked + per-expert GGUF/MatMulNBits tensors. Gate/up rows are interleaved for + ``swiglu_fusion=1``; shared-expert tensors are left untouched. + """ + quantization = config.quantization + if quantization is None: + raise ValueError("Fused QMoE packing requires quantization settings") + + result: dict[str, torch.Tensor] = {} + fused: dict[str, dict[str, dict[str, torch.Tensor]]] = {} + per_expert: dict[str, dict[int, dict[str, dict[str, torch.Tensor]]]] = {} + for key, value in state_dict.items(): + match = _FUSED_EXPERT_RE.match(key) + if match is not None: + fused.setdefault(match["prefix"], {}).setdefault( + match["projection"], {} + )[match["kind"]] = value + continue + match = _PER_EXPERT_RE.match(key) + if match is not None: + expert = int(match["expert"]) + per_expert.setdefault(match["prefix"], {}).setdefault(expert, {}).setdefault( + match["projection"], {} + )[match["kind"]] = value + continue + result[key] = value + + prefixes = set(fused) | set(per_expert) + for prefix in prefixes: + if prefix in fused and prefix in per_expert: + raise ValueError(f"Mixed fused and per-expert weights under {prefix}") + if prefix in fused: + gate_up = _prepare_fused_projection( + fused[prefix].get("gate_up_proj"), + quantization.quant_method, + quantization.bits, + quantization.group_size, + ) + down = _prepare_fused_projection( + fused[prefix].get("down_proj"), + quantization.quant_method, + quantization.bits, + quantization.group_size, + ) + fc1 = { + kind: _interleave_gate_up(value) + for kind, value in gate_up.items() + } + fc2 = down + else: + experts = per_expert[prefix] + expected = set(range(config.num_local_experts or 0)) + if set(experts) != expected: + raise ValueError( + f"{prefix} has expert indices {sorted(experts)}, " + f"expected {sorted(expected)}" + ) + fc1, fc2 = _stack_per_expert_projections(experts) + + _store_qmoe_projection( + result, + f"{prefix}.moe.fc1", + fc1, + symmetric=quantization.sym, + ) + _store_qmoe_projection( + result, + f"{prefix}.moe.fc2", + fc2, + symmetric=quantization.sym, + ) + return result + + +def _prepare_fused_projection( + tensors: dict[str, torch.Tensor] | None, + quant_method: str, + bits: int, + block_size: int, +) -> dict[str, torch.Tensor]: + if tensors is None: + raise ValueError("Missing fused routed-expert projection") + if "qweight" in tensors: + preprocess = ( + preprocess_gptq_weights if quant_method == "gptq" else preprocess_awq_weights + ) + num_experts = tensors["qweight"].shape[0] + processed: dict[str, list[torch.Tensor]] = {} + for expert in range(num_experts): + expert_state = { + f"projection.{kind}": value[expert] + for kind, value in tensors.items() + if kind in {"qweight", "qzeros", "scales"} + } + for key, value in preprocess( + expert_state, bits=bits, group_size=block_size + ).items(): + processed.setdefault(key.rsplit(".", 1)[-1], []).append(value) + return {kind: torch.stack(values) for kind, values in processed.items()} + required = {"weight", "scales"} + if not required.issubset(tensors): + raise ValueError(f"Missing fused projection tensors: {required - tensors.keys()}") + return tensors + + +def _stack_per_expert_projections( + experts: dict[int, dict[str, dict[str, torch.Tensor]]], +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + fc1: dict[str, list[torch.Tensor]] = {} + fc2: dict[str, list[torch.Tensor]] = {} + for expert in sorted(experts): + projections = experts[expert] + for name in ("gate_proj", "up_proj", "down_proj"): + if name not in projections: + raise ValueError(f"Expert {expert} is missing {name}") + kinds = set(projections["gate_proj"]) | set(projections["up_proj"]) + for kind in kinds: + if kind not in projections["gate_proj"] or kind not in projections["up_proj"]: + raise ValueError(f"Expert {expert} gate/up {kind} tensors are incomplete") + gate = projections["gate_proj"][kind] + up = projections["up_proj"][kind] + fc1.setdefault(kind, []).append( + torch.stack((gate, up), dim=1).flatten(0, 1) + ) + for kind, value in projections["down_proj"].items(): + fc2.setdefault(kind, []).append(value) + return ( + {kind: torch.stack(values) for kind, values in fc1.items()}, + {kind: torch.stack(values) for kind, values in fc2.items()}, + ) + + +def _interleave_gate_up(value: torch.Tensor) -> torch.Tensor: + if value.shape[1] % 2: + raise ValueError(f"gate_up projection has odd output size {value.shape[1]}") + intermediate = value.shape[1] // 2 + return value.reshape(value.shape[0], 2, intermediate, *value.shape[2:]).transpose( + 1, 2 + ).flatten(1, 2) + + +def _store_qmoe_projection( + result: dict[str, torch.Tensor], + stem: str, + tensors: dict[str, torch.Tensor], + *, + symmetric: bool, +) -> None: + weight = tensors["weight"] + if weight.ndim == 4: + weight = weight.flatten(-2) + result[f"{stem}_experts_weights"] = weight + result[f"{stem}_scales"] = tensors["scales"].float() + zero_points = tensors.get("zero_points") + if not symmetric: + if zero_points is None: + raise ValueError(f"Asymmetric QMoE projection {stem} requires zero-points") + result[f"{stem}_experts_zero_points"] = zero_points diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index c48bcdfc..fdb9245f 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -7,13 +7,24 @@ import onnx_ir as ir import pytest +import torch +from mobius._builder import build_from_module +from mobius._configs import QuantizationConfig from mobius._testing import ( + count_op_type, create_test_builder, create_test_input, make_config, ) -from mobius.components._moe import MoELayer, SparseMixerGate, TopKGate +from mobius.components._moe import ( + FusedQuantizedMoE, + MoELayer, + SparseMixerGate, + TopKGate, + pack_fused_quantized_moe_weights, +) +from mobius.models.deepseek import DeepSeekMoEGate, DeepSeekV3CausalLMModel class TestTopKGate: @@ -118,3 +129,256 @@ def test_moe_layer_forward_with_sparse_mixer_gate(self): result = layer(op, hidden) builder._adapt_outputs([result], "") assert graph.num_nodes() > 0 + + +def test_fused_qmoe_wires_asymmetric_zero_points(): + config = make_config( + hidden_size=32, + intermediate_size=16, + moe_intermediate_size=16, + num_local_experts=4, + num_experts_per_tok=2, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gptq", + sym=False, + ), + ) + layer = FusedQuantizedMoE(config, DeepSeekMoEGate(config)) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 2, 32]) + builder._adapt_outputs([layer(op, hidden)], "") + + qmoe = next(node for node in graph if node.op_type == "QMoE") + assert qmoe.inputs[11] is not None + assert qmoe.inputs[12] is not None + assert qmoe.inputs[3] is not None + assert qmoe.inputs[3].producer().op_type == "Cast" + assert qmoe.inputs[3].dtype == ir.DataType.FLOAT + assert layer.fc1_experts_zero_points.shape == ir.Shape([4, 32, 1]) + assert layer.fc2_experts_zero_points.shape == ir.Shape([4, 32, 1]) + + +def test_deepseek_fused_qmoe_graph_has_one_node_per_moe_layer(): + config = make_config( + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=16, + num_hidden_layers=2, + first_k_dense_replace=1, + num_local_experts=4, + num_experts_per_tok=2, + n_shared_experts=1, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gptq", + sym=True, + ), + ) + graph = build_from_module(DeepSeekV3CausalLMModel(config), config)["model"].graph + + assert count_op_type(graph, "QMoE") == 1 + assert not any( + value is not None + and value.name is not None + and ".moe.experts." in value.name + for node in graph + for value in node.inputs + ) + shared_prefix = "model.layers.1.mlp.shared_experts." + assert ( + sum( + node.op_type == "MatMulNBits" + and any( + value is not None + and value.name is not None + and shared_prefix in value.name + for value in node.inputs + ) + for node in graph + ) + == 3 + ) + + +def test_expert_major_packing_matches_static_64_expert_top6_reference(): + torch.manual_seed(0) + experts, top_k = 64, 6 + hidden_size, intermediate_size, block_size = 32, 16, 16 + gate_codes = torch.randint(0, 16, (experts, intermediate_size, hidden_size)) + up_codes = torch.randint(0, 16, (experts, intermediate_size, hidden_size)) + down_codes = torch.randint(0, 16, (experts, hidden_size, intermediate_size)) + gate_scales = torch.rand(experts, intermediate_size, hidden_size // block_size) + up_scales = torch.rand(experts, intermediate_size, hidden_size // block_size) + down_scales = torch.rand(experts, hidden_size, intermediate_size // block_size) + state_dict = {} + for expert in range(experts): + prefix = f"model.layers.0.mlp.experts.{expert}" + for name, codes, scales in ( + ("gate_proj", gate_codes, gate_scales), + ("up_proj", up_codes, up_scales), + ("down_proj", down_codes, down_scales), + ): + state_dict[f"{prefix}.{name}.weight"] = _pack_matmul_nbits( + codes[expert], block_size + ) + state_dict[f"{prefix}.{name}.scales"] = scales[expert] + + config = make_config( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_intermediate_size=intermediate_size, + num_local_experts=experts, + num_experts_per_tok=top_k, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=block_size, + quant_method="gguf", + sym=True, + ), + ) + packed = DeepSeekV3CausalLMModel(config).preprocess_weights(state_dict) + fc1_weight = packed["model.layers.0.mlp.moe.fc1_experts_weights"] + fc1_scales = packed["model.layers.0.mlp.moe.fc1_scales"] + fc2_weight = packed["model.layers.0.mlp.moe.fc2_experts_weights"] + fc2_scales = packed["model.layers.0.mlp.moe.fc2_scales"] + assert fc1_weight.shape == (experts, 2 * intermediate_size, hidden_size // 2) + assert fc2_weight.shape == (experts, hidden_size, intermediate_size // 2) + + packed_fc1 = _dequant(fc1_weight, fc1_scales, block_size) + packed_fc2 = _dequant(fc2_weight, fc2_scales, block_size) + static_fc1 = torch.stack( + ( + _dequant_codes(gate_codes, gate_scales, block_size), + _dequant_codes(up_codes, up_scales, block_size), + ), + dim=2, + ).flatten(1, 2) + static_fc2 = _dequant_codes(down_codes, down_scales, block_size) + + hidden = torch.randn(3, hidden_size) + router_probs = torch.softmax(torch.randn(3, experts), dim=-1) + weights, selected = router_probs.topk(top_k, dim=-1) + weights /= weights.sum(dim=-1, keepdim=True) + expected = _static_moe( + hidden, weights, selected, static_fc1, static_fc2, intermediate_size + ) + actual = _static_moe( + hidden, weights, selected, packed_fc1, packed_fc2, intermediate_size + ) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + +def test_fused_gptq_checkpoint_tensors_pack_to_qmoe_layout(): + experts, hidden_size, intermediate_size, block_size = 4, 32, 16, 16 + gate_up_codes = torch.randint( + 0, 16, (experts, 2 * intermediate_size, hidden_size) + ) + down_codes = torch.randint(0, 16, (experts, hidden_size, intermediate_size)) + gate_up_scales = torch.rand( + experts, 2 * intermediate_size, hidden_size // block_size + ) + down_scales = torch.rand( + experts, hidden_size, intermediate_size // block_size + ) + state_dict = { + "model.layers.0.mlp.experts.gate_up_proj.qweight": _to_gptq_qweight( + gate_up_codes + ), + "model.layers.0.mlp.experts.gate_up_proj.scales": gate_up_scales.transpose( + -1, -2 + ), + "model.layers.0.mlp.experts.down_proj.qweight": _to_gptq_qweight( + down_codes + ), + "model.layers.0.mlp.experts.down_proj.scales": down_scales.transpose(-1, -2), + } + config = make_config( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_intermediate_size=intermediate_size, + num_local_experts=experts, + num_experts_per_tok=2, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=block_size, + quant_method="gptq", + sym=True, + ), + ) + + packed = pack_fused_quantized_moe_weights(state_dict, config) + fc1 = packed["model.layers.0.mlp.moe.fc1_experts_weights"] + scales = packed["model.layers.0.mlp.moe.fc1_scales"] + expected_codes = gate_up_codes.reshape( + experts, 2, intermediate_size, hidden_size + ).transpose(1, 2).flatten(1, 2) + expected_scales = gate_up_scales.reshape( + experts, 2, intermediate_size, hidden_size // block_size + ).transpose(1, 2).flatten(1, 2) + + assert fc1.shape == (experts, 2 * intermediate_size, hidden_size // 2) + torch.testing.assert_close( + _dequant(fc1, scales, block_size), + _dequant_codes(expected_codes, expected_scales, block_size), + ) + + +def _pack_matmul_nbits(codes: torch.Tensor, block_size: int) -> torch.Tensor: + low = codes[..., 0::2].to(torch.uint8) + high = codes[..., 1::2].to(torch.uint8) + return (low | (high << 4)).reshape( + codes.shape[0], codes.shape[1] // block_size, block_size // 2 + ) + + +def _to_gptq_qweight(codes: torch.Tensor) -> torch.Tensor: + shifts = torch.arange(8, dtype=torch.int64) * 4 + packed = torch.sum( + codes.to(torch.int64).reshape(*codes.shape[:-1], -1, 8) << shifts, + dim=-1, + ).to(torch.int32) + return packed.transpose(-1, -2).contiguous() + + +def _dequant_codes( + codes: torch.Tensor, scales: torch.Tensor, block_size: int +) -> torch.Tensor: + blocks = codes.shape[-1] // block_size + values = codes.reshape(*codes.shape[:-1], blocks, block_size).float() - 8.0 + return (values * scales.unsqueeze(-1)).flatten(-2) + + +def _dequant( + packed: torch.Tensor, scales: torch.Tensor, block_size: int +) -> torch.Tensor: + codes = torch.stack((packed & 0x0F, packed >> 4), dim=-1).flatten(-2) + return _dequant_codes(codes, scales, block_size) + + +def _static_moe( + hidden: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + fc1: torch.Tensor, + fc2: torch.Tensor, + intermediate_size: int, +) -> torch.Tensor: + result = torch.zeros_like(hidden) + for expert in range(fc1.shape[0]): + projected = hidden @ fc1[expert].T + activated = torch.nn.functional.silu(projected[:, 0::2]) + activated *= projected[:, 1::2] + assert activated.shape[-1] == intermediate_size + expert_output = activated @ fc2[expert].T + weight = ( + routing_weights + * (selected_experts == expert).to(routing_weights.dtype) + ).sum(dim=-1, keepdim=True) + result += expert_output * weight + return result diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index 849eaa2a..2ad1b147 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -769,13 +769,17 @@ def _load_quantized_state_dict( raw_expert = raw_flat[ expert_idx * blocks_per_expert : (expert_idx + 1) * blocks_per_expert ] + repacked = None if source_params == (target_bits, target_block_size): - repacked = repack_gguf_tensor( + candidate = repack_gguf_tensor( raw_expert.ravel().view(np.uint8), qtype_val, (n_out, k_in), ) - else: + target_has_zero_points = not target_symmetric + if (candidate.zero_points is not None) == target_has_zero_points: + repacked = candidate + if repacked is None: _require_supported_requantization( bits=target_bits, block_size=target_block_size, diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index 98b0a02e..8917f594 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -496,6 +496,65 @@ def _repack(raw, qtype, shape): assert "model.layers.0.mlp.experts.1.gate_proj.weight" in result +def test_load_quantized_experts_requantizes_when_zero_point_presence_differs( + monkeypatch, +): + from types import SimpleNamespace + + from gguf import GGMLQuantizationType + + from mobius.integrations.gguf import _repacker, _tencent_q1_0, _tensor_mapping + from mobius.integrations.gguf._builder import _load_quantized_state_dict + + requantized = [] + + def _repack(*_args, **_kwargs): + return SimpleNamespace( + weight=np.zeros((2, 1, 16), dtype=np.uint8), + scales=np.ones((2, 1), dtype=np.float32), + zero_points=None, + ) + + def _requantize(values, **_kwargs): + requantized.append(values) + return SimpleNamespace( + weight=np.zeros((2, 1, 16), dtype=np.uint8), + scales=np.ones((2, 1), dtype=np.float32), + zero_points=np.zeros((2, 1), dtype=np.uint8), + ) + + monkeypatch.setattr(_tencent_q1_0, "is_tencent_q1_0_layout", lambda _model: False) + monkeypatch.setattr( + _tensor_mapping, + "map_gguf_to_hf_names", + lambda _name, _arch: "model.layers.0.mlp.experts.gate_proj.weight", + ) + monkeypatch.setattr(_repacker, "repack_quant_params", lambda _qtype: (4, 32)) + monkeypatch.setattr(_repacker, "repack_gguf_tensor", _repack) + monkeypatch.setattr(_repacker, "repack_dequantized_tensor", _requantize) + + raw = np.arange(8, dtype=np.uint8) + model = SimpleNamespace( + _tensor_index={"experts": object()}, + tensor_items_raw=lambda: [ + ("experts", raw, GGMLQuantizationType.Q4_0, (2, 2, 2)) + ], + dequantize_raw_tensor=lambda *_args: np.ones((2, 2), dtype=np.float32), + ) + config = SimpleNamespace( + quantization=SimpleNamespace(bits=4, group_size=32, sym=False), + num_attention_heads=1, + num_key_value_heads=1, + model_type="glm_moe_dsa", + ) + + result = _load_quantized_state_dict(model, "glm-dsa", SimpleNamespace(named_modules=list), config) + + assert len(requantized) == 2 + assert "model.layers.0.mlp.experts.0.gate_proj.zero_points" in result + assert "model.layers.0.mlp.experts.1.gate_proj.zero_points" in result + + def test_load_quantized_glm_fuses_split_kv_b(monkeypatch): from types import SimpleNamespace diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index 0a3f400e..32923171 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -29,6 +29,7 @@ make_quantized_linear_factory, ) from mobius.components._deepseek_mla import DeepSeekMLA +from mobius.components._moe import pack_fused_quantized_moe_weights from mobius.models.base import CausalLMModel if TYPE_CHECKING: @@ -398,7 +399,9 @@ def __init__(self, config: ArchitectureConfig): # Detect MLA vs standard attention use_mla = config.qk_nope_head_dim is not None and config.qk_nope_head_dim > 0 - LayerClass = DeepSeekMLADecoderLayer if use_mla else _DeepSeekStandardDecoderLayer # noqa: N806 + LayerClass = ( # ruff:ignore[non-lowercase-variable-in-function] + DeepSeekMLADecoderLayer if use_mla else _DeepSeekStandardDecoderLayer + ) # Build layers: dense for first k, MoE for rest first_k = config.first_k_dense_replace @@ -490,17 +493,28 @@ def preprocess_weights( - MoE expert weights: HF stores all experts in a single fused tensor experts.gate_up_proj: (n_experts, 2*intermediate, hidden) experts.down_proj: (n_experts, hidden, intermediate) - These are split into per-expert ONNX weights: + Static MoE exports split these into per-expert ONNX weights: moe.experts.{i}.gate_proj.weight: (intermediate, hidden) moe.experts.{i}.up_proj.weight: (intermediate, hidden) moe.experts.{i}.down_proj.weight: (hidden, intermediate) + Fused quantized exports instead pack expert-major FC1/FC2 tensors for QMoE. """ renamed = {} + routed_experts = {} + qc = self.config.quantization + use_fused_qmoe = ( + self.config.fused_quantized_moe + and qc is not None + and qc.quant_method != "none" + ) for key, value in state_dict.items(): new_key = key # Remap MoE layer names: mlp.gate.* → mlp.moe.gate.* new_key = new_key.replace(".mlp.gate.", ".mlp.moe.gate.") + if use_fused_qmoe and ".mlp.experts." in new_key: + routed_experts[new_key] = value + continue # HF stores all routed experts in fused tensors: # layers.N.mlp.experts.gate_up_proj (n_experts, 2*mid, hidden) @@ -522,4 +536,7 @@ def preprocess_weights( renamed[new_key] = value # Handle weight tying - return super().preprocess_weights(renamed) + processed = super().preprocess_weights(renamed) + if use_fused_qmoe: + processed.update(pack_fused_quantized_moe_weights(routed_experts, self.config)) + return processed From 6f2a52a9d20d401169c42777504020f466c72f8b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 14:55:10 +0000 Subject: [PATCH 17/23] fix(lint): restore Ruff noqa suppressions Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/components/_moe.py | 2 +- src/mobius/models/deepseek.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 6bc5b2aa..7d8bc24a 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -110,7 +110,7 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): # Renormalize selected weights to sum to 1 (prevents vanishing gradients) weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True) routing_weights = op.Div(routing_weights, op.Add(weight_sum, 1e-9)) - if self.routed_scaling_factor != 1.0: # ruff:ignore[float-equality-comparison] + if self.routed_scaling_factor != 1.0: # noqa: RUF069 routing_weights = op.Mul(routing_weights, self.routed_scaling_factor) return routing_weights, selected_experts diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index 32923171..b0773fab 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -399,7 +399,7 @@ def __init__(self, config: ArchitectureConfig): # Detect MLA vs standard attention use_mla = config.qk_nope_head_dim is not None and config.qk_nope_head_dim > 0 - LayerClass = ( # ruff:ignore[non-lowercase-variable-in-function] + LayerClass = ( # noqa: N806 DeepSeekMLADecoderLayer if use_mla else _DeepSeekStandardDecoderLayer ) From bd88fa8438da04b920c728659e4a5f002f4b8726 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 15:39:42 +0000 Subject: [PATCH 18/23] fix(deepseek): align low-precision router dtypes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/components/_deepseek_mla_test.py | 26 +++++++++++++++++++++ src/mobius/components/_moe_test.py | 26 +++++++++++++++++++++ src/mobius/models/deepseek.py | 16 ++++++------- 3 files changed, 60 insertions(+), 8 deletions(-) diff --git a/src/mobius/components/_deepseek_mla_test.py b/src/mobius/components/_deepseek_mla_test.py index 3946ec41..508d457e 100644 --- a/src/mobius/components/_deepseek_mla_test.py +++ b/src/mobius/components/_deepseek_mla_test.py @@ -5,8 +5,10 @@ from __future__ import annotations +import onnx_ir as ir import pytest +from mobius._builder import build_from_module from mobius._testing import ( count_op_type, create_test_builder, @@ -14,6 +16,7 @@ make_config, ) from mobius.components._deepseek_mla import DeepSeekMLA +from mobius.models.deepseek import DeepSeekV3CausalLMModel # DeepSeek MLA config: uses low-rank KV compression with separate # nope (non-positional) and rope (rotary) head dimensions. @@ -31,6 +34,29 @@ ) +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_low_precision_export_matches_attention_auxiliary_dtypes(dtype): + config = _mla_config( + dtype=dtype, + intermediate_size=128, + num_hidden_layers=1, + num_key_value_heads=4, + head_dim=8, + first_k_dense_replace=1, + num_local_experts=None, + ) + graph = build_from_module(DeepSeekV3CausalLMModel(config), config)["model"].graph + + rotary_nodes = [node for node in graph if node.op_type == "RotaryEmbedding"] + assert len(rotary_nodes) == 2 + for node in rotary_nodes: + assert [value.dtype for value in node.inputs] == [dtype, dtype, dtype] + + attention = next(node for node in graph if node.op_type == "Attention") + assert attention.inputs[3].dtype == dtype + assert attention.inputs[3].dtype == attention.inputs[0].dtype + + def _mla_config(**overrides): """Create a test config for DeepSeek MLA.""" kw = {**_MLA_DEFAULTS, **overrides} diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index fdb9245f..21ffd907 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -77,6 +77,32 @@ def test_gate_top_k_1(self): assert graph.num_nodes() > 0 +@pytest.mark.parametrize("route_for_qmoe", [False, True]) +def test_deepseek_gate_casts_both_router_matmul_inputs_to_float32(route_for_qmoe): + config = make_config( + dtype=ir.DataType.FLOAT16, + num_local_experts=4, + num_experts_per_tok=2, + ) + gate = DeepSeekMoEGate(config) + builder, op, graph = create_test_builder() + hidden = create_test_input( + builder, + "hidden", + [1, 2, config.hidden_size], + dtype=ir.DataType.FLOAT16, + ) + + outputs = gate.route_for_qmoe(op, hidden) if route_for_qmoe else gate(op, hidden) + builder._adapt_outputs(list(outputs), "") + + router_matmul = next(node for node in graph if node.op_type == "MatMul") + assert [value.dtype for value in router_matmul.inputs] == [ + ir.DataType.FLOAT, + ir.DataType.FLOAT, + ] + + class TestMoELayer: def test_moe_layer_has_gate(self): config = make_config(num_local_experts=4, num_experts_per_tok=2) diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index b0773fab..253dc573 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -83,12 +83,7 @@ def __init__(self, config: ArchitectureConfig): self.e_score_correction_bias = nn.Parameter([self.num_experts]) def forward(self, op: OpBuilder, hidden_states: ir.Value): - # Compute routing logits: hidden @ weight^T - weight_t = op.Transpose(self.weight, perm=[1, 0]) - router_logits = op.MatMul( - op.Cast(hidden_states, to=1), - weight_t, # Cast to float32 - ) + router_logits = self._router_logits(op, hidden_states) # Score computation depends on scoring function if self.scoring_func == "sigmoid": @@ -137,8 +132,7 @@ def route_for_qmoe(self, op: OpBuilder, hidden_states: ir.Value): combine weights (normalized sigmoid * routed_scaling_factor) match the per-expert path bit-for-bit. """ - weight_t = op.Transpose(self.weight, perm=[1, 0]) - router_logits = op.MatMul(op.Cast(hidden_states, to=1), weight_t) + router_logits = self._router_logits(op, hidden_states) if self.scoring_func == "sigmoid": scores = op.Sigmoid(router_logits) @@ -162,6 +156,12 @@ def route_for_qmoe(self, op: OpBuilder, hidden_states: ir.Value): return scores_for_choice, routing_weights, selected_experts + def _router_logits(self, op: OpBuilder, hidden_states: ir.Value): + """Compute routing logits in float32 for stable expert selection.""" + hidden_states = op.Cast(hidden_states, to=1) + weight_t = op.Cast(op.Transpose(self.weight, perm=[1, 0]), to=1) + return op.MatMul(hidden_states, weight_t) + def _group_topk_selection(self, op, scores_for_choice): """Group-based expert selection: pick topk_group groups first.""" experts_per_group = self.num_experts // self.n_group From 4816c20e4335fc20cb9e2222653a56d33c966a8c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 24 Jul 2026 06:59:42 +0000 Subject: [PATCH 19/23] Lower GLM DSA sparse-attention to pkg.nxrt::IndexShare The GLM Deepseek-Sparse-Attention (DSA) mask path was exported as a dense TopK -> ScatterElements sparse-bias island feeding op.Attention. Under the native runtime's single-token decode, the attention mask is frozen to the fixed KV capacity (for CUDA-graph capture) while the indexer/KV caches are exposed at the growing logical length, so the dense full-length bias adds were not broadcast-compatible (Add ... shapes [[1,1,S],[1,1,capacity]]). Emit pkg.nxrt::IndexShare instead: it gathers only the indexer's selected key positions, so causality is carried by the (already causal) top-k indices and no dense additive mask is emitted. Also slice the indexer's own causal bias down to the score key length so its Add stays broadcast-safe under both the logical (ORT) and fixed-capacity (native) exposures, and sort the top-k indices ascending to satisfy IndexShare's strictly-increasing requirement. MLA's narrower value head is zero-padded up to qk_head_dim for the kernel (no Pad op, which the CUDA EP lacks) and sliced back off the output; the present key/value stay in native BNSH so the packed KV cache is unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/glm_moe_dsa.py | 113 +++++++++++++++++++++++++------ src/mobius/tasks/_base.py | 2 +- 2 files changed, 93 insertions(+), 22 deletions(-) diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index c9fb1b53..bb14534f 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -121,11 +121,25 @@ def select( weights = op.Cast(weights, to=ir.DataType.FLOAT) weights = op.Mul(weights, self.index_n_heads**-0.5) scores = op.ReduceSum(op.Mul(scores, op.Unsqueeze(weights, [3])), [2], keepdims=False) - scores = op.Add(scores, op.Cast(op.Squeeze(attention_bias, [1]), to=ir.DataType.FLOAT)) - + # The additive causal mask is sized to the attention key axis, which the + # native decoder exposes at fixed KV capacity while the indexer key cache + # grows with the logical prefix. Slice the mask down to the score key + # length so the add stays broadcast-compatible under both the logical + # (ORT) and fixed-capacity (native single-token decode) exposures. key_length = op.Shape(scores, start=2, end=3) + causal = op.Squeeze(attention_bias, [1]) + causal = op.Slice(causal, op.Constant(value_ints=[0]), key_length, op.Constant(value_ints=[2])) + scores = op.Add(scores, op.Cast(causal, to=ir.DataType.FLOAT)) + k = op.Min(key_length, op.Constant(value_ints=[self.index_topk])) _, indices = op.TopK(scores, k, axis=-1, largest=1, sorted=0, _outputs=2) + # ``pkg.nxrt::IndexShare`` requires the selected key positions to be + # strictly increasing per row. TopK returns them in score order, so sort + # the chosen positions ascending. The runtime TopK kernel is f32-only, so + # round-trip through float before restoring the Int64 index dtype. + indices = op.Cast(indices, to=ir.DataType.FLOAT) + indices, _ = op.TopK(indices, k, axis=-1, largest=0, sorted=1, _outputs=2) + indices = op.Cast(indices, to=ir.DataType.INT64) return indices def forward( @@ -211,18 +225,80 @@ def _pack_present( ) return op.Unsqueeze(key, [1]), op.Unsqueeze(value, [1]) - def _sparse_bias( + def _index_share_attention( self, op: OpBuilder, - indices: ir.Value, - attention_bias: ir.Value, - ) -> ir.Value: - shape = op.Shape(op.Squeeze(attention_bias, [1])) - masked = op.Expand(op.Constant(value_float=float(self.dtype.min)), shape) - updates = op.Mul(op.Cast(indices, to=ir.DataType.FLOAT), 0.0) - sparse = op.ScatterElements(masked, indices, updates, axis=-1, reduction="none") - sparse = op.Cast(op.Unsqueeze(sparse, [1]), to=self.dtype) - return op.Add(attention_bias, sparse) + q: ir.Value, + key: ir.Value, + value: ir.Value, + main_past: tuple | None, + topk_indices: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + """Device-resident selected-token attention via ``pkg.nxrt::IndexShare``. + + Replaces the dense ``ScatterElements`` sparse-mask + ``Attention`` lowering + with the frozen ``IndexShare`` op, which gathers only the indexer's + selected key positions. Causality is carried entirely by the (already + causal) ``topk_indices``, so no dense additive mask island — and thus no + logical-length-vs-fixed-capacity broadcast — is emitted here. + + ``IndexShare`` requires a homogeneous head size across query/key/value. + MLA's value head (``v_head_dim``) is narrower than the query/key head + (``qk_head_dim``), so the value is zero-padded up to ``qk_head_dim`` for + the kernel and the padding columns are sliced back off the output. The + returned present key/value stay in their native (unpadded) BNSH layout so + the packed KV cache is unchanged. + """ + q_bnsh = op.Transpose( + op.Reshape(q, [0, 0, self.num_heads, self.qk_head_dim]), perm=[0, 2, 1, 3] + ) + cur_key = op.Transpose( + op.Reshape(key, [0, 0, self.num_heads, self.qk_head_dim]), perm=[0, 2, 1, 3] + ) + cur_value = op.Transpose( + op.Reshape(value, [0, 0, self.num_heads, self.v_head_dim]), perm=[0, 2, 1, 3] + ) + present_key = cur_key + present_value = cur_value + if main_past is not None: + present_key = op.Concat(main_past[0], cur_key, axis=2) + present_value = op.Concat(main_past[1], cur_value, axis=2) + + value_pad = self.qk_head_dim - self.v_head_dim + padded_value = present_value + if value_pad > 0: + # Zero-pad the value head up to ``qk_head_dim`` without ``Pad`` (which + # the CUDA EP has no handler for). Build a device-resident zero block + # of the value's own dtype and concatenate it onto the head axis. + pad_shape = op.Concat( + op.Shape(present_value, start=0, end=3), + op.Constant(value_ints=[value_pad]), + axis=0, + ) + zeros = op.Expand(op.Cast(op.Constant(value_float=0.0), to=self.dtype), pad_shape) + padded_value = op.Concat(present_value, zeros, axis=3) + + selected_indices = op.Unsqueeze(topk_indices, [1]) + attn_output = op.IndexShare( + q_bnsh, + present_key, + padded_value, + None, + None, + selected_indices, + num_heads=self.num_heads, + kv_num_heads=self.num_heads, + scale=self.scaling, + _domain="pkg.nxrt", + _outputs=1, + ) + if value_pad > 0: + attn_output = op.Slice(attn_output, [0], [self.v_head_dim], [3]) + attn_output = op.Reshape( + op.Transpose(attn_output, perm=[0, 2, 1, 3]), + [0, 0, self.num_heads * self.v_head_dim], + ) + return attn_output, present_key, present_value def forward( self, @@ -313,18 +389,13 @@ def forward( "Shared GLM DSA layers require top-k indices from a preceding full layer" ) - sparse_bias = self._sparse_bias(op, topk_indices, attention_bias) - attn_output, present_key, present_value = op.Attention( + attn_output, present_key, present_value = self._index_share_attention( + op, q, key, value, - attn_mask=sparse_bias, - past_key=main_past[0] if main_past is not None else None, - past_value=main_past[1] if main_past is not None else None, - q_num_heads=self.num_heads, - kv_num_heads=self.num_heads, - scale=self.scaling, - _outputs=3, + main_past, + topk_indices, ) attn_output = self.o_proj(op, attn_output) present = self._pack_present( diff --git a/src/mobius/tasks/_base.py b/src/mobius/tasks/_base.py index b09276fc..38b5e938 100644 --- a/src/mobius/tasks/_base.py +++ b/src/mobius/tasks/_base.py @@ -117,7 +117,7 @@ def _make_graph( [], nodes=[], name=name, - opset_imports={"": OPSET_VERSION, "com.microsoft": 1}, + opset_imports={"": OPSET_VERSION, "com.microsoft": 1, "pkg.nxrt": 1}, ) return graph, GraphBuilder(graph) From 7453e2c116b2feadc0fb8cc83c348644436e301b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 24 Jul 2026 07:19:12 +0000 Subject: [PATCH 20/23] test(glm-dsa): assert IndexShare lowering instead of dense attention Co-authored-by: GitHub Copilot --- src/mobius/models/glm_moe_dsa_test.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/mobius/models/glm_moe_dsa_test.py b/src/mobius/models/glm_moe_dsa_test.py index 990fc3e1..72ec95e2 100644 --- a/src/mobius/models/glm_moe_dsa_test.py +++ b/src/mobius/models/glm_moe_dsa_test.py @@ -97,7 +97,7 @@ def test_indexer_types_length_must_match_layers(self): with pytest.raises(ValueError, match="expected 4, got 1"): _indexer_types(config) - def test_builds_full_attention_mla_moe_graph(self): + def test_builds_index_share_mla_moe_graph(self): config = _tiny_glm_moe_dsa_config() package = build_from_module( GlmMoeDsaCausalLMModel(config), @@ -106,8 +106,9 @@ def test_builds_full_attention_mla_moe_graph(self): ) graph = package["model"].graph - assert count_op_type(graph, "Attention") == config.num_hidden_layers - assert count_op_type(graph, "ScatterElements") >= 1 + assert count_op_type(graph, "IndexShare") == config.num_hidden_layers + assert count_op_type(graph, "Attention") == 0 + assert count_op_type(graph, "ScatterElements") == 0 assert count_op_type(graph, "Sigmoid") >= config.num_hidden_layers - 1 assert any( "layers.0.self_attn.indexer.wk.weight" in name for name in graph.initializers @@ -121,8 +122,9 @@ def test_builds_full_attention_mla_moe_graph(self): assert set(package) == {"model", "mtp"} mtp_graph = package["mtp"].graph - assert count_op_type(mtp_graph, "Attention") == 1 - assert count_op_type(mtp_graph, "ScatterElements") == 1 + assert count_op_type(mtp_graph, "IndexShare") == 1 + assert count_op_type(mtp_graph, "Attention") == 0 + assert count_op_type(mtp_graph, "ScatterElements") == 0 assert {value.name for value in mtp_graph.outputs} == { "mtp_hidden", "present.0.key", From 1c3396fb3f5037e2ced24c70d14323a94ba765a6 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 24 Jul 2026 11:35:36 +0000 Subject: [PATCH 21/23] fix(export): address PR 404 review feedback Forward fused GPTQ g_idx tensors, omit symmetric 8-bit zero points, and pass string paths to ONNX Runtime sessions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/components/_moe.py | 13 +++++++------ src/mobius/components/_moe_test.py | 13 ++++++++++++- src/mobius/integrations/gguf/_repacker.py | 10 ++++++---- src/mobius/integrations/gguf/_repacker_test.py | 14 ++++++++++++++ src/mobius/models/llada_test.py | 2 +- src/mobius/models/unet_parity_test.py | 8 ++++---- 6 files changed, 44 insertions(+), 16 deletions(-) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 7d8bc24a..e8e044df 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -431,7 +431,7 @@ def _route(self, op: OpBuilder, hidden_states: ir.Value): _FUSED_EXPERT_RE = re.compile( r"^(?P.*\.mlp)\.experts\." r"(?Pgate_up_proj|down_proj)\." - r"(?Pqweight|qzeros|weight|scales|zero_points)$" + r"(?Pqweight|qzeros|g_idx|weight|scales|zero_points)$" ) @@ -530,11 +530,12 @@ def _prepare_fused_projection( num_experts = tensors["qweight"].shape[0] processed: dict[str, list[torch.Tensor]] = {} for expert in range(num_experts): - expert_state = { - f"projection.{kind}": value[expert] - for kind, value in tensors.items() - if kind in {"qweight", "qzeros", "scales"} - } + expert_state = {} + for kind in {"qweight", "qzeros", "scales", "g_idx"} & tensors.keys(): + value = tensors[kind] + expert_state[f"projection.{kind}"] = ( + value if kind == "g_idx" and value.ndim == 1 else value[expert] + ) for key, value in preprocess( expert_state, bits=bits, group_size=block_size ).items(): diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index 21ffd907..494490f7 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -299,7 +299,8 @@ def test_expert_major_packing_matches_static_64_expert_top6_reference(): torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) -def test_fused_gptq_checkpoint_tensors_pack_to_qmoe_layout(): +@pytest.mark.parametrize("shared_g_idx", [False, True]) +def test_fused_gptq_checkpoint_tensors_pack_to_qmoe_layout(caplog, shared_g_idx): experts, hidden_size, intermediate_size, block_size = 4, 32, 16, 16 gate_up_codes = torch.randint( 0, 16, (experts, 2 * intermediate_size, hidden_size) @@ -311,6 +312,12 @@ def test_fused_gptq_checkpoint_tensors_pack_to_qmoe_layout(): down_scales = torch.rand( experts, hidden_size, intermediate_size // block_size ) + gate_up_g_idx = torch.arange(hidden_size, dtype=torch.int32) // block_size + gate_up_g_idx[0] = 1 + down_g_idx = torch.arange(intermediate_size, dtype=torch.int32) // block_size + if not shared_g_idx: + gate_up_g_idx = gate_up_g_idx.repeat(experts, 1) + down_g_idx = down_g_idx.repeat(experts, 1) state_dict = { "model.layers.0.mlp.experts.gate_up_proj.qweight": _to_gptq_qweight( gate_up_codes @@ -318,10 +325,12 @@ def test_fused_gptq_checkpoint_tensors_pack_to_qmoe_layout(): "model.layers.0.mlp.experts.gate_up_proj.scales": gate_up_scales.transpose( -1, -2 ), + "model.layers.0.mlp.experts.gate_up_proj.g_idx": gate_up_g_idx, "model.layers.0.mlp.experts.down_proj.qweight": _to_gptq_qweight( down_codes ), "model.layers.0.mlp.experts.down_proj.scales": down_scales.transpose(-1, -2), + "model.layers.0.mlp.experts.down_proj.g_idx": down_g_idx, } config = make_config( hidden_size=hidden_size, @@ -353,6 +362,8 @@ def test_fused_gptq_checkpoint_tensors_pack_to_qmoe_layout(): _dequant(fc1, scales, block_size), _dequant_codes(expected_codes, expected_scales, block_size), ) + assert "desc_act models with non-trivial g_idx" in caplog.text + assert not any(key.endswith(".g_idx") for key in packed) def _pack_matmul_nbits(codes: torch.Tensor, block_size: int) -> torch.Tensor: diff --git a/src/mobius/integrations/gguf/_repacker.py b/src/mobius/integrations/gguf/_repacker.py index 7f369c9b..946c4f0b 100644 --- a/src/mobius/integrations/gguf/_repacker.py +++ b/src/mobius/integrations/gguf/_repacker.py @@ -470,20 +470,22 @@ def repack_dequantized_tensor( block_max = np.maximum(blocks.max(axis=-1), 0.0) if symmetric: scales = np.maximum(-block_min / 128.0, block_max / 127.0) - zero_points = np.full_like(scales, 128, dtype=np.uint8) + zero_points_arr = np.full_like(scales, 128, dtype=np.uint8) else: scales = (block_max - block_min) / 255.0 safe_scales = np.where(scales != 0, scales, 1.0) - zero_points = np.clip(np.rint(-block_min / safe_scales), 0, 255).astype(np.uint8) + zero_points_arr = np.clip( + np.rint(-block_min / safe_scales), 0, 255 + ).astype(np.uint8) safe_scales = np.where(scales != 0, scales, 1.0) quants = np.rint(blocks / safe_scales[:, :, None]) - quants += zero_points[:, :, None] + quants += zero_points_arr[:, :, None] weight = np.clip(quants, 0, 255).astype(np.uint8) weight = np.where(scales[:, :, None] != 0, weight, 0).astype(np.uint8) return RepackedTensor( weight=weight, scales=scales.astype(np.float32), - zero_points=zero_points, + zero_points=None if symmetric else zero_points_arr, block_size=block_size, bits=bits, ) diff --git a/src/mobius/integrations/gguf/_repacker_test.py b/src/mobius/integrations/gguf/_repacker_test.py index 7cbc5531..a7b04850 100644 --- a/src/mobius/integrations/gguf/_repacker_test.py +++ b/src/mobius/integrations/gguf/_repacker_test.py @@ -682,6 +682,20 @@ def test_asymmetric_q8_round_trip_bound(self): assert result.weight.shape == (3, 3, 32) assert result.zero_points.shape == (3, 3) + def test_symmetric_q8_omits_zero_points(self): + values = np.array([[-128.0, 0.0, 127.0] + [0.0] * 29], dtype=np.float32) + result = repack_dequantized_tensor( + values, bits=8, block_size=32, symmetric=True + ) + + assert result.zero_points is None + dequantized = ( + result.weight.astype(np.float32) - 128.0 + ) * result.scales[:, :, None] + np.testing.assert_allclose( + dequantized.reshape(values.shape), values, atol=result.scales.max() * 0.51 + ) + # ---- Q1_0 tests ---- diff --git a/src/mobius/models/llada_test.py b/src/mobius/models/llada_test.py index ff0872ac..01686e4a 100644 --- a/src/mobius/models/llada_test.py +++ b/src/mobius/models/llada_test.py @@ -173,7 +173,7 @@ def _build_onnx_session(config: ArchitectureConfig, state: dict[str, torch.Tenso onnx_ir.save(model, model_path) # onnxruntime reads the model fully at construction, so the temp file # can be removed as soon as the session exists. - return ort.InferenceSession(model_path) + return ort.InferenceSession(str(model_path)) def test_llada_matches_torch_reference(): diff --git a/src/mobius/models/unet_parity_test.py b/src/mobius/models/unet_parity_test.py index 5dd7b348..893df1ec 100644 --- a/src/mobius/models/unet_parity_test.py +++ b/src/mobius/models/unet_parity_test.py @@ -67,7 +67,7 @@ def test_cross_attention_block_matches_diffusers(): with tempfile.TemporaryDirectory() as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) actual = session.run(None, {"hidden": hidden.numpy(), "context": context.numpy()})[0] assert np.abs(actual - expected).max() < 1e-4 @@ -120,7 +120,7 @@ def test_unet_matches_diffusers(): with tempfile.TemporaryDirectory() as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) actual = session.run( None, { @@ -191,7 +191,7 @@ def test_unet_sd1x_mixed_block_types_matches_diffusers(): with tempfile.TemporaryDirectory() as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) actual = session.run( None, { @@ -278,7 +278,7 @@ def test_unet_lora_gate_parity(): with tempfile.TemporaryDirectory() as temp_dir: model_path = Path(temp_dir) / "model.onnx" onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) + session = ort.InferenceSession(str(model_path)) feed = { "sample": sample.numpy(), "timestep": timestep.numpy().astype(np.int64), From 5ba283ff9b41c5bd4b9afcf55fe5f4c0cf3f71ab Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 29 Jul 2026 17:38:00 +0000 Subject: [PATCH 22/23] Guard fused QMoE CUDA routing Keep fused QMoE activation inputs in the model dtype and fail fast for CUDA exports that require router_weights, since ORT CUDA QMoE currently ignores input 14. Add regression tests for activation dtype and the CUDA guard.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/components/_moe.py | 20 ++++++++++-- src/mobius/components/_moe_test.py | 50 ++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 64e644fa..fd0b49df 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -14,6 +14,7 @@ import torch from onnxscript import OpBuilder, nn +from mobius._build_context import ep_capabilities from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius._weight_utils import preprocess_awq_weights, preprocess_gptq_weights from mobius.components._mlp import MLP @@ -279,6 +280,7 @@ def __init__( config: ArchitectureConfig, gate: nn.Module | None = None, linear_class: type | None = None, + enable_qmoe: bool = True, ): super().__init__() assert config.num_local_experts is not None @@ -296,7 +298,11 @@ def __init__( if config.moe_intermediate_size is not None else config ) - if self._qmoe_quantization is not None and hasattr(self.gate, "qmoe_routing"): + if ( + enable_qmoe + and self._qmoe_quantization is not None + and hasattr(self.gate, "qmoe_routing") + ): self.experts = None self._init_qmoe_parameters(expert_config) else: @@ -529,12 +535,20 @@ def __init__( ) def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + if ep_capabilities().name == "cuda": + raise ValueError( + "FusedQuantizedMoE is disabled for CUDA because ORT CUDA QMoE " + "currently ignores router_weights (input 14), which is required " + "for GLM/DeepSeek group-limited routing. Use the decomposed " + "MatMulNBits export path for CUDA." + ) hidden = self._hidden # QMoE requires 2-D router_probs, so flatten [B, S, H] -> [rows, H]. orig_shape = op.Shape(hidden_states) flat = op.Reshape(hidden_states, op.Constant(value_ints=[-1, hidden])) - # QMoE input/router_probs must be float32. + # Router math must be float32, but QMoE activation input stays in the + # model dtype: CUDA QMoE kernels are registered for fp16/bf16 inputs. flat_f32 = op.Cast(flat, to=1) scores_for_choice, routing_weights, selected_experts = self._route(op, flat_f32) @@ -545,7 +559,7 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: aggregation = op.ScatterElements(zeros, selected_experts, routing_weights, axis=-1) moe_out = op.QMoE( - flat_f32, # 0: input + flat, # 0: input scores_for_choice, # 1: router_probs (selection logits) self.fc1_experts_weights, # 2 op.Cast(self.fc1_scales, to=1), # 3 diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index c587ff1e..b19c03e1 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -186,6 +186,30 @@ def test_fused_qmoe_wires_asymmetric_zero_points(): assert layer.fc2_experts_zero_points.shape == ir.Shape([4, 32, 1]) +def test_fused_qmoe_keeps_activation_input_model_dtype(): + config = make_config( + hidden_size=32, + intermediate_size=16, + moe_intermediate_size=16, + num_local_experts=4, + num_experts_per_tok=2, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gptq", + sym=True, + ), + ) + layer = FusedQuantizedMoE(config, DeepSeekMoEGate(config)) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 2, 32], dtype=ir.DataType.FLOAT16) + builder._adapt_outputs([layer(op, hidden)], "") + + qmoe = next(node for node in graph if node.op_type == "QMoE") + assert qmoe.inputs[0].dtype == ir.DataType.FLOAT16 + assert qmoe.inputs[1].dtype == ir.DataType.FLOAT + + def test_deepseek_fused_qmoe_graph_has_one_node_per_moe_layer(): config = make_config( hidden_size=32, @@ -230,6 +254,32 @@ def test_deepseek_fused_qmoe_graph_has_one_node_per_moe_layer(): ) +def test_deepseek_fused_qmoe_rejects_cuda(): + config = make_config( + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=16, + num_hidden_layers=2, + first_k_dense_replace=1, + num_local_experts=4, + num_experts_per_tok=2, + n_shared_experts=1, + fused_quantized_moe=True, + quantization=QuantizationConfig( + bits=4, + group_size=16, + quant_method="gguf", + sym=True, + ), + ) + with pytest.raises(ValueError, match="CUDA.*ignores router_weights"): + build_from_module( + DeepSeekV3CausalLMModel(config), + config, + execution_provider="cuda", + ) + + def test_expert_major_packing_matches_static_64_expert_top6_reference(): torch.manual_seed(0) experts, top_k = 64, 6 From 66d99023f64110775ca338a0c9475ecf346d8172 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 30 Jul 2026 05:03:53 +0000 Subject: [PATCH 23/23] fix: merge main, fix RUF043 raw string pattern, reformat glm_moe_dsa.py --- src/mobius/components/_moe_test.py | 2 +- src/mobius/models/glm_moe_dsa.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index b19c03e1..a6faab6a 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -272,7 +272,7 @@ def test_deepseek_fused_qmoe_rejects_cuda(): sym=True, ), ) - with pytest.raises(ValueError, match="CUDA.*ignores router_weights"): + with pytest.raises(ValueError, match=r"CUDA.*ignores router_weights"): build_from_module( DeepSeekV3CausalLMModel(config), config, diff --git a/src/mobius/models/glm_moe_dsa.py b/src/mobius/models/glm_moe_dsa.py index bb14534f..451d4337 100644 --- a/src/mobius/models/glm_moe_dsa.py +++ b/src/mobius/models/glm_moe_dsa.py @@ -128,7 +128,9 @@ def select( # (ORT) and fixed-capacity (native single-token decode) exposures. key_length = op.Shape(scores, start=2, end=3) causal = op.Squeeze(attention_bias, [1]) - causal = op.Slice(causal, op.Constant(value_ints=[0]), key_length, op.Constant(value_ints=[2])) + causal = op.Slice( + causal, op.Constant(value_ints=[0]), key_length, op.Constant(value_ints=[2]) + ) scores = op.Add(scores, op.Cast(causal, to=ir.DataType.FLOAT)) k = op.Min(key_length, op.Constant(value_ints=[self.index_topk]))