From 79a167233785ea99d000ef6691beca0ff3ba56e2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 11:42:08 +0000 Subject: [PATCH 01/11] feat(chatglm): import fused GPTQ checkpoints Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_configs/_base.py | 21 +++++-- src/mobius/_configs_test.py | 33 +++++++++++ src/mobius/models/chatglm.py | 64 ++++++++++++++++------ src/mobius/models/chatglm_test.py | 91 +++++++++++++++++++++++++++++++ 4 files changed, 187 insertions(+), 22 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 33dd17aa..4385a392 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -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) @@ -549,6 +549,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: rope_config = RoPEConfig( rope_type="default", rope_theta=_IMPLICIT_ROPE_DEFAULTS[model_type], + partial_rotary_factor=0.5 if model_type == "chatglm" else 1.0, ) # Some hierarchical models (Segformer, Swin) use plural list attrs @@ -568,6 +569,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: num_hidden_layers = ( getattr(config, "num_hidden_layers", None) or getattr(config, "n_layers", None) + or getattr(config, "num_layers", None) or getattr(config, "num_encoder_blocks", None) or 0 ) @@ -590,12 +592,18 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: config.head_dim if (hasattr(config, "head_dim") and config.head_dim is not None) else getattr(config, "d_kv", None) + or getattr(config, "kv_channels", None) or _as_int(hidden_size) // _as_int(num_attention_heads) ), num_attention_heads=_as_int(num_attention_heads), num_key_value_heads=_as_int( getattr(config, "num_key_value_heads", None) or getattr(config, "n_kv_heads", None) + or ( + getattr(config, "multi_query_group_num", None) + if getattr(config, "multi_query_attention", False) + else None + ) or num_attention_heads ), num_hidden_layers=_as_int(num_hidden_layers), @@ -643,7 +651,9 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: or 1e-6 ), attn_qkv_bias=( - getattr( + getattr(config, "add_qkv_bias", False) + or getattr(config, "add_bias_linear", False) + or getattr( config, "attention_bias", getattr( @@ -677,7 +687,8 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: ) ), attn_o_bias=( - getattr( + getattr(config, "add_bias_linear", False) + or getattr( config, "attention_bias", getattr( @@ -722,7 +733,8 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: ), attn_qk_norm_full=(model_type in ("flex_olmo", "olmoe", "olmo2", "olmo3")), mlp_bias=( - getattr( + getattr(config, "add_bias_linear", False) + or getattr( config, "use_mlp_bias", getattr( @@ -765,6 +777,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: max_position_embeddings=( getattr(config, "max_position_embeddings", None) or getattr(config, "max_sequence_length", None) + or getattr(config, "seq_length", None) or 0 ), tie_word_embeddings=( diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index 219355a1..da97a82e 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -188,6 +188,39 @@ class FakeQwen3Config: config = ArchitectureConfig.from_transformers(FakeQwen3Config()) assert config.attn_qk_norm is True + def test_from_transformers_chatglm4_legacy_fields(self): + class FakeChatGLMConfig: + model_type = "chatglm" + num_attention_heads = 32 + multi_query_attention = True + multi_query_group_num = 2 + num_layers = 40 + hidden_size = 4096 + ffn_hidden_size = 13696 + kv_channels = 128 + padded_vocab_size = 151552 + vocab_size = 151552 + seq_length = 131072 + layernorm_epsilon = 1.5625e-7 + add_bias_linear = False + add_qkv_bias = True + pad_token_id = 151329 + tie_word_embeddings = False + + config = ArchitectureConfig.from_transformers(FakeChatGLMConfig()) + + assert config.num_hidden_layers == 40 + assert config.num_key_value_heads == 2 + assert config.head_dim == 128 + assert config.intermediate_size == 13696 + assert config.hidden_act == "silu" + assert config.max_position_embeddings == 131072 + assert config.partial_rotary_factor == pytest.approx(0.5) + assert config.rope_interleave is True + assert config.attn_qkv_bias is True + assert config.attn_o_bias is False + assert config.mlp_bias is False + def test_from_transformers_rope_scaling(self): class FakeConfig: model_type = "llama" diff --git a/src/mobius/models/chatglm.py b/src/mobius/models/chatglm.py index ff6a4b66..77e9eccb 100644 --- a/src/mobius/models/chatglm.py +++ b/src/mobius/models/chatglm.py @@ -5,29 +5,57 @@ import torch -from mobius._weight_utils import rename_weight_keys -from mobius.models.base import CausalLMModel +from mobius._weight_utils import split_fused_qkv +from mobius.models.base import FusedGateUpCausalLMModel -class ChatGLMCausalLMModel(CausalLMModel): - """ChatGLM model with partial rotary (0.5 factor) and MLP name remapping. +class ChatGLMCausalLMModel(FusedGateUpCausalLMModel): + """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 tensors along their output dimension. Keeping gate/up fused + allows packed tensors to load without dequantization. """ 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 + + 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..79948d95 --- /dev/null +++ b/src/mobius/models/chatglm_test.py @@ -0,0 +1,91 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import torch + +from mobius._configs import ArchitectureConfig, QuantizationConfig +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_qkv_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) + 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), + } + _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" in result + 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 From ce61f462d30b57a473119d622f98a221f582b81f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 12:00:47 +0000 Subject: [PATCH 02/11] perf(chatglm): pre-split packed gate-up projections Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/chatglm.py | 21 +++++++--- src/mobius/models/chatglm_test.py | 64 ++++++++++++++++++++++++++++++- 2 files changed, 78 insertions(+), 7 deletions(-) diff --git a/src/mobius/models/chatglm.py b/src/mobius/models/chatglm.py index 77e9eccb..bf47963a 100644 --- a/src/mobius/models/chatglm.py +++ b/src/mobius/models/chatglm.py @@ -5,18 +5,19 @@ import torch -from mobius._weight_utils import split_fused_qkv -from mobius.models.base import FusedGateUpCausalLMModel +from mobius._weight_utils import split_fused_qkv, split_gate_up_proj +from mobius.models.base import CausalLMModel -class ChatGLMCausalLMModel(FusedGateUpCausalLMModel): +class ChatGLMCausalLMModel(CausalLMModel): """ChatGLM model with partial rotary and fused projection imports. 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 tensors along their output dimension. Keeping gate/up fused - allows packed tensors to load without dequantization. + 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( @@ -58,4 +59,14 @@ def preprocess_weights( 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 index 79948d95..3399c40e 100644 --- a/src/mobius/models/chatglm_test.py +++ b/src/mobius/models/chatglm_test.py @@ -6,6 +6,8 @@ import torch from mobius._configs import ArchitectureConfig, QuantizationConfig +from mobius._weight_utils import preprocess_gptq_weights +from mobius.components._mlp import MLP from mobius.models.chatglm import ChatGLMCausalLMModel @@ -23,7 +25,7 @@ def _add_gptq(state_dict: dict[str, torch.Tensor], stem: str, k: int, n: int) -> state_dict[f"{stem}.{suffix}"] = value -def test_glm4_gptq_checkpoint_remap_and_fused_qkv_split(): +def test_glm4_gptq_checkpoint_remap_and_fused_projection_split(): config = ArchitectureConfig( model_type="chatglm", vocab_size=64, @@ -42,6 +44,7 @@ def test_glm4_gptq_checkpoint_remap_and_fused_qkv_split(): 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), @@ -49,6 +52,9 @@ def test_glm4_gptq_checkpoint_remap_and_fused_qkv_split(): "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, @@ -84,8 +90,62 @@ def test_glm4_gptq_checkpoint_remap_and_fused_qkv_split(): 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" 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:]) From 00309bece3de6c145e4f87164274994e3f21b0e6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 13:40:43 +0000 Subject: [PATCH 03/11] fix(gqa): retain attention for unequal KV heads Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../rewrite_rules/_group_query_attention.py | 27 +++++++++++++++++-- .../_group_query_attention_test.py | 24 +++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) 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..425ae94d 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -209,6 +209,30 @@ def test_preserves_non_matching_model(self): assert vision_counts.get("Attention", 0) == vision_attn_before assert vision_counts.get("GroupQueryAttention", 0) == 0 + @pytest.mark.arch_validation + def test_deepseek_v2_lite_unequal_kv_heads_retain_attention(self): + """CUDA export must not rewrite unequal-dimension MLA K/V to GQA.""" + with pytest.warns(UserWarning, match="GQA fusion expected"): + pkg = build( + "deepseek-ai/DeepSeek-V2-Lite", + dtype="f16", + execution_provider="cuda", + load_weights=False, + ) + model = pkg["model"] + counts = count_ops(model) + + assert counts.get("Attention", 0) == 27 + assert counts.get("GroupQueryAttention", 0) == 0 + + attention_nodes = [node for node in model.graph if node.op_type == "Attention"] + assert len(attention_nodes) == 27 + for node in attention_nodes: + past_key = node.inputs[4] + past_value = node.inputs[5] + assert past_key is not None and past_key.shape[-1] == 192 + assert past_value is not None and past_value.shape[-1] == 128 + def test_fallback_attention_to_gqa_no_rope(self): """AttentionToGQA fallback fires when applied in isolation (do_rotary=0). From 85fb2c6098ce9430b6f056ad7028fb223e3a1b8c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 07:00:16 -0700 Subject: [PATCH 04/11] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/models/chatglm_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mobius/models/chatglm_test.py b/src/mobius/models/chatglm_test.py index 3399c40e..be95d13d 100644 --- a/src/mobius/models/chatglm_test.py +++ b/src/mobius/models/chatglm_test.py @@ -7,7 +7,7 @@ from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius._weight_utils import preprocess_gptq_weights -from mobius.components._mlp import MLP +from mobius.components import MLP from mobius.models.chatglm import ChatGLMCausalLMModel From 7e823269d86f0e180b24a2a7b907d62dcae60b48 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 14:03:19 +0000 Subject: [PATCH 05/11] feat(moe): emit expert-major int4 QMoE Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_builder.py | 5 +- src/mobius/_weight_utils.py | 76 +++++++-- src/mobius/_weight_utils_test.py | 38 +++++ src/mobius/components/_deepseek_mla.py | 15 +- src/mobius/components/_moe.py | 146 +++++++++++++++- src/mobius/components/_moe_test.py | 158 ++++++++++++++++++ src/mobius/models/deepseek.py | 128 +++++++++----- .../_group_query_attention_test.py | 66 +++++++- 8 files changed, 564 insertions(+), 68 deletions(-) diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index a080f36a..4638ecd4 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -91,13 +91,14 @@ def _cast_module_dtype(module: nn.Module, dtype: ir.DataType) -> None: caches), the underlying data is also cast. Only recasts parameters that are currently FLOAT (float32). Integer - parameters and non-float types are left unchanged. + parameters, non-float types, and parameters marked ``_keep_float32`` are + left unchanged. """ if dtype == ir.DataType.FLOAT: return torch_dtype = tensor_adapters.to_torch_dtype(dtype) for param in module.parameters(): - if param.dtype != ir.DataType.FLOAT: + if param.dtype != ir.DataType.FLOAT or getattr(param, "_keep_float32", False): continue param.type = ir.TensorType(dtype) if param.const_value is not None: diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index 57f6c4c3..f7b198b9 100644 --- a/src/mobius/_weight_utils.py +++ b/src/mobius/_weight_utils.py @@ -557,19 +557,22 @@ def vlm_vision_weights( def _reshape_packed_qweight(value: torch.Tensor, blob_size: int) -> torch.Tensor: """Transpose and reshape a packed qweight tensor for MatMulNBits. - Converts ``[K_packed, N]`` int32 → ``[N, n_blocks, blob_size]`` uint8. + Converts ``[..., K_packed, N]`` int32 to + ``[..., N, n_blocks, blob_size]`` uint8. """ - transposed = value.t().contiguous() - n = transposed.shape[0] + transposed = value.transpose(-1, -2).contiguous() + prefix = transposed.shape[:-2] + n = transposed.shape[-2] packed = transposed.view(torch.uint8) - n_blocks = packed.shape[1] // blob_size - return packed.reshape(n, n_blocks, blob_size) + n_blocks = packed.shape[-1] // blob_size + return packed.reshape(*prefix, n, n_blocks, blob_size) def _reshape_packed_qzeros(value: torch.Tensor, bits: int, n_blocks: int) -> torch.Tensor: """Transpose and unpack packed qzeros for MatMulNBits. - Converts ``[n_groups_packed, N]`` int32 → ``[N, zp_cols]`` uint8 + Converts ``[..., n_groups_packed, N]`` int32 to + ``[..., N, zp_cols]`` uint8 where ``zp_cols = ceil(n_blocks * bits / 8)``. For 4-bit this packs two zero-point values per byte, matching ORT's expectation. @@ -578,11 +581,56 @@ def _reshape_packed_qzeros(value: torch.Tensor, bits: int, n_blocks: int) -> tor bits: Quantization bit-width (4 or 8). n_blocks: Actual number of quantization blocks (``ceil(K / block_size)``). """ - transposed = value.t().contiguous() - n = transposed.shape[0] - flat_uint8 = transposed.flatten().view(torch.uint8).reshape(n, -1) + transposed = value.transpose(-1, -2).contiguous() + prefix = transposed.shape[:-2] + n = transposed.shape[-2] + flat_uint8 = transposed.reshape(-1).view(torch.uint8).reshape(*prefix, n, -1) zp_cols = math.ceil(n_blocks * bits / 8) - return flat_uint8[:, :zp_cols] + return flat_uint8[..., :zp_cols] + + +def pack_qmoe_expert_weights( + state_dict: dict[str, torch.Tensor], + *, + target_moe_path: str = ".mlp.moe", +) -> dict[str, torch.Tensor]: + """Map fused expert-major quantized tensors to native QMoE parameters.""" + packed: dict[str, torch.Tensor] = {} + projections = { + ".mlp.experts.gate_up_proj.weight": ( + f"{target_moe_path}.fc1_experts_weights", + True, + ), + ".mlp.experts.gate_up_proj.scales": ( + f"{target_moe_path}.fc1_scales", + False, + ), + ".mlp.experts.gate_up_proj.zero_points": ( + f"{target_moe_path}.fc1_experts_zero_points", + False, + ), + ".mlp.experts.down_proj.weight": ( + f"{target_moe_path}.fc2_experts_weights", + True, + ), + ".mlp.experts.down_proj.scales": ( + f"{target_moe_path}.fc2_scales", + False, + ), + ".mlp.experts.down_proj.zero_points": ( + f"{target_moe_path}.fc2_experts_zero_points", + False, + ), + } + for key, value in state_dict.items(): + for source, (target, flatten_blocks) in projections.items(): + if source in key: + key = key.replace(source, target) + if flatten_blocks: + value = value.flatten(-2) + break + packed[key] = value + return packed def preprocess_gptq_weights( @@ -645,12 +693,12 @@ def preprocess_gptq_weights( raise ValueError( f"Missing {qw_key} — qweight must be present alongside qzeros for {key}" ) - k = state_dict[qw_key].shape[0] * 32 // bits + k = state_dict[qw_key].shape[-2] * 32 // bits n_blocks = math.ceil(k / group_size) result[new_key] = _reshape_packed_qzeros(value, bits, n_blocks) elif key.endswith(".scales"): - result[key] = value.t().contiguous() + result[key] = value.transpose(-1, -2).contiguous() else: result[key] = value @@ -792,7 +840,7 @@ def preprocess_awq_weights( raise ValueError( f"Missing {qw_key} — qweight must be present alongside qzeros for {key}" ) - k = state_dict[qw_key].shape[0] * 32 // bits + k = state_dict[qw_key].shape[-2] * 32 // bits n_blocks = math.ceil(k / group_size) # AWQ zero points have an implicit +1 offset; subtract it # before unpacking so MatMulNBits sees the raw value. @@ -810,7 +858,7 @@ def preprocess_awq_weights( result[new_key] = zp_int16.clamp(min=0).to(torch.uint8) elif key.endswith(".scales"): - result[key] = value.t().contiguous() + result[key] = value.transpose(-1, -2).contiguous() else: result[key] = value diff --git a/src/mobius/_weight_utils_test.py b/src/mobius/_weight_utils_test.py index 14b05e57..4c1e6dda 100644 --- a/src/mobius/_weight_utils_test.py +++ b/src/mobius/_weight_utils_test.py @@ -510,6 +510,44 @@ def test_scales_transposed(self): result = preprocess_gptq_weights(sd, bits=self.BITS, group_size=self.GROUP_SIZE) assert result["q_proj.scales"].shape == (self.N, self.N_GROUPS) + def test_expert_major_tensors_preserve_expert_axis(self): + num_experts = 64 + sd = { + "experts.qweight": torch.randint( + 0, + 255, + (num_experts, self.K_PACKED, self.N), + dtype=torch.int32, + ), + "experts.qzeros": torch.randint( + 0, + 255, + (num_experts, max(1, self.N_GROUPS_PACKED), self.N), + dtype=torch.int32, + ), + "experts.scales": torch.randn(num_experts, self.N_GROUPS, self.N), + } + result = preprocess_gptq_weights( + sd, bits=self.BITS, group_size=self.GROUP_SIZE + ) + + assert result["experts.weight"].shape == ( + num_experts, + self.N, + self.N_GROUPS, + self.BLOB_SIZE, + ) + assert result["experts.zero_points"].shape == ( + num_experts, + self.N, + self.N_GROUPS * self.BITS // 8, + ) + assert result["experts.scales"].shape == ( + num_experts, + self.N, + self.N_GROUPS, + ) + def test_non_gptq_keys_pass_through(self): t = torch.randn(4, 8) sd = {"model.embed_tokens.weight": t, "lm_head.weight": t.clone()} 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..6d636eab 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -6,16 +6,14 @@ from __future__ import annotations import dataclasses -from typing import TYPE_CHECKING +import math +import onnx_ir as ir from onnxscript import OpBuilder, nn -from mobius._configs import ArchitectureConfig +from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius.components._mlp import MLP -if TYPE_CHECKING: - import onnx_ir as ir - class TopKGate(nn.Module): """Standard top-k expert routing gate. @@ -38,6 +36,13 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): routing_weights = op.Softmax(routing_weights, axis=-1) return routing_weights, selected_experts + def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): + """Return selection and aggregation tensors for QMoE.""" + weight_t = op.Transpose(self.weight, perm=[1, 0]) + router_logits = op.MatMul(hidden_states, weight_t) + router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) + return router_logits, None, True, 1.0 + class SoftmaxTopKGate(nn.Module): """Softmax-first top-k expert routing gate (Qwen3-Next style). @@ -68,6 +73,14 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): routing_weights = op.Div(routing_weights, weight_sum) return routing_weights, selected_experts + def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): + """Return selection and aggregation tensors for QMoE.""" + weight_t = op.Transpose(self.weight, perm=[1, 0]) + router_logits = op.MatMul(hidden_states, weight_t) + router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) + routing_probs = op.Softmax(router_logits, axis=-1) + return routing_probs, routing_probs, self.norm_topk_prob, 1.0 + class SigmoidTopKGate(nn.Module): """Sigmoid-first top-k expert routing gate (GLM4-MoE style). @@ -109,6 +122,19 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): routing_weights = op.Mul(routing_weights, self.routed_scaling_factor) return routing_weights, selected_experts + def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): + """Return selection and aggregation tensors for QMoE.""" + weight_t = op.Transpose(self.weight, perm=[1, 0]) + router_logits = op.MatMul(hidden_states, weight_t) + router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) + routing_probs = op.Sigmoid(router_logits) + return ( + routing_probs, + routing_probs, + self.norm_topk_prob, + self.routed_scaling_factor, + ) + class SparseMixerGate(nn.Module): """Sparsemixer-style routing gate (used by PhiMoE). @@ -196,6 +222,7 @@ def __init__(self, config: ArchitectureConfig, gate: nn.Module | None = None): assert config.num_experts_per_tok is not None self.num_experts = config.num_local_experts self.top_k = config.num_experts_per_tok + self._qmoe_quantization = _supported_qmoe_quantization(config.quantization) if gate is not None: self.gate = gate else: @@ -206,9 +233,102 @@ 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)]) + if self._qmoe_quantization is not None and hasattr(self.gate, "qmoe_routing"): + self.experts = None + self._init_qmoe_parameters(expert_config) + else: + self.experts = nn.ModuleList( + [MLP(expert_config) for _ in range(self.num_experts)] + ) + + def _init_qmoe_parameters(self, expert_config: ArchitectureConfig) -> None: + quantization = self._qmoe_quantization + assert quantization is not None + hidden_size = expert_config.hidden_size + intermediate_size = expert_config.intermediate_size + block_size = quantization.group_size + bits = quantization.bits + fc1_out = 2 * intermediate_size + if hidden_size % block_size or intermediate_size % block_size: + raise ValueError( + "QMoE dimensions must be divisible by the quantization group size" + ) + + self.fc1_experts_weights = nn.Parameter( + [self.num_experts, fc1_out, hidden_size * bits // 8], + dtype=ir.DataType.UINT8, + ) + self.fc1_scales = nn.Parameter( + [self.num_experts, fc1_out, hidden_size // block_size] + ) + self.fc1_scales._keep_float32 = True + self.fc2_experts_weights = nn.Parameter( + [self.num_experts, hidden_size, intermediate_size * bits // 8], + dtype=ir.DataType.UINT8, + ) + self.fc2_scales = nn.Parameter( + [self.num_experts, hidden_size, intermediate_size // block_size] + ) + self.fc2_scales._keep_float32 = True + if quantization.sym: + self.fc1_experts_zero_points = None + self.fc2_experts_zero_points = None + else: + self.fc1_experts_zero_points = nn.Parameter( + [ + self.num_experts, + fc1_out, + math.ceil((hidden_size // block_size) * bits / 8), + ], + dtype=ir.DataType.UINT8, + ) + self.fc2_experts_zero_points = nn.Parameter( + [ + self.num_experts, + hidden_size, + math.ceil((intermediate_size // block_size) * bits / 8), + ], + dtype=ir.DataType.UINT8, + ) + + def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value): + quantization = self._qmoe_quantization + assert quantization is not None + router_probs, router_weights, normalize, output_scale = self.gate.qmoe_routing( + op, hidden_states + ) + result = op.QMoE( + hidden_states, + router_probs, + self.fc1_experts_weights, + self.fc1_scales, + None, + self.fc2_experts_weights, + self.fc2_scales, + None, + None, + None, + None, + self.fc1_experts_zero_points, + self.fc2_experts_zero_points, + None, + router_weights, + activation_type="swiglu", + normalize_routing_weights=int(normalize), + k=self.top_k, + expert_weight_bits=quantization.bits, + block_size=quantization.group_size, + swiglu_fusion=2, + _domain="com.microsoft", + ) + if output_scale != 1.0: # noqa: RUF069 + result = op.Mul(result, op.CastLike(output_scale, result)) + return result def forward(self, op: OpBuilder, hidden_states: ir.Value): + if self.experts is None: + return self._qmoe_forward(op, hidden_states) + routing_weights, selected_experts = self.gate(op, hidden_states) result = None @@ -226,3 +346,17 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): result = op.Add(result, contribution) return result + + +def _supported_qmoe_quantization( + quantization: QuantizationConfig | None, +) -> QuantizationConfig | None: + """Return quantization settings when they match the native QMoE ABI.""" + if ( + quantization is None + or quantization.bits != 4 + or quantization.float_zero_point + or quantization.quant_method not in {"gptq", "awq"} + ): + return None + return quantization diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index c48bcdfc..0f3e9ab7 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -7,12 +7,16 @@ import onnx_ir as ir import pytest +import torch +from mobius._builder import _cast_module_dtype +from mobius._configs import QuantizationConfig from mobius._testing import ( create_test_builder, create_test_input, make_config, ) +from mobius._weight_utils import pack_qmoe_expert_weights, preprocess_gptq_weights from mobius.components._moe import MoELayer, SparseMixerGate, TopKGate @@ -118,3 +122,157 @@ 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_int4_moe_emits_expert_major_qmoe(self): + config = make_config( + hidden_size=64, + intermediate_size=32, + moe_intermediate_size=32, + num_local_experts=64, + num_experts_per_tok=6, + quantization=QuantizationConfig( + bits=4, + group_size=32, + quant_method="gptq", + sym=False, + ), + ) + layer = MoELayer(config) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 8, 64]) + result = layer(op, hidden) + builder._adapt_outputs([result], "") + + qmoe_nodes = [node for node in graph if node.op_type == "QMoE"] + assert len(qmoe_nodes) == 1 + qmoe = qmoe_nodes[0] + assert qmoe.attributes["k"].value == 6 + assert qmoe.attributes["block_size"].value == 32 + assert qmoe.attributes["swiglu_fusion"].value == 2 + assert layer.fc1_experts_weights.shape == ir.Shape([64, 64, 32]) + assert layer.fc1_scales.shape == ir.Shape([64, 64, 2]) + assert layer.fc1_experts_zero_points.shape == ir.Shape([64, 64, 1]) + assert layer.fc2_experts_weights.shape == ir.Shape([64, 64, 16]) + assert layer.fc2_scales.shape == ir.Shape([64, 64, 1]) + assert layer.fc2_experts_zero_points.shape == ir.Shape([64, 64, 1]) + assert sum(node.op_type == "MatMul" for node in graph) == 1 # router only + + _cast_module_dtype(layer, ir.DataType.FLOAT16) + assert layer.gate.weight.dtype == ir.DataType.FLOAT16 + assert layer.fc1_scales.dtype == ir.DataType.FLOAT + assert layer.fc2_scales.dtype == ir.DataType.FLOAT + + def test_expert_major_packing_matches_static_64_expert_top6_reference(self): + """Packed QMoE math matches the existing loop-over-experts semantics.""" + torch.manual_seed(0) + num_experts, top_k = 64, 6 + hidden_size, intermediate_size, block_size = 32, 16, 16 + fc1_codes = torch.randint( + 0, 16, (num_experts, 2 * intermediate_size, hidden_size) + ) + fc2_codes = torch.randint( + 0, 16, (num_experts, hidden_size, intermediate_size) + ) + fc1_scales = torch.rand( + num_experts, 2 * intermediate_size, hidden_size // block_size + ) + fc2_scales = torch.rand( + num_experts, hidden_size, intermediate_size // block_size + ) + raw = { + "model.layers.1.mlp.experts.gate_up_proj.qweight": _to_gptq_qweight( + fc1_codes + ), + "model.layers.1.mlp.experts.gate_up_proj.scales": fc1_scales.transpose( + -1, -2 + ), + "model.layers.1.mlp.experts.down_proj.qweight": _to_gptq_qweight( + fc2_codes + ), + "model.layers.1.mlp.experts.down_proj.scales": fc2_scales.transpose( + -1, -2 + ), + } + packed = pack_qmoe_expert_weights( + preprocess_gptq_weights(raw, bits=4, group_size=block_size) + ) + fc1_packed = packed["model.layers.1.mlp.moe.fc1_experts_weights"] + fc2_packed = packed["model.layers.1.mlp.moe.fc2_experts_weights"] + assert fc1_packed.shape == (num_experts, 2 * intermediate_size, hidden_size // 2) + assert fc2_packed.shape == (num_experts, hidden_size, intermediate_size // 2) + + packed_fc1 = _dequant_qmoe(fc1_packed, fc1_scales, block_size) + packed_fc2 = _dequant_qmoe(fc2_packed, fc2_scales, block_size) + static_fc1 = _dequant_codes(fc1_codes, fc1_scales, block_size) + static_fc2 = _dequant_codes(fc2_codes, fc2_scales, block_size) + + hidden = torch.randn(3, hidden_size) + router_probs = torch.softmax(torch.randn(3, num_experts), dim=-1) + selected_weights, selected_experts = router_probs.topk(top_k, dim=-1) + selected_weights /= selected_weights.sum(dim=-1, keepdim=True) + static = _static_moe( + hidden, + selected_weights, + selected_experts, + static_fc1, + static_fc2, + intermediate_size, + ) + fused_reference = _static_moe( + hidden, + selected_weights, + selected_experts, + packed_fc1, + packed_fc2, + intermediate_size, + ) + torch.testing.assert_close(fused_reference, static, atol=1e-5, rtol=1e-5) + + +def _to_gptq_qweight(codes: torch.Tensor) -> torch.Tensor: + """Pack output-major int4 codes into GPTQ's [E, K/8, N] int32 layout.""" + 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_qmoe( + packed: torch.Tensor, scales: torch.Tensor, block_size: int +) -> torch.Tensor: + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), 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_idx in range(fc1.shape[0]): + projected = hidden @ fc1[expert_idx].T + activated = torch.nn.functional.silu(projected[:, :intermediate_size]) + activated *= projected[:, intermediate_size:] + expert_output = activated @ fc2[expert_idx].T + weight = ( + routing_weights + * (selected_experts == expert_idx).to(routing_weights.dtype) + ).sum(dim=-1, keepdim=True) + result += expert_output * weight + return result diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index ed835417..62f27f1f 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -9,12 +9,13 @@ from __future__ import annotations import dataclasses -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 pack_qmoe_expert_weights from mobius.components import ( MLP, Attention, @@ -26,10 +27,23 @@ initialize_rope, ) from mobius.components._deepseek_mla import DeepSeekMLA +from mobius.components._quantized_linear import make_quantized_linear_factory from mobius.models.base import CausalLMModel -if TYPE_CHECKING: - import onnx_ir as ir + +def _linear_factory(config: ArchitectureConfig): + quantization = config.quantization + if quantization is None or quantization.quant_method == "none": + return None + zero_point_dtype = ( + config.dtype if quantization.float_zero_point else ir.DataType.UINT8 + ) + return make_quantized_linear_factory( + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, + zero_point_dtype=zero_point_dtype, + ) class DeepSeekMoEGate(nn.Module): @@ -64,26 +78,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 - ) - - # Score computation depends on scoring function - if self.scoring_func == "sigmoid": - scores = op.Sigmoid(router_logits) # (B*S, num_experts) - # Add correction bias for expert selection (V3) - scores_for_choice = op.Add(scores, self.e_score_correction_bias) - else: - # Softmax scoring (V2) - scores = op.Softmax(router_logits, axis=-1) - scores_for_choice = scores - - # Expert selection: group-based or simple TopK - if self.n_group > 1 and self.topk_method != "greedy": - scores_for_choice = self._group_topk_selection(op, scores_for_choice) + scores, scores_for_choice = self._routing_scores(op, hidden_states) # Select top-k experts k_val = op.Constant(value_ints=[self.top_k]) @@ -103,6 +98,38 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return routing_weights, selected_experts + def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): + """Return distinct expert-selection and aggregation scores for QMoE.""" + scores, scores_for_choice = self._routing_scores(op, hidden_states) + return ( + scores_for_choice, + scores, + self.norm_topk_prob, + float(self.routed_scaling_factor), + ) + + def _routing_scores(self, op: OpBuilder, hidden_states: ir.Value): + weight_t = op.Transpose(self.weight, perm=[1, 0]) + router_logits = op.MatMul( + op.Cast(hidden_states, to=1), + op.Cast(weight_t, to=1), + ) + + # Score computation depends on scoring function + if self.scoring_func == "sigmoid": + scores = op.Sigmoid(router_logits) # (B*S, num_experts) + # Add correction bias for expert selection (V3) + scores_for_choice = op.Add(scores, self.e_score_correction_bias) + else: + # Softmax scoring (V2) + scores = op.Softmax(router_logits, axis=-1) + scores_for_choice = scores + + # Expert selection: group-based or simple TopK + if self.n_group > 1 and self.topk_method != "greedy": + scores_for_choice = self._group_topk_selection(op, scores_for_choice) + return scores, scores_for_choice + 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 @@ -152,12 +179,13 @@ class DeepSeekMLADecoderLayer(nn.Module): def __init__(self, config: ArchitectureConfig, is_moe: bool = False): super().__init__() - self.self_attn = DeepSeekMLA(config) + linear_class = _linear_factory(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) @@ -199,12 +227,13 @@ class _DeepSeekStandardDecoderLayer(nn.Module): def __init__(self, config: ArchitectureConfig, is_moe: bool = False): super().__init__() - self.self_attn = Attention(config) + linear_class = _linear_factory(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,7 +273,9 @@ 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 @@ -252,7 +283,9 @@ def __init__(self, config: ArchitectureConfig, gate: nn.Module): # 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 +296,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) @@ -389,6 +429,12 @@ def preprocess_weights( moe.experts.{i}.down_proj.weight: (hidden, intermediate) """ renamed = {} + use_qmoe = ( + self.config.quantization is not None + and self.config.quantization.bits == 4 + and self.config.quantization.quant_method in {"gptq", "awq"} + and not self.config.quantization.float_zero_point + ) for key, value in state_dict.items(): new_key = key @@ -399,14 +445,14 @@ def preprocess_weights( # layers.N.mlp.experts.gate_up_proj (n_experts, 2*mid, hidden) # layers.N.mlp.experts.down_proj (n_experts, hidden, mid) # Split into per-expert weights for our ModuleList. - if new_key.endswith(".mlp.experts.gate_up_proj"): + if not use_qmoe and new_key.endswith(".mlp.experts.gate_up_proj"): prefix = new_key[: -len(".mlp.experts.gate_up_proj")] mid = value.shape[1] // 2 for i in range(value.shape[0]): renamed[f"{prefix}.mlp.moe.experts.{i}.gate_proj.weight"] = value[i, :mid] renamed[f"{prefix}.mlp.moe.experts.{i}.up_proj.weight"] = value[i, mid:] continue - if new_key.endswith(".mlp.experts.down_proj"): + if not use_qmoe and new_key.endswith(".mlp.experts.down_proj"): prefix = new_key[: -len(".mlp.experts.down_proj")] for i in range(value.shape[0]): renamed[f"{prefix}.mlp.moe.experts.{i}.down_proj.weight"] = value[i] @@ -414,5 +460,9 @@ def preprocess_weights( renamed[new_key] = value - # Handle weight tying - return super().preprocess_weights(renamed) + # Handle weight tying and GPTQ/AWQ conversion before flattening the + # expert-major MatMulNBits blobs into the QMoE ABI. + processed = super().preprocess_weights(renamed) + if use_qmoe: + processed = pack_qmoe_expert_weights(processed) + return processed diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index 425ae94d..cae259f3 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -3,13 +3,16 @@ from __future__ import annotations +import dataclasses + +import onnx_ir as ir import pytest from onnxscript.rewriter import rewrite from onnxscript.rewriter._rewrite_rule import RewriteRuleSet from mobius import build from mobius._builder import build_from_module -from mobius._configs import ArchitectureConfig, Gemma2Config +from mobius._configs import ArchitectureConfig, Gemma2Config, QuantizationConfig from mobius._registry import registry from mobius._testing.ort_inference import OnnxModelSession from mobius.rewrite_rules import group_query_attention_rules, pack_qkv_for_gqa_rules @@ -233,6 +236,67 @@ def test_deepseek_v2_lite_unequal_kv_heads_retain_attention(self): assert past_key is not None and past_key.shape[-1] == 192 assert past_value is not None and past_value.shape[-1] == 128 + @pytest.mark.arch_validation + def test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer(self): + """The int4 graph keeps shared experts dense and routed experts fused.""" + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained( + "deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True + ) + config = dataclasses.replace( + ArchitectureConfig.from_transformers(hf_config), + dtype=ir.DataType.FLOAT16, + quantization=QuantizationConfig( + bits=4, + group_size=128, + quant_method="gptq", + sym=True, + ), + ) + with pytest.warns(UserWarning, match="GQA fusion expected"): + model = build_from_module( + registry.get("deepseek_v3")(config), + config, + execution_provider="cuda", + )["model"] + counts = count_ops(model) + + assert counts.get("QMoE", 0) == 26 + assert counts.get("GroupQueryAttention", 0) == 0 + assert counts.get("Attention", 0) == 27 + for node in (node for node in model.graph if node.op_type == "QMoE"): + assert node.inputs[1] is not None + assert node.inputs[1].dtype == ir.DataType.FLOAT + assert node.inputs[3] is not None + assert node.inputs[3].dtype == ir.DataType.FLOAT + assert node.inputs[6] is not None + assert node.inputs[6].dtype == ir.DataType.FLOAT + assert node.inputs[14] is not None + assert node.inputs[14].dtype == ir.DataType.FLOAT + + input_names = [ + value.name + for node in model.graph + for value in node.inputs + if value is not None and value.name is not None + ] + assert not any(".moe.experts." in name for name in input_names) + for layer_idx in range(1, 27): + prefix = f"model.layers.{layer_idx}.mlp.shared_experts." + shared_matmuls = [ + node + for node in model.graph + if node.op_type == "MatMulNBits" + and any( + value is not None + and value.name is not None + and prefix in value.name + for value in node.inputs + ) + ] + assert len(shared_matmuls) == 3 + def test_fallback_attention_to_gqa_no_rope(self): """AttentionToGQA fallback fires when applied in isolation (do_rotary=0). From bb70029e2bd901612af03edfc6afbd5ba58e20c3 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 23 Jul 2026 14:12:10 +0000 Subject: [PATCH 06/11] test: cover MLA GQA dimension guard Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_configs/_base.py | 6 +- .../_group_query_attention_test.py | 63 ++++++++++++------- 2 files changed, 44 insertions(+), 25 deletions(-) diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 4385a392..d1634e8e 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 and 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 """ @@ -295,7 +295,7 @@ def _extract_vision_config(config, parent_config, model_type: str) -> dict: hooks live under :mod:`mobius._configs.per_model` and are registered with :mod:`mobius._configs._extractors` at import time. """ - from mobius._configs import per_model # noqa: F401 - side-effect import + from mobius._configs import per_model # ruff:ignore[unused-import] - side effect from mobius._configs._extractors import extract_vision_config as _dispatch return _dispatch(config, parent_config, model_type) @@ -308,7 +308,7 @@ def _extract_audio_config(config, parent_config, model_type: str) -> dict: hooks live under :mod:`mobius._configs.per_model` and are registered with :mod:`mobius._configs._extractors` at import time. """ - from mobius._configs import per_model # noqa: F401 - side-effect import + from mobius._configs import per_model # ruff:ignore[unused-import] - side effect from mobius._configs._extractors import extract_audio_config as _dispatch return _dispatch(config, parent_config, model_type) diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index cae259f3..3535d899 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -39,6 +39,28 @@ pad_token_id=0, ) +# Tiny DeepSeek MLA config for structural K/V head-dimension coverage. +_DEEPSEEK_MLA_CONFIG = ArchitectureConfig( + hidden_size=32, + intermediate_size=64, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=16, + num_hidden_layers=1, + vocab_size=64, + max_position_embeddings=32, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + dtype=ir.DataType.FLOAT16, + q_lora_rank=16, + kv_lora_rank=16, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=8, +) + # Tiny qwen3 config: has QK norm, weights NOT packable _QWEN3_CONFIG = ArchitectureConfig( hidden_size=64, @@ -212,29 +234,26 @@ def test_preserves_non_matching_model(self): assert vision_counts.get("Attention", 0) == vision_attn_before assert vision_counts.get("GroupQueryAttention", 0) == 0 - @pytest.mark.arch_validation - def test_deepseek_v2_lite_unequal_kv_heads_retain_attention(self): - """CUDA export must not rewrite unequal-dimension MLA K/V to GQA.""" - with pytest.warns(UserWarning, match="GQA fusion expected"): - pkg = build( - "deepseek-ai/DeepSeek-V2-Lite", - dtype="f16", - execution_provider="cuda", - load_weights=False, - ) - model = pkg["model"] - counts = count_ops(model) + @pytest.mark.parametrize( + ("v_head_dim", "expected_attention", "expected_gqa"), + [(8, 1, 0), (16, 0, 1)], + ) + def test_mla_gqa_fusion_requires_equal_kv_head_dimensions( + self, v_head_dim, expected_attention, expected_gqa + ): + """GQA fusion declines unequal MLA K/V dimensions and accepts equal ones.""" + config = dataclasses.replace(_DEEPSEEK_MLA_CONFIG, v_head_dim=v_head_dim) + model = build_from_module(registry.get("deepseek_v3")(config), config)["model"] + attention = next(node for node in model.graph if node.op_type == "Attention") + assert attention.inputs[4] is not None + assert attention.inputs[4].shape[-1] == 16 + assert attention.inputs[5] is not None + assert attention.inputs[5].shape[-1] == v_head_dim - assert counts.get("Attention", 0) == 27 - assert counts.get("GroupQueryAttention", 0) == 0 - - attention_nodes = [node for node in model.graph if node.op_type == "Attention"] - assert len(attention_nodes) == 27 - for node in attention_nodes: - past_key = node.inputs[4] - past_value = node.inputs[5] - assert past_key is not None and past_key.shape[-1] == 192 - assert past_value is not None and past_value.shape[-1] == 128 + rewrite(model, pattern_rewrite_rules=group_query_attention_rules()) + counts = count_ops(model) + assert counts.get("Attention", 0) == expected_attention + assert counts.get("GroupQueryAttention", 0) == expected_gqa @pytest.mark.arch_validation def test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer(self): From c5daf876286bb293b89101bf6c91a3fa3b4a25a3 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 24 Jul 2026 17:12:30 +0000 Subject: [PATCH 07/11] docs(moe): explain why TopKGate.qmoe_routing passes raw logits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Copilot review feedback questioning whether TopKGate.qmoe_routing should pass raw router logits (instead of probabilities) as the QMoE `router_probs` input. This is intentional and correct: with `router_weights=None` and `normalize_routing_weights=1`, com.microsoft::QMoE reproduces `TopKGate.forward` exactly — it selects the top-k experts by logit value (softmax is monotonic) and, because `router_weights` is omitted, computes aggregation weights as softmax over the k selected logits, i.e. `Softmax(TopK(router_logits))`. Verified against the ORT QMoE kernels: the CPU default path in contrib_ops/cpu/moe/moe_quantization_cpu.cc computes `exp(logit - max) / sum` over the selected experts, and the CUDA fused/unfused paths in contrib_ops/cuda/moe/ft_moe/moe_kernel.cu apply full softmax then renormalize over the selected k when normalize=1 (the softmax denominator cancels). The Softmax/Sigmoid gates instead pass their activated probabilities via `router_weights` to bypass the op's internal softmax-over-selected. No routing math changes; adds a clarifying docstring only. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/components/_moe.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 6d636eab..7e873143 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -37,7 +37,27 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return routing_weights, selected_experts def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): - """Return selection and aggregation tensors for QMoE.""" + """Return selection and aggregation tensors for QMoE. + + Raw router logits are intentionally passed as ``router_probs`` with + ``router_weights=None`` and ``normalize_routing_weights=1``. In this + configuration com.microsoft::QMoE reproduces ``TopKGate.forward`` + exactly: it selects the top-k experts by logit value (softmax is + monotonic, so argmax over logits == argmax over softmax) and, because + ``router_weights`` is not supplied, computes the aggregation weights as + a softmax over the *k selected* logits. That is precisely + ``Softmax(TopK(router_logits))``. + + The other gates (Softmax/Sigmoid) must instead pass their already + activated probabilities as ``router_weights`` to bypass the op's + internal softmax-over-selected (otherwise QMoE would apply softmax on + top of an existing softmax/sigmoid). See the QMoE ``router_weights`` + semantics in ORT (contrib_ops/.../moe_quantization_cpu.cc and + ft_moe/moe_kernel.cu): when ``router_weights`` is omitted the kernel's + default path is ``softmax(top_k(router_probs))``; when provided it + gathers the given weights at the selected experts and (optionally) + renormalizes them. + """ weight_t = op.Transpose(self.weight, perm=[1, 0]) router_logits = op.MatMul(hidden_states, weight_t) router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) From 5b0903c5bc1c5289bbaaa180df57bdd61801b485 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 24 Jul 2026 17:12:50 +0000 Subject: [PATCH 08/11] test(gqa): make DeepSeek-V2-Lite int4 QMoE guard run in offline CI `test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer` was marked `@pytest.mark.arch_validation` but lives under `src/`, so it ran in no CI job: main CI excludes `arch_validation` and nightly L2 only scans `tests/arch_validation_test.py`. It also downloaded the config from HuggingFace (`AutoConfig.from_pretrained("deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True)`). Make it a self-contained regression guard: build a tiny DeepSeek-V2-Lite-shaped model from an inline `ArchitectureConfig` (MLA attention, first_k_dense_replace, softmax routing, shared experts) with no network access, and drop the `arch_validation` marker so it runs in the offline main CI. hidden_size and moe_intermediate_size are multiples of the quantization group_size (128) as QMoE requires; small MLA/dense dims are fine for MatMulNBits (ceil-padded K). The invariant is unchanged: one com.microsoft::QMoE per routed MoE layer, shared experts stay dense MatMulNBits, QMoE float routing inputs (1/3/6/14), and no `.moe.experts.` initializers leak into the graph. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../_group_query_attention_test.py | 78 +++++++++++++------ 1 file changed, 56 insertions(+), 22 deletions(-) diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index 3535d899..6d95812d 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -61,6 +61,48 @@ v_head_dim=8, ) +# Synthetic DeepSeek-V2-Lite-shaped config for the int4 QMoE regression guard. +# Mirrors the real architecture (MLA attention + first_k_dense_replace, softmax +# routing, shared experts) at tiny dimensions so the test runs fully offline with +# no HuggingFace download. hidden_size and moe_intermediate_size are multiples of +# the quantization group_size (128) because QMoE requires divisibility, while the +# small MLA/dense dimensions are fine for MatMulNBits (which ceil-pads the K axis). +# Layer 0 is a dense MLP; layers 1..2 are routed MoE layers -> 2 QMoE, 3 Attention. +_DEEPSEEK_V2_LITE_INT4_CONFIG = ArchitectureConfig( + hidden_size=128, + intermediate_size=256, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=16, + num_hidden_layers=3, + vocab_size=128, + max_position_embeddings=32, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + dtype=ir.DataType.FLOAT16, + q_lora_rank=16, + kv_lora_rank=16, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=8, + num_local_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=128, + n_shared_experts=1, + first_k_dense_replace=1, + scoring_func="softmax", + norm_topk_prob=False, + routed_scaling_factor=1.0, + quantization=QuantizationConfig( + bits=4, + group_size=128, + quant_method="gptq", + sym=True, + ), +) + # Tiny qwen3 config: has QK norm, weights NOT packable _QWEN3_CONFIG = ArchitectureConfig( hidden_size=64, @@ -255,24 +297,18 @@ def test_mla_gqa_fusion_requires_equal_kv_head_dimensions( assert counts.get("Attention", 0) == expected_attention assert counts.get("GroupQueryAttention", 0) == expected_gqa - @pytest.mark.arch_validation def test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer(self): - """The int4 graph keeps shared experts dense and routed experts fused.""" - from transformers import AutoConfig + """The int4 graph keeps shared experts dense and routed experts fused. - hf_config = AutoConfig.from_pretrained( - "deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True - ) - config = dataclasses.replace( - ArchitectureConfig.from_transformers(hf_config), - dtype=ir.DataType.FLOAT16, - quantization=QuantizationConfig( - bits=4, - group_size=128, - quant_method="gptq", - sym=True, - ), - ) + Self-contained regression guard: builds a tiny DeepSeek-V2-Lite-shaped + model from an inline config (no HuggingFace download) so it runs in the + offline main CI. Layer 0 is a dense MLP and layers 1..2 are routed MoE + layers, so exactly one com.microsoft::QMoE is emitted per routed MoE + layer while the shared experts remain dense MatMulNBits. + """ + config = _DEEPSEEK_V2_LITE_INT4_CONFIG + num_layers = config.num_hidden_layers + num_moe_layers = num_layers - config.first_k_dense_replace with pytest.warns(UserWarning, match="GQA fusion expected"): model = build_from_module( registry.get("deepseek_v3")(config), @@ -281,9 +317,9 @@ def test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer(self): )["model"] counts = count_ops(model) - assert counts.get("QMoE", 0) == 26 + assert counts.get("QMoE", 0) == num_moe_layers assert counts.get("GroupQueryAttention", 0) == 0 - assert counts.get("Attention", 0) == 27 + assert counts.get("Attention", 0) == num_layers for node in (node for node in model.graph if node.op_type == "QMoE"): assert node.inputs[1] is not None assert node.inputs[1].dtype == ir.DataType.FLOAT @@ -301,16 +337,14 @@ def test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer(self): if value is not None and value.name is not None ] assert not any(".moe.experts." in name for name in input_names) - for layer_idx in range(1, 27): + for layer_idx in range(config.first_k_dense_replace, num_layers): prefix = f"model.layers.{layer_idx}.mlp.shared_experts." shared_matmuls = [ node for node in model.graph if node.op_type == "MatMulNBits" and any( - value is not None - and value.name is not None - and prefix in value.name + value is not None and value.name is not None and prefix in value.name for value in node.inputs ) ] From cf3e35f06b93874ae7655f5b7a03f381d3f74dd6 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 24 Jul 2026 17:30:33 +0000 Subject: [PATCH 09/11] docs(moe): correct QMoE routing semantics Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/components/_moe.py | 30 +++++++++++++----------------- 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 7e873143..29cf4c12 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -40,23 +40,19 @@ def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): """Return selection and aggregation tensors for QMoE. Raw router logits are intentionally passed as ``router_probs`` with - ``router_weights=None`` and ``normalize_routing_weights=1``. In this - configuration com.microsoft::QMoE reproduces ``TopKGate.forward`` - exactly: it selects the top-k experts by logit value (softmax is - monotonic, so argmax over logits == argmax over softmax) and, because - ``router_weights`` is not supplied, computes the aggregation weights as - a softmax over the *k selected* logits. That is precisely - ``Softmax(TopK(router_logits))``. - - The other gates (Softmax/Sigmoid) must instead pass their already - activated probabilities as ``router_weights`` to bypass the op's - internal softmax-over-selected (otherwise QMoE would apply softmax on - top of an existing softmax/sigmoid). See the QMoE ``router_weights`` - semantics in ORT (contrib_ops/.../moe_quantization_cpu.cc and - ft_moe/moe_kernel.cu): when ``router_weights`` is omitted the kernel's - default path is ``softmax(top_k(router_probs))``; when provided it - gathers the given weights at the selected experts and (optionally) - renormalizes them. + ``router_weights=None`` and ``normalize_routing_weights=1``. QMoE + selects the top-k by raw value (monotonic activations preserve that + selection) and, on this path, gives the selected logits softmax + weights. This matches ``TopKGate.forward``: + ``Softmax(TopK(router_logits))`` on both CPU and CUDA EPs. + + ORT's CPU QMoE honors Input 14 (``router_weights``) by gathering it at + the selected experts, but CUDA QMoE ignores Input 14 and always uses + softmax-top-k on ``router_probs``; see + ``contrib_ops/cpu/moe/moe_quantization_cpu.cc`` and + ``contrib_ops/cuda/moe/moe_quantization.cc``. Thus ``None`` is correct + here on both EPs, while activated probabilities used by Softmax/Sigmoid + gates are correct on CPU but double-activated on CUDA. """ weight_t = op.Transpose(self.weight, perm=[1, 0]) router_logits = op.MatMul(hidden_states, weight_t) From a387ee2ce7beb63b29b030e236c0ae5072479798 Mon Sep 17 00:00:00 2001 From: Rico Date: Fri, 24 Jul 2026 23:47:03 +0000 Subject: [PATCH 10/11] fix(glm4-gptq): resolve lint failures and QMoE routing consistency - Replace invalid '# ruff:ignore[unused-import]' directives in _configs/_base.py with the repo-standard '# noqa: F401' for the per_model side-effect imports (fixes RUF102 + F401 Lint failures). - Apply ruff-format to PR-touched files (_moe_test, _weight_utils_test, deepseek, _moe) that had line-length reflow violations. - SoftmaxTopKGate/SigmoidTopKGate.qmoe_routing now pass raw float32 logits as the QMoE router_probs (input 1) and the activated probabilities as router_weights (input 14). CUDA QMoE ignores router_weights and applies softmax-top-k on router_probs, so passing activated probs there double-activated (softmax-of- softmax / softmax-of-sigmoid); logits reproduce forward on both EPs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_configs/_base.py | 4 +-- src/mobius/_weight_utils_test.py | 4 +-- src/mobius/components/_moe.py | 38 +++++++++++++++++++------- src/mobius/components/_moe_test.py | 43 ++++++++---------------------- src/mobius/models/deepseek.py | 4 +-- 5 files changed, 43 insertions(+), 50 deletions(-) diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index d1634e8e..68e5d3a5 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -295,7 +295,7 @@ def _extract_vision_config(config, parent_config, model_type: str) -> dict: hooks live under :mod:`mobius._configs.per_model` and are registered with :mod:`mobius._configs._extractors` at import time. """ - from mobius._configs import per_model # ruff:ignore[unused-import] - side effect + from mobius._configs import per_model # noqa: F401 - imported for registration side effect from mobius._configs._extractors import extract_vision_config as _dispatch return _dispatch(config, parent_config, model_type) @@ -308,7 +308,7 @@ def _extract_audio_config(config, parent_config, model_type: str) -> dict: hooks live under :mod:`mobius._configs.per_model` and are registered with :mod:`mobius._configs._extractors` at import time. """ - from mobius._configs import per_model # ruff:ignore[unused-import] - side effect + from mobius._configs import per_model # noqa: F401 - imported for registration side effect from mobius._configs._extractors import extract_audio_config as _dispatch return _dispatch(config, parent_config, model_type) diff --git a/src/mobius/_weight_utils_test.py b/src/mobius/_weight_utils_test.py index 4c1e6dda..a4e4b9e9 100644 --- a/src/mobius/_weight_utils_test.py +++ b/src/mobius/_weight_utils_test.py @@ -527,9 +527,7 @@ def test_expert_major_tensors_preserve_expert_axis(self): ), "experts.scales": torch.randn(num_experts, self.N_GROUPS, self.N), } - result = preprocess_gptq_weights( - sd, bits=self.BITS, group_size=self.GROUP_SIZE - ) + result = preprocess_gptq_weights(sd, bits=self.BITS, group_size=self.GROUP_SIZE) assert result["experts.weight"].shape == ( num_experts, diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 29cf4c12..cf4f5226 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -90,12 +90,24 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return routing_weights, selected_experts def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): - """Return selection and aggregation tensors for QMoE.""" + """Return selection and aggregation tensors for QMoE. + + ``router_probs`` (QMoE input 1) is the raw float32 logits and the + pre-softmaxed probabilities are passed as ``router_weights`` (input + 14). CUDA QMoE ignores ``router_weights`` and applies softmax-top-k on + ``router_probs``, so feeding logits reproduces ``forward`` exactly + (``Softmax`` then top-k then renormalize); feeding pre-softmaxed probs + here would double-softmax on CUDA. CPU QMoE selects the top-k over + ``router_probs`` (softmax is monotonic, so logits give the same + selection) and gathers ``router_weights`` at the selected experts, + renormalizing when ``normalize_routing_weights=1`` -- which equals + ``forward``'s renormalized top-k of the full softmax. + """ weight_t = op.Transpose(self.weight, perm=[1, 0]) router_logits = op.MatMul(hidden_states, weight_t) router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) routing_probs = op.Softmax(router_logits, axis=-1) - return routing_probs, routing_probs, self.norm_topk_prob, 1.0 + return router_logits, routing_probs, self.norm_topk_prob, 1.0 class SigmoidTopKGate(nn.Module): @@ -139,13 +151,23 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return routing_weights, selected_experts def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): - """Return selection and aggregation tensors for QMoE.""" + """Return selection and aggregation tensors for QMoE. + + The sigmoid-activated probabilities are passed as ``router_weights`` + (QMoE input 14) while the raw float32 logits are passed as + ``router_probs`` (input 1). CPU QMoE selects the top-k over + ``router_probs`` (sigmoid is monotonic, so logits give the same + selection) and gathers ``router_weights`` at the selected experts, + renormalizing per ``normalize_routing_weights``. Passing logits (rather + than the sigmoid probs) as ``router_probs`` avoids a softmax-of-sigmoid + on CUDA, which ignores ``router_weights``. + """ weight_t = op.Transpose(self.weight, perm=[1, 0]) router_logits = op.MatMul(hidden_states, weight_t) router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) routing_probs = op.Sigmoid(router_logits) return ( - routing_probs, + router_logits, routing_probs, self.norm_topk_prob, self.routed_scaling_factor, @@ -253,9 +275,7 @@ def __init__(self, config: ArchitectureConfig, gate: nn.Module | None = None): self.experts = None self._init_qmoe_parameters(expert_config) else: - self.experts = nn.ModuleList( - [MLP(expert_config) for _ in range(self.num_experts)] - ) + self.experts = nn.ModuleList([MLP(expert_config) for _ in range(self.num_experts)]) def _init_qmoe_parameters(self, expert_config: ArchitectureConfig) -> None: quantization = self._qmoe_quantization @@ -274,9 +294,7 @@ def _init_qmoe_parameters(self, expert_config: ArchitectureConfig) -> None: [self.num_experts, fc1_out, hidden_size * bits // 8], dtype=ir.DataType.UINT8, ) - self.fc1_scales = nn.Parameter( - [self.num_experts, fc1_out, hidden_size // block_size] - ) + self.fc1_scales = nn.Parameter([self.num_experts, fc1_out, hidden_size // block_size]) self.fc1_scales._keep_float32 = True self.fc2_experts_weights = nn.Parameter( [self.num_experts, hidden_size, intermediate_size * bits // 8], diff --git a/src/mobius/components/_moe_test.py b/src/mobius/components/_moe_test.py index 0f3e9ab7..16c3f8df 100644 --- a/src/mobius/components/_moe_test.py +++ b/src/mobius/components/_moe_test.py @@ -167,31 +167,15 @@ def test_expert_major_packing_matches_static_64_expert_top6_reference(self): torch.manual_seed(0) num_experts, top_k = 64, 6 hidden_size, intermediate_size, block_size = 32, 16, 16 - fc1_codes = torch.randint( - 0, 16, (num_experts, 2 * intermediate_size, hidden_size) - ) - fc2_codes = torch.randint( - 0, 16, (num_experts, hidden_size, intermediate_size) - ) - fc1_scales = torch.rand( - num_experts, 2 * intermediate_size, hidden_size // block_size - ) - fc2_scales = torch.rand( - num_experts, hidden_size, intermediate_size // block_size - ) + fc1_codes = torch.randint(0, 16, (num_experts, 2 * intermediate_size, hidden_size)) + fc2_codes = torch.randint(0, 16, (num_experts, hidden_size, intermediate_size)) + fc1_scales = torch.rand(num_experts, 2 * intermediate_size, hidden_size // block_size) + fc2_scales = torch.rand(num_experts, hidden_size, intermediate_size // block_size) raw = { - "model.layers.1.mlp.experts.gate_up_proj.qweight": _to_gptq_qweight( - fc1_codes - ), - "model.layers.1.mlp.experts.gate_up_proj.scales": fc1_scales.transpose( - -1, -2 - ), - "model.layers.1.mlp.experts.down_proj.qweight": _to_gptq_qweight( - fc2_codes - ), - "model.layers.1.mlp.experts.down_proj.scales": fc2_scales.transpose( - -1, -2 - ), + "model.layers.1.mlp.experts.gate_up_proj.qweight": _to_gptq_qweight(fc1_codes), + "model.layers.1.mlp.experts.gate_up_proj.scales": fc1_scales.transpose(-1, -2), + "model.layers.1.mlp.experts.down_proj.qweight": _to_gptq_qweight(fc2_codes), + "model.layers.1.mlp.experts.down_proj.scales": fc2_scales.transpose(-1, -2), } packed = pack_qmoe_expert_weights( preprocess_gptq_weights(raw, bits=4, group_size=block_size) @@ -239,17 +223,13 @@ def _to_gptq_qweight(codes: torch.Tensor) -> torch.Tensor: return packed.transpose(-1, -2).contiguous() -def _dequant_codes( - codes: torch.Tensor, scales: torch.Tensor, block_size: int -) -> torch.Tensor: +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_qmoe( - packed: torch.Tensor, scales: torch.Tensor, block_size: int -) -> torch.Tensor: +def _dequant_qmoe(packed: torch.Tensor, scales: torch.Tensor, block_size: int) -> torch.Tensor: low = packed & 0x0F high = packed >> 4 codes = torch.stack((low, high), dim=-1).flatten(-2) @@ -271,8 +251,7 @@ def _static_moe( activated *= projected[:, intermediate_size:] expert_output = activated @ fc2[expert_idx].T weight = ( - routing_weights - * (selected_experts == expert_idx).to(routing_weights.dtype) + routing_weights * (selected_experts == expert_idx).to(routing_weights.dtype) ).sum(dim=-1, keepdim=True) result += expert_output * weight return result diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index 62f27f1f..3d1713aa 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -35,9 +35,7 @@ def _linear_factory(config: ArchitectureConfig): quantization = config.quantization if quantization is None or quantization.quant_method == "none": return None - zero_point_dtype = ( - config.dtype if quantization.float_zero_point else ir.DataType.UINT8 - ) + zero_point_dtype = config.dtype if quantization.float_zero_point else ir.DataType.UINT8 return make_quantized_linear_factory( bits=quantization.bits, block_size=quantization.group_size, From 98153fd4f2dcc8008cf50a976b9abfebc239295d Mon Sep 17 00:00:00 2001 From: Rico Date: Sat, 25 Jul 2026 00:45:16 +0000 Subject: [PATCH 11/11] docs(moe): document MoELayer's conditional QMoE vs loop-over-experts dispatch MoELayer now switches to a fused com.microsoft::QMoE path (experts=None) when the quantization config matches the native QMoE ABI and the gate implements qmoe_routing, otherwise it uses the loop-over-experts MLP dispatch. Update the class docstring to describe both paths so the weight/preprocess expectations are clear when debugging. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/components/_moe.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index cf4f5226..71ae1f59 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -250,8 +250,16 @@ class MoELayer(nn.Module): Routes each token to top-k experts via a gating mechanism, applies each expert MLP, and accumulates weighted results. - Uses loop-over-experts dispatch: each expert processes all tokens, - then results are masked and weighted. + Two dispatch paths are selected at construction time: + + - Loop-over-experts (default): each expert ``MLP`` processes all + tokens, then results are masked and weighted. Used when the model + is unquantized or the gate has no ``qmoe_routing`` hook. + - Fused ``com.microsoft::QMoE`` (``experts=None``): used when the + quantization config matches the native QMoE ABI + (:func:`_supported_qmoe_quantization`) and the gate implements + ``qmoe_routing``. Expert weights are packed into quantized + ``fc1``/``fc2`` parameters instead of per-expert ``MLP`` modules. """ def __init__(self, config: ArchitectureConfig, gate: nn.Module | None = None):